Research map / Federated learning and privacy / Communication-efficient FL
Thinking Forward: Memory-Efficient Federated Finetuning of Language Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Spry, a federated LLM finetuning algorithm that splits LoRA trainable layers across clients and has each client estimate gradients with forward-mode AD, cutting memory versus backpropagation and improving accuracy over zero-order federated methods.
Problem
Federated finetuning of LLMs with backpropagation needs excessive memory, mainly for intermediate activations, which overwhelms edge devices; existing zero-order alternatives converge slowly and lose accuracy.
Approach
Assign each client a small subset of trainable LoRA layers so forward-mode AD (a single jvp per batch with random perturbations) gives gradient estimates closer to the true gradient; server aggregates the per-layer weights and applies FedYogi.
Key results
- Spry reduces peak memory by 1.4-7.1x vs backpropagation; Llama2-7B peak memory falls from 33.9GB to 6.2GB with LoRA.
- Spry achieves 5.15-13.50% higher generalized accuracy than the best zero-order federated method across datasets, within 0.60-6.16% of backpropagation methods.
- Spry converges 1.15-1.59x faster than FwdLLM+ and 6.16-20.28x faster than Baffle+.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Spry instead of the strongest zero-order baseline | RoBERTa Large and LLMs · Heterogeneous classification and SQuADv2 tasks | generalized accuracy / F1 score: 5.15--13.50% | Helped |
| Use Spry instead of backpropagation-based federated optimizers | RoBERTa Large and LLMs · Heterogeneous classification and SQuADv2 tasks | generalized accuracy / F1 score: 0.60--6.16% | Hurt |
| Use Spry forward-mode AD instead of backpropagation | RoBERTa Large, Llama2-7B, OPT6.7B, OPT13B · Training peak memory | peak memory consumption: 27.90--86.26% | Helped |
| Use Spry instead of zero-order methods | RoBERTa Large and LLMs · Time to convergence | time to convergence: 1.15-1.59×, 6.16-20.28×, and 1.34-2.98× | Helped |
| Increase communication frequency from per-epoch to per-iteration (ablation) | Spry · Accuracy performance | accuracy: 4.47% | Helped |
| Increase perturbations per batch from K=1 to K=100 (ablation) | Spry · Generalized test accuracy | accuracy: 1.1% | Helped |
| Increase perturbations per batch from K=1 to K=10 (ablation) | Spry · Convergence | FL rounds to steady state: 200; 500 | Helped |
| Increase participating clients from C=10 to C=50 to C=100 (ablation) | Spry on SST2 · SST2 | accuracy: 85.14%; 86.56%; 88.08% | Helped |
| Increase participating clients from C=10 to C=50 to C=100 (ablation) | Spry on SST2 · SST2 | rounds to approximately 85% accuracy: 500, 450, and 150 rounds | Helped |
| Split trainable layers across clients in FedAvg (ablation) | FedAvgSplit · Federated finetuning | accuracy: 2.60-10.00% | Hurt |
| Remove weight splitting from forward-mode AD (ablation) | FedFgd · RoBERTa Large with LoRA | convergence | Hurt |
| Use LoRA instead of IA3, BitFit, or classifier-only finetuning (ablation) | Spry · Language-model finetuning | accuracy: 10.60-16.53% | Helped |
Surprising findings
- Naive forward-gradient FL (FedFgd) fails to converge as trainable weight count grows (e.g., RoBERTa Large with LoRA, 1.15M weights).
- Increasing the number of perturbations per batch K has little effect on final accuracy, mainly affecting convergence speed.
Limitations
- Forward-mode AD jvp computation in PyTorch is slower than backpropagation for billion-scale models, so no speedup there.
- Memory still higher than zero-order methods (1.54-1.96x) due to parallel evaluation and perturbed weight copies.
- Heterogeneous data introduces estimation bias; the method does not directly address data heterogeneity.
Experimental setup
Datasets: AG News SST2 Yelp Yahoo SNLI MNLI SQuADv2 MultiRC
Benchmarks: AG News SST2 Yelp Yahoo SNLI MNLI SQuADv2 MultiRC
Metrics: generalized accuracy personalized accuracy exact match F1 peak memory consumption time to convergence
Baselines: FedAvg FedYogi FedSgd FwdLLM FedMeZO Baffle
Builds on: LoRA FedAvg FedYogi MeZO forward gradients
Closest papers by meaning
- FlowerTune: A Cross-Domain Benchmark for Federated Fine-Tuning of Large Language Models (NeurIPS 2025 · similarity 0.90)
- FLoRA: Federated Fine-Tuning Large Language Models with Heterogeneous Low-Rank Adaptations (NeurIPS 2024 · similarity 0.89)
- Federated Full-Parameter Tuning of Billion-Sized Language Models with Communication Cost under 18 Kilobytes (ICML 2024 · similarity 0.89)
- FedLLM-Bench: Realistic Benchmarks for Federated Learning of Large Language Models (NeurIPS 2024 · similarity 0.89)
- Splitting with Importance-aware Updating for Heterogeneous Federated Learning with Large Language Models (ICML 2025 · similarity 0.88)
- FedBPT: Efficient Federated Black-box Prompt Tuning for Large Language Models (ICML 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.