Training Dynamics of In-Context Learning in Linear Attention

Understand

While attention-based models have demonstrated the remarkable ability of in-context learning (ICL), the theoretical understanding of how these models acquired this ability through gradient descent training is still preliminary.

• Towards answering this question, we study the gradient descent dynamics of multi-head linear self-attention trained for in-context linear regression.
• We examine two parametrizations of linear self-attention: one with the key and query weights merged as a single matrix (common in theoretical studies), and one with separate key and query matrices (closer to practical settings).
• For the merged parametrization, we show that the training dynamics has two fixed points and the loss trajectory exhibits a single, abrupt drop.