[论文解读] Linear Transformers Are Secretly Fast Weight Programmers

Imanol Schlag, Kazuki Irie+1 more|arXiv (Cornell University)|Feb 22, 2021

Neural Networks and Applications被引用 24

一句话总结

本文建立了线性化自注意力机制与20世纪90年代的快速权重编程器（FWPs）之间的形式等价性，揭示了线性Transformer通过键和值的加法外积隐式学习快速权重记忆。该文提出一种受delta规则启发的更新机制，实现动态内存编辑，并在过容量设置下提升学习性能，实验证明其在语言建模与机器翻译任务中达到最先进性能，且计算开销极低。

ABSTRACT

We show the formal equivalence of linearised self-attention mechanisms and fast weight controllers from the early '90s, where a ``slow" neural net learns by gradient descent to program the ``fast weights" of another net through sequences of elementary programming instructions which are additive outer products of self-invented activation patterns (today called keys and values). Such Fast Weight Programmers (FWPs) learn to manipulate the contents of a finite memory and dynamically interact with it. We infer a memory capacity limitation of recent linearised softmax attention variants, and replace the purely additive outer products by a delta rule-like programming instruction, such that the FWP can more easily learn to correct the current mapping from keys to values. The FWP also learns to compute dynamically changing learning rates. We also propose a new kernel function to linearise attention which balances simplicity and effectiveness. We conduct experiments on synthetic retrieval problems as well as standard machine translation and language modelling tasks which demonstrate the benefits of our methods.

研究动机与目标

• 揭示线性化自注意力机制与20世纪90年代一类模型——快速权重编程器（FWPs）之间的形式等价性。
• 识别并解决由纯加法外积更新导致的线性Transformer内存容量限制问题。
• 提出一种受delta规则启发的新颖可微编程指令，使FWP能够通过梯度下降动态修正键到值的映射。
• 提出一种新的、简单且高效的核函数，用于线性化注意力机制，平衡计算效率与模型性能。
• 在合成检索任务与标准NLP基准（包括机器翻译与语言建模）上实证验证所提方法。

提出的方法

• 本文证明线性自注意力机制在形式上等价于快速权重编程器（FWPs），其中慢速网络通过自生成的键与值的加法外积序列学习编程快速权重。
• 提出一种类似delta规则的更新规则，作为纯加法外积的替代方案，使FWP能够通过梯度下降学习修正当前键到值的映射。
• 该方法将标准加法更新替换为一种可微指令，其形式类似于误差校正的delta规则，从而实现自适应内存编辑。
• 提出一种新的核函数，用于线性化softmax注意力机制，兼顾简洁性与有效性，以实现高效的序列建模。
• 模型采用截断的时间反向传播进行训练，跨段落保留快速权重记忆，同时将梯度计算限制在段内范围内。
• 实验使用自定义CUDA内核以确保效率，并与标准线性Transformer、Performer及Transformer-XL在合成与真实任务上进行对比。

实验结果

研究问题

• RQ1 线性化自注意力机制与20世纪90年代的快速权重编程器（FWPs）之间是否存在形式等价性？
• RQ2 在纯加法外积更新下，线性Transformer的内存容量存在哪些局限性，特别是在过容量场景下？
• RQ3 受delta规则启发的更新规则能否提升FWP在过容量设置下动态编辑内存内容的能力？
• RQ4 所提出的用于线性化注意力的核函数在简洁性与有效性方面相较于现有方法表现如何？
• RQ5 改进后的FWP机制能否在标准语言建模与机器翻译基准上达到最先进性能？

主要发现

• 采用所提delta规则更新的Delta网络在Wikitext-103测试集上达到27.4的困惑度，优于基线线性Transformer，并与状态更优的Transformer-XL相当，但参数量显著更小。
• Delta网络在无限上下文长度下仍保持稳定性能，而基线线性Transformer因采用朴素求和更新规则，在过容量条件下性能崩溃。
• 模型推理速度达63 K词/秒，显存占用13 GB，优于标准PyTorch Transformer（33 K词/秒，17 GB），且与Performer（57 K词/秒）性能相当。
• 所提核函数实现了更高效且更简洁的线性注意力，平衡了计算效率与模型性能。
• 从FWP视角出发，模型能够学习动态学习率并选择性管理键-值关联，显著提升长上下文场景下的鲁棒性。
• 该方法在合成检索任务与真实世界NLP基准上均表现出强大泛化能力，证实了通过梯度学习的编程指令实现动态内存编辑的优势。