简介

我们正处于基础模型 (Foundation Models, FMs) 的时代。从聊天机器人到代码助手，大型语言模型 (LLMs) 在零样本和少样本学习中展现了惊人的能力。然而，当前的 AI 生态系统存在一个巨大的摩擦点: 隐私 。

大多数强大的模型都驻留在中心化的数据中心。要在敏感的私有数据 (如医疗记录、法律文件或个人聊天记录) 上微调这些模型，用户通常不得不将数据上传到云端。这简直是一场隐私噩梦，而联邦学习 (Federated Learning, FL) 正是为了解决这个问题而生。FL 允许模型在分布式设备 (客户端) 上进行训练，而数据从未离开过设备。

但也存在一个问题: 基础模型非常庞大。即便是参数量在个位数 (十亿级) 的较小“端侧”变体 (ODFMs) ，其计算负担也很重。由于内存和电池的限制，在智能手机上微调整个模型通常是不可能的。

此外，并非所有设备都是生而平等的。在现实世界的网络中，你可能拥有一台高端平板电脑、一部全新的旗舰手机和一部老旧的廉价智能手机，它们都试图为同一个模型做贡献。这就是所谓的系统异构性 。

在这篇文章中，我们将深入探讨一篇提出该问题新颖解决方案的研究论文: HETLORA (异构 LoRA) 。 作者介绍了一种方法，允许不同能力的设备使用不同“秩 (rank) ”的模型复杂度进行训练，并智能地聚合结果以构建一个强大的全局模型。

背景: 微调的必要性

在剖析解决方案之前，我们必须了解为什么不能直接在设备上使用预训练模型。虽然巨型模型 (如 GPT-4 或 PaLM-L) 具有惊人的零样本能力，但较小的“端侧”模型 (ODFMs) 如果没有特定的微调，往往表现不佳。

如下表所示，与通过联邦学习进行完全微调的模型相比，像 PaLM 2 (XXS 和 XS 尺寸) 这样的小型模型在零样本或少样本设置下的性能显着下降。

表 1: PaLM 2 在零样本、少样本和完整联邦微调下的困惑度。

困惑度 (越低越好) 的差异是惊人的。为了在边缘设备上获得可用的性能，微调不是一种奢侈，而是一种必需。

参数效率的挑战

微调通用模型的所有参数 (全量微调) 在计算上是非常昂贵的。为了解决这个问题，研究人员使用了参数高效微调 (PEFT) 。 目前最流行的方法是 LoRA (低秩自适应) 。

什么是 LoRA？

LoRA 不会重新训练预训练模型的巨大权重矩阵 \(\mathbf{W}_0\)，而是冻结原始权重并注入可训练的低秩分解矩阵。

在数学上，更新 \(\Delta \mathbf{W}\) 表示为两个较小矩阵 \(\mathbf{B}\) 和 \(\mathbf{A}\) 的乘积:

\[ \Delta \mathbf{W} = \mathbf{B} \mathbf{A} \]

其中 \(\mathbf{B} \in \mathbb{R}^{d \times r}\) 且 \(\mathbf{A} \in \mathbb{R}^{r \times l}\)。变量 \(r\) 即为秩。

LoRA 的神奇之处在于 \(r\) 可以非常小。如下表所示，使用 LoRA 可以将可训练参数的数量减少到原始模型大小的一小部分 (通常小于 1%) 。

表 2: 与原始模型大小相比，LoRA 参数的百分比。

问题: 同构 LoRA 的权衡

在使用 LoRA 的标准联邦学习设置中，服务器通常强制执行同构秩 (Homogeneous Rank) 。 这意味着每部智能手机，无论是强大的 iPhone 15 还是老旧的 Android 设备，都必须使用完全相同的秩 \(r\) (例如 \(r=16\)) 。

这种“一刀切”的方法造成了一个艰难的权衡:

  1. 高秩: 如果每个人都使用高秩 (更多参数) ，模型学得快，但倾向于在本地数据上过拟合 。 这也给较弱的设备带来了负担。
  2. 低秩: 如果每个人都使用低秩 (较少参数) ，模型计算量小且稳定，但收敛非常慢。

