Official Review · ICLR 2026

This paper investigates how multi-layer softmax-attention transformers perform in-context linear regression under the challenging setting where the covariate distribution (specifically, the covariance matrix) varies across sequences. The authors demonstrate that depth is critical for robust performance, with multi-layer models substantially outperforming single-layer ones. Through a novel probing methodology, they show that the transformer learns to implement a specific variant of gradient descent, characterized by consistent diagonal structures in the model's weight matrices. The authors further extend their findings to show that with a chain-of-thought-style intermediate step, the transformer can solve in-context instrumental variable (IV) regression. 1. The paper moves beyond the common assumption of a fixed covariate distribution, studying a more realistic and challenging scenario with varying covariance across sequences. This directly addresses the critical issue of out-of-distribution generalization in ICL. 2. The paper provides a clear, empirical mechanistic interpretation of the transformer's operation. The identification of consistent diagonal patterns in the Key-Query (KQ) and Output-Value (OV) circuits across layers and heads is a strong contribution, offering a simplified and interpretable view of the learned algorithm. 3. The study is thorough, with extensive experiments and ablations (e.g., varying depth, width, sequence length). The successful application of the interpreted mechanism to the more complex IV regression task strengthens the generalizability and practical relevance of the findings. 1. The contribution may be perceived as incremental. The paper focuses on softmax transformers but relies on a large-sequence-length regime, which simplifies the analysis and makes the behavior analogous to that of linear attention transformers. A more direct discussion of the necessity and advantages of softmax attention in this specific setting would strengthen the claim of novelty. 2. The parameters in Algorithm 1, such as ηx and ηy, are central to the interpretation. It is unclear whether these are precisely extracted from a trained transformer's weights, and whether the learning rates used in standard gradient descent is optimal for each step/layer in fig3. 3. The algorithm is noted to be non-convergent when applied beyond the trained depth. While it empirically outperforms standard GD within that depth, the paper lacks a theoretical explanation for whythis specific, divergent algorithm is more effective than a convergent one for the in-context learning task 1. Given the reliance on the large-sequence-length regime for analysis, how do the conclusions hold for very short context lengths where the softmax attention's non-linearity is more critical? 2. How exactly are the step sizes ηx and ηy in Algorithm 1 extracted from the learned model parameters? Are they consistent across layers in the looped transformer, and is there an int