Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Finite-Time Analysis of Gradient Descent for Shallow Transformers

Enes Arda, Semih Cayci, Atilla Eryilmaz

TL;DR

The paper analyzes training dynamics of a shallow, multi‑head Transformer in the kernel regime under projected gradient descent, addressing nonconvex optimization with a finite‑time, nonasymptotic lens. It preserves attention nonlinearity, derives a Transformer NTK decomposition including a nonzero attention component, and uses transportation mappings to connect initialization to a kernel‑regime RKHS. The results show that width scales only logarithmically with the sample size $n$, the optimization error is independent of sequence length $T$, and a clear memory‑time trade‑off exists due to full context attention. Empirical teacher–student experiments validate the predicted scaling laws and demonstrate favorable comparisons to IndRNNs on long‑range dependency tasks, highlighting practical implications for training efficiency and memory usage.

Abstract

Understanding why Transformers perform so well remains challenging due to their non-convex optimization landscape. In this work, we analyze a shallow Transformer with $m$ independent heads trained by projected gradient descent in the kernel regime. Our analysis reveals two main findings: (i) the width required for nonasymptotic guarantees scales only logarithmically with the sample size $n$, and (ii) the optimization error is independent of the sequence length $T$. This contrasts sharply with recurrent architectures, where the optimization error can grow exponentially with $T$. The trade-off is memory: to keep the full context, the Transformer's memory requirement grows with the sequence length. We validate our theoretical results numerically in a teacher-student setting and confirm the predicted scaling laws for Transformers.

Finite-Time Analysis of Gradient Descent for Shallow Transformers

TL;DR

The paper analyzes training dynamics of a shallow, multi‑head Transformer in the kernel regime under projected gradient descent, addressing nonconvex optimization with a finite‑time, nonasymptotic lens. It preserves attention nonlinearity, derives a Transformer NTK decomposition including a nonzero attention component, and uses transportation mappings to connect initialization to a kernel‑regime RKHS. The results show that width scales only logarithmically with the sample size , the optimization error is independent of sequence length , and a clear memory‑time trade‑off exists due to full context attention. Empirical teacher–student experiments validate the predicted scaling laws and demonstrate favorable comparisons to IndRNNs on long‑range dependency tasks, highlighting practical implications for training efficiency and memory usage.

Abstract

Understanding why Transformers perform so well remains challenging due to their non-convex optimization landscape. In this work, we analyze a shallow Transformer with independent heads trained by projected gradient descent in the kernel regime. Our analysis reveals two main findings: (i) the width required for nonasymptotic guarantees scales only logarithmically with the sample size , and (ii) the optimization error is independent of the sequence length . This contrasts sharply with recurrent architectures, where the optimization error can grow exponentially with . The trade-off is memory: to keep the full context, the Transformer's memory requirement grows with the sequence length. We validate our theoretical results numerically in a teacher-student setting and confirm the predicted scaling laws for Transformers.
Paper Structure (56 sections, 16 theorems, 155 equations, 2 figures)

This paper contains 56 sections, 16 theorems, 155 equations, 2 figures.

Table of Contents

  1. INTRODUCTION
  2. RELATED WORK
  3. Neural Tangent Kernel (NTK) analysis.
  4. Asymptotic theory for Transformers.
  5. Nonasymptotic theory for Transformers.
  6. Notation.
  7. PROBLEM SETUP
  8. Data
  9. Transformer Architecture
  10. Algorithms for Empirical Risk Minimization
  11. OVERVIEW OF MAIN RESULTS
  12. TRANSFORMER NTK ANALYSIS
  13. Transformer Gradient Analysis
  14. Initialization
  15. Transformer NTK
  16. ...and 41 more sections

Key Result

Lemma 1

Fix $i\in[m]$ and write $a_i:=a(X;W_i)$. Then where and $J_s(z) \;=\; \operatorname{diag}\!(\sigma_s(z)) - \sigma_s(z)\sigma_s(z)^\top$ is the Jacobian matrix of softmax.

Figures (2)

  • Figure 1: Width scaling across three metrics (log–log). Shaded bands show mean $\pm$95% CI across seeds. Fitted slopes are consistent with the predicted $m^{-1/2}$ decay.
  • Figure 2: Performance comparison of IndRNN and Transformer architectures with varying time lag $L$.

Theorems & Definitions (39)

  • Remark 1
  • Remark 2
  • Remark 3
  • Lemma 1: Gradients
  • Remark 4
  • Remark 5
  • Remark 6
  • Definition 1
  • Definition 2
  • Remark 7
  • ...and 29 more