作者在下图中通过经验证明了这种权衡。请注意，高秩 (绿色，\(r=50\)) 的困惑度下降很快，但随后停滞或恶化 (过拟合) ，而最低秩 (红色，\(r=1\)) 虽然慢但很稳定。

图 3: 不同秩下同构 LoRA 的性能。

我们需要一个兼顾两者的解决方案: 既有高秩的速度，又有低秩的稳定性，同时适应不同能力的设备。

核心方法: HETLORA

作者提出了 HETLORA , 这是一个支持异构 LoRA 配置的框架。这允许不同的客户端根据其系统能力和数据复杂性使用不同的秩 (\(r_k\)) 进行训练。

图 1: LoRA 异构秩部署概览。

如上图所示，HETLORA 允许网络混合不同的秩 (\(r_1 < r_2 < r_3\)) 。服务器管理一个具有最大秩 (\(r_{max}\)) 的全局模型，但单个设备只下载和训练它们能处理的部分。

HETLORA 包含三个不同的阶段: 通过截断进行分发、带有秩自剪枝的本地训练以及稀疏加权聚合 。 让我们逐一解析。

1. 通过截断进行分发

服务器维护一个具有最大秩 (\(r_{max}\)) 的全局 LoRA 模块。当客户端连接时，它们不一定要下载整个模块。

如果客户端的带宽或计算能力有限，服务器会对矩阵进行截断 。 对于需要秩 \(r_k\) 的客户端，服务器发送:

  • \(\mathbf{B}_{:, :r_k}\) (B 的前 \(r_k\) 列)
  • \(\mathbf{A}_{:r_k, :}\) (A 的前 \(r_k\) 行)

这个简单的步骤确保了弱设备不会被迫下载或计算它们无法处理的高秩参数的梯度。

2. 带有秩自剪枝的本地训练

客户端收到其 LoRA 模块后，开始在本地私有数据上进行训练。然而，作者引入了一个巧妙的调整: 秩自剪枝 (Rank Self-Pruning) 。

有时，一个客户端可能被分配了一个高秩 (例如 \(r=50\)) ，因为它有一个强大的处理器，但它的数据可能非常简单，不需要那么高的复杂度。在简单数据上使用高秩会导致过拟合 (噪声) 。

为了解决这个问题，本地损失函数包含一个正则化项，用于惩罚高秩的幅度。

\[ \text{Loss} + \lambda \| \text{Higher Ranks of B} \| \| \text{Higher Ranks of A} \| \]

在训练过程中，如果这些高秩的权重缩小到某个阈值以下 (由衰减因子 \(\gamma\) 控制) ，客户端实际上会对自己进行“剪枝”。它意识到，“我不需要秩 50；秩 30 就够了。”然后它会将较小的更新发送回服务器。

这充当了本地噪声过滤器，防止客户端上传过拟合的参数。

3. 稀疏加权聚合

这是 HETLORA 在数学上最关键的部分。服务器现在接收到来自不同秩的更新。

  • 客户端 A 发送秩为 5 的更新。
  • 客户端 B 发送秩为 50 的更新。

如何将它们结合起来？

一个朴素的方法是将小矩阵零填充 (Zero-pad) 以匹配大矩阵的大小，然后对它们进行平均。然而，在这里简单的平均是危险的。高秩客户端通常会引入更多的噪声 (正如我们在过拟合问题中看到的那样) 。如果我们简单地对它们进行平均，充满噪声的高秩更新可能会淹没稳定的低秩特征。

作者提出了稀疏加权聚合 。

图 2: 零填充和稀疏加权聚合方法概览。

该过程如下 (参考上图) :

  1. 零填充 (a): 小秩更新被填充零以匹配全局最大秩。
  2. 计算稀疏权重: 服务器为每个客户端计算一个权重 \(p_k\)。该权重基于更新的 Frobenius 范数 (幅度) 。
  • 其直觉依赖于奇异值分解 (SVD) 理论。包含丰富信息的更新往往具有独特的奇异值模式。
  • 作者使用重构矩阵 (\(\Delta \mathbf{W} = \mathbf{B}\mathbf{A}\)) 的 Frobenius 范数作为计算成本较低的代理指标，而不是运行昂贵的 SVD。
  1. 加权聚合 (b): 服务器使用这些计算出的权重聚合更新。 \[ \overline{\mathbf{B}} = \sum p_k \mathbf{B}_k \] 这确保了稳定的、信息丰富的更新 (通常来自低秩) 被优先考虑，而来自高秩客户端的充满噪声、稀疏的更新则被降低权重。

