Finite-Time Analysis of Gradient Descent for Shallow Transformers

AI summary

Authors: Enes Arda (The Ohio State University), Semih Cayci (RWTH Aachen University), Atilla Eryilmaz (The Ohio State University) arXiv: 2601.16514v2 [cs.LG], 02 Apr 2026

Overview

Research area: Optimization theory for deep learning — specifically, nonasymptotic convergence analysis of gradient descent for Transformer architectures, sitting at the intersection of learning theory, kernel methods (NTK), and attention-based sequence models.

Technical level: Advanced. The paper is written for readers comfortable with neural tangent kernels, overparameterization arguments, RKHS theory, and convergence-rate proofs. The informal theorem statement and the "methodology in plain English" framing make the conclusions accessible, but the body relies heavily on matrix calculus and concentration inequalities.

Scope in one sentence: The paper proves finite-time, nonasymptotic training guarantees for a shallow multi-head Transformer with genuine softmax attention trained by projected gradient descent, showing that the required width grows only logarithmically in the sample size and that the optimization error does not depend on the input sequence length.

What This Paper Is About

Transformers achieve state-of-the-art results across language, vision, and multimodal tasks, but a sharp nonasymptotic understanding of their training dynamics is still missing. Two obstacles stand out: self-attention is highly nonlinear because each output coordinate depends on all inputs through a softmax reweighting, and the training landscape is nonconvex, which complicates guarantees for gradient-based methods beyond infinite-width limits.

This paper's goal is to analyze a shallow Transformer with m independent heads trained by projected gradient descent in the kernel regime, while deliberately preserving the real softmax nonlinearity rather than removing or linearizing it. The authors ask two concrete questions: how wide does such a network need to be to guarantee training convergence, and how does that training cost scale with the length of the input sequence?

Key Contributions

1. Preserving attention nonlinearity with multiple heads. Unlike prior analyses that remove or linearize softmax, or that force degeneracy via a β = m⁻¹ scaling, this analysis retains the genuine attention nonlinearity and permits independent heads. The head-independent design is shown to be equivalent to a block-diagonal feed-forward weight matrix with one neuron per head.

2. Bypassing positive definiteness of the NTK/Gram matrix. Where the original NTK literature requires strict positive definiteness of a Gram matrix to drive convergence, this analysis does not. The authors work instead in a projected near-initialization regime where the linearization error is explicitly controlled.

3. Logarithmic overparameterization. The width required for nonasymptotic training guarantees scales only logarithmically with the sample size n. The authors state that, to their knowledge, this had not previously been shown for Transformers.

4. Optimization error independent of input sequence length T. The convergence bound does not degrade with the sequence length T. The authors contrast this sharply with RNN analyses, where the dependence on T can be exponential under long-term dependencies. The stated trade-off is memory: retaining full context implies the Transformer's memory footprint grows with T due to attention over all tokens.

5. Experimental validation. Controlled teacher–student experiments validate the predicted scaling laws, and the paper contrasts Transformers with independently recurrent neural networks (IndRNNs) in their handling of long-term dependencies and memory complexity.

Main Findings

• Three-term error decomposition. After τ steps of ProjGD, the averaged iterate satisfies, with probability at least 1 − δ − δ′ over the random initialization:

  L̂_n(φ̄^(τ)) ≲ D²/√τ [optimization error] + √(D log(n/δ)/m) [approximation error] + √(D³/m) [linearization error], where D := d + log(m/δ′).

  For ProjSGD, the same bound holds for E[L̂_n(φ̄^(τ)) | φ⁽⁰⁾] with the same rates.

• Non-asymptotic in both time and width. The training loss admits a finite-time optimization–approximation–linearization decomposition with the canonical 1/√τ optimization rate and 1/√m width dependence. This decomposition holds for any τ, m ≥ 1, with no hidden asymptotic regime.

• Logarithmic dependence on sample size. The sample size enters only through the approximation error. To make that term at most ε, it suffices to take m ≳ ε⁻² log(n/δ). In particular, m needs to grow only logarithmically with n.

• No explicit sequence-length dependence. Attention gradients are uniformly bounded in T via a softmax-covariance structure, yielding an optimization error independent of the sequence length. The key object is the softmax-weighted covariance matrix M(X, Wᵢ) = Σₜ αₜ Xₜ Xₜᵀ − μμᵀ, which satisfies tr M(X, Wᵢ) ≤ 1 and hence ‖M(X, Wᵢ)‖_F ≤ 1 uniformly in T. The authors note that for RNNs the analogous gradient contains an explicit sum over T recurrent steps, leading to constants that can grow with T.

• The attention block of the NTK is nonzero. The Transformer NTK decomposes as K = K_c + K_u + K_w, where the attention-weight component K_w is nonzero. This follows from not using the degenerate m⁻¹ scaling; in contrast, the NTK study in Wu et al. (2023), Lemma 1, yields a vanishing attention component.

• Deterministic kernel with independent heads. With multiple independent heads, the empirical NTK converges to the deterministic kernel K, which the authors describe as solving the limit kernel problem of single-head attention raised by Hron et al. (2020).

• Comparison with Wu et al. (2023). That work proves global convergence via a well-conditioned last-layer feature/Gram matrix under the regime d ≥ n, achieving a linear rate with polynomial overparameterization in n (their Proposition 1). This paper instead works in a projected near-initialization NTK regime without assuming a positive definite Gram matrix, obtains the 1/√τ rate with only logarithmic width dependence on n, has optimization constants explicitly independent of T, and retains a nontrivial attention block in the NTK.

