PaperlayerView exact revision on arXiv

Source-linked AI summary

FedSubMuon: Communication-Efficient Federated LLM Fine-Tuning via Structured Subspace Muon

Shaolong Chen, Youming Tao, Shuzhen Chen, Falko Dressler, Qingqing Ye, Di Wang

arXiv:2609.06073v1cs.LGcs.AI

TL;DR

Communication cost constrains cross-device federated LLM fine-tuning, and existing Muon approaches either communicate full layer-size updates or change the optimized object when applied to LoRA factors. FedSubMuon applies Muon to compact coefficients in structured subspaces, while FedSubMuon-GT adds projected-gradient basis tracking. FedSubMuon-GT ranks first on four of five dataset–model pairs, while FedSubMuon performs best under matched communication budgets.

  • Problem

    Cross-device federated LLM fine-tuning is communication-constrained: full-matrix federated Muon communicates layer-size updates, while applying Muon to LoRA factors changes the optimized object.

    Source: §1 · p. 1

  • Method

    FedSubMuon applies Muon to compact coefficient matrices in shared structured subspaces, and FedSubMuon-GT uses projected gradients to adapt tracked subspace bases.

    Source: §1 · p. 2, §3.1 · p. 2

  • Results

    FedSubMuon-GT ranks first on four of five dataset–model pairs, while FedSubMuon performs best on both evaluated backbones under matched communication budgets.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Structured subspace coefficients provide an effective path for bringing Muon matrix optimization to federated LLM fine-tuning.

    Source: §5 · p. 9

    Federated learning alone provides no formal privacy guarantee, and model updates, coefficient matrices, or gradient-related information may leak information under strong attacks.

    Source: Limitations · p. 10

Abstract

from arXiv · show

Federated fine-tuning adapts large language models (LLMs) to decentralized client data, but its scalability in cross-device training is often limited by the high communication cost. Muon is an optimizer that improves optimization performance by orthogonalizing momentum for matrix-valued parameters. Existing federated Muon methods demonstrate the benefit of matrix-aware optimization in federated learning, but still require transmitting full layer-size updates and optimizer state. A natural way to reduce communication is to directly apply Muon to LoRA factors, but this changes the optimized object and weakens Muon's matrix-aware update geometry. We propose FedSubMuon, a communication-efficient federated Muon fine-tuning method that optimizes compact coefficient matrices within shared structured subspaces. This design keeps Muon on a single matrix-valued trainable object, while reducing the client upload to compact coefficient matrices. We further introduce FedSubMuon-GT, an accuracy-oriented extension that uses projected gradients to adapt tracked subspace bases toward task-relevant gradient directions. Experiments on instruction tuning and mathematical reasoning show that FedSubMuon-GT achieves the best overall accuracy on four of five dataset-model pairs, while FedSubMuon performs best under all matched communication budgets. On Dolly-15K, the closest communication baseline requires 5.5 times and 1.4 times more total communication on Llama-1B and Qwen-4B, respectively.

1 Introduction

Federated LLM fine-tuning is constrained by communication cost, while existing federated Muon methods still communicate full layer-size updates and applying Muon to LoRA factors changes the optimized object. FedSubMuon uses structured subspace coefficients to retain matrix-aware optimization with compact communication, and FedSubMuon-GT adds basis tracking for accuracy.

1 Introduction

  • Full-matrix federated Muon methods require transmitting full layer-size parameter updates, so they do not remove the fine-tuning communication bottleneck.Muon orthogonalizes momentum before updating hidden-layer weights.Source: §1 · p. 1
  • Applying Muon to LoRA factors changes the optimized object, and the cited Table 1 reports no performance improvement from replacing the optimizer in federated LoRA baselines.The LoRA factors are coupled low-rank factors, not a parametrization-independent update of the underlying matrix.Source: §1 · p. 1, Table 1 · p. 1
  • FedSubMuon transmits compact coefficient matrices in a shared matrix subspace, while clients optimize those matrices with Muon.At refresh rounds, the server commits aggregated subspace updates and samples a new coefficient matrix.Source: §1 · p. 2
  • FedSubMuon-GT ranks first on four of five dataset–model pairs, while FedSubMuon performs best under matched communication budgets on both evaluated backbones.Experiments cover instruction tuning and mathematical reasoning with Llama and Qwen backbones.Source: §1 · p. 2
  • The paper evaluates instruction tuning and mathematical reasoning across 1–8B LLMs to study accuracy–communication trade-offs.Its stated contributions include FedSubMuon-GT, which uses projected gradients to adapt tracked subspace bases toward task-relevant directions.Source: §1 · p. 2

