Title: Interpreting Multi-Layer Transformers for In-Context Linear Regression with Varying Covariance
Primary area: interpretability and explainable AI
Keywords: In-context learning; Linear regression; IV regression; Transformer

Abstract: We study how multi-layer softmax-attention transformers perform in-context linear regression, focusing one the challenging setting where the covariate distribution varies across sequences. We show that multi-layer transformers substantially outperform single-layer models, demonstrating that depth is critical for robust generalization. Through a novel probing methodology based on Gaussian test data, we reveal that the transformer approximates linear regression by implementing a variant version of gradient descent algorithm. This algorithm's parameters are dependent on model depth and data distribution, but insensitive to model width or sequence length, and correspond to a consistent diagonal structure in the learned weight matrices. Building on this insight, we show that by incorporating a chain-of-thought-style intermediate step, the transformer can solve in-context instrumental variable (IV) regression, achieving performance comparable to a two-stage least squares estimator. Our findings provide new evidence that depth enables transformers to learn sophisticated in-context algorithms, bridging the gap between empirical performance and interpretable algorithmic behavior.

TL;DR: We investigate how multi-layer softmax attention model learns to perform in-context linear regression with covariates distribution varies across sequences, revealing an intersting GD-like algorithm learned by the model.