这一次，谷歌 DeepMind 在基础模型方面又有了新动作。

我们知道，循环神经网络（RNN）在深度学习和自然语言处理研究的早期发挥了核心作用，并在许多应用中取得了实功，包括谷歌第一个端到端机器翻译系统。不过近年来，深度学习和 NLP 都以 Transformer 架构为主，该架构融合了多层感知器（MLP）和多头注意力（MHA）。

Transformer 已经在实践中实现了比 RNN 更好的性能，并且在利用现代硬件方面也非常高效。基于 Transformer 的大语言模型在从网络收集的海量数据集上进行训练，取得了显著的成功。

纵然取得了很大的成功，但 Transformer 架构仍有不足之处，比如由于全局注意力的二次复杂性，Transformer 很难有效地扩展到长序列。此外，键值（KV）缓存随序列长度线性增长，导致 Transformer 在推理过程中变慢。这时，循环语言模型成为一种替代方案，它们可以将整个序列压缩为固定大小的隐藏状态，并迭代更新。但若想取代 Transformer，新的 RNN 模型不仅必须在扩展上表现出相当的性能，而且必须实现类似的硬件效率。

在谷歌 DeepMind 近日的一篇论文中，研究者提出了 RG-LRU 层，它是一种新颖的门控线性循环层，并围绕它设计了一个新的循环块来取代多查询注意力（MQA）。

他们使用该循环块构建了两个新的模型，一个是混合了 MLP 和循环块的模型 Hawk，另一个是混合了 MLP 与循环块、局部注意力的模型 Griffin。

• 论文标题：Griffin: Mixing Gated Linear Recurrences with Local Attention for Efficient Language Models
• 论文链接：https://arxiv.org/pdf/2402.19427.pdf

研究者表示，Hawk 和 Griffin 在 held-out 损失和训练 FLOPs 之间表现出了幂律缩放，最高可以达到 7B 参数，正如之前在 Transformers 中观察到的那样。其中 Griffin 在所有模型规模上实现了比强大 Transformer 基线略低的 held-out 损失。

研究者针对一系列模型规模、在 300B tokens 上对 Hawk 和 Griffin 进行了过度训练，结果显示，Hawk-3B 在下游任务的性能上超越了 Mamba-3B，尽管训练的 tokens 数量只有后者的一半。Griffin-7B 和 Griffin-14B 的性能与 Llama-2 相当，尽管训练的 tokens 数量只有后者的 1/7。

此外，Hawk 和 Griffin 在 TPU-v3 上达到了与 Transformers 相当的训练效率。由于对角 RNN 层受内存限制，研究者使用了 RG-LRU 层的内核来实现这一点。

同时在推理过程中，Hawk 和 Griffin 都实现比 MQA Transformer 更高的吞吐量，并在采样长序列时实现更低的延迟。当评估的序列比训练中观察到的更长时，Griffin 的表现比 Transformers 更好，并且可以有效地从训练数据中学习复制和检索任务。不过当在未经微调的情况下在复制和精确检索任务上评估预训练模型时，Hawk 和 Griffin 的表现不如 Transformers。

共同一作、DeepMind 研究科学家 Aleksandar Botev 表示，混合了门控线性循环和局部注意力的模型 Griffin 保留了 RNN 的所有高效优势和 Transformer 的表达能力，最高可以扩展到 14B 参数规模。

来源：https://twitter.com/botev_mg/status/1763489634082795780

Griffin 模型架构

Griffin 所有模型都包含以下组成部分：(i) 一个残差块，(ii) 一个 MLP 块，(iii) 一个时间混合块。所有模型的 (i) 和 (ii) 都是相同的，但时间混合块有三个：全局多查询注意（MQA）、局部（滑动窗口）MQA 和本文提出的循环块。作为循环块的一部分，研究者使用了真实门控线性循环单元（RG-LRU）—— 一种受线性循环单元启发的新型循环层。

如图 2（a）所示，残差块定义了 Griffin 模型的全局结构，其灵感来自 pre-normTransformer。在嵌入输入序列后，研究者将其通过