2 Preliminaries

The preliminaries describe federated fine-tuning of selected weight matrices, with client data kept private and updates aggregated at the server. They introduce full-matrix Muon and low-dimensional random-subspace optimization as background for the paper’s method.

2 Preliminaries

  • Federated fine-tuning keeps client data private while selected clients train and upload adapter parameters for server aggregation.The setup uses a frozen pretrained backbone and client-specific local objectives.Source: §2 · p. 2
  • Muon updates hidden-layer weight matrices by orthogonalizing momentum; Newton–Schulz iterations can efficiently approximate the resulting direction.The optimizer is presented for full-matrix updates using the global gradient.Source: §2 · p. 2
  • Random-subspace optimization restricts training to a low-dimensional coordinate, represented by a projection matrix and a trainable subspace coordinate.The projection can be resampled periodically, changing the subspace used by subsequent updates.Source: §2 · p. 2

3 Method

FedSubMuon trains compact coefficient matrices in shared subspaces while applying Muon to a matrix-valued object. Its refresh mechanisms commit accumulated updates and, in FedSubMuon-GT, steer bases using projected gradients.

3.1 Overview

  • With a frozen backbone, FedSubMuon uses shared random-seed-generated bases so clients train compact coefficient matrices with Muon instead of full layer matrices.Clients synchronize their accumulators by replaying missing commits before local training.Source: §3.1 · pp. 2–3

3.2 Client Synchronization and Subspace Coefficient Update

  • Clients project layer gradients onto the current left and right bases, apply Muon to the resulting r × r coefficient matrix, and upload it for weighted aggregation.Momentum remains local and is reinitialized when a client is selected for a new round.Source: §3.2 · p. 3

3.3 Commit-and-Reset Refresh

  • At refresh rounds, the server commits the active coefficient update to the global accumulator, resets the coefficient matrix, and samples a new basis for subsequent updates.Because accumulated updates can span different bases, inactive clients replay logged commits when selected again.Source: §3.3 · p. 4

3.4 FEDSUBMUON-GT: Gradient-Tracked Subspace Refresh

  • FedSubMuon-GT replaces random refresh with projected-gradient tracking that rotates existing bases toward useful, underrepresented gradient directions.Clients send projected gradients at refresh rounds rather than full gradient matrices; the basis update follows the current coefficient commit.Source: §3.4 · pp. 4–5

3.5 Communication Analysis

  • FedSubMuon sends r × r coefficient matrices per client in both directions, while FedSubMuon-GT adds tracked-basis downlink traffic and amortized projected-gradient uplink traffic.FedSubMuon-GT's tracked-basis downlink scales as O((dout + din)r), and projected-gradient uplink adds O((dout + din)r/τ) when refreshes occur every τ rounds; the committed accumulator remains an ambient matrix that can be merged into the backbone at deployment.Source: §3.5 · p. 5

4 Experiments

The experiments test structured subspace optimization across instruction-tuning and reasoning tasks, model scales, and communication budgets. Results compare downstream accuracy, training progress, efficiency, and component choices.

4.1 Setup

  • Evaluation spans instruction tuning and reasoning on Llama 3.2 1B, Qwen3 4B, and Qwen3 8B, using ROUGE-L or exact match.Natural Instructions, Dolly-15K, GSM8K, and MATH are evaluated; training uses early stopping and task-dependent maximum round schedules.Source: §4.1 · p. 6

4.2 Performance

  • FEDSUBMUON-GT ranks first on four of five dataset–model pairs and second on the remaining pair, while FEDSUBMUON remains competitive.The evaluation spans instruction tuning and mathematical reasoning across the reported datasets and model scales.Source: §4.2 · p. 7

4.3 Performance Analysis under Limited Communication

  • FEDSUBMUON is best in all three matched-budget settings, even against baselines given rank sweeps; reducing LoRA rank alone is not a reliable substitute.The comparison uses 220 MB for Natural Instructions with Llama 1B, 264 MB for Dolly-15K with Llama 1B, and 3457 MB for Dolly-15K with Qwen 4B.Source: §4.3 · p. 7, Table 3 · p. 7

