引言

“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\) 在主导项中不会与自身相乘。

比较 LLaMA 与线性模型的模型参数和 FLOPs 清单。

核心发现: Scaling Laws 确实存在

这篇论文的主要贡献是证明了 线性复杂度模型遵循与 Transformer 相同的幂律缩放趋势。

Scaling Laws 允许研究人员根据训练所用的计算预算来预测模型的损失 (错误率) 。研究人员遵循了 “Chinchilla” 方法论 (Hoffmann et al., 2022) ,训练了数十个模型,以找到模型大小和数据集大小之间的最佳权衡。

缩放的可视化

下图是论文中最关键的可视化图表。它绘制了训练损失与计算预算 (PFLOPs-days) 的关系。

  • 左列 (损失 vs. 计算量) : 所有四种架构 (LLaMA 和三种线性模型) 在双对数坐标图上都显示为一条直线。这证实了随着计算量的增加,模型的性能会以可预测的速率提升。
  • 中间列 (最佳模型大小) : 这显示了在给定预算下,你的模型应该有多大。
  • 右列 (最佳 Token 数) : 这显示了你应该使用多少数据。

四种架构的训练曲线拟合,展示了损失、模型大小和 Token 数与计算量的关系。

关于效率的结论

值得注意的是,研究发现线性复杂度模型表现出与传统 Transformer 相似的缩放能力 。 事实上,在相同的计算预算下,线性模型通常能实现 更低 的训练损失。

下表总结了从实验中得出的数学幂律。线性模型的系数 (\(\alpha\) 和 \(\beta\)) 与 LLaMA 极具竞争力。

Scaling Laws 总结,展示了损失、参数和语料库大小之间的关系。

超越损失: 下游任务表现

低训练损失固然好,但模型真的理解语言吗?为了测试这一点,研究人员在下游任务上评估了模型,包括 常识推理 (CSR)验证困惑度 (Validation Perplexity)

熟练度与知识

结果令人惊讶。在许多情况下,当按计算量归一化后, 线性模型的表现优于 LLaMA

  • 困惑度: 像 HGRN2 和 cosFormer2 这样的线性模型在 WikiText-2 等数据集上实现了更低的困惑度 (更好的预测) 。
  • 推理: 在 HellaSwag 和 PIQA 等基准测试中,线性模型在 7B 参数规模下始终表现出卓越的推理能力。

下图跟踪了不同模型规模的性能。你可以看到线性模型 (TNL、HGRN2、cosFormer2) 在准确率上呈上升趋势,在困惑度上呈下降趋势,经常紧贴或穿过 LLaMA 的曲线。

CSR 和困惑度等不同基准测试的比较性能。

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

包括 CSR、困惑度和检索在内的下游任务基准测试。

阿喀琉斯之踵: 检索任务

如果线性模型速度更快且同样聪明,为什么我们没有在所有地方都使用它们?该研究揭示了一个重大限制: 检索 (Retrieval) 。

研究人员使用“大海捞针” (Needle in a Haystack, NIAH) 基准测试对模型进行了测试。在这个任务中,一条特定的信息 (针) 被隐藏在一段长文本 (大海) 中,模型必须将其检索出来。

“回忆”问题

虽然 LLaMA (Softmax 注意力) 擅长回顾特定的 token,无论它们在过去多远的地方,但线性模型却很挣扎。因为线性模型将上下文压缩成固定大小的循环状态,信息会随着时间的推移被“稀释”。

这个公式展示了 Softmax 注意力机制。注意它是如何显式计算查询 (\(Q\)) 和键 (\(K\)) 之间的交互的:

展示 Softmax 注意力计算的公式。

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

展示线性循环计算的公式。

视觉证据: 热力图

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

LLaMA 7B 的大海捞针热力图,主要显示绿色的成功检索。

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

HGRN2 7B 的大海捞针热力图,显示混合结果及红色失败区域。

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

TNL 7B 的大海捞针热力图,显示其表现挣扎。

研究人员得出的结论是,虽然线性模型在一般语言建模和推理方面表现出色,但它们缺乏 Transformer 天生具备的“翻阅书籍” (Going Through a Book, GTB) 能力——即重新扫描精确先前输入的能力。

架构细微差别: 形状很重要

该论文强调了与传统 Scaling Laws 的另一个有趣分歧: 宽高比敏感性 (Aspect Ratio Sensitivity) 。

对于像 LLaMA 这样的 Transformer,模型的具体形状 (有多深 vs 有多宽) 通常并不重要,只要总参数量相同即可。然而,线性模型对此要敏感得多。

如下表所示,如果将隐藏层维度推得太高 (从而减少层数) ,会导致像 cosFormer2 这样的线性模型的检索性能崩溃,而 LLaMA 则保持相对稳定。

宽高比和模型容量的基准测试,展示了线性模型的敏感性。

结论与启示

这项研究为高效 LLM 的未来提供了坚实的基础。主要结论如下:

  1. 可预测性: 我们现在可以自信地使用已建立的幂律来扩展线性模型。
  2. 效率至上: 对于一般的语言生成和推理,线性模型提供了一种比 Transformer 更快、更高效的替代方案,且不牺牲质量。
  3. 检索差距: 剩下的主要障碍是从长上下文中进行精确的信息检索。线性模型压缩历史,而 Transformer 则精确地保留它。

其意义是重大的。对于像聊天机器人、代码生成和创意写作这样“主旨”和推理比照相式记忆更重要的应用,线性模型 (如 HGRN2 或 TNL) 已经准备好进入大联盟了。然而,对于需要从海量文档中精确引用的任务,传统的 Transformer 目前仍然占据统治地位。