PaperlayerView exact revision on arXiv

Source-linked AI summary

How Much Regularization Survives Averaging? Update Masking in Federated Learning

Wenhao Yan, Fu Kuroda, Yucheng Jin, Zhenke Chen

arXiv:2608.23286v2cs.LGcs.AI

TL;DR

Federated learning seeks flat minima for generalization across non-IID clients, but update masking has not served as an implicit federated regularizer. The paper analyzes masking under FedAvg and finds that independent masks dilute its curvature penalty by cohort size, while synchronized masks recover it according to inverse gradient diversity; in CIFAR-10, recovery is typically weak and costly to retain.

  • Problem

    Federated learning needs generalization across non-IID client distributions, while update-noise regularization from centralized training has not been carried over as an implicit federated regularizer.

  • Method

    The paper derives how FedAvg changes masking's variance-induced curvature penalty under independent and synchronized client masks, then measures recovery across local-training configurations.

  • Results

    Independent masks attenuate the penalty by cohort size, while synchronized masks recover it by inverse gradient diversity; on CIFAR-10, the factor is 1.19 versus a ceiling of 10, rising to 8.96 with full-batch training.

  • Takeaways & Limitations

    Minibatch sampling noise, rather than data heterogeneity, holds recovery near one, and configurations recovering the regularization train too poorly to use.

  • Takeaways & Limitations

    The analysis relies on a small-update second-order expansion, uses a Euclidean norm instead of the exact Hessian-weighted norm, and covers one architecture and dataset.

Abstract

from arXiv · show

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 as an implicit regularizer. 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 thousandfold leaves it between 1.17 and 1.50. The configurations keeping the regularization train far too poorly to use.

I. INTRODUCTION

Federated learning on non-IID client data seeks flat minima for cross-distribution generalization, but update-noise regularization has not transferred as an implicit federated regularizer. The paper analyzes why, within FedAvg's client-and-server training setup.

  • I. INTRODUCTION: FL must generalize across data distributions that clients never jointly observe.Existing approaches pursue flat minima and import sharpness-aware minimization from centralized training.
  • I. INTRODUCTION: Update masking induces an implicit curvature penalty through variance while discarding and rescaling random parameter blocks.The paper distinguishes this update-dynamics route from explicitly perturbing the loss function.
  • I. INTRODUCTION: The paper argues that the absence of masking as a federated implicit regularizer is explained by averaging rather than oversight.This frames the study's central problem.
  • I. INTRODUCTION: Under independent per-client masks, FedAvg attenuates masking's curvature penalty by exactly the cohort size.The paper derives the resulting penalty for this configuration.
  • I. INTRODUCTION: Synchronizing masks restores the penalty by the reciprocal of cohort gradient diversity, reaching the cohort size at peak but falling below one when updates cancel.The contribution summary states that synchronization can then hurt model performance.

B. Update Masking

Block-wise masking creates unbiased noisy updates whose variance induces a curvature penalty, while gradient diversity measures how client effort survives aggregation. The reciprocal diversity can range from cohort-size amplification to less than one under cancellation.

  • B. Update Masking: Block-wise masking replaces a dense update with independently sampled, rescaled parameter blocks.The 1/p factor keeps the masked update unbiased, and the masking granularity determines which curvature terms survive.
  • B. Update Masking: Masking's variance produces a curvature penalty whose structure loses cross-coordinate terms under independent block sampling.The paper focuses on masking whole parameter tensors.
  • C. Connection with Gradient Diversity: The reciprocal gradient diversity 1/∆S measures how much individual client effort survives aggregation.It equals Kc when clients move together and falls below one when their updates conflict.
  • C. Connection with Gradient Diversity: 1/∆S ≤ Kc, with equality when all client updates are equal; it can fall below one when updates point in conflicting directions and largely cancel.Gradient diversity itself is at least 1/Kc, but neither quantity has a finite bound on the opposite side.
  • C. Connection with Gradient Diversity: FedExP's server step size is exactly related to gradient diversity for any number of local steps.Its rearranged step size is (Kc/2) · ∆S, so the server takes larger steps when the cohort disagrees.

III. ANALYSIS

