谷歌爆改Transformer，“无限注意力”让1B小模型读完10部小说

明敏发自凹非寺

量子位 | 公众号 QbitAI

谷歌大改 Transformer，"无限" 长度上下文来了。

现在，1B 大模型上下文长度可扩展到 1M（100 万 token，大约相当于 10 部小说），并能完成 Passkey 检索任务。

8B 大模型在 500K 上下文长度 的书籍摘要任务中，拿下最新 SOTA。

这就是谷歌最新提出的 Infini-attention 机制（无限注意力）。

它能让 Transformer 架构大模型在有限的计算资源里处理无限长的输入，在内存大小上实现 114 倍压缩比。

什么概念？

就是在内存大小不变的情况下，放进去 114 倍多的信息。好比一个存放 100 本书的图书馆，通过新技术能存储 11400 本书了。

这项最新成果立马引发学术圈关注，大佬纷纷围观。

加之最近 DeepMind 也改进了 Transformer 架构，使其可以动态分配计算资源，以此提高训练效率。

有人感慨，基于最近几个新进展，感觉大模型越来越像一个包含高度可替换、商品化组件的软件栈了。

该论文核心提出了一种新机制 Infini-attention。

它通过将压缩记忆（compressive memory）整合到线性注意力机制中，用来处理无限长上下文。

压缩记忆允许模型在处理新输入时保留和重用之前的上下文信息。它通过固定数量的参数来存储和回忆信息，而不是随着输入序列长度的增加而增加参数量，能减少内存占用和计算成本。

线性注意力机制不同于传统 Transformer 中的二次方复杂度注意力机制，它能通过更小的计算开销来检索和更新长期记忆。

在 Infini-attention 中，旧的 KV 状态（{KV}s-1）被存储在压缩记忆中，而不是被丢弃。

通过将查询与压缩记忆中存储的键值进行匹配，模型就可以检索到相关的值。

PE 表示位置嵌入，用于给模型提供序列中元素的位置信息。

对比来看 Transformer-XL，它只缓存最后一段 KV 状态，在处理新的序列段时就会丢弃旧的键值对，所以它只能保留最近一段的上下文信息。

对比几种不同 Transformer 模型可处理上下文的长度和内存占用情况。

Infini-attention 能在内存占用低的情况下，有效处理非常长的序列。

Infini-attention 在训练后，分化出了两种不同类型的注意力头，它们协同处理长期和短期上下文信息。

专门化的头（Specialized heads）：这些头在训练过程中学习到了特定的功能，它们的门控得分（gating score）接近 0 或 1。这意味着它们要么通过局部注意力机制处理当前的上下文信息，要么从压缩记忆中检索信息。
混合头（Mixer heads）：这些头的门控得分接近 0.5，它们的作用是将当前的上下文信息和长期记忆内容聚合到单一的输出中。