Finite-Time Analysis of Gradient Descent for Shallow Transformers
This provides theoretical insights into Transformer optimization for researchers in machine learning, though it is incremental as it focuses on a simplified model.
The paper tackles the challenge of understanding Transformers' performance by analyzing a shallow Transformer trained with projected gradient descent, finding that the required width scales logarithmically with sample size and optimization error is independent of sequence length, unlike recurrent architectures.
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.