Mixture-of-Subspaces in Low-Rank Adaptation (MoSLoRA)
Chen Kai BOSS

LoRA 把“全量微调”压缩成一个低秩更新,但它也引入了一个直觉上的限制:单个低秩子空间未必足以覆盖任务所需的多样变化;而 MoE 虽然能增加容量,却常带来路由开销、训练不充分与推理不可合并等工程代价。 MoSLoRA 的思路更像是把“多专家”做成“多子空间”:把 LoRA 的更新拆成多个小子空间,再用一个可学习的 Mixer 在这些子空间之间做融合,从而在几乎不增加推理复杂度的前提下提升表达能力,并尽量保留 LoRA 的可合并性。下面我会按“动机—结构—训练/推理行为—实验结论”的顺序拆解这篇工作,重点看它的 Mixer 设计到底带来了哪些收益,以及它与 LoRA/MoE 的边界在哪里。

背景介绍

LoRA 方法简介: LoRA( Low-Rank Adaptation)是一种用于大型语言模型的参数高效微调方法,其核心思想是通过添加低秩分支(由两个低秩矩阵组成)来近似权重的更新,从而显著减少需要微调的参数数量。 Hu 等人 (2022) 的研究表明, LoRA 在不影响模型性能的前提下,将参数更新的数量从全连接层的 降低到 ,其中 $ r d_1 d_2$。

MoE 方法简介: Mixture-of-Experts (MoE) 是一种通过引入多个专家模型,并通过一个门控机制 (gate router) 来动态选择部分专家参与计算的方法。 MoE 的设计目标是在增加模型容量的同时,保持计算效率。通常情况下, MoE 的 gate router 会根据输入生成专家的选择权重,选取 Top-K 个专家参与计算,从而在不显著增加计算开销的前提下提升模型的性能 (Fedus et al., 2022a)。

研究动机: 早期的研究尝试将 LoRA 与 MoE 结合,但这些尝试通常是将 LoRA 作为 MoE 的专家模块,嵌入到 MoE 结构中。这样的设计有几个问题:

  1. 缺乏理论动机:直接将 LoRA 作为专家模块并没有明确的理论支持,缺少对两种方法结合后新特性的分析。
  2. 影响 LoRA 的可合并性:由于 MoE 的 gate router 引入了选择机制,使得这些 LoRA 模块在推理阶段无法合并回原始权重中,增加了推理延迟。
  3. 训练效率低: MoE 的选择机制导致部分参数得不到充分训练,影响了整体效率。

为了解决这些问题,作者提出了反其道而行之的设计思路,即将 MoE 的思想融入到 LoRA 中,而不是将 LoRA 作为 MoE 的组件。

具体细节

MoSLoRA 的核心思想是通过引入一个可学习的 Mixer 矩阵来融合多个子空间,而不是采用 MoE 的 gate 机制来选择专家。具体来说, MoSLoRA 将 LoRA 中的低秩矩阵进一步分解为多个子空间,并通过一个可学习的矩阵来对这些子空间进行加权组合。这样一来,每个子空间的输出都参与了最终的结果计算,既避免了传统 MoE 选择机制带来的推理延迟问题,又充分利用了 LoRA 的低秩优势。

  1. 去掉 Gate 机制: MoE 的 gate router 机制虽然能够选择部分专家进行计算,但在 LoRA 的框架下并不适用。原因是 LoRA 的核心优势在于其低秩结构能够实现参数高效的模型微调,而引入 gate 机制会破坏这一优势。因此, MoSLoRA 选择去掉 gate 机制,直接在所有子空间上进行加权求和,确保所有子空间的输出都能被利用。
  2. 引入可学习的 Mixer 矩阵: 为了更灵活地融合子空间信息, MoSLoRA 引入了一个可学习的 Mixer 矩阵,而不是固定的组合权重。这样做的好处是,模型可以在训练过程中学习到最佳的子空间组合方式,从而提升模型的适应能力和表现力。
  3. 子空间的多样性与组合: 在传统的 LoRA 中,低秩矩阵的秩 通常较小,为了进一步提升模型的表示能力, MoSLoRA 将低秩矩阵分解为多个子空间,每个子空间的秩更小(例如,将一个秩为 的矩阵分解为两个秩为 的子矩阵)。这种分解方式类似于多头注意力机制中的多头结构,通过并行处理不同的子空间并在最终阶段进行组合,能够更好地捕捉数据中的多样性信息。
  4. 实验与性能: 作者在多种下游任务(如常识推理、视觉指令调优、图像生成等)上对 MoSLoRA 进行了实验验证,结果表明,与传统的 LoRA 和其他 PEFT 方法相比, MoSLoRA 在多个基准上均取得了更好的性能,特别是在细粒度子空间视角下,其能够更灵活地融合信息,提升模型的复杂信息建模能力。

