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

学学看科技 2024-03-05 14:52:59

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

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

0 阅读:0

学学看科技

简介:感谢大家的关注