Large-Step Training Dynamics of a Two-Factor Linear Transformer Model

2026-05-20

AI Summary

This work investigates the finite-step training dynamics of linear Transformers under large-step gradient descent, revealing that high learning rates not only accelerate convergence but can also alter the structure of attractors. By constructing a single-prompt training problem amenable to exact reduction, the authors normalize it into a two-factor product map governed by an effective step-size parameter μ. Combining gradient flow analysis, dynamical systems theory, and Chebyshev elliptic invariants, they demonstrate that for 0 < μ < 2, the system exhibits a structural bifurcation: training trajectories may converge to periodic orbits or bounded chaotic regimes rather than the unique in-context linear regression solution. The study further identifies, for the first time, an explicit Chebyshev ellipse delineating regions of transverse attraction and repulsion, offering theoretical insights relevant to mini-batch gradient descent.

Abstract

Gradient-flow analyses show that simplified linear transformers can learn the in-context linear-regression algorithm, but they do not explain the finite-step behavior of gradient descent at large learning rates. Motivated by empirical work on high-learning-rate transformer instabilities and by the cubic-map phase diagram for quadratic regression, we study an exactly reducible one-prompt linear-transformer training problem. After normalization, the dynamics reduce to a two-factor product map with an effective step-size parameter \(μ\). On the balanced slice, this map recovers the known scalar cubic transition from monotone convergence to catapult convergence, periodic and chaotic bounded nonconvergence, and divergence. We then analyze the full two-dimensional system and show that, for \(0<μ<2\), it has an explicit invariant Chebyshev ellipse separating forward-invariant regions; this ellipse carries off-balanced chaotic dynamics but is transversely repelling, while balanced scalar attractors can be transversely attracting. These results show that large constant learning rates can change the training attractor of the learned transformer rather than merely accelerating convergence: beyond sharp stability thresholds, finite-step training may settle into cycles, bounded chaos, or divergence instead of a single in-context linear-regression solution. We also discuss the consequences for mini-batch gradient descent based training methods.

Problem

• linear transformer
• large learning rate
• training dynamics
• gradient descent
• chaotic behavior

Innovation

• large learning rate
• linear transformer
• training dynamics
• chaotic convergence
• Chebyshev ellipse