• Initialization produces zero output at start. The symmetric random initialization (with Wᵢ and Uᵢ duplicated for i and i + m/2, and cᵢ negated) ensures f(X, φ⁽⁰⁾) = 0 for all X in the input domain, which the authors state is critical to their error bounds.

• Experimental setting. Numerical validation is done in a teacher–student setting, and Transformers are compared against IndRNNs on an autoregressive task. The paper also reports experiments trained without projection. The truncated content provided does not report specific numerical values, dataset sizes, or error figures from these experiments.

Methodology in Plain English

The authors study a deliberately simplified but structurally faithful model: a single attention layer with m heads, each head attending over a sequence of T tokens using a single non-trainable query vector, followed by a pointwise activation and a linear combination of head outputs. This reduces the standard (W_Q, W_K, W_V) parameterization by collapsing the query and key matrices into one matrix Wᵢ, absorbing the value/output matrices into the later linear map Uᵢ, and absorbing the softmax scaling β = 1/√m into the initialization.

Rather than proving convergence by showing that a Gram matrix is strictly positive definite — the standard NTK recipe — the authors take a different route. They restrict all training iterates to a small neighborhood Ωρ around the random initialization by projecting after each gradient step. Inside this neighborhood, two things become controllable: the error from replacing the network by its first-order Taylor expansion at initialization (the linearization error), and the local Lipschitz and smoothness constants of each head's output with respect to its parameters.

To handle arbitrary target functions, they use transportation mappings: for a target function in the NTK's reproducing kernel Hilbert space, a mapping v translates the initial weights to a nearby point where the network better approximates that target. This gives an approximation bound whose dependence on the sample size is only logarithmic. The function class of interest is the norm-constrained set F_ν̄, and ‖ν̄‖₂ serves as the complexity measure for the regression problem.

Finally, to remove the sequence-length dependence, they observe that the only place T enters the gradient expressions is through a softmax-weighted covariance matrix, which is uniformly bounded in Frobenius norm by 1. This feeds into T-independent Lipschitz and smoothness constants and ultimately into a T-independent optimization error.

Why This Matters

Impact on research. The paper sharpens the gap between theory for Transformers and theory for recurrent networks. It shows that the "curse of sequence length" that afflicts RNN optimization bounds is not intrinsic to sequence models — it comes from the recurrent parameter-sharing structure, not from processing long sequences per se. It also weakens the overparameterization assumptions required for Transformer training guarantees, moving from polynomial in n (with a d ≥ n condition) to logarithmic in n, and it removes the positive-definiteness assumption on the Gram/NTK matrix that anchors much of the earlier literature. For readers of the Wu et al. (2023) line of work, it is a direct and explicit point of comparison.

Real-world applications (these are implications, not claims made by the paper):

• Long-context language and document models, where the ability to reason about long sequences without training-cost blowup is a central practical bottleneck.
• Time-series and sensor-sequence forecasting, including the autoregressive setting the authors use experimentally, where sequence length can be very large.
• Efficient architecture design, since the analysis gives a principled reason to favor independent-head, block-diagonal multi-head designs as a structural regularizer.
• Training-recipe selection, because the projected-gradient analysis suggests that mechanisms keeping parameters near initialization — weight decay, explicit norm regularization, or early stopping — can play the same analytical role.

Industry relevance. The result that width need only grow logarithmically with sample size speaks directly to model-sizing decisions: it is an argument against the intuition that more data always requires proportionally more parameters to guarantee trainability. The explicit memory trade-off — sequence-length-independent optimization error, but memory growing with T — frames a design tension that practitioners already face when choosing between attention and recurrent architectures for long sequences.

Future Directions

• Deepening the architecture. The analysis covers a shallow, single-attention-layer Transformer with scalar pooled outputs. Extending the guarantees to deeper stacks, multi-layer attention, or non-pooled sequence outputs is not addressed.
• Beyond the kernel regime. The guarantees hold in a projected near-initialization regime. Whether the T-independence and logarithmic width scaling survive feature learning away from initialization is an open question the paper does not resolve.
• Generalization, not just optimization. The bounds are on the empirical training loss (with an approximation term tied to the target function class). Generalization guarantees to unseen data are not established here.
• Removing the projection device. Projection onto Ωρ is described as an analytical device rather than a practical algorithm. The paper reports experiments trained without projection, but a formal treatment of alternative mechanisms such as weight decay or early stopping is left open.
• The memory side of the trade-off. The paper identifies memory growth with sequence length as the cost of T-independent optimization error. Improving that memory scaling while preserving the convergence guarantees is a natural next target.

Target Audience

This paper is aimed at optimization and learning-theory researchers working on Transformers, neural tangent kernels, and overparameterization; at readers who follow the Wu et al. (2023) and Cayci and Eryilmaz (2024a) lines of work and want a direct technical comparison; and at theoretically inclined machine learning practitioners who want a rigorous account of how width and sequence length affect the trainability of attention models. A solid background in kernel methods, matrix calculus, and concentration inequalities is assumed; readers without it will still extract value from the informal theorem statement and the stated scaling laws, but the proofs and Lemmas require advanced preparation.

Authors’ 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 compare Transformers with recurrent architectures on an autoregressive task.