FedSubMuon: 通过结构化子空间Muon的通信高效联邦LLM微调

# 通过结构化子空间Muon实现通信高效的联邦LLM微调
来源：https://arxiv.org/html/2609.06073

Shaolong Chen  
附属机构：阿卜杜拉国王科技大学 (KAUST)  
附属机构：可验证的负责任AI与数据分析 (PRADA) 实验室  
附属机构：伦敦帝国理工学院  

Shuzhen Chen  
附属机构：中国海洋大学  

Falko Dressler  
附属机构：柏林工业大学 (TU Berlin)  

Qingqing Ye  
附属机构：香港理工大学  

Di Wang††thanks:通讯作者\.  
附属机构：阿卜杜拉国王科技大学 (KAUST)  
附属机构：可验证的负责任AI与数据分析 (PRADA) 实验室

###### 摘要

联邦微调将大语言模型 (LLMs) 适配到去中心化的客户端数据，但其跨设备训练的可扩展性常受高通信成本的限制。Muon是一种通过正交化动量来优化矩阵值参数，从而提升优化性能的优化器。现有的联邦Muon方法展示了矩阵感知优化在联邦学习中的益处，但仍然需要传输完整层大小的更新和优化器状态。一种自然的减少通信的方法是直接将Muon应用于LoRA因子，但这改变了优化目标并削弱了Muon的矩阵感知更新几何特性。我们提出了FedSubMuon，一种通信高效的联邦Muon微调方法，该方法在共享的结构化子空间内优化紧凑的系数矩阵。该设计使Muon保持在单一矩阵值的可训练对象上，同时将客户端上传量减少为紧凑的系数矩阵。我们进一步引入了FedSubMuon-GT，这是一种以准确性为导向的扩展，它使用投影梯度来调整跟踪的子空间基，使其趋向任务相关的梯度方向。在指令微调和数学推理任务上的实验表明，FedSubMuon-GT在五个数据集-模型组合中的四个上取得了最佳的整体准确性，而FedSubMuon在所有匹配的通信预算下表现最佳。在Dolly-15K数据集上，最接近的通信基线在Llama-1B和Qwen-4B上分别需要5.5倍和1.4倍的总通信量。

## 1 引言

联邦微调在不将原始样本传输到中央服务器的情况下，利用去中心化的客户端数据适配大语言模型 (LLMs)。然而，在跨设备设置中，这种适配常受到通信成本的限制：在每轮通信中，服务器将可训练参数发送给选定的客户端，客户端返回本地更新进行聚合（Zhang et al., 2024；Liu et al., 2025a）。因此，通信更新的大小成为联邦LLM微调的核心设计约束。

表 1：在Llama 1B和Natural Instructions数据集上，LoRA因子联邦基线的优化器控制。每个单元格报告两次运行的平均ROUGE-L%。Muon是一种矩阵感知优化器，通过在更新隐层权重前正交化动量来提升优化性能（Jordan, 2024）。现有工作已开始将Muon适配到联邦学习中，包括用于矩阵正交化联邦优化的FedMuon变体（Liu et al., 2025b；Takezawa et al., 2026；Zhang and Gao, 2025；Wang et al., 2026）。这一研究方向是相关的，因为利用模型更新的矩阵结构已被证明可以在一系列联邦设置中提升优化性能。然而，现有的FedMuon工作并未解决微调的通信瓶颈：将Muon应用于完整权重矩阵需要传输完整层大小的参数更新。

一个自然的替代方案是在像LoRA这样的参数高效联邦微调方法中应用Muon，但这改变了优化的对象。LoRA通过两个耦合的低秩因子表示更新，先前的工作表明，将Muon应用于这些因子并不等同于对底层矩阵进行与参数化无关的更新（Bogachev et al., 2026）。与这一问题一致，表1显示，在联邦LoRA基线中用Muon替换优化器并未提升性能。这些观察促使我们提出一种联邦Muon微调方法，该方法在仅传输紧凑的可训练对象的同时，保留Muon的矩阵更新几何特性。

我们提出了FedSubMuon，一种用于结构化子空间联邦Muon微调的通信高效框架。FedSubMuon通过传输共享矩阵子空间中的紧凑系数矩阵来减少通信。客户端使用Muon优化系数矩阵，并将其上传进行聚合。在刷新轮次，服务器将聚合的子空间更新提交到全局累加器中，并采样新的系数矩阵。接着，我们引入了FedSubMuon-GT，它在刷新轮次使用投影梯度来更新跟踪的子空间基。虽然FedSubMuon强调通信效率，但FedSubMuon-GT通过添加基跟踪通信来提高准确性。