实验与结果

研究人员在两个现实任务上使用 PaLM 2 模型 (XXS 和 XS 尺寸) 评估了 HETLORA:

  1. 多轮次聊天 (MSC) : 生成对话回复。
  2. Reddit 摘要: 总结用户发布的帖子。

HETLORA 是否击败了基线？

结果表明，HETLORA 显着优于朴素方法。

首先，让我们看看如果我们在没有特殊聚合或剪枝的情况下使用异构秩会发生什么 (朴素异构 LoRA) 。

图 4: 在没有秩剪枝且仅使用简单平均聚合的情况下，HETLoRA 的性能。

如上图 4 所示，朴素异构设置仍然遭受过拟合 (困惑度下降，然后回升) 。这证实了简单地混合秩是不够的；你需要剪枝和加权聚合机制。

当应用完整的 HETLORA 方法时，结果要强得多。下图将 HETLORA (绿线) 与秩为 5 (蓝色) 和秩为 50 (紫色) 的同构 LoRA 以及全量微调 (红色) 进行了比较。

图 5: 同构 LoRA、异构 LoRA 和全量微调的性能比较。

图 5 的关键要点:

  • 速度: HETLORA 的学习速度比低秩基线快得多。
  • 稳定性: 与高秩基线 (紫色) 不同，HETLORA 的困惑度没有飙升；它保持稳定。
  • 性能: 它接近全量微调 (红色虚线) 的性能，而仅训练了一小部分参数。

定量比较

下表提供了最终的确切数字。我们可以看到，与同构 LoRA (“HOMLoRA”) 相比，HETLORA 在聊天数据上实现了更好 (更低) 的困惑度，在 Reddit 数据上实现了更好 (更高) 的 RougeL 分数。

表 3: 最终 RougeL 分数和困惑度比较。

值得注意的是，请看 Recon+SVD 这一行。这代表了另一种策略，即服务器重构完整矩阵并使用 SVD 重新分发。HETLORA 击败了这种方法，证明聚合分解后的矩阵 (\(\mathbf{B}\) 和 \(\mathbf{A}\)) 比重构完整权重能更好地保留跨客户端关系。

效率提升

在设备上使用 LoRA 的主要目标之一是降低通信成本。下图显示了达到特定目标性能所需的通信参数比例。

图 6: 相对于全量微调的通信参数比例。

HETLORA (绿色条/点) 所需的通信量明显少于全量微调。在 Reddit 任务中，同构 LoRA (最小秩) 未能达到目标精度 (标记为 X) ，而 HETLORA 仅用极小的开销就成功了。

剪枝的作用

最后，作者进行了消融实验，以观察自剪枝步骤的重要性。他们测试了不同的 \(\gamma\) 值 (衰减因子) 。

表 4: 衰减因子 gamma 影响的消融实验。

结果表明 \(\gamma = 0.99\) 是最佳点。

  • \(\gamma = 1\) (无剪枝) : 性能较差，因为没有过滤掉高秩噪声。
  • \(\gamma = 0.85\) (激进剪枝) : 性能下降，因为模型丢弃了太多有用的信息。
  • \(\gamma = 0.99\): 在降噪和信息保留之间取得了平衡。

结论与启示

论文《用于端侧基础模型联邦微调的异构 LoRA》为在现实环境中部署大模型提供了一个稳健的解决方案。

通过承认异构性是一种特性，而非缺陷 , HETLORA 实现了:

  1. 包容性: 老旧设备可以通过低秩参与训练。
  2. 性能: 系统获得了高秩的学习速度，且没有不稳定性。
  3. 效率: 与全量微调相比，它显着减少了通信和计算量。

这项研究弥合了基础模型的巨大潜力与现实世界隐私/硬件限制之间的鸿沟。随着移动硬件继续以不同的速度发展，像 HETLORA 这样的技术对于创建完全在边缘端运行的个性化、私密 AI 助手将至关重要。

对于对联邦学习感兴趣的学生来说，这篇论文强调了一个重要趋势: 从“统一”算法转向灵活、自适应的系统，以尊重分布式计算的多样性本质。