引言
“Transformer” 架构已成为近期人工智能大爆发的代名词。从 ChatGPT 到 LLaMA, Softmax 注意力 (Softmax Attention) 机制驱动着这些模型理解和生成人类语言的能力。然而,这种强大的能力伴随着巨大的代价: 二次复杂度 (Quadratic Complexity) 。
简单来说,当文本 (上下文) 的长度翻倍时,标准 Transformer 的计算成本会变为原来的四倍。这种 \(O(N^2)\) 的复杂度在处理书籍、代码库或长对话历史时造成了巨大的瓶颈。
于是, 线性复杂度模型 (Linear Complexity Models) 应运而生。这些架构旨在以 \(O(N)\) 的复杂度处理文本——这意味着如果你将文本长度翻倍,成本也仅仅翻倍。它们承诺实现 AI 领域的圣杯: 拥有 Transformer 的性能,同时具备循环神经网络 (RNN) 的效率。
但一直存在一个挥之不去的问题: 这些线性模型具备可扩展性吗 (Do these linear models scale) ? 多亏了 Scaling Laws (缩放定律) ,我们知道如何预测标准 Transformer 随着规模变大时的性能,但对于线性替代方案,我们还没有这种确定性。
在这篇文章中,我们将深入探讨一篇至关重要的研究论文,该论文确立了 线性复杂度语言模型的 Scaling Laws 。 研究人员在 3000 亿 token 上训练了从 7000 万到 70 亿参数不等的模型,以找出线性模型是否真的能与巨头们抗衡。
竞争者: 二次方 vs. 线性
为了理解这种比较,我们必须首先介绍本研究中评估的架构。研究人员将标准的 LLaMA 基线与三种不同的线性架构进行了对比。
1. 基线: LLaMA (Softmax 注意力)
这代表了当前的标准。它使用传统的注意力机制,其中每个词都会关注其他所有词。它功能强大,但计算量大。
2. 挑战者 (线性复杂度)
该研究考察了三种“高效”架构,它们修改了注意力或记忆的工作方式以实现线性速度:
- TNL (TransNormerLLM): 一种使用“数据无关衰减 (data-independent decay) ”的线性注意力模型。它使用滑动窗口方法配合闪电注意力 (Lightning Attention) 来处理记忆。
- HGRN2 (Hierarchically Gated RNN): 一种使用“数据相关衰减 (data-dependent decay) ”的现代 RNN。它具有状态扩展机制,可以在不激增参数数量的情况下增加其循环记忆的大小。
- cosFormer2: 一种不带衰减的线性注意力模型。它使用基于余弦的重加权机制来突出重要信息。
根本的区别在于它们如何计算 token 之间的关系。如下表所示,虽然 LLaMA 的 FLOPs (浮点运算次数) 与 \(n\) (序列长度) 呈二次项关系,但线性模型保持了一种关系,即序列长度 \(n\) 在主导项中不会与自身相乘。

核心发现: Scaling Laws 确实存在
这篇论文的主要贡献是证明了 线性复杂度模型遵循与 Transformer 相同的幂律缩放趋势。
Scaling Laws 允许研究人员根据训练所用的计算预算来预测模型的损失 (错误率) 。研究人员遵循了 “Chinchilla” 方法论 (Hoffmann et al., 2022) ,训练了数十个模型,以找到模型大小和数据集大小之间的最佳权衡。
缩放的可视化
下图是论文中最关键的可视化图表。它绘制了训练损失与计算预算 (PFLOPs-days) 的关系。
- 左列 (损失 vs. 计算量) : 所有四种架构 (LLaMA 和三种线性模型) 在双对数坐标图上都显示为一条直线。这证实了随着计算量的增加,模型的性能会以可预测的速率提升。
- 中间列 (最佳模型大小) : 这显示了在给定预算下,你的模型应该有多大。
- 右列 (最佳 Token 数) : 这显示了你应该使用多少数据。

关于效率的结论
值得注意的是,研究发现线性复杂度模型表现出与传统 Transformer 相似的缩放能力 。 事实上,在相同的计算预算下,线性模型通常能实现 更低 的训练损失。
下表总结了从实验中得出的数学幂律。线性模型的系数 (\(\alpha\) 和 \(\beta\)) 与 LLaMA 极具竞争力。

