[论文解读] Repeat After Me: Transformers are Better than State Space Models at Copying

Samy Jelassi, David Brandfonbrener+2 more|arXiv (Cornell University)|Feb 1, 2024

Simulation Techniques and ApplicationsDecision Sciences被引用 3

一句话总结

该论文表明，由于注意力机制能够存储和检索n-gram，Transformers在长输入序列中复制和检索信息方面优于广义状态空间模型（GSSMs）。理论上，两层Transformer可以复制指数级长度的序列，而GSSMs则因其固定大小的潜在状态而受到根本限制，其编码的位数无法超过其维度。

ABSTRACT

Transformers are the dominant architecture for sequence modeling, but there is growing interest in models that use a fixed-size latent state that does not depend on the sequence length, which we refer to as "generalized state space models" (GSSMs). In this paper we show that while GSSMs are promising in terms of inference-time efficiency, they are limited compared to transformer models on tasks that require copying from the input context. We start with a theoretical analysis of the simple task of string copying and prove that a two layer transformer can copy strings of exponential length while GSSMs are fundamentally limited by their fixed-size latent state. Empirically, we find that transformers outperform GSSMs in terms of efficiency and generalization on synthetic tasks that require copying the context. Finally, we evaluate pretrained large language models and find that transformer models dramatically outperform state space models at copying and retrieving information from context. Taken together, these results suggest a fundamental gap between transformers and GSSMs on tasks of practical interest.

研究动机与目标

• 探究尽管GSSMs具有高效性，是否能在从输入上下文中复制和检索信息方面与Transformers相媲美。
• 分析由于固定大小的潜在状态，GSSMs在表示长序列时的根本性限制。
• 评估Transformers与GSSMs在合成复制任务上的泛化能力与效率。
• 在零样本设置下，测试预训练模型（如Pythia与Mamba）在上下文检索与复制任务中的表现。
• 理解架构差异对下游语言模型能力的实际影响，超越困惑度指标。

提出的方法

• 使用GSSMs与Transformers的形式化定义，对字符串复制任务进行理论分析，证明GSSMs无法复制超过其潜在状态大小允许的序列。
• 构建一个两层Transformer，通过注意力机制实现n-gram存储与检索，使其能够复制长度为注意力头数量指数级的序列。
• 在合成数据集上进行实证评估，模型被训练以复制逐渐增长长度的序列，同时测量训练效率与长度泛化能力。
• 使用Hard-Alibi位置编码，以提升Transformers在训练上下文窗口之外的泛化能力。
• 在基于上下文的查找任务上评估预训练模型（Pythia与Mamba）的表现，测量不同段落长度下的F1分数。
• 分析训练后Transformers中学习到的注意力模式，以确认其复制机制依赖于n-gram查找机制。

(a) Copying: training efficiency.

实验结果

研究问题

• RQ1 Transformer模型能否复制长度为自身规模指数级的序列？如果可以，它是如何实现的？
• RQ2 由于固定大小的潜在状态，GSSMs在复制任务中的根本性表示限制是什么？
• RQ3 在合成复制任务中，Transformers与GSSMs在训练效率和向更长序列泛化方面的表现如何比较？
• RQ4 预训练的Transformer是否在零样本上下文检索与复制任务中优于预训练的GSSM？
• RQ5 GSSMs固定的内存容量在多大程度上会降低其在具有长输入上下文的自然语言任务中的性能？

主要发现

• 两层Transformer能够复制长度为注意力头数量指数级的序列，展示了通过n-gram存储实现上下文检索的强大机制。
• GSSMs在复制任务中存在根本性限制：它们无法准确复制比特数超过其潜在状态大小的序列。
• 在合成复制任务中，Transformers的训练速度显著快于GSSMs，且在更长序列上的泛化能力也远超GSSMs，而Mamba在超过训练长度后性能迅速下降。
• 在零样本上下文查找任务中，基于Transformer的预训练模型Pythia表现优于基于GSSM的Mamba，且随着段落长度增加，Mamba的F1分数下降更为显著。
• Transformers中学习到的复制机制依赖于基于注意力的n-gram查找，证实了理论构造的正确性。
• 尽管Mamba在Pile数据集上的困惑度低于Pythia，但在上下文密集型任务中表现更差，表明训练损失本身并不能预测模型在检索与复制任务中的表现。

(b) Copying: length generalization