[논문 리뷰] Transformers Trained via Gradient Descent Can Provably Learn a Class of Teacher Models

Chenyang Zhang, Qingyue Zhao외 2명|arXiv (Cornell University)|2026. 03. 24.

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, 희소 토큰 선택, 그룹-희소 예측기를 교사 모델로 포함하는 통합 이중선형 구조 프레임워크를 정의한다.
• 이 설정에서 그라디언트 하강으로 학습된 한 층 트랜스포머에 대한 수렴 및 일반화 보장을 확립한다.
• 합성 및 실 데이터 실험을 통해 이론이 관찰된 학습 역학 및 주의 패턴과 일치함을 입증한다.

제안 방법

• 이중선형 구조를 갖고 다양한 구현(CNN, GCN, STS, GSLP)로 구성된 f*(X) = sigma(V* X S*) 형태의 교사 모델을 정의한다.
• 위치-전용 자기 주의를 갖는 간소화된 한 층 트랜스포머를 채택한다: TF(Z; WV; WKQ) = sigma(WV X S) 여기서 S는 학습된 어텐션 점수이다.
• WV와 WKQ를 0으로 초기화하고 Gaussian 입력 X에 대해 모집단 손실에 대해 그래디언트 하강으로 학습하고 반복 업데이트(3.3)-(3.4)를 도출한다.
• 해당 논문은 실제 매개변수(V*, S*)로의 수렴을 이론적으로 분석하고 초과 손실 및 매개변수 수렴에 대해 엄밀한 상한을 제공한다(정리 3.1).
• 완화된 모멘트 가정 하에서 분포 외 일반화 경계로 확장한다(정리 3.2).
• CNN, GCN, STS, GSLP 과제 및 MNIST 기반 설정을 포함한 합성 및 실 데이터 실험으로 검증한다.

실험 결과

연구 질문

• RQ1 한 층의 트랜스포머가 그래디언트 하강으로 학습될 때, 이중선형 구조를 가진 광범위한 교사 모델의 모든 매개변수 블록을 복구할 수 있는가?
• RQ2 학습된 주의 점수와 값 행렬의 수렴 속도는 어떠하며, 초과 손실은 반복에 따라 어떻게 스케일하는가?
• RQ3 학습 데이터 분포를 넘어서는 분포 외 데이터에 대해 학습된 트랜스포머가 일반화되는가?
• RQ4 이론적 결과가 CNN, GCN, 희소 토큰 선택, 그룹-희소 선형 예측기 교사 모델에서 어떻게 드러나는가?
• RQ5 경험적 실험이 예측된 매개변수 및 손실 역학과 주의 패턴을 반영하는가?

주요 결과

• 그래디언트 하강으로 학습된 한 층 트랜스포머가 교사의 값 행렬 V*와 소프트맥스 점수 S*를 정밀한 수렴 보장과 함께 복구할 수 있다.
• 주의 점수는 초기 상태에서ground-truth S*로 수렴하는 속도가 ||S(T)−S*||F = Theta(D^{5/2} / (||V*|| sqrt(eta T)))이다.
• 값 행렬은 ||WV^(T)−V*||F = Theta(D^2 sqrt(K/(eta T)))의 속도로 ground-truth V*로 수렴한다.
• 초과 손실 L(WV^(T); WKQ^(T)) − L_opt은 Theta(K D^4 /(eta T))로 상한과 하한 모두에 의해 제시된다.
• 이 프레임워크는 평균 풀링을 사용하는 CNN, 규칙 그래프상 GCN, 희소 토큰 선택, 그룹-희소 선형 예측기를 포함한 다양한 교사 모델을 포괄하며, 합성 실험에서 수렴하는 손실과 정렬된 주의 패턴으로 이론을 확인한다.
• OOD 일반화 경계는 학습된 트랜스포머의 OOD 손실이 두 번째 모멘트가 제한된 경우 교사의 OOD 손실보다 epsilon 이내임을 보여주어 견고한 일반화를 확립한다.