[论文解读] Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection

Yu Bai, Fan Chen+3 more|arXiv (Cornell University)|Jun 7, 2023

Adversarial Robustness in Machine Learning被引用 9

一句话总结

该论文提出一个全面理论，表明变换器可以实现广泛的上下文学习（in-context learning）算法类（如最小二乘、岭回归、Lasso、GLMs），并且能够进行上下文算法选择以适应不同任务，具有理论保证和经验验证。

ABSTRACT

Neural sequence models based on the transformer architecture have demonstrated remarkable \emph{in-context learning} (ICL) abilities, where they can perform new tasks when prompted with training and test examples, without any parameter update to the model. This work first provides a comprehensive statistical theory for transformers to perform ICL. Concretely, we show that transformers can implement a broad class of standard machine learning algorithms in context, such as least squares, ridge regression, Lasso, learning generalized linear models, and gradient descent on two-layer neural networks, with near-optimal predictive power on various in-context data distributions. Using an efficient implementation of in-context gradient descent as the underlying mechanism, our transformer constructions admit mild size bounds, and can be learned with polynomially many pretraining sequences. Building on these ``base'' ICL algorithms, intriguingly, we show that transformers can implement more complex ICL procedures involving \emph{in-context algorithm selection}, akin to what a statistician can do in real life -- A \emph{single} transformer can adaptively select different base ICL algorithms -- or even perform qualitatively different tasks -- on different input sequences, without any explicit prompting of the right algorithm or task. We both establish this in theory by explicit constructions, and also observe this phenomenon experimentally. In theory, we construct two general mechanisms for algorithm selection with concrete examples: pre-ICL testing, and post-ICL validation. As an example, we use the post-ICL validation mechanism to construct a transformer that can perform nearly Bayes-optimal ICL on a challenging task -- noisy linear models with mixed noise levels. Experimentally, we demonstrate the strong in-context algorithm selection capabilities of standard transformer architectures.

研究动机与目标

• 证明变换器在不更新参数的情况下可以在上下文中实现标准ML算法。
• 提供端到端的理论，阐明表达能力、上下文预测和预训练样本复杂度。
• 引入并分析上下文算法选择机制（预ICL测试和后ICL验证）。
• 证明单一变换器可以在不同任务和数据分布之间自适应选择基础ICL算法。

提出的方法

• 为上下文岭回归和最小二乘构建显式的基于变换器的实现。
• 扩展到上下文广义线性模型和凸风险最小化。
• 在变换器内开发高效的上下文梯度下降机制。
• 证明两种算法选择机制：后ICL验证和前ICL测试。
• 给出预训练结果，表明变换器可以从多序列中学习多种ICL任务。
• 从理论和经验角度验证并证明单一变换器在嘈杂设置下接近贝叶斯最优的ICL。

实验结果

研究问题

• RQ1 变换器是否能够在上下文中实现广泛的标准ML算法（如最小二乘、岭回归、Lasso、GLMs）？
• RQ2 单一变换器通过哪些机制在序列之间选择不同的基础ICL算法？
• RQ3 如何界定实现准确上下文学习所需的层/头部复杂度和权重范数？
• RQ4 在多样任务的训练下，变换器在具有挑战性噪声的设置中是否接近贝叶斯最优ICL？
• RQ5 预训练如何影响变换器的上下文学习和算法选择能力？

主要发现

• 变换器可以在上下文中逼近岭回归和最小二乘，且具有可证明的误差保证。
• 变换器可以在上下文中实现GLMs的凸风险最小化以及两层网络的梯度下降。
• 存在两种普适的算法选择机制：前ICL测试和后ICL验证，能够实现自适应的任务/算法选择。
• 在带有混合噪声水平的嘈杂线性模型上，单一变换器通过后ICL验证实现接近贝叶斯最优的ICL。
• 预训练变换器有助于从多序列训练中学习多种ICL任务。
• 实验结果支持在代表性任务上对上下文算法选择能力的有力证据。