[论文解读] Transformers Trained via Gradient Descent Can Provably Learn a Class of Teacher Models

Chenyang Zhang, Qingyue Zhao+2 more|arXiv (Cornell University)|Mar 24, 2026

Advanced Graph Neural Networks被引用 0

一句话总结

论文表明，采用简化的位置仅注意力的一层变换器通过梯度下降可严格学习广义的教师模型族，在收敛到教师参数且具备良好的分布外泛化的同时，达到上界和下界一致的匹配。

ABSTRACT

Transformers have achieved great success across a wide range of applications, yet the theoretical foundations underlying their success remain largely unexplored. To demystify the strong capacities of transformers applied to versatile scenarios and tasks, we theoretically investigate utilizing transformers as students to learn from a class of teacher models. Specifically, the teacher models covered in our analysis include convolution layers with average pooling, graph convolution layers, and various classic statistical learning models, including a variant of sparse token selection models [Sanford et al., 2023, Wang et al., 2024] and group-sparse linear predictors [Zhang et al., 2025]. When learning from this class of teacher models, we prove that one-layer transformers with simplified "position-only'' attention can successfully recover all parameter blocks of the teacher models, thus achieving the optimal population loss. Building upon the efficient mimicry of trained transformers towards teacher models, we further demonstrate that they can generalize well to a broad class of out-of-distribution data under mild assumptions. The key in our analysis is to identify a fundamental bilinear structure shared by various learning tasks, which enables us to establish unified learning guarantees for these tasks when treating them as teachers for transformers.

研究动机与目标

• 通过理论保证理解变换器在从一类教师模型学习时的机理。
• 定义一个统一的双线性结构框架，覆盖 CNN、GCN、稀疏标记选择，以及组稀疏预测器等作为教师模型。
• 在该设定下，证明一层变换器经梯度下降训练的收敛性和泛化性。
• 通过合成数据与真实数据实验，验证理论与训练动态和注意力模式的一致性。

提出的方法

• 定义形式为 f* (X) = sigma(V* X S*) 的教师模型，具有双线性结构及多种实例化（CNN、GCN、STS、GSLP）。
• 采用简化的一层变换器，使用仅位置自注意力：TF(Z; WV; WKQ) = sigma(WV X S)，其中 S 为学习得到的注意力分数。
• 对总体损失进行梯度下降训练，WV 与 WKQ 零初始化，输入 X 为高斯分布，推导迭代更新式（3.3)-(3.4）。
• 从理论上分析对真实分量 (V*, S*) 的收敛，并给出过剩损失和参数收敛的严格界（定理 3.1）。
• 在温和矩的假设下扩展到分布外（OOD）泛化界限（定理 3.2）。
• 通过合成数据与真实数据实验进行验证，包括 CNN、GCN、STS、GSLP 任务，以及基于 MNIST 的设置。

实验结果

研究问题

• RQ1 一层变换器在用梯度下降学习时，是否能够恢复具有双线性结构的一类广义教师模型的所有参数块？
• RQ2 学习得到的注意力分数和值矩阵的收敛速率是多少，过剩损失如何随迭代次数缩放？
• RQ3 学习得到的变换器是否能泛化到超出训练分布的分布外数据？
• RQ4 理论结果在 CNN、GCN、稀疏标记选择和组稀疏线性预测器等教师模型上如何体现？
• RQ5 实证实验是否反映出预测的参数与损失动态及注意力模式？

主要发现

• 对一层变换器进行梯度下降训练可以在严格的收敛保证下恢复教师的值矩阵 V* 和 softmax 分数 S*。
• 注意力分数以 ||S(T)−S*||F = Theta(D^{5/2} / (||V*|| sqrt(eta T))) 的速率收敛到真实值 S*。
• 值矩阵以 ||WV^(T)−V*||F = Theta(D^2 sqrt(K/(eta T))) 的速率收敛到真实值 V*。
• 过剩损失 L(WV^(T); WKQ^(T)) − L_opt 在上下界之间被 Theta(K D^4 /(eta T)) 夹住。
• 该框架覆盖多样的教师模型，包括带平均池化的 CNN、常规图上的 GCN、稀疏标记选择，以及组稀疏线性预测器，并通过合成实验确认理论：损失收敛且注意力模式对齐。
• OOD 泛化界限显示，在二阶矩有界的条件下，训练后的变换器的 OOD 损失与教师的 OOD 损失相差不超过一个 epsilon，建立了鲁棒的泛化性。