High-Dimensional Learning Dynamics of Attention-Indexed Models

2026-09-03

AI Summary

研究通过分析高维度下的注意力索引模型，揭示了注意力机制训练动态特性，并提出使用有限截断系统近似无限矩阵矩层次以优化学习过程。

Abstract

Attention mechanisms are central to modern foundation models, yet their training dynamics remain poorly understood, especially when the attention matrices have extensive rank. In this work, we study attention-indexed models, a broad framework that can represent multi-layer and multi-head attention architectures. First, we show that, in a suitable high-dimensional limit, the population-loss landscape is characterized by a finite set of trace order parameters. In contrast, online stochastic gradient descent (SGD) is governed by an infinite hierarchy of matrix moments, which we show can be exponentially well-approximated by a finite truncated system. Second, this framework reveals that attention parameterization itself can act as an architectural implicit bias. Direct optimization of an attention matrix $S\in\mathbb{R}^{d\times d}$ can remain trapped in an uninformative state. Tied attention ($S=WW^\top$) induces an automatic symmetry-breaking mechanism and yields weak recovery in $Θ(d^2\log d)$ samples. For untied attention, $S=UV^\top$, we uncover a fast-slow mechanism: the pre-activation mean first evolves on a fast timescale, while the overlaps evolve on a slower one. Weak recovery on the $Θ(d^2\log d)$ scale occurs when the state selected by the fast dynamics breaks the initial symmetry.

Problem

• attention mechanisms
• training dynamics
• high-dimensional
• foundation models

Innovation

• High-Dimensional Learning Dynamics
• Attention-Indexed Models
• Stochastic Gradient Descent (SGD)
• Trace Order Parameters
• Implicit Bias