Looped Transformers可自动学习多步梯度下降,实现零样本推理。
Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning?
- 用循环权重共享的多层Transformer模拟多步梯度下降算法。
- 理论证明训练后全局最优解对应自适应预条件梯度下降,收敛快。
- 首次为多层Transformer在上下文学习中提供理论支持,适合研究者参考。
Transformer在无需微调的情况下展现出强大的推理与少样本学习能力,普遍认为其源于权重在单次前向传播中隐式模拟多步算法(如梯度下降)的能力。近期研究从表达性角度表明Transformer能表示此类算法,但对更根本的可学习性问题——尤其在多层模型中——仍知之甚少。本文针对线性上下文回归任务,分析了具有权重共享的循环多层Transformer(looped Transformer),证明其总体训练损失的全局最小值对应于一种自适应数据分布的预条件梯度下降算法。同时,通过建立新的梯度主导性条件,证明了梯度流在非凸损失面上仍具快速收敛性。这是首个针对该设置的多层Transformer理论分析。合成实验验证了理论发现。
原文摘要 · Abstract (English)
The remarkable capability of Transformers to do reasoning and few-shot learning, without any fine-tuning, is widely conjectured to stem from their ability to implicitly simulate a multi-step algorithms -- such as gradient descent -- with their weights in a single forward pass. Recently, there has been progress in understanding this complex phenomenon from an expressivity point of view, by demonstrating that Transformers can express such multi-step algorithms. However, our knowledge about the more fundamental aspect of its learnability, beyond single layer models, is very limited. In particular, can training Transformers enable convergence to algorithmic solutions? In this work we resolve this for in-context linear regression with linear looped Transformers -- a multi-layer model with weight sharing that is conjectured to have an inductive bias to learn fix-point iterative algorithms. More specifically, for this setting we show that the global minimizer of the population training loss implements multi-step preconditioned gradient descent, with a preconditioner that adapts to the data distribution. Furthermore, we show a fast convergence for gradient flow on the regression loss, despite the non-convexity of the landscape, by proving a novel gradient dominance condition. To our knowledge, this is the first theoretical analysis for multi-layer Transformer in this setting. We further validate our theoretical findings through synthetic experiments.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。