FedRot-LoRA: Mitigating Rotational Misalignment in Federated LoRA¶

会议: ICML 2026
arXiv: 2602.23638
代码: https://github.com/haoran-zh/FedRot-LoRA (有)
领域: 联邦学习 / 参数高效微调 / LoRA
关键词: 联邦学习, LoRA, Procrustes 对齐, 旋转不变性, 子空间对齐

一句话总结¶

本文指出联邦 LoRA 中朴素 factor-wise 平均的真正"敌人"是旋转不变性导致的潜在子空间错位，提出在客户端用正交 Procrustes 求解出旋转矩阵 \(R_i^t\) 对齐 \(A,B\) 因子后再聚合，理论与实验都证明能显著降低聚合误差且不增加通信开销。

研究背景与动机¶

领域现状：LoRA 把权重更新表示为 \(\Delta W = BA\)，\(B\in\mathbb{R}^{d\times r}, A\in\mathbb{R}^{r\times d}\)，参数量大幅下降，是 LLM 联邦微调最自然的载体 (FedIT、FFA-LoRA、FlexLoRA 等都基于此)。

现有痛点：理想聚合应该是 \(\Delta W_{ideal}=\tfrac1N\sum B_i A_i\)，但这个 rank 一般 \(>r\)，没法继续保持低秩；退而求其次的 factor-wise 平均 \(\Delta W_{naive}=\bar B \bar A\) 又会引入交叉项 \(B_i A_j\)，造成训练不稳定。已有解法走三条路：参数冻结 (FFA-LoRA、RoLoRA) 损失表达力；SVD 投影 (FlexLoRA) 计算昂贵；高通信传残差 (FedEx-LoRA) 违背 LoRA 初衷。

核心矛盾：现有分析只看到"算子不可交换"的代数原因，忽视了 LoRA 因子分解本身的旋转不变性——对任意正交 \(R\in\mathbb{R}^{r\times r}\)，\((B_i R)(R^\top A_i)=B_i A_i\)。这意味着语义等价的更新会被表示在不同的潜在子空间里，朴素平均时这些错位子空间相互"破坏性干涉"，比代数项造成的误差更大。

本文目标：在不增加通信、不冻参数、不做高维 SVD 的前提下，把"旋转误差"这一被低估的因素显式消掉。

切入角度：既然旋转 \(R\) 不改变语义，那就主动选一个让所有客户端因子对齐到共同参考的 \(R_i^t\)；这个 \(R\) 在正交群上有 \(r(r-1)/2\) 自由度，刚好能对齐子空间，又能用 Procrustes 闭式解高效求得。

核心 idea：用前一轮全局 \(\bar A^{t-1}/\bar B^{t-1}\) 当 reference，每个客户端解一个 Procrustes 问题得到旋转 \(R_i^t\)，再交替地对齐 \(A\) 或 \(B\)，最后插值出"软旋转"避免早期 reference 噪声过大。

方法详解¶

整体框架¶

每通信轮 \(t\)：① 服务器广播 \((\bar A^{t-1}, \bar B^{t-1})\)；② 客户端本地训练得到 \((A_i^t, B_i^t)\)；③ 若 \(t\) 是奇数对齐 \(A\)、偶数对齐 \(B\)，解 Procrustes 得 \(R_i^{t,*}\)；④ 用 \(\lambda\) 在恒等矩阵与 \(R_i^{t,*}\) 之间插值得软旋转 \(R_{i,\text{soft}}^t\)；⑤ 应用 \(\tilde A_i^t=(R_{i,\text{soft}}^t)^\top A_i^t,\;\tilde B_i^t=B_i^t R_{i,\text{soft}}^t\) 后上传聚合。

关键设计¶