超越损失: 下游任务表现
低训练损失固然好,但模型真的理解语言吗?为了测试这一点,研究人员在下游任务上评估了模型,包括 常识推理 (CSR) 和 验证困惑度 (Validation Perplexity) 。
熟练度与知识
结果令人惊讶。在许多情况下,当按计算量归一化后, 线性模型的表现优于 LLaMA 。
- 困惑度: 像 HGRN2 和 cosFormer2 这样的线性模型在 WikiText-2 等数据集上实现了更低的困惑度 (更好的预测) 。
- 推理: 在 HellaSwag 和 PIQA 等基准测试中,线性模型在 7B 参数规模下始终表现出卓越的推理能力。
下图跟踪了不同模型规模的性能。你可以看到线性模型 (TNL、HGRN2、cosFormer2) 在准确率上呈上升趋势,在困惑度上呈下降趋势,经常紧贴或穿过 LLaMA 的曲线。

对 70 亿参数大关的得分进行的详细细分显示,HGRN2 和 TNL 在一般语言任务中是特别强有力的竞争者。

阿喀琉斯之踵: 检索任务
如果线性模型速度更快且同样聪明,为什么我们没有在所有地方都使用它们?该研究揭示了一个重大限制: 检索 (Retrieval) 。
研究人员使用“大海捞针” (Needle in a Haystack, NIAH) 基准测试对模型进行了测试。在这个任务中,一条特定的信息 (针) 被隐藏在一段长文本 (大海) 中,模型必须将其检索出来。
“回忆”问题
虽然 LLaMA (Softmax 注意力) 擅长回顾特定的 token,无论它们在过去多远的地方,但线性模型却很挣扎。因为线性模型将上下文压缩成固定大小的循环状态,信息会随着时间的推移被“稀释”。
这个公式展示了 Softmax 注意力机制。注意它是如何显式计算查询 (\(Q\)) 和键 (\(K\)) 之间的交互的:

相比之下,线性循环通过更新运行状态来工作。它不会重新扫描历史记录;它只是更新记忆:

视觉证据: 热力图
当可视化时,差异是明显的。下面是 LLaMA 7B 在“标准模式” (检索+理解) 下的热力图。绿色表示成功 (得分 10) ,红色表示失败。LLaMA 大部分是绿色的,成功地在各种深度和长度上检索到了信息。

现在,看看同一任务上的 HGRN2 7B (一个线性模型) 。虽然它在某些区域 (绿色斑块) 表现不错,但在很多地方出现了明显的红色区域,无法检索到信息,尤其是在任务变得更加复杂时。

同样,与 Transformer 基线相比, TNL 7B 显示出在整个上下文窗口中保持一致检索能力的挣扎。

研究人员得出的结论是,虽然线性模型在一般语言建模和推理方面表现出色,但它们缺乏 Transformer 天生具备的“翻阅书籍” (Going Through a Book, GTB) 能力——即重新扫描精确先前输入的能力。
架构细微差别: 形状很重要
该论文强调了与传统 Scaling Laws 的另一个有趣分歧: 宽高比敏感性 (Aspect Ratio Sensitivity) 。
对于像 LLaMA 这样的 Transformer,模型的具体形状 (有多深 vs 有多宽) 通常并不重要,只要总参数量相同即可。然而,线性模型对此要敏感得多。
如下表所示,如果将隐藏层维度推得太高 (从而减少层数) ,会导致像 cosFormer2 这样的线性模型的检索性能崩溃,而 LLaMA 则保持相对稳定。

结论与启示
这项研究为高效 LLM 的未来提供了坚实的基础。主要结论如下:
- 可预测性: 我们现在可以自信地使用已建立的幂律来扩展线性模型。
- 效率至上: 对于一般的语言生成和推理,线性模型提供了一种比 Transformer 更快、更高效的替代方案,且不牺牲质量。
- 检索差距: 剩下的主要障碍是从长上下文中进行精确的信息检索。线性模型压缩历史,而 Transformer 则精确地保留它。
其意义是重大的。对于像聊天机器人、代码生成和创意写作这样“主旨”和推理比照相式记忆更重要的应用,线性模型 (如 HGRN2 或 TNL) 已经准备好进入大联盟了。然而,对于需要从海量文档中精确引用的任务,传统的 Transformer 目前仍然占据统治地位。
](https://deep-paper.org/en/paper/2406.16690/images/cover.png)