How Much Regularization Survives Averaging? Update Masking in Federated Learning
2026-08-24 • Machine Learning
Machine LearningArtificial Intelligence
AI summaryⓘ
The authors study a method to help federated learning work better when data varies a lot across users, aiming for solutions that work well for everyone (flat minima). They explore a known technique called masking that adds noise to updates and helps avoid sharp minima but find it behaves differently when applied to federated learning. They show that if each user uses a different mask, the beneficial effect is reduced by the number of users, but using the same mask partially restores it depending on how diverse the users' data gradients are. Their experiments on image classification indicate that typical setups reduce this effect a lot, and those that enhance it train poorly.
federated learningnon-IID dataflat minimasharpness-aware minimizationmaskingparameter updatesgradient diversityfederated averagingCIFAR-10noise regularization
Authors
Wenhao Yan, Fu Kuroda, Yucheng Jin, Zhenke Chen
Abstract
Federated learning on non-IID data seeks flat minima to generalize across clients, and existing methods borrow sharpness-aware minimization from centralized training. There is a second way to reach flat minima, in which the regularization comes for free from noise added to the parameter updates, and it has never been carried over to the federated setting. We show the reason. Masking charges the optimizer for moving in sharp directions. We prove that when each client draws its own mask, federated averaging weakens that charge by exactly the cohort size, and that giving every client the same mask brings it back by a factor equal to the inverse gradient diversity of the cohort. In our experiment setting on CIFAR-10, that factor is 1.19 out of a possible 10. Turning off minibatch sampling raises it to 8.96, while changing data heterogeneity a hundredfold leaves it between 1.17 and 1.50. The configurations keeping the regularization train far too poorly to use.