Transformers learn variable-order Markov chains in-context
作者: Ruida Zhou, Chao Tian, Suhas Diggavi
分类: cs.LG, cs.IT
发布日期: 2024-10-07
💡 一句话要点
研究Transformer上下文学习变阶马尔可夫链能力，并提出CTW算法的Transformer构造。
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 上下文学习 变阶马尔可夫链 Transformer 数据压缩 上下文树加权
📋 核心要点
- 现有研究主要集中在Transformer上下文学习固定阶马尔可夫链（FOMC），忽略了自然语言更适合用变阶马尔可夫链（VOMC）建模的特性。
- 论文将语言建模视为数据压缩，利用上下文树加权（CTW）等成熟算法作为基线，研究Transformer在上下文中学习VOMC的能力。
- 实验表明Transformer能有效压缩VOMC，且性能对层数不敏感。论文还构建了能模拟CTW算法的Transformer结构，并验证了其有效性。
📝 摘要（中文）
大型语言模型展现了卓越的上下文学习（ICL）能力。然而，Transformer如何实现ICL，尤其是在复杂场景下，仍然不清楚。为了解决这个问题，一些研究探讨了Transformer如何在上下文中学习固定阶马尔可夫链（FOMC），但自然语言更适合用变阶马尔可夫链（VOMC），即上下文树（CTs）建模。本文将语言建模视为一种数据压缩形式，研究VOMC的ICL，并关注小字母表和低阶VOMC。这种视角使我们能够利用成熟的压缩算法，如上下文树加权（CTW）和部分匹配预测（PPM）算法作为基线，其中CTW对于一类CTW先验是贝叶斯最优的。实验观察到：1) Transformer确实可以学习在上下文中压缩VOMC，而PPM表现不佳；2) Transformer的性能对层数不敏感，即使是两层Transformer也能很好地进行上下文学习；3) 在非CTW先验上训练和测试的Transformer可以显著优于CTW算法。为了解释这些现象，我们分析了Transformer的注意力图，并提取了两种机制，并基于此提供了两种Transformer构造：1) 一种具有D+2层的构造，可以精确地模拟最大阶数为D的CT的CTW算法；2) 一种使用前馈网络进行概率混合的两层Transformer。与FOMC设置的一个区别是，计数机制似乎起着重要作用。我们实现了这些合成Transformer层，并表明这种混合Transformer可以匹配Transformer的ICL性能，更有趣的是，尽管参数集大大减少，但其中一些可以表现得更好。
🔬 方法详解
问题定义：现有研究主要关注Transformer在上下文中学习固定阶马尔可夫链（FOMC）的能力，但自然语言更适合用变阶马尔可夫链（VOMC）建模。现有方法难以有效处理VOMC的上下文学习，且缺乏对Transformer内部机制的深入理解。
核心思路：论文将语言建模问题转化为数据压缩问题，并借鉴了数据压缩领域的上下文树加权（CTW）算法。通过分析Transformer的注意力机制，提取出Transformer学习VOMC的关键机制，并基于此构建具有特定结构的Transformer，以模拟CTW算法。
技术框架：论文首先通过实验观察Transformer在VOMC上的上下文学习能力，并与PPM等算法进行比较。然后，分析Transformer的注意力图，提取出两种关键机制：一是用于模拟CTW算法的机制，二是用于概率混合的机制。基于这些机制，论文提出了两种Transformer构造：一种是D+2层的Transformer，用于精确模拟CTW算法；另一种是两层Transformer，利用前馈网络进行概率混合。
关键创新：论文的关键创新在于：1) 将语言建模问题与数据压缩问题联系起来，利用CTW算法作为基线；2) 通过分析Transformer的注意力图，揭示了Transformer学习VOMC的关键机制；3) 基于这些机制，构建了具有特定结构的Transformer，能够有效模拟CTW算法，并取得了良好的性能。
关键设计：论文的关键设计包括：1) 使用小字母表和低阶VOMC，简化了问题；2) 设计了D+2层的Transformer结构，其中D是上下文树的最大阶数，该结构能够精确模拟CTW算法；3) 设计了两层Transformer结构，利用前馈网络进行概率混合；4) 强调了计数机制在VOMC学习中的重要作用。
🖼️ 关键图片
📊 实验亮点
实验结果表明，Transformer能够有效地学习在上下文中压缩VOMC，并且性能对层数不敏感。更重要的是，论文构建的混合Transformer能够在参数量大大减少的情况下，匹配甚至超过标准Transformer的ICL性能，验证了论文提出的机制的有效性。
🎯 应用场景
该研究成果可应用于提升语言模型的上下文学习能力，尤其是在需要处理变长上下文依赖关系的场景中，例如自然语言理解、文本生成、对话系统等。通过理解Transformer学习VOMC的机制，可以设计更高效、更可控的语言模型。
📄 摘要（原文）
Large language models have demonstrated impressive in-context learning (ICL) capability. However, it is still unclear how the underlying transformers accomplish it, especially in more complex scenarios. Toward this goal, several recent works studied how transformers learn fixed-order Markov chains (FOMC) in context, yet natural languages are more suitably modeled by variable-order Markov chains (VOMC), i.e., context trees (CTs). In this work, we study the ICL of VOMC by viewing language modeling as a form of data compression and focus on small alphabets and low-order VOMCs. This perspective allows us to leverage mature compression algorithms, such as context-tree weighting (CTW) and prediction by partial matching (PPM) algorithms as baselines, the former of which is Bayesian optimal for a class of CTW priors. We empirically observe a few phenomena: 1) Transformers can indeed learn to compress VOMC in-context, while PPM suffers significantly; 2) The performance of transformers is not very sensitive to the number of layers, and even a two-layer transformer can learn in-context quite well; and 3) Transformers trained and tested on non-CTW priors can significantly outperform the CTW algorithm. To explain these phenomena, we analyze the attention map of the transformers and extract two mechanisms, on which we provide two transformer constructions: 1) A construction with $D+2$ layers that can mimic the CTW algorithm accurately for CTs of maximum order $D$, 2) A 2-layer transformer that utilizes the feed-forward network for probability blending. One distinction from the FOMC setting is that a counting mechanism appears to play an important role. We implement these synthetic transformer layers and show that such hybrid transformers can match the ICL performance of transformers, and more interestingly, some of them can perform even better despite the much-reduced parameter sets.