How do Transformers perform In-Context Autoregressive Learning?
Michael E. Sander Raja Giryes Taiji Suzuki Mathieu Blondel Gabriel Peyré
Transformers have achieved state-of-the-art performance in language modeling
tasks. However, the reasons behind their tremendous success are still unclear.
In this paper, towards a better understanding, we train a Transformer model on
a simple next token prediction task, where sequences are generated as a
first-order autoregressive process $s_{t+1} = W s_t$. We show how a trained
Transformer predicts the next token by first learning $W$ in-context, then
applying a prediction mapping. We call the resulting procedure in-context
autoregressive learning. More precisely, focusing on commuting orthogonal
matrices $W$, we first show that a trained one-layer linear Transformer
implements one step of gradient descent for the minimization of an inner
objective function, when considering augmented tokens. When the tokens are not
augmented, we characterize the global minima of a one-layer diagonal linear
multi-head Transformer. Importantly, we exhibit orthogonality between heads and
show that positional encoding captures trigonometric relations in the data. On
the experimental side, we consider the general case of non-commuting orthogonal
matrices and generalize our theoretical findings.