In-Context Linear Regression Demystified: Training Dynamics and Mechanistic Interpretability of Multi-Head Softmax Attention

Published in arXiv preprint arXiv:2503.12734, 2025

We study how multi-head softmax attention learns in-context linear regression, characterizing the training dynamics and the kernel estimator a one-layer Transformer converges to.

arXiv

Recommended citation: Jianliang He, Xintian Pan, Siyu Chen, Zhuoran Yang. "In-Context Linear Regression Demystified: Training Dynamics and Mechanistic Interpretability of Multi-Head Softmax Attention." arXiv preprint arXiv:2503.12734, 2025.
Download Paper