Asymmetric EMA for Federated LoRA: Feed Back A, Not B
Abstract
Federated LoRA can suffer from majority‑class collapse in the high‑heterogeneity regime (Dirichlet ), and whether a run collapses often depends on the random seed. The two LoRA factors multiply, so per‑round errors in either factor compound from one round to the next. Applying a time‑axis exponential moving average (EMA) over rounds is a natural remedy, but the choice of which factor to feed back is not obvious: feeding a factor back keeps only a fraction of its per‑round displacement, while not feeding it back lets the gap between the EMA and the training iterate accumulate across rounds. We show that both effects are exact server‑side identities that hold for any LoRA rank, local optimizer, and number of local steps, and that this common cost acts asymmetrically. Drift is most damaging for the direction factor , whereas the displacement squeeze is most damaging for the scale factor . We propose Stabilized A-Feedback (), which feeds only the direction factor back through the EMA, keeps the scale factor on the path, and adds one line of code with no new hyperparameters. Under a rank‑1, single‑step SGD model we quantify four consequences: a residual error that does not vanish as the round budget grows when ‑feedback is absent, standard SGD‑style convergence of , an optimization‑bias penalty when is fed back that cannot be recovered by retuning the learning rate within the stability constraint, and a heterogeneity crossover below which wins. Across four federated LoRA methods and six tasks (four GLUE and two MedMNIST; 24 five‑seed cells), outperforms both and EMA in 22 cells and reduces the MNLI seed standard deviation from up to to at most ; ablations that hold the deployed model fixed separate training‑time ‑feedback from deployment‑time averaging. We also report where the method stops: at 355M outperforms on all three seeds but its ordering against EMA is seed‑dependent, and stacking with FedSA‑LoRA is structurally compatible but not shown to be additive.
Then back it, or bet against it.
Related papers
Open the market on this paper to see 7 more related papers.