正交 Procrustes 对齐 (Alternating Factor Alignment):
- 功能：把每个客户端的 LoRA 因子旋转到全局参考所在的子空间，消除因式分解的旋转歧义。
- 核心思路：奇数轮解 \(\min_{R}\|R^\top A_i^t - A_{ref}\|_F^2,\;\text{s.t.}\;R^\top R=I,\det R>0\)，这是经典 Procrustes 问题，对相关矩阵 \(M=A_{ref}(A_i^t)^\top\) 做 SVD \(M=U\Sigma V^\top\) 后闭式解 \(R_i^{t,*}=V\cdot\text{diag}(1,\dots,1,\det(UV^\top))\cdot U^\top\)。偶数轮换成 \(B\) 的对齐，相关矩阵改为 \(M=(B_{ref})^\top B_i^t\)。计算复杂度 \(\mathcal{O}(d r^2+r^3)\)，远小于 FlexLoRA 的 \(\mathcal{O}(d^3)\)。
- 设计动机：作者证明 (Theorem 4.1) 标量缩放只有 1 个自由度，无法消除子空间错位；非约束的可逆矩阵又会 ill-conditioned。正交矩阵在两者之间取得最优平衡——足够灵活 (\(r(r-1)/2\) 自由度) 又保持良态。
交替对齐 (Alternating \(A\) and \(B\)):
- 功能：避免重复对齐同一因子时另一侧不可控漂移。
- 核心思路：奇数轮固定 \(A\) 的语义到 \(A_{ref}\)、\(B\) 跟随补偿；偶数轮反之。每轮只解一次 SVD，但全局看两个因子轮流被"校准"。
- 设计动机：消融表明，只对齐 \(B\) 时性能掉很多 (SST-2: 0.879 vs 0.954)，原因是 \(B\) 初始化范数小，早期对齐信号弱；交替策略保证两个因子的子空间都被定期校准，互相约束。
软旋转插值 (Soft Rotation):
- 功能：早期训练阶段 reference 噪声大，硬旋转会做出剧烈过度修正。
- 核心思路：构造 \(R'=(1-\lambda)I+\lambda R_i^{t,*}\) 后再投影回正交群得 \(R_{i,\text{soft}}^t\)；\(\lambda=0\) 退化为 FedIT，\(\lambda=1\) 是硬 Procrustes。Lemma A.1 证明 \(\|R_{\text{soft}}-I\|_F\le 2\lambda\|R-I\|_F\)，即软旋转的修正幅度被 \(\lambda\) 线性 bound。
- 设计动机：早期全局模型还没收敛，把客户端用力拉过去反而会破坏个性化收敛轨迹；\(\lambda\in[0.2,0.8]\) 在实验中都比硬对齐好，最优常在 0.4-0.6。

损失函数 / 训练策略¶

保留标准 FedIT 训练流程，仅在客户端上传前插入旋转步骤。论文给出了非凸下的收敛分析 (Theorem 4.4)，把误差分解为初始 gap + 累积聚合误差 \(\|E^t\|_F^2\) + \(\mathcal{O}(\eta)\)；Theorem 4.8 进一步证明对齐后的误差上界比 naive 严格更紧，紧性 gain 为 \(\Gamma(\lambda)=(c_0-\tfrac{4\sqrt\tau\kappa\eta G_B}{\delta_A})\lambda - 4\kappa^2\lambda^2\tau\)，给出了 \(\lambda\) 的可行区间。

实验关键数据¶

主实验¶

在 RoBERTa-Large 上跑 GLUE 五任务，rank=4，三个客户端规模 \(N\in\{3,10,50\}\)。

任务/规模	FedIT	FFA-LoRA	RoLoRA	FedRot-LoRA
MNLI (\(N=3\))	0.866	0.862	0.868	0.876
RTE (\(N=3\))	0.840	0.830	0.854	0.868
GLUE Avg (\(N=50\))	0.768	0.772	0.824	0.873
GSM8K (Llama 3-8B)	0.429	0.436	0.344	0.444
HumanEval pass@1	0.288	0.385	0.295	0.409

聚合误差降低尤其显著：MNLI 上 FedIT 误差 \(3.98\times10^{-3}\)，FedRot-LoRA 只有 \(1.48\times10^{-4}\)，整整一个数量级。

消融实验¶

配置	MNLI Acc
不对齐 (FedIT)	0.866
随机旋转	0.318
标量缩放对齐	0.865
只对齐 \(A\)	0.861
只对齐 \(B\)	0.862
交替 \(A/B\) (Full)	0.876
reference = \(W^{t-2}\)	0.866
reference = \(W^{t-1}\) (默认)	0.876

关键发现¶

标量缩放 (一维旋转) 在高维 LoRA 中几乎无效——证明 rank>1 时必须做子空间级 (正交) 对齐而不是范数调整。
随机旋转把性能砸到 0.318，说明对齐"方向有意义"才有效；这反过来排除了"FedRot-LoRA 的提升只是因为引入额外随机性"这种 trivial 解释。
越异质 (Dirichlet \(h=0.5\)、client 数 \(N=50\)) 优势越大；近 IID (\(h=100\)) 时 baseline 也都不错，FedRot-LoRA 仍能稳定领先 1-2 个点。

亮点与洞察¶

把"代数 trick 当作根因"翻成"几何 invariance 当作根因"：之前的解法都在补救 \((B_i-B_j)(A_i-A_j)\) 这个代数项，本文指出真正的破坏者是潜在子空间的相对旋转，这种 reframe 让解法从"冻结/投影"变成"对齐"，思想转变非常优雅。
Procrustes 是 LoRA 联邦聚合的完美匹配工具——闭式解、\(\mathcal{O}(r^3)\) 复杂度、保正交、保语义、不增加通信，几乎所有 desiderata 一次满足。
软旋转 \(\lambda\) 这个细节非常重要：硬对齐在某些 setting 反而掉点，作者用 \(\lambda\in[0,1]\) 插值控制对齐强度，是把理论上的"最优正交矩阵"工程化的关键。这种"训练初期保守、后期自信"的思路可以迁移到其他 federated/distributed scheme。

局限与展望¶

论文假设 reference 用前一轮全局模型，"如果客户端被选中频率不均"或"通信丢包"情况下 \(W^{t-1}\) 可能严重过时，作者也实验了 \(W^{t-2}\) 会掉点。
\(\lambda\) 是超参数，论文给的可行区间依赖 \(c_0,\delta_A,\kappa,\tau\) 等难以估计的常数，实际只能 grid search。
没有验证 rank>24 或 \(N>50\) 的极端规模——超大 rank 时 SVD 的 \(\mathcal{O}(r^3)\) 会逐渐显著。
没和 FedSA-LoRA、FedEx-LoRA 在生成任务上直接对比，附录 Table 7 也只是综述性比较。

评分¶

新颖性: ⭐⭐⭐⭐⭐ "旋转不变性是真正根因"这一观察很犀利，把已有零散尝试统一在一个清晰几何框架下。
实验充分度: ⭐⭐⭐⭐⭐ GLUE 五任务 × 三规模 × 五 rank + GSM8K + HumanEval + 多种消融 + 标量缩放/随机旋转对照，覆盖很全。
写作质量: ⭐⭐⭐⭐⭐ Figure 1-2 把动机讲得非常直观，理论部分 Theorem 4.4 + 4.8 配以 Corollary 4.9 给出 \(\lambda\) 可行域，结构清晰。
价值: ⭐⭐⭐⭐ 直接可用于现有联邦 LoRA 框架，只需替换聚合步骤；对工业部署很友好。