代码实现

为了在实践中实现 MoSLoRA,首先需要对标准的 LoRA 方法进行扩展,将其权重矩阵的低秩分解替换为多个子空间的组合。以下是实现 MoSLoRA 的一些关键步骤和代码片段。

  1. 定义子空间:我们首先将 LoRA 中的权重矩阵分解为多个子空间矩阵,这些子空间的秩通常更小。例如,将一个原始的低秩矩阵分解为两个更小的低秩子空间。

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    import torch
    import torch.nn as nn

    class LoRABase(nn.Module):
    def __init__(self, original_matrix, rank=4):
    super(LoRABase, self).__init__()
    self.rank = rank
    self.low_rank_A = nn.Parameter(torch.randn(original_matrix.size(0), rank))
    self.low_rank_B = nn.Parameter(torch.randn(rank, original_matrix.size(1)))

    def forward(self, x):
    return x @ self.low_rank_A @ self.low_rank_B
  2. 引入可学习的 Mixer 矩阵:接下来,我们定义一个可学习的 Mixer 矩阵,用于对各子空间的输出进行加权组合。这个矩阵的大小取决于子空间的数量。

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    class MoSLoRA(nn.Module):
    def __init__(self, original_matrix, num_subspaces=2, rank=4):
    super(MoSLoRA, self).__init__()
    self.subspaces = nn.ModuleList([LoRABase(original_matrix, rank=rank//num_subspaces) for _ in range(num_subspaces)])
    self.mixer = nn.Parameter(torch.randn(num_subspaces, 1))

    def forward(self, x):
    outputs = [subspace(x) for subspace in self.subspaces]
    weighted_outputs = [output * self.mixer[i] for i, output in enumerate(outputs)]
    return sum(weighted_outputs)
  3. 训练和推理:在训练过程中,模型会学习最佳的 Mixer 矩阵参数,以优化子空间的组合权重。在推理过程中,可以直接使用学习到的 Mixer 矩阵来组合子空间的输出,从而得到最终的结果。

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    original_matrix = torch.randn(100, 100)
    model = MoSLoRA(original_matrix, num_subspaces=2, rank=4)

    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    for epoch in range(100):
    optimizer.zero_grad()
    output = model(input_data)
    loss = criterion(output, target_data)
    loss.backward()
    optimizer.step()

以上代码演示了 MoSLoRA 的基本实现流程。通过这种方式,可以将原始 LoRA 的低秩矩阵分解为多个子空间,并通过可学习的 Mixer 矩阵进行优化组合。

基于场景的创新

多模态任务

在处理多模态任务时, MoSLoRA 的子空间组合机制确实具有优势。不同模态的数据(如文本、图像、视频等)通常有不同的特征空间,传统方法需要单独微调每个模态的模型或者使用巨大的联合模型来处理所有模态。而 MoSLoRA 通过引入多个子空间并使用可学习的 Mixer 矩阵对这些子空间进行加权组合,可以更灵活地适应多模态数据。这种方法可以通过以下步骤实现:

  1. 定义子空间:针对每种模态的数据特征,定义不同的子空间。例如,文本数据可以使用一个低秩子空间,图像数据使用另一个子空间。
  2. 训练子空间:使用各自模态的数据分别训练这些子空间。在此阶段,可以选择只训练与当前模态相关的参数,而冻结其他模态的参数,减少训练时间和资源消耗。
  3. 融合子空间:引入可学习的 Mixer 矩阵,动态调整各子空间的权重,使得模型能够根据任务需求和输入数据的模态类型,优化不同子空间的组合方式。
  4. 推理阶段:在推理阶段,使用训练好的 Mixer 矩阵,根据输入数据的模态类型,自动调整各子空间的权重,进行推理计算。

领域适应

在跨领域任务中, MoSLoRA 的多子空间策略同样适用,因为它允许为不同领域定制专门的子空间,并通过训练学习最佳组合方式。例如,在金融和医疗两个完全不同的领域,可以设置两个独立的子空间,一个针对金融数据优化,另一个针对医疗数据优化。 MoSLoRA 在训练过程中会学习如何最优地融合这些子空间,以便在混合数据集或跨领域应用中取得最佳效果。

  1. 数据预处理:将数据划分为不同领域,并根据领域特征对数据进行标准化和处理。
  2. 子空间构建:为每个领域创建单独的子空间,并初始化这些子空间的参数。
  3. 联合训练:在训练时,使用每个领域的训练数据,分别优化其相关的子空间,同时训练可学习的 Mixer 矩阵,使得各子空间在联合任务中的表现得到优化。
  4. 适应性调整:通过训练过程中的损失函数和性能指标反馈,动态调整各子空间的权重和 Mixer 矩阵的参数。

个性化推荐系统

对于个性化推荐系统, MoSLoRA 可以利用其子空间来捕获不同用户的偏好。通过多个子空间表示用户的不同行为模式,并在推理阶段根据用户的历史行为动态调整 Mixer 矩阵的权重,可以更精确地生成个性化推荐。

  1. 用户行为建模:收集用户的历史行为数据,如浏览记录、点击记录、购买记录等。
  2. 子空间学习:根据不同的用户行为模式,定义多个子空间,并分别训练这些子空间以捕获特定的行为特征。
  3. 动态 Mixer 矩阵:在推理阶段,使用用户的历史行为数据,动态调整 Mixer 矩阵的权重,使得推荐模型能够根据当前用户的偏好,选择合适的子空间组合进行推理。
  4. 实时推荐:在推理时,通过动态调整的 Mixer 矩阵生成个性化推荐列表,提高推荐的相关性和用户满意度。

实现方案与潜在问题

实现方案:基于上述不同的场景, MoSLoRA 的实现需要针对特定任务进行细致的配置和优化。首先,需要确定适合的子空间数量和类型,其次是选择合适的初始化策略(如使用 Kaiming 初始化来提高模型收敛速度),最后是动态调整 Mixer 矩阵的训练策略(可以采用 Adam 或 RMSprop 等优化算法)。

潜在问题:在实际应用中,可能会面临以下几个问题:

  • 模型复杂性增加:引入多个子空间和 Mixer 矩阵会增加模型的复杂性,可能导致训练时间延长和内存消耗增加。
  • 子空间选择困难:需要精心选择子空间的数量和类型,以确保它们能够有效地捕获数据特征。
  • 融合策略优化:如何有效地训练和优化 Mixer 矩阵,以实现最佳的子空间组合,也是一个挑战。

针对这些问题,可以通过以下方式解决:

  • 模型压缩技术:使用模型压缩技术,如剪枝和量化,减少模型参数量。
  • 自动化子空间选择:利用自动化机器学习技术(如 AutoML),帮助选择最佳的子空间数量和类型。
  • 优化算法调整:尝试不同的优化算法和学习率策略,寻找最适合的训练设置。
  • 本文标题:Mixture-of-Subspaces in Low-Rank Adaptation (MoSLoRA)
  • 本文作者:Chen Kai
  • 创建时间:2025-02-18 14:30:00
  • 本文链接:https://www.chenk.top/Mixture-of-Subspaces-in-Low-Rank-Adaptation-MoSLoRA/
  • 版权声明:本博客所有文章除特别声明外,均采用 BY-NC-SA 许可协议。转载请注明出处!
 评论