将MoE塞到LoRA：一篇文章的诞生

©PaperWeekly 原创 · 作者 | Taki

单位 | 香港大学

研究方向 | LLM efficiency, trustworthy

Nothing will work unless you do. ——Maya Angelou

本文主要介绍一篇论文是怎么诞生。文章基本信息如下：

论文标题：

Mixture-of-Subspaces in Low-Rank Adaptation

论文链接：

https://arxiv.org/pdf/2406.11909

代码链接：

https://github.com/wutaiqiang/MoSLoRA

简介：在传统的 LoRA 中加入一个 Mixer 矩阵，进行混个不同子空间的信息。设计非常简单：

最初的想法

说来也是巧合，之前有很多的文章尝试将 LoRA 和 MoE 结合起来，他们基本上都是把 LoRA 当做 MoE 的 Expert，然后塞到 MoE 结构之中，之前也介绍过一些，如文章 1 2 3 4。这些文章无疑都是将 LoRA 看作 MoE 的 expert，一来缺乏动机，二来影响了 LoRA 的可合并性，三来 训练还慢。

闲来与同事聊天，同事说没见过有文章把 MoE 塞到 LoRA 里面，我当时愣了一下。啊？MoE 塞到 LoRA 里面，意思是说把 MoE 的那种 gate+多专家去做 LoRA 的 lora_A 和 lora_B ?

最直观的设计就是：

▲ 有点抽象，但稍微知道点 MoE 和 LoRA 的应该都能懂

其实想出这种设计还是很直接的，毕竟 lora 和 MoE 都是很成熟，很简单的设计。

先不谈有没有动机，反正水文章嘛，都能找到点。就说这个设计，其实有点不合适，为什么呢？

核心就在于 Gate 这玩意，MoE 是希望尽可能训多点参数但计算量不要大太多，因此整了多个Expert 选用一部分并设计了 Gate Router 的机制。但是，LoRA 本身参数量就不大，且rank 大又不一定效果好，堆这个参数属实没必要。

此外，LoRA 的好处就在于可以 Merge 回原来的权重，infer 的时候 0 延迟。这个 Router Gate 因为和输入 x 耦合，因此没法 merge 回去了。这就带来了推理延迟。

去掉Gate，直接上

有了上面的分析，下一步自然就是要去掉 Gate 了。为了确保能合并，因此所有的 expert 都得用，此时就变成了：

▲ 仿佛在拼积木

