线性（包括RNN系列），不再需要被困在少数B参数的范围内。
一套方法可以对现有的各种量级的模型进行线性化，最高可达Llama 3.1 405B，只需要十几块显卡就可以在两天内完成！
这就是斯坦福、麻省理工等科研机构推出的低秩线性变换（Low-rank with ）。
论文及代码：
应用，可以实现从传统（）到线性的无缝转移，
并且转换后只需要非常低成本的微调（LoRA），0.2%的参数更新就可以恢复精度。与类似的线性注意力模型或方法相比，5-shot MMLU 直接提升了约 20 个点！
换句话说，LLM的计算复杂度从二次降低到线性，而几乎不失去大型模型的语言能力。
线性问题已被前人详细记录。然而，这是我们第一次真正能够做大做强。
特别实用的是，它以最小的开销实现了接近原始模型的性能。
线性化转换只需要两步：
首先，用线性形式替换原来的部分，并使用简单的MSE损失来训练新的参数来近似注意力；
然后通过低成本微调（LoRA）进一步提高模型的准确性。
为了实现可扩展性，作者采用了更精细的“逐块”训练，将LLM的每k层视为一个块，在块内联合训练注意力，以提高层次注意力匹配。
如上图所示，羊驼（Llama）可以看作是多只小刺猬叠在一起。每个小刺猬都有独特的线性化参数，并且可以相互独立地进行训练。
加速法学硕士
为了避免昂贵的培训费用，研究人员一直在两个方面进行探索：
做得又快又快
Mamba、RWKV、Hawk等高效工具不断出现。
而对流行的LLM进行线性化的工作也让我们眼前一亮。
然而，线性化的 LLM 通常伴随着模型质量的显着下降。你甚至可以通过MMLU测试分数，或者模型中传统块的比例来猜测模型是否具有传统架构。
另外，从实用角度来看，只有获得大型量产级模型，线性路径才能真正与传统看齐。
初步知识
我们先打好基础：为什么要线性化？
正常注意力可以表示为下图上方的公式：
因为 ，我们只能先计算Q乘以K，导致中间缓冲量和计算量随着序列长度的平方而增加；
线性化就是设计两个函数进行逼近，从而将公式转化为如下形式。
此时，Q和K不再需要捆绑在一起。可以先计算K乘以V。这种顺序的改变会导致中间缓冲和计算量随着向量长度的平方而增加，与序列长度呈线性关系。
这就是线性化的意义，让你不再害怕长序列带来的压力。
开始线性化
在这篇文章中，作者的主要想法是在线性化中添加三个简单的概念：
1. ()：可学习（线性）注意力 2. Low-rank：低秩适应 3. Layer-wise：分层优化
首先训练线性注意力来模拟和替代注意力。这种“注意力转移”的灵感来自于作者之前的作品：。
论文地址：
如何设计优雅复杂的功能来近似注意力？
作者说：与其让人类费事，不如让AI自己学习！
与仅使用可学习的线性注意力相比，作者将其推广为可学习的线性注意力和+滑动窗口。
研究人员将线性和注意力统一在一层中，并训练一些额外的参数来近似整体的注意力。
对于N个token的序列，前W个token用于计算注意力，最后NW个token用于计算线性注意力，然后将这些值组合起来。
在 中，作者通过 KL 散度训练特征图来匹配注意力权重，而本文则在注意力层的输出上使用 MSE 损失。
这绕过了一个限制：需要将所有注意力权重实例化为监督目标。
相反，我们可以用来计算注意力输出，并将线性化注意力的内存消耗保持在 O(N)。
只需将这些特征图插入到每个现有的注意力中即可创建线性化的 LLM。冻结所有其他权重并仅训练这些特征图，对于 7B LLM 仅需要调整 0.2% 的参数。
低阶
以前的线性化工作通常需要相对昂贵的端到端训练阶段。
但在这里，只需将低秩适应 (LoRA) 应用于 QKVO 注意力权重即可恢复模型的性能。
冻结其他所有内容，仅训练 LoRA 权重，以最小化 LLM 输出在某些自然语言数据上的下一个标记预测损失。
逐层
大多数情况下，只需要上述两个步骤。但对于 Llama 3.1 405B 大小的型号来说，它仍然需要一些工作。
通过简单地联合优化所有层，可以在 7B 到 70B 参数范围内成功线性化 LLM，但是当作为一个整体进行训练时，后面层的 MSE 将大于前面层的 MSE。
随着模型变得更大、更深，MSE 进行了升级，以微调 Llama 3.1 405B 的实际问题。
为此，研究人员采用了更精细的逐块训练，将 Llama 3.1 405B 划分为多个 k 层块，并仅在每个块内联合训练注意力。
当使用一些线性化数据并行训练所有模块时，只需为每个块预先计算 LLM 的隐藏状态。
可以调整 k 以平衡并行训练的速度与用于预计算和将隐藏状态保存到磁盘的内存。无需花哨的成本模型即可实现 50M 代币的线性化：
当k=1时，需要2字节×126层×50M令牌×16384（大小）=200TB的磁盘空间来存储隐藏状态。
当k = 9时，磁盘空间需求将减少到22TB，并且每个块（9层）仍然可以在单个GPU上并行训练。
——后者显然更加友好，因此作者将 Llama 3.1 405B 的 126 层拆分为 14 个 9 层块，并在 14 个 GPU 上并行线性化注意力。整个过程只用了5个小时。然后使用 LoRA 将它们拼接在一起以获得最终模型。
实验结果
质量恢复
下表给出了 6 个流行的 LLM 评估任务的结果。
与最近的一些线性化方法相比，不同任务和不同LLM的质量和训练效率都有显着提高。
虽然只训练了0.2%的模型参数（40M token），但线性化后与原模型的性能差距平均缩小了80%以上，token到模型的效率提高了500~2500倍。
在7B量级上，优于所有线性注意力（包括RNN系列）模型：Mamba、RWKV、Hawk、和。
挑战405B大型号
最后，作者将线性化扩展到 Llama 3.1 70B 和更大的 405B 模型。
首先是与以前的线性化方法相比，质量显着提高。通过控制相同的线性+滑动窗口层，对于Llama 3.1 70B，在5-shot MMLU上的精度提高了39点，对于Llama 3.1 405B，精度也提高了38.3点。
二是培训效率的提升。在单个 H100 上对 Llama 3.1 70B 进行线性化仅需 18 小时，而对 Llama 3.1 405B 进行线性化所需的时间比之前针对 8B 模型的方法要少。
参考：
