谷歌发布 RNN 新架构，同等规模强于Mamba

3月1日，在谷歌 DeepMind 近日的一篇论文中，研究者提出了 RG-LRU 层，它是一种新颖的门控线性循环层，并围绕它设计了一个新的循环块来取代多查询注意力（MQA）。他们使用该循环块构建了两个新的模型，一个是混合了MLP和循环块的模型Hawk，另一个是混合了MLP与循环块、局部注意力的模型Griffin。

https://arxiv.org/pdf/2402.19427.pdf

针对一系列模型规模、在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 层的内核来实现这一点。