将 MOE 塞到 LoRA: 一篇文章的诞生

Nothing will work unless you do. --Maya Angelou

本文主要介绍一篇论文是怎么诞生。

文章的基本信息是：

标题：Mixture-of-Subspaces in Low-Rank Adaptation

链接：

代码：

简介：在传统的 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 ?

最直观的设计就是：

其实想出这种设计还是很直接的，毕竟 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'。从这个视角去看，这种多 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 (2024/5/23 挂出)还是蛮像的，他们的论文是从 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) 不可合并 这些性质或者特点。

后记

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