这是用户在 2024-7-2 23:13 为 https://zhuanlan.zhihu.com/p/704761512 保存的双语快照页面,由 沉浸式翻译 提供双语支持。了解如何保存?

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

229 人赞同了该文章
发布于 2024-06-21 23:30・IP 属地中国香港 ,编辑于 2024-06-22 12:41・IP 属地中国香港
目录
收起
最初的想法
去掉Gate,直接上
'多头注意力'的视角
阶段性结果,但还不够
混合矩阵的引入
回到MoE的视角
后记

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维度的向量,拆分,逐位相加成一半长度,然后复制,再拼接)用矩阵来看,是什么呢?

一番分析下来,不难得到相当于中间加了一个固定的蝴蝶矩阵因子(关于蝴蝶矩阵因子,可以参考:weld.stanford.edu/2019/ )。

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

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

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

注1:这种