The analysis restates masking's centralized curvature penalty and explains its mechanism through update variance and local sharpness. The penalty charges movement into narrow directions, with coefficient 1−p over 2p.

  • III. ANALYSIS: The analysis extends the centralized masking penalty from one worker's optimizer step to round-level client updates averaged by the server.This establishes the object whose survival under federated aggregation is analyzed.
  • III. ANALYSIS: Masking's excess over the unmasked step originates entirely in variance because the masked update is unbiased.The analysis expresses this excess through the Hessian and masked-update covariance.
  • III. ANALYSIS: The quadratic form measures how sharply the loss curves along each masked block's update direction.It is large in a narrow valley and small along a flat one.
  • III. ANALYSIS: Independent block masks leave only diagonal Hessian blocks in expectation, so masking charges the optimizer for moving in sharp directions.This is the paper's curvature interpretation of the induced regularization.
  • III. ANALYSIS: The coefficient (1−p)/2p sets the masking price, vanishing as p approaches 1 and increasing as more update is discarded.The coefficient directly controls the strength of the curvature penalty.

A. Independent Masks

Independent client masks preserve each client's local masking penalty but dilute the penalty evaluated at the aggregated model. This dilution is exactly a further factor of the cohort size because averaging is linear while the penalty is quadratic.

  • A. Independent Masks: The natural independent-mask configuration lets each client draw its own mask.The paper then evaluates the resulting penalty at the aggregated model.
  • A. Independent Masks: Independent masks across clients and blocks define the aggregated curvature penalty Rind.This is the configuration formalized in Proposition 2.
  • A. Independent Masks: Client-mask independence makes aggregated-update covariance scale with 1/Kc^2 while summing Kc individual contributions.Each contribution corresponds to a client's masking penalty.
  • A. Independent Masks: The resulting aggregated penalty is the cohort average of single-worker penalties, attenuated by an additional factor Kc.Each client still incurs its full penalty along its own local trajectory.
  • A. Independent Masks: Averaging contributes a squared coefficient to the quadratic penalty but only Kc terms to the sum.This explains the exact cohort-size dilution.

B. Synchronized Masks

Sharing one mask across clients prevents masking noise from averaging away, restoring the curvature penalty on the aggregated update. The recovery factor is the reciprocal of cohort gradient diversity and depends on the cohort’s update alignment.

  • B. Synchronized Masks: Broadcasting one mask seed per round removes the averaging that dilutes independently masked updates.Every client applies the same block mask, so the aggregated update retains the covariance of a single masked update.
  • B. Synchronized Masks: The synchronized penalty acts on the aggregated direction rather than each client’s local trajectory.This targets the update direction that the federated server actually uses for the model update.
  • B. Synchronized Masks: Synchronizing the mask recovers the curvature penalty by 1/∆S, the reciprocal gradient diversity of the cohort.The comparison assumes the Hessian is isotropic within blocks so quadratic forms reduce to squared norms.
  • B. Synchronized Masks: Recovery reaches at most the cohort size and vanishes when cohort updates are mutually orthogonal.When updates cancel sufficiently, 1/∆S < 1, so synchronized masking can perform worse than independent masking.

D. Signal-Noise Decomposition

The signal-noise decomposition explains recovery-factor values by separating a shared update component from client-specific noise. Shared updates add coherently, whereas independent noise adds incoherently, making the factor a signal-to-noise ratio in disguise.

  • D. Signal-Noise Decomposition: Each client update is decomposed into a shared component and client-specific noise with zero mean.The noise terms are assumed independent across clients and to have common variance.
  • D. Signal-Noise Decomposition: The shared component adds coherently across clients, while noise adds incoherently and grows more slowly with cohort size.This difference produces the gap between synchronized and independently masked aggregation.
  • D. Signal-Noise Decomposition: The recovery factor approaches Kc when the shared-to-noise ratio is large and approaches 1 when it is small.The paper identifies the factor as a signal-to-noise ratio in disguise.
  • D. Signal-Noise Decomposition: Server-side statistics can separate the shared component from client noise using quantities already formed during aggregation.The estimators require Kc squared norms of client updates and no additional communication.
  • D. Signal-Noise Decomposition: The estimators assume independent client noise with comparable magnitude, which is considered reasonable for clients holding similar amounts of data.They distinguish a low recovery factor caused by a weak shared direction from one caused by large client noise.

