ICML 2024 | 揭示非线形Transformer在上下文学习中学习和泛化的机制

上下文学习 (in-context learning, 简写为 ICL) 已经在很多 LLM 有关的应用中展现了强大的能力，但是对其理论的分析仍然比较有限。人们依然试图理解为什么基于 Transformer 架构的 LLM 可以展现出 ICL 的能力。

近期，一个来自美国伦斯勒理工大学和 IBM 研究院的团队从优化和泛化理论的角度分析了带有非线性注意力模块 (attention) 和多层感知机 (MLP) 的 Transformer 的 ICL 能力。他们特别从理论端证明了单层 Transformer 首先在 attention 层根据 query 选择一些上下文示例，然后在 MLP 层根据标签嵌入进行预测的 ICL 机制。该文章已收录在 ICML 2024。

• 论文题目：How Do Nonlinear Transformers Learn and Generalize in In-Context Learning?
• 论文地址：https://arxiv.org/pdf/2402.15607

背景介绍

上下文学习 in context learning (ICL)

ICL 理论工作的进展

近期的很多理论工作都是基于 [1] 所提出的研究框架，即人们可以直接使用 prompt 的格式来对 Transformer 进行训练 (这一步也可以理解为在模拟一种简化的 LLM 预训练模式)，从而使得模型具有 ICL 能力。已有的理论工作聚焦于模型的表达能力 (expressive power) 的角度 [2]。他们发现，人们能够找到一个有着 “完美” 的参数的 Transformer 可以通过前向运算执行 ICL，甚至隐含地执行梯度下降等经典机器学习算法。但是这些工作无法回答为什么 Transformer 可以被训练成这样 “完美” 的，具有 ICL 能力的参数。因此，还有一些工作试图从 Transformer 的训练或泛化的角度理解 ICL 机制 [3,4]。不过，受制于分析 Transformer 结构的复杂性，这些工作目前止步于研究线性回归任务，而所考虑的模型通常会略去 Transformer 中的非线形部分。

本文从优化和泛化理论的角度分析了带有非线性 attention 和 MLP 的 Transformer 的 ICL 能力和机制：

• 基于一个简化的分类模型，本文具体量化了数据的特征如何影响了一层单头 Transformer 的域内 (in-domain) 和域外 (out-of-domain, OOD) 的 ICL 泛化能力。
• 本文进一步阐释了 ICL 是如何通过被训练的 Transformer 来实现了。
• 基于被训练的 Transformer 的特点，本文还分析了在 ICL 推断的时候使用基于幅值的模型剪枝 (magnitude-based pruning) 的可行性。

理论部分

问题描述

训练网络为一个单层单头 Transformer：

预训练过程是求解一个对所有训练任务的经验风险最小化 (empirical risk minimization)。损失函数使用的是适合二分类问题的 Hinge loss，训练算法是随机梯度下降。

本文定义了两种 ICL 泛化的情况。一个是 in-domain 的，即泛化的时候测试数据的分布和训练数据一样，注意这个情况里面测试任务不必和训练任务一样，即这里已经考虑了对未见任务 (unseen task) 的泛化。另一个是 out-of-domain 的，即测试、训练数据分布不一样。

本文还涉及了在 ICL 推断的时候进行 magnitude-based pruning 的分析，这里的剪枝方式是指对于训练得到的中的各个神经元，根据其幅值大小，进行从小到大的删除。

对数据和任务的构建

理论结果

首先，对于 in-domain 的情况，本文先给了一个 condition 3.2 来规定训练任务需要满足的条件，即训练任务需要覆盖所有的 IDR pattern 和标签。然后 in-domain 的结果如下：

这里表明：1，训练任务的数量只需要在全部任务中占比达到满足 condition 3.2 的小比例，我们就可以对 unseen task 实现很好的泛化；2，跟当前任务相关的 IDR pattern 在 prompt 中的比例越高，就可以以更少的训练数据，训练迭代次数，以及更短的 training/testing prompt 实现理想的泛化。

接下来是 out-of-domain 泛化的结果。

然后，本文给出了带有 magnitude-based pruning 的 ICL 泛化结果。

ICL 机制

通过对预训练过程的刻画，本文得到了单层单头非线性 Transformer 做 ICL 的内在机制，这一部分在原文的 Section 4。该过程可以用下图表示。

简而言之，attention 层会选择和 query 的 ODR/IDR pattern 一样的上下文，赋予它们几乎全部 attention 权重，然后 MLP 层会重点根据 attention 层输出中的标签嵌入来作出最后的分类。

总结

本文讲解了在 ICL 当中，非线性 Transformer 的训练机制，以及对于新任务和分布偏移数据的泛化能力。理论结果对于设计 prompt 选择算法和 LLM 剪枝算法有一定实际意义。

参考文献

[1] Garg, et al., Neurips 2022. "What can transformers learn in-context? a case study of simple function classes."

[2] Von Oswald et al., ICML 2023. "Transformers learn in-context by gradient descent."

[3] Zhang et al., JMLR 2024. "Trained transformers learn linear models in-context."

[4] Huang et al., ICML 2024. "In-context convergence of transformers."

文章来自于微信公众号“机器之心”，作者 “机器之心”