有了这个设计以后，同时又出现了一些 concern：虽然说 infer 的时候，大家都可以合并到原始权重，都是 0 延迟。但是训练的时候，比如我这个图画的，训的参数是之前的 3 倍多。(在当今这个大环境下，怕是要被审稿人喷）

所以说，要说公平，那就不能设置为 r，每个模块还是得设置成 r/k，上图的 case 对应的就是 r/3，这样训练的参数没变，同时 infer 都是 0 延迟。

这也就是论文里面的 two-subspace-mixing 的方法的由来。

'多头注意力'的视角

既然把每个专家设置成了 r/k 的大小，这玩意就很像是多头注意力了，有 维度切分+并行操作+最后合并 的操作。这不禁让我思考，这和多头注意力有什么关系？原始的 LoRA 能不能等价拆开？

说到拆开，有两个量可以拆，一个是 rank，一个是输入的维度 d。若是直接说多头，可能大家想到的都是直接把 d 拆开，而不是把 rank 拆开。那么这两种拆开我们都可以分析：

i）把 d 拆开的视角：

▲ 一如既往的抽象，熟悉矩阵运算的应该能一眼看明白

图中展示的是 d 拆分 2 个 d/2 的 case，为了好理解，我刻意画了矩阵视角。从矩阵运算角度来看，在 d 维度切分以后，相当于过了两个 A，求和，然后再过两个 B，最后拼接在一起。这三个视角都是等价的。

说实话，要说改进这个，真就没啥好改的。

ii）把 r 拆开的视角：

▲ 这个视角去看，就挺好的，也比较简介

类似地，也可以将 rank 去拆开。上图展示了将 rank 拆开成两个子块的过程。可以看出，等价于两条支路，每个支路 rank=r/2，最后求和。明显比上面的拆分 d 的方法更优雅。

在这个视角下，一个很简单的改进，就呼之欲出了：

思路很简单的，其实就是将中间的平行支路扭在一起，从公式的角度来看，从A1B1+A2B2 变成了 (A1+A2)(B1+B2)=A1B1+A2B2+A1B2+A2B1.

这么来，相当于多了两项。暂且称这个为扭麻花方案吧。

阶段性结果，但还不够

有了上面的分析，那么就开始做实验了：

微调 LLaMA3 做 commonsense reasoning，发现还是有提高的。

不过，这么做还有个问题，那就是 代码效率其实不高。划几条并行的线然后扭个麻花很简单，但是实现起来得看怎么去实现。我初始化了两个 expert 依次去 forward，因此计算效率不高。当然，也可以学习 MHA 的代码，先整个 infer，然后再拆分向量（相当于 A1 和 A2 两个线性层拼在一起 forward，得到结果后再将向量拆开）。

这就启发了另外一个思考，也就是说，这一通操作有很多的线性层的拆分与合并操作，我们之前的分析都是从 linear 层的拆分合并去考虑的，没有考虑向量的拆分合并操作。向量角度等价于：

▲ 核心在于中间的 r 向量进行一系列操作（切分，求和，复制）

之前提到的扭麻花操作，等价于中间的 r 维度的向量，拆分，逐位相加成一半长度，然后复制，再拼接，获得最终的 r'。从这个视角去看，这种多 expert 的扭麻花本质就是在 r 维的向量上加一套组合拳。

混合矩阵的引入

既然是加一套组合拳，这个组合拳 （r 维度的向量，拆分，逐位相加成一半长度，然后复制，再拼接）用矩阵来看，是什么呢？

一番分析下来，不难得到相当于中间加了一个固定的蝴蝶矩阵因子（关于蝴蝶矩阵因子，可以参考：https://weld.stanford.edu/2019/06/13/butterfly/ ）。

既然如此，那么有没有可能模仿 Tri Dao 的做法， 引入一堆蝴蝶矩阵因子？想想还是没啥必要，因为 lora 本身计算量不大，无需这样的拆分，其次就是延迟可能变大很多（此外，调研发现，蝴蝶矩阵序列在 OFT 系列里面是有应用的，也就是 BOFT）。

不往蝴蝶矩阵序列走，另外一个直观的想法就是把这个矩阵升级为可学习的矩阵了。我在论文中把这个矩阵称为 Mixer 矩阵，那么：

原始的 LoRA 相当于使用固定的单位矩阵做 Mixer，中间的扭麻花方案相当于插入固定的蝴蝶因子矩阵做 Mixer，论文里升级为可学习的 Mixer，且矩阵全部元素可学习，也就是所提出的 MoSLoRA 方法。

注1：这种形式和 AdaLoRA 还是蛮像的，不过 AdaLoRA 中间是一个 SVD 分解的特征值，且前后矩阵都加上了正交化约束。 

注2：我在写论文的时候，发现了 Arxiv 有个优秀的同期工作 FLoRA: Low-Rank Core Space for N-dimension，他们的论文是从 Tucker 分解的角度去切入的，思路很巧妙，也很优雅，感兴趣的也可以看看他们文章和解读。

回到MoE的视角

回到 MoE 的视角去看，也就是回到论文最开始的图：

我们可以简单地将 Mixer 理解为 MoE 的 Gate 生成的 weight，此外这个 Gate 有几个特性：

1. 这个 weight 和输入无关，进而确保可合并性

2. 这个 weight 是稠密的，意思是所有的 expert 都用上，而不是 MoE 的那种选取 top-k

3. 原始的 vanilla LoRA 可以看作是 这个 Mixer 矩阵固定为单位矩阵。

看到这，还可以看明白另外一件事，也就是：

【多个并行的 LoRA 分支 选 top-k个输出 最后求和】这种常规 LoRA+MoE 设计，本质上相当于 Mixer 具备：i）每行都是同一个元素 ii）部分行全行为 0 iii）非 0 行的元素由输入来确定 iv）不可合并 这些性质或者特点。

后记

写到这里，其实也把整个思维的推进过程都说清楚了。当然，论文不可能这么写，太冗长且难以理解。博文尚且没几个人有耐心看完，更别说审稿人了。不过整个的思考过程还是收获很多的，可能一个东西刚开始想的时候复杂，换个角度以后，竟然会如此简单。

补充证明

直觉上来说，中间插一个 W，如果把 AW 合并看作 A'，那岂不是和直接学 A‘B 效果是一样的？

其实，并不是一回事，就算是初始化等价，不代表后续优化的路径是一致的。正如重参数化，虽然看起来是等价的，但是学的结果就是不一样。这个角度去看，Mixer 也可以看作是重参数化分支的形式：

其中 I 是固定的不学习的矩阵。这样就相当于原始 LoRA 的旁边加了一个并行分支，和 RegVGG 等重参数化一致了。

当然，这里也给出一个【后续优化的路径是不一致的】的简单证明：

https://github.com/wutaiqiang/MoSLoRA/blob/main/MoSLoRA_proof.pdf

也可以直接看图：

只有当 W 是固定的正交矩阵，才是等价的，不然就算初始化一致，优化过程也会有差异。

在 MoSLoRA 中，W 是可学习的，且我们分析了初始化对结果的影响。

更多阅读

#投 稿 通 道#

 让你的文字被更多人看到 

如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。

总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。 

PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。

📝 稿件基本要求：

• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注 

• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题

• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算

📬 投稿通道：

• 投稿邮箱：hr@paperweekly.site 

• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者

• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿

△长按添加PaperWeekly小编

🔍

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

·

·

·