4.5 Communication and Computational Efficiency

  • FLoRG requires 5.5× more communication on Llama 1B and 1.4× more on Qwen 4B than FEDSUBMUON.FEDSUBMUON-GT has the strongest overall accuracy but spends additional communication on basis refreshes and lazy catch-up packages; it is fastest in wall-clock time and second-best in GPU memory.Source: Table 4 · p. 8, §4.5 · p. 8

4.4 Training Loss under Rounds and Communication

  • On GSM8K, FEDSUBMUON-GT reaches 0.396 final validation loss versus FEDSUBMUON at 0.397; on MATH, FEDSUBMUON ends at 0.577 and FEDSUBMUON-GT at 0.578, with both below baselines.Across round-based curves, both converge faster than baselines; in communication plots, FEDSUBMUON has the strongest lower-left frontier, while FEDSUBMUON-GT spends extra basis-refresh communication.Source: §4.4 · p. 8

4.6 Ablation Studies

  • Ablations favor moderate rank, five-round basis refresh, random initialization, and Muon for optimizing the coefficient matrix.Rank 64 is best for both variants; refreshing every five rounds scores 30.09 versus 28.31 at τ = 1 and 28.66 at τ = 10. Random initialization scores 30.09, and Muon outperforms AdamW and SGD.Source: Figure 3 · p. 9, §4.6 · p. 9

5 Conclusion

FEDSUBMUON applies matrix-aware Muon optimization to compact structured-subspace coefficients, reducing client uploads. FEDSUBMUON-GT adds projected-gradient basis tracking and achieves the strongest overall accuracy.

5 Conclusion

  • FEDSUBMUON moves matrix-aware optimization from full weight matrices to compact structured-subspace coefficients, reducing client uploads.Muon remains applied to a single matrix-valued trainable object.Source: §5 · p. 9
  • FEDSUBMUON-GT adds projected-gradient basis tracking and achieves the strongest overall accuracy, while FEDSUBMUON is the strongest communication-efficient variant.Source: §5 · p. 9
  • The experiments support structured subspace coefficients as an effective path for federated Muon optimization in LLM fine-tuning.Source: §5 · p. 9

Limitations

The evaluation is limited to controlled settings and selected attention projection matrices, and federated learning alone does not provide a formal privacy guarantee.

Limitations

  • Experiments are controlled rather than real-world cross-device deployments, with larger models and more varied client conditions left for future study.The paper identifies heterogeneous hardware, unstable networks, stragglers, and stricter memory constraints as practical factors that may affect synchronization, basis refresh, and aggregation.Source: Limitations · pp. 9–10
  • Experiments adapt selected attention projection matrices, leaving other modules, broader adaptation scopes, and backbone families unexplored.These settings may exhibit different accuracy–communication trade-offs.Source: Limitations · p. 10
  • Federated learning keeps raw examples decentralized but provides no formal privacy guarantee, and updates or gradient-related information may leak under strong attacks.The paper recommends combining the method with privacy and security mechanisms such as secure aggregation and differential privacy for sensitive data.Source: Limitations · p. 10

Ethical Considerations

The paper reports that its experiments use public benchmarks without collecting personal data or conducting human-subject studies, and situates its work among three research lines.

Ethical Considerations

  • All experiments use public benchmark datasets; the authors collect no new personal data and conduct no human-subject studies.They identify no specific additional risks beyond those commonly associated with federated learning and LLM fine-tuning.Source: Ethical Considerations · p. 10

A Related Work

  • The related-work discussion situates the paper among communication-efficient federated fine-tuning, random subspace optimization, and matrix-aware optimizers.Source: A Related Work · p. 12

A.1 Federated Fine-Tuning

The related-work discussion distinguishes FEDSUBMUON from prior federated fine-tuning, random-subspace, and matrix-aware optimization approaches by how it represents and optimizes updates.

A.1 Federated Fine-Tuning

  • FEDSUBMUON is closest to seeded-subspace methods but uses an ambient-space matrix accumulator and applies Muon to matrix-valued coefficients rather than scalar coordinates.The discussion describes LoRA and seeded-subspace alternatives, including FLoRG and FedKRSO.Source: §A.1 · p. 12

A.2 Random Subspace Optimization

  • FEDSUBMUON treats its subspace coefficient matrix as a per-round variable; committed updates accumulate in the ambient matrix as refreshed bases redirect future search.The stated design preserves previously accumulated gradient information.Source: §A.2 · pp. 12–13

