AEGIS seals gradient leaks in federated LLM fine-tuning

Problem

The paper studies how gradients shared during federated fine-tuning of transformer language models leak private training text. It shows that transformer gradients expose distinct algebraic channels that reveal token identity: a low-rank subspace in the attention output-projection gradient, a row-norm sparsity pattern in the embedding gradient that directly signals which tokens are present, and a recoverable low-rank signal in the MLP expansion gradient. State-of-the-art analytical attacks exploit these channels to recover near-exact token sets in seconds, undermining the privacy assumptions of collaborative training.

Approach

The authors introduce AEGIS, a client-side, architecture-agnostic defence that closes all three analytical channels via three backward-pass operations applied after backpropagation and before the optimiser step. The operations are: freeze attention projection parameters so their gradients are identically zero (neutralising the attention subspace channel); replace the sparse embedding-gradient matrix with a dense, calibrated Gaussian uniformisation that removes the zero-vs-nonzero row indicator while retaining a controlled fraction of the learning signal; and apply an analogous per-block Gaussian flooding to the MLP expansion gradients to submerge any low-rank SVD signal beneath the random-matrix noise bulk. The same masked gradients are used locally and exported to the server, preserving training while denying a clean export. Formal guarantees are provided for the attention freeze and embedding uniformisation; the MLP flooding is motivated by random-matrix design criteria and validated empirically. Evaluation covers 11 models across four architecture families and six datasets, and includes adaptive adversaries that retarget the MLP channel.

Key Findings

• AEGIS collapses analytical token recovery to near zero: across 11 models and six datasets, DAGER ROUGE-1 falls from undefended levels of at least 0.87 (and typically ≈0.98) to at most 0.005, a relative reduction of at least 99% against the considered attacks.
• Utility is preserved or improved: test perplexity for defended models is at or below the undefended baseline in the reported experiments, with some checkpoints showing modest improvements in held-out PPL and unchanged classification metrics.
• Triple-channel closure is necessary: an ablation that froze attention and flooded embeddings but omitted MLP flooding left a recoverable MLP channel on larger decoders; adding MLP uniformisation restored near-zero recovery without harming utility.
• Robustness versus attack classes: AEGIS eliminates closed-form analytical attacks that rely on the three channels; optimisation-based gradient-matching attacks are degraded but not fully defeated, indicating complementary mechanisms are required to address those.
• Low overhead: the defence adds modest runtime cost (measured step-time multipliers ≈1.10–1.22×) and can reduce peak memory and dense payload sizes by freezing attention gradients, though embedding flooding makes the exported gradient dense.

Limitations

AEGIS does not provide formal differential-privacy (ε,δ) guarantees, and it does not fully defeat optimisation-based inversion attacks, robust-PCA denoisers, multi-step aggregation across rounds, colluding clients, or an active malicious server that modifies weights. Evaluation is limited to a single-client, single-step FedSGD threat model with one fine-tuning seed per cell, and formal lower bounds for the MLP denoising problem are left to future work. Flooding increases dense gradient payloads for large vocabularies, which may affect some deployments.

Implications

Offensive implications are clear: without structural defences, an adversary that observes cleartext client gradients can recover private tokens quickly by exploiting attention, embedding and MLP gradient channels. AEGIS substantially raises the bar for such attackers by removing the analytical shortcuts; an attacker must instead resort to more expensive or different strategies such as iterative gradient-matching, aggregation across rounds, active model manipulation, or advanced denoising methods. However, those alternative avenues remain feasible and warrant continued adversarial development and combined defences (for example composition with DP-SGD or secure aggregation) to defend against adaptive, persistent, or stronger threat models.