谷歌Infini 间接裁减到有限长

谷歌又加大招了，颁布下一代 Transformer 模型 Infini-Transformer。

Infini-Transformer 引入了一种有效的方法，可以将基于 Transformer 的大型言语模型 (LLM) 裁减到有限长输入，而不参与内存和计算需求。经常使用该技术，钻研者成功将一个 1B 的模型高低文长度提高到 100 万；运行到 8B 模型上，模型能处置 500K 的书籍摘要义务。

自 2017 年开创性钻研论文《Attention is All You Need》问世以来，Transformer 架构就不时主导着生成式人工自动畛域。而谷歌对 Transformer 的优化设计最近比拟频繁，几天前，他们降级了 Transformer 架构，颁布 Mixture-of-Depths（MoD），扭转了以往 Transformer 计算形式。没过几天，谷歌又放出了这项新钻研。

专一 AI 畛域的钻研者都了解内存的关键性，它是自动的基石，可认为 LLM 提供高效的计算。但是，Transformer 和基于 Transformer 的 LLM 因为留意力机制的固有个性，即 Transformer 中的留意力机制在内存占用和计算期间上都体现出二次复杂性。例如，关于批大小为 512、高低文长度为 2048 的 500B 模型，留意力键 - 值 (KV) 形态的内存占用为 3TB。但理想上，规范 Transformer 架构有时须要将 LLM 裁减到更长的序列（如 100 万 token），这就带来渺小的内存开支，并且随着高低文长度的参与，部署老本也在参与。

基于此，谷歌引入了一种有效的方法，其关键组成部分是一种称为 Infini-attention（有限留意力）的新留意力技术。不同于传统的 Transformer 经常使用部分留意力摈弃旧片段，为新片段监禁内存空间。Infini-attention 参与了紧缩内存（compressive memory），可以将经常使用后的旧片段存储到紧缩内存中，输入时集聚合高低文消息以及紧缩内存中的消息，因此模型可以检索完整的高低文历史。

该方法使 Transformer LLM 在有限内存的状况下裁减到有限长高低文，并以流的形式处置极长的输入启动计算。

试验标明，该方法在长高低白话语建模基准测试中的性能优于基线，同时内存参数缩小了 100 倍以上。当经常使用 100K 序列长度启动训练时，该模型成功了更好的困惑度。此外该钻研发现，1B 模型在 5K 序列长度的密钥实例上启动了微调，处置了 1M 长度的疑问。最后，论文展现了具备 Infini-attention 的 8B 模型经过继续的预训练和义务微调，在 500K 长度的书籍摘要义务上到达了新的 SOTA 结果。

本文奉献总结如下：

方法引见

Infini-attention 使 Transformer LLM 能够经过有限的内存占用和计算有效地处置有限长的输入。如下图 1 所示，Infini-attention 将紧缩记忆融入到个别的留意力机制中，并在单个 Transformer 块中构建了掩码部分留意力和常年线性留意力机制。

对 Transformer 留意力层启动这种巧妙但关键的修正可以经过继续的预训练和微调将现有 LLM 的高低文窗口裁减到有限长。

Infini-attention 驳回规范留意力计算的一切键、值和查问形态，以启动常年记忆坚固（memory consolidation）和检索，并将留意力的旧 KV 形态存储在紧缩内存中，而不是像规范留意力机制那样摈弃它们。在处置后续序列时，Infini-attention 经常使用留意查问形态从内存中检索值。为了计算最终的高低文输入，Infini-attention 聚合了常年记忆检索值和部分留意力高低文。

如下图 2 所示，钻研团队比拟了基于 Infini-attention 的 Infini-Transformer 和 Transformer-XL。与 Transformer-XL 相似，Infini-Transformer 对 segment 序列启动操作，并计算每个 segment 中的规范因果点积留意力高低文。因此，点积留意力计算在某种意义上是部分的。

但是，部分留意力在处置下一个 segment 时会摈弃前一个 segment 的留意力形态，但 Infini-Transformer 复用旧的 KV 留意力形态，以经过紧缩存储来保养整个高低文历史。因此，Infini-Transformer 的每个留意力层都具备全局紧缩形态和部分细粒度形态。

与多头留意力（MHA）相似，除了点积留意力之外，Infini-attention 还为每个留意力层保养 H 个并行紧缩内存（H 是留意力头的数量）。

下表 1 列出了几种模型依据模型参数和输入 segment 长度，定义的高低文内存占用和有效高低文长度。Infini-Transformer 允许具备有限内存占用的有限高低文窗口。

试验

该钻研在长高低白话语建模、长度为 1M 的密钥高低文块检索和 500K 长度的书籍摘要义务上评价了 Infini-Transformer 模型，这些义务具备极长的输入序列。关于言语建模，钻研者选用从头开局训练模型，而关于密钥和书籍摘要义务，钻研者驳回不时预训练 LLM 的形式，以证实 Infini-attention 即插即用的长高低文顺应才干。

长高低白话语建模。表 2 结果标明 Infini-Transformer 优于 Transformer-XL 和 Memorizing Transformers 基线，并且与 Memorizing Transformer 模型相比，存储参数缩小了 114 倍。

密钥义务。表 3 为 Infini-Transformer 在 5K 长度输入上启动微调后，处置了高达 1M 高低文长度的密钥义务。试验中输入 token 的范围从 32K 到 1M，关于每个测试子集，钻研者控制密钥的位置，使其位于输入序列的扫尾、两边或开头左近。试验报告了零样本准确率和微调准确率。在对 5K 长度输入启动 个步骤的微调后，Infini-Transformer 处置了高达 1M 高低文长度的义务。

摘要义务。表 4 将 Infini-Transformer 与专门为摘要义务构建的编码器 - 解码器模型启动了比拟。结果标明 Infini-Transformer 逾越了之前最佳结果，并且经过处置书中的整个文本在 BookSum 上成功了新的 SOTA。

钻研者还在图 4 中绘制了 BookSum 数据验证宰割的总体 Rouge 分数。依据折线趋向标明，随着输入长度的参与，Infini-Transformers 提高了摘要性能目的。

原文链接:​ ​​ ​