Transformers Trained via Gradient Descent Can Provably Learn a Class of Teacher Models

TL;DR

This paper provides a theoretical analysis showing that one-layer transformers trained via gradient descent can learn a broad class of teacher models, including convolutional and statistical models, with provable guarantees.

Contribution

It introduces a theoretical framework demonstrating transformers' ability to learn diverse teacher models with simplified attention mechanisms and provides generalization guarantees.

Findings

Transformers can recover teacher model parameters with optimal loss.

One-layer transformers with position-only attention can learn various teacher models.

Transformers generalize well to out-of-distribution data under mild conditions.

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.