单层线性自注意力下梯度下降一步被证明是最优的上下文学习器
机器学习
2023-07-10 v1
摘要
近期工作对上下文学习进行了实证分析，并表明在合成线性回归任务上训练的 transformer 能够学会实现岭回归（即贝叶斯最优预测器），前提是容量充足 [Akyürek et al., 2023]，而具有线性自注意力且无 MLP 层的单层 transformer 会学会在最小二乘线性回归目标上实现一步梯度下降（GD）[von Oswald et al., 2022]。然而，这些观察背后的理论仍知之甚少。我们从理论上研究具有单层线性自注意力的 transformer，其在合成含噪线性回归数据上训练。首先，我们从数学上证明，当协变量从标准高斯分布中抽取时，使预训练损失最小的单层 transformer 将实现最小二乘线性回归目标上的一步 GD。接着，我们发现将协变量和权向量的分布改为非各向同性高斯分布对所学算法有强烈影响：预训练损失的全局极小元现在实现一步 GD。然而，如果仅改变响应的分布，则对所学算法影响不大：即使响应来自更一般的 函数族，预训练损失的全局极小元仍实现最小二乘线性回归目标上的一步 GD。
引用
@article{arxiv.2307.03576,
title = {One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention},
author = {Arvind Mahankali and Tatsunori B. Hashimoto and Tengyu Ma},
journal= {arXiv preprint arXiv:2307.03576},
year = {2023}
}