A.3 Muon and Matrix-Aware Optimizers

  • Federated Muon methods operate on full weight matrices, while FEDSUBMUON applies Muon to an r × r coefficient matrix and reduces the payload to O(r^2) per layer.Applying Muon directly to LoRA factors is not equivalent to optimizing the underlying update matrix because the factors are non-unique coordinates of a low-rank product (Bogachev et al., 2026).Source: §A.3 · p. 13

B Algorithms

The section identifies algorithms summarizing FEDSUBMUON and FEDSUBMUON-GT.

B Algorithms

C Hyperparameter Settings

The experiments vary client participation by dataset and use fixed local-training settings. FedSubMuon and FedSubMuon-GT use coefficient-matrix learning rates, with an additional subspace learning rate for FedSubMuon-GT.

Experiment settings

  • Natural Instructions uses 738 clients with 3% participation; Dolly-15K uses 200 clients with 5%, and GSM8K and MATH use 100 clients with 10%.Each participating client uses batch size 1; clients run 100 local steps on Natural Instructions and one local epoch per round on the other datasets.Source: C Hyperparameter Settings · p. 13

Learning rates

  • Baseline learning rates are searched over {10^-4, 5 × 10^-5, 10^-5, 5 × 10^-6}; the final rate is 10^-5 for FedKRSO and 5 × 10^-5 for other baselines.All methods use a LoRA scaling factor of 1.Source: C Hyperparameter Settings · p. 13
  • FedSubMuon uses coefficient-matrix learning rates of 2 × 10^-3 on Natural Instructions and 5 × 10^-3 on the other datasets.FedSubMuon-GT additionally uses subspace learning rates of 10^-2 on Natural Instructions and 5 × 10^-2 on the other datasets.Source: C Hyperparameter Settings · p. 13

Algorithm settings

  • At refresh rounds, FedSubMuon-GT forms projected gradients and updates tracked subspace bases, while clients apply local Muon steps to coefficient matrices.The algorithm also specifies server client sampling, synchronization, and aggregation of client coefficient matrices.Source: C Hyperparameter Settings · p. 14

D Dataset Processing and Prompt Templates

The preprocessing and prompt formats differ across datasets and backbones. Dolly-15K and Natural Instructions use dataset-specific federated partitions, while GSM8K and MATH use official splits for training and evaluation.

Dolly-15K

  • Dolly-15K examples are normalized into instruction, optional context, response, and category fields; seven categories train clients, while summarization is reserved for evaluation.The default client partition is Dirichlet with concentration 0.5; Llama uses Alpaca-style prompts and Qwen uses single-turn chat templates.Source: D Dataset Processing and Prompt… · p. 14

Natural Instructions

  • Natural Instructions examples contain a task definition, instance input, and first reference output; retained training tasks become clients after filtering and subsampling.Tasks with fewer than 20 examples are removed, retained training tasks are subsampled to 20%, and evaluation tasks with more than 20 instances are subsampled to max(20, 2%).Source: D Dataset Processing and Prompt… · p. 14

Mathematical reasoning

  • GSM8K uses its official training split for federated training and official test split for evaluation and test, measuring exact-match accuracy on final answers.MATH divides its official training split into federated training and a 500-example development set, then uses the official test split; accuracy is exact match on extracted boxed answers.Source: D Dataset Processing and Prompt… · pp. 14–15

E Additional Ablation Studies

The ablations vary client participation and data heterogeneity on Dolly-15K with Llama 1B. FedSubMuon-GT leads at partial participation, while FedSubMuon leads at both ends of the tested heterogeneity range.

Client participation ratio

  • 30.79 at q = 0.1 and 29.75 at q = 0.5 are FedSubMuon-GT’s best results under low and medium participation; most methods perform similarly at full participation.Relative to FedSubMuon, basis refresh adds 0.94 at q = 0.1 and 0.69 at q = 0.5.Source: Figure 4 · p. 15, E Additional Ablation Studies · p. 15

Client-data distribution

  • 31.90 at α = 0.1 and 30.70 at α = 1.0 are FedSubMuon’s scores, exceeding the strongest baseline by 1.33 and 0.30 ROUGE-L, respectively.At α = 0.5, FedSubMuon-GT reaches 30.79 and improves over the best baseline by 0.33 ROUGE-L.Source: Figure 4 · p. 15, E Additional Ablation Studies · p. 15

Related papers

Loading 2609.06073v1…