新智元报道
编辑:alan
【新智元导读】 近日,来自斯坦福、MIT等机构的研究人员推出了低秩线性转换方法,让传统注意力无缝转移到线性注意力,仅需0.2%的参数更新即可恢复精度,405B大模型两天搞定!
生产级大模型应用线性注意力的方法,来了。
线性Attention(包括RNN系列),再也不用困在几B参数的范围内娱乐了。
一套方法,即可线性化现有各种量级的Transformer模型,上至Llama 3.1 405B,也只需要十来张显卡在两天内搞定!
这就是斯坦福、MIT等科研机构推出的低秩线性转换LoLCATs(Low-rank Linear Conversion with Attention Transfer)。
论文与代码:https://github.com/HazyResearch/lolcats
应用LoLCATs,可以实现传统注意力(softmax)到线性注意力的无缝转移,
且转换后仅需开销很低的微调(LoRA),0.2%的参数更新即可恢复精度,对比同类的线性注意力模型或方法, 5-shot MMLU直接提高了20分左右!
也就是说,在几乎不损失Transformer大模型语言能力的基础上,将LLM的计算复杂度从二次方降到了线性。
线性Attention一事,前人之述备矣,然则,能够真正做大做强,还是第一次。
尤其具有实用价值的是,LoLCATs实现了极小的开销和接近原始模型的性能。
LoLCATs的线性化转换只需两个步骤:
首先使用线性Attention的形式替换原始Attention部分,并利用简单的MSE损失训练新增的参数,以近似softmax注意力;
然后通过低成本的微调(LoRA)来进一步提高模型的精度。
为了实现可扩展性,作者采用更精细的「block by block」训练,将LLM的每k层看成一个block,尽在块内联合训练注意力,以提高分层注意力匹配。
就如上图所表示的那样,一个羊驼(Llama)可以看成多个小刺猬叠在一起,每个小刺猬拥有独特的用于线性化的参数,并且相互之间可以独立训练。
LoLCATS 加速 LLM
为了避免昂贵的训练成本,研究者们一直在不断探索两个方面:
make models fast 与 create fast models
诸如Mamba、RWKV、TransNormer、Hawk、 Griffin和 StripedHyena等高效的subquadratic models不断出现,
而关于将流行的LLM线性化的工作也让我们眼前一亮。
但是线性化LLM往往伴随着模型质量的显著降低,你甚至能通过MMLU的测试分数猜出一个模型是不是传统的Attention架构,或者传统Attention块在模型中的占比。
另外,从实用的角度讲,只有拿下了生产级别的大模型,线性化的道路才能真正与传统Transformer平分秋色。
预备知识
先打基础:为什么要线性化?
正常的softmax注意力可以表示为下图上面的公式:
由于softmax的缘故,只能先算Q乘K,导致中间缓存和计算量随序列长度的平方增长;
线性化就是设计俩函数来近似softmax,从而把公式转化成下面的形式。
此时Q和K不需要绑在一起了,就可以先算K乘V,这个顺序的改变导致中间缓存和计算量随向量长度的平方增长,而相对于序列长度是线性关系。
这就是线性化的意思,这样的Attention也就不惧怕长序列带来的压力了。
开始线性化
本文中,作者的主要想法是向线性化Transformer中添加三个简单的概念:
Learnable Attentions
首先训练线性注意力来模拟和替换softmax注意力。这种「注意力转移」的灵感来自作者之前的一篇工作:Hedgehog。
论文地址:https://arxiv.org/pdf/2402.04347
如何设计设计精妙复杂的函数来近似softmax注意力?
作者表示:与其让人类煞费苦心,不如交给AI自己去学!
相比于Hedgehog中只使用可学习的线性注意力,作者在LoLCATs中,将其推广为可学习的线性注意力和 + 滑动窗口。
研究人员将线性和softmax注意力统一在一个层中,训练一些新增的参数以从整体上近似softmax注意力。
对于N个token的序列,前W个token用于计算softmax注意力,后N-W个token用于计算线性注意力,然后将这些值组合。
在Hedgehog中,作者通过KL散度来训练特征图以匹配注意力权重,而本文改为在注意力层的输出上使用MSE 失。
这绕过了Hedgehog的一个限制:需要将所有注意力权重实例化为监督目标。
相反,LoLCATs可以使用FlashAttention来计算softmax注意力输出,并将线性化注意力的内存消耗保持在O(N)。
只需将这些特征图插入到每个现有的注意力中,即可创建线性化的 LLM。冻结所有其他权重,只训练这些特征图,对于7B的LLM来说,只需要调整0.2%的参数。
Low-rank Adaptation
之前的线性化工作,通常需要一个比较昂贵的端到端训练阶段。
但在LoLCATs这里,可以通过简单地将低秩适应(LoRA)应用于注意力的QKVO权重来恢复模型的性能。
冻结所有其他内容,只训练LoRA权重,在某些自然语言数据上,最大限度地减少LLM输出的next-token预测损失。
Layer-wise Optimization
大多数情况下,只需要以上两步就搞定了。但对于像Llama 3.1 405B这种规模的模型来说,还需要努力一下。
通过简单地联合优化所有层,可以成功地线性化7B到70B参数范围的LLM,但整体训练时,后面层的MSE会比前面的层更大。
当模型变得更大更深时,MSE升级为了微调Llama 3.1 405B的真正问题。
为此,研究人员使用了更精细的逐块训练,将Llama 3.1 405B分成多个k层块,并仅在每个块内联合训练注意力。
当使用一些线性化数据并行训练所有模块时,只需为每个块预先计算LLM的隐藏状态。
可以调节k来平衡并行训练的速度与预计算的内存,并将隐藏状态保存到磁盘。不需要花哨的成本模型,对于50M token的线性化来说:
k = 1时,需要2字节 × 126层 × 50M token × 16384(hidden size)= 200TB的磁盘空间来存储隐藏状态。
而k = 9时,磁盘空间的需求将减少为22TB,这时仍然能在单个GPU上并行训练每个块(9层)。
——后者显然更友好一点,所以作者将Llama 3.1 405B的126层拆分为14个9层块,在14个GPU上并行进行注意力的线性化,过程仅需5个小时。然后用LoRA将它们全部拼接在一起,就得到了最终模型。
实验结果
质量恢复
下表给出了6个流行的LLM评估任务的结果。
与最近的一些线性化方法相比,LoLCATs显著提高了不同任务和不同LLM的质量和训练效率。
尽管只训练了0.2% 的模型参数(40M token),LoLCATs将线性化与原始模型的性能差距平均缩小了80%以上,token to model的效率提高了500~2500倍。
在7B这个量级上,LoLCATs优于所有的线性注意力(包括RNN系列)模型:Mamba、RWKV、TransNormer、Hawk、 Griffin和 StripedHyena。
挑战405B大模型
最后,作者使用LoLCATs将线性化扩展到Llama 3.1 70B和更大的405B模型。
与之前的线性化方法相比,首先是质量上的显著改进。通过控制相同的线性 + 滑动窗口层,对于Llama 3.1 70B,在5-shot MMLU上的精度实现了39点的提升,对于Llama 3.1 405B,同样实现了38.3分的改进。
其次是训练效率的提高,在单个8x80GB H100上线性化Llama 3.1 70B仅需18个小时,而线性化Llama 3.1 405B所花费的时间比之前用于8B模型的方法还要少。
参考资料:
https://x.com/simran_s_arora/status/1845909074774475125