A. Setup

The CIFAR-10 study sweeps local batch size and local epochs under FedAvg, revealing that minibatch sampling—not simply the number of local steps—largely determines recovery-factor collapse.

  • Experimental setup: The experiment trains ResNet-18 with GroupNorm on CIFAR-10 across 100 Dirichlet-partitioned clients, sampling 10 clients per round for 200 rounds.Each client performs τ local SGD epochs, and the server uses an n_k-weighted update mean.
  • Experimental setup: The sweep varies batch sizes 32, 128, and full, local epochs τ ∈ {1, 2, 5}, and separately Dirichlet concentration α ∈ {100, 1, 0.1}.The local-step range is approximately T = 1 to T = 80, with batch size and local-step count coupled.
  • Recovery-factor results: 8.96, 4.13, and 3.96 are the recovery factors for the three full-batch configurations, whereas every minibatch configuration lies between 1.10 and 1.36.The results form a cliff with an empty middle range rather than a smooth gradient.
  • Recovery-factor results: 3.96 versus 1.36 shows that five full-batch steps can recover far more regularization than four steps at batch size 128, despite similar accuracy.The corresponding accuracies are 0.493 and 0.449, so training progress does not explain the gap.
  • Recovery-factor results: Increasing τ from 1 to 5 changes recovery from 1.24 to 1.10 at batch size 32, but from 8.96 to 3.96 at full batch.The mechanism is most sensitive when gradient noise is smallest.
  • Recovery-factor results: 4.5% of rounds fall below one for batch size 128 with τ = 1, while that regime disappears in the three noisiest cells.Below one, synchronized masks perform worse than independent masks.

C. Decomposition of the Recovery Factor

The recovery-factor decomposition separates shared cohort motion from client-specific noise and shows that sampling noise, rather than training progress alone, drives the collapse.

  • Signal and noise: Noise variance spans 4.9×10^-4 to 6.0×10^-1 across configurations, while the shared component spans only 6.2×10^-4 to 7.3×10^-3.More local steps increase both shared motion and accumulated noise, but noise grows much faster.
  • Separating training progress: 0.493 accuracy for full batch with τ = 5 lies between 0.449 and 0.534 for batch size 128 with τ = 1 and 2, yet its recovery factor is three times higher.Its noise variance is more than twofold lower than both neighboring configurations.
  • Accuracy trade-off: The recovery-factor-versus-accuracy plot contains a large empty middle gap, so moving toward better recovery costs substantial accuracy without a smooth trade-off.Only one configuration is dominated, and most regularization remains absent even after the costly move.

E. Effect of Data Heterogeneity

Changing data heterogeneity over three orders of magnitude barely changes the recovery factor, while the analysis bounds the interpretation to its approximation, norm, and evaluation scope.

  • Heterogeneity results: At α = 100 and α = 1, the unweighted recovery factors are indistinguishable, while α = 0.1 rises to 1.50.The α = 0.1 partition also produces lower accuracy, partly confounding the comparison with training progress.
  • Aggregation conventions: Under n_k-weighted aggregation, all three heterogeneity settings remain near 1.18.The two aggregation conventions diverge mainly at α = 0.1, where shard sizes are highly uneven.
  • Heterogeneity results: Over three orders of magnitude in α, the recovery factor ranges only from 1.17 to 1.50, compared with 1.19 to 8.96 when the local protocol changes.The paper therefore identifies heterogeneity as not the main factor holding recovery down.
  • Limitations: The analysis assumes small updates, uses a Euclidean norm because the global Hessian is unavailable, and treats the penalty as indicative rather than exact for multi-step minibatch training.The exact corollary requires the Hessian-weighted norm.
  • Limitations: A recovery factor of 1.19 measures surviving regularization, not whether masking would improve performance, and all measurements use one architecture and dataset.The paper explicitly separates survival of the penalty from its practical value.
  • Relation to prior methods: Mean-shifting perturbations such as sharpness-aware minimization pass through averaging undiminished, unlike zero-mean independent perturbations such as masking and additive gradient noise.This provides a supported explanation for why some flatness-seeking methods transfer more successfully to federated learning.
Loading 2608.23286v2…