在Llama和Qwen骨干网络上进行的指令微调和数学推理实验展示了这一权衡。FedSubMuon-GT取得了最强的整体准确性，在五个数据集-模型组合中的四个上排名第一。在匹配的通信预算下，FedSubMuon在两种评估的骨干网络上都表现最佳，这表明简单地降低LoRA基线的秩并不是所提出的子空间系数参数化的可靠替代方案。

我们的贡献包括：
- 我们提出了FedSubMuon，它通过传输紧凑的系数矩阵而非完整权重矩阵，将联邦Muon适配到通信高效的微调中。
- 我们开发了FedSubMuon-GT，这是一种以性能为导向的扩展，它使用投影梯度来调整跟踪的子空间基，使其适应任务相关的梯度方向。
- 我们在1-8B的LLM上，对指令微调和数学推理任务评估了所提出的方法，展示了在标准和匹配预算设置下精度-通信权衡的改进。

## 2 前提知识

#### 联邦微调\。
我们考虑在N个客户端上，对冻结的预训练骨干网络\(W^{0}\)进行跨设备联邦微调，其中隐私数据保留在客户端（McMahan et al., 2017）。客户端\(i\)拥有一个本地目标函数\(f_{i}\)，全局目标函数为
\[
\min_{\Delta W} f(\Delta W) := \sum_{i=1}^{N} p_i f_i(W^{0} + \Delta W),
\]
其中\(\mathcal{L}\)表示所选权重矩阵的索引集，\(\Delta W = \{\Delta W_{\ell}\}_{\ell \in \mathcal{L}}\)是这些矩阵上的更新，\(p_{i}\)是客户端\(i\)的聚合权重。在第\(t\)轮，我们记当前模型为\(W^{t} := W^{0} + \Delta W^{t}\)。一种流行的联邦微调策略是使用低秩矩阵适配所选的投影矩阵（Hu et al., 2022）；选定的客户端训练并上传适配器参数供服务器聚合（Zhang et al., 2023b；Zhang et al., 2024）。

#### Muon优化器\。
Muon是一种矩阵感知优化器，它通过对动量进行正交化来更新隐层权重矩阵（Jordan, 2024）。考虑某一层\(\ell\)的完整矩阵更新。对于全局梯度\(G_{\ell}^{t} = \nabla_{W_{\ell}} f(W^{t})\)和矩阵动量\(M_{\ell}^{t}\)，完整的Muon将使用
\[
M_{\ell}^{t+1} = (1 - \beta) M_{\ell}^{t} + \beta G_{\ell}^{t},
\]
\[
W_{\ell}^{t+1} = W_{\ell}^{t} - \eta \, \operatorname{Muon}(M_{\ell}^{t+1}),
\]
其中，对于紧凑的SVD分解\(A = \widetilde{U} \Sigma \widetilde{V}^{\top}\)，\(\operatorname{Muon}(A) = \widetilde{U} \widetilde{V}^{\top}\)；在实践中，Newton–Schulz (NS) 方法为此方向提供了有效的近似。

#### 随机子空间\。
随机子空间优化将训练限制在完整参数空间内的一个低维坐标中（Li et al., 2018；Aghajanyan et al., 2021）。对于层\(\ell\)，一个通用的向量化随机子空间为
\[
\operatorname{vec}(\Delta W_{\ell}) = P_{\ell} b_{\ell},
\]
其中\(P_{\ell} \in \mathbb{R}^{d_{\mathrm{out}} d_{\mathrm{in}} \times k}\)是一个随机投影矩阵，\(b_{\ell} \in \mathbb{R}^{k}\)是可训练的子空间坐标。投影矩阵可以定期重新采样（Yang et al., 2026；Rajabi et al., 2025）；切换\(P_{\ell}\)会改变后续更新所使用的低维子空间。

## 3 方法

### 3.1 概述

FedSubMuon是一种用于联邦Muon微调的结构化子空间方法。骨干网络保持冻结，而服务器在选定的权重矩阵上维护一个全局更新累加器。客户端使用从共享随机种子生成的固定基来训练紧凑的子空间系数矩阵。这种参数化将Muon应用于具有矩阵结构的系数更新，同时减少了每轮可训练参数数量和通信负载。

图1展示了FedSubMuon的关键工作流程。服务器采样客户端，发送当前的随机种子、活跃的子空间系数矩阵和任何缺失的提交记录给客户端。选定的客户端首先重放这些提交，以将其本地累加器与全局更新累加器同步。然后，它从随机种子重新生成共享基，在紧凑的子空间系数矩阵上运行本地Muon更新，并上传更新后的矩阵进行服务器聚合。

刷新轮次将累积的模型更新与后续的子空间更新分开。在刷新轮次，服务器将完整的模型更新提交到全局更新累加器中，将聚合的系数矩阵和相应的种子追加到提交日志中，重置系数矩阵，并采样新的随机种子用于新基。如果当前通信轮不是刷新轮次，服务器保持全局更新累加器和随机种子不变，并将聚合值设置为下一轮通信的活跃子空间系数矩阵。

附录B提供了FedSubMuon和FedSubMuon-GT的完整流程描述，分别对应算法1和2。

### 3.2 客户端同步与子空间系数更新

**图 1：FedSubMuon工作流程。** 我们首先描述刷新分支分离前的一个通信轮次。选定的客户端首先通过重放缺失的提交来同步其本地累加器；由于这些记录是在刷新轮次创建的，我们将在3.3节给出重放规则。同步之后，所有选定的客户端从相同的子空间系数矩阵\(X^{t}\)开始本地训练，并使用相同的基。

对于选定的层\(\ell\)，当前轮使用的更新分为两部分：同步的本地累加器\(\Delta W_{i,\ell}\)和活跃的子空间系数矩阵\(X_{i,\ell}\)。前者存储已提交的矩阵更新，而后者是当前轮唯一的可训练对象。在FedSubMuon中，当前的基从随机种子\(s^{t}\)生成。对于每个选定的层，共享种子确定性地生成高斯矩阵\(\Omega_{U,\ell}^{t} \in \mathbb{R}^{d_{\mathrm{out}} \times r}\)和\(\Omega_{V,\ell}^{t} \in \mathbb{R}^{d_{\mathrm{in}} \times r}\)，其元素为独立同分布的\(\mathcal{N}(0,1)\)，且
\[
U_{\ell}^{t} = \operatorname{qf}(\Omega_{U,\ell}^{t}), \qquad V_{\ell}^{t} = \operatorname{qf}(\Omega_{V,\ell}^{t}),
\]
其中\(\operatorname{qf}(\cdot)\)返回具有固定确定性符号约定的瘦型QR正交因子。

在当前基固定的情况下，客户端\(i\)仅优化子空间系数矩阵\(X_{i,\ell} \in \mathbb{R}^{r \times r}\)。用于训练的有效本地层为
\[
W_{i,\ell}(X_{i,\ell}) = W_{\ell}^{0} + \Delta W_{i,\ell} + U_{\ell}^{t} X_{i,\ell} (V_{\ell}^{t})^{\top}.
\]
(1)
式(1)是本地训练模型：冻结的骨干网络\(W_{\ell}^{0}\)，固定的已同步累加器\(\Delta W_{i,\ell}\)，固定的子空间基\((U_{\ell}^{t}, V_{\ell}^{t})\)，只有\(X_{i,\ell}\)被更新。在本地优化开始时，客户端初始化
\[
X_{i,\ell}^{t,0} = X_{\ell}^{t}, \qquad M_{i,\ell}^{t,0} = 0.
\]
(2)
对于本地步\(e\)，它计算关于系数矩阵的随机梯度，
\[
\widehat{G}_{i,\ell}^{t,e} = \nabla_{X_{i,\ell}} f_i(W_i^{t}(X_i); \xi_i^{t,e}) \big|_{X_i = X_i^{t,e}}.
\]
(3)
根据链式法则，这是当前层梯度投影到当前左右基上的结果：当完整层梯度写作\(\nabla_{W_{i,\ell}} f_i\)时，相应的系数梯度具有形式\((U_{\ell}^{t})^{\top} (\nabla_{W_{i,\ell}} f_i) V_{\ell}^{t}\)。因此，式(3)更新的是\(r \times r\)的系数矩阵，而非完整层大小的矩阵。

客户端随后更新本地动量并将Muon应用于此小矩阵：
\[
M_{i,\ell}^{t,e+1} = (1 - \beta) M_{i,\ell}^{t,e} + \beta \widehat{G}_{i,\ell}^{t,e},
\]
\[
X_{i,\ell}^{t,e+1} = X_{i,\ell}^{t,e} - \eta \, \operatorname{Muon}(M_{i,\ell}^{t,e+1}).
\]
此本地更新对应图1中的步骤3。经过\(E\)个本地步后，客户端设置\(X_{i,\ell}^{t+1} = X_{i,\ell}^{t,E}\)。每当客户端被选中参与新轮次时，动量会被重新初始化，因此它保持为本地的\(r \times r\)优化器状态。

客户端上传\(\{X_{i,\ell}^{t+1}\}_{\ell \in \mathcal{L}}\)，服务器使用归一化的轮次权重聚合上传的系数矩阵，\(\bar{X}_{\ell}^{t+1} = \sum_{i \in \mathcal{S}_t} a_i^t X_{i,\ell}^{t+1}, \ell \in \mathcal{L}\)。由于所有选定的客户端都从相同的活跃系数矩阵\(X_{\ell}^{t}\)开始，此聚合对相同当前子空间内的紧凑系数更新进行平均。

### 3.3 提交-重置刷新

聚合后，服务器检查...