单层 Transformer 可在上下文学习中可证地学习单近邻规则
机器学习
2024-11-19 v1 人工智能
最优化与控制
摘要
Transformer 近年来取得了巨大成功。有趣的是，Transformer 展现出了特别强的上下文学习能力——即使不进行微调，它们仍能仅基于特定任务的提示很好地解决未见任务。本文在理论框架下研究单层 Transformer 学习最经典的非参数估计器之一——单近邻预测规则的能力，该框架中提示包含一系列带标签的训练数据和无标签的测试数据。我们证明，尽管使用梯度下降训练时损失函数是非凸的，但单个 softmax 注意力层仍能成功学习表现得像一个单近邻分类器。我们的结果给出了 Transformer 如何被训练以实现非参数机器学习算法的一个具体例子，并揭示了 softmax 注意力在 Transformer 模型中的作用。
引用
@article{arxiv.2411.10830,
title = {One-Layer Transformer Provably Learns One-Nearest Neighbor In Context},
author = {Zihao Li and Yuan Cao and Cheng Gao and Yihan He and Han Liu and Jason M. Klusowski and Jianqing Fan and Mengdi Wang},
journal= {arXiv preprint arXiv:2411.10830},
year = {2024}
}