Transformers Handle Endogeneity in In-Context Linear Regression
Authors: Haodong Liang
Krishna Balasubramanian
Lifeng Lai
Venue: ICLR 2025 Poster (2025)
Keywords: endogeneity
instrumental variables
transformers
bi-level gradient descent
Primary area: learning theory
TL;DR: Transformers handle endogeneity in in-context linear regression by implementing gradient-based instrumental variable regression.

Abstract: We explore the capability of transformers to address endogeneity in in-context linear regression. Our main finding is that transformers inherently possess a mechanism to handle endogeneity effectively using instrumental variables (IV). First, we demonstrate that the transformer architecture can emulate a gradient-based bi-level optimization procedure that converges to the widely used two-stage least squares (2SLS) solution at an exponential rate. Next, we propose an in-context pretraining scheme and provide theoretical guarantees showing that the global minimizer of the pre-training loss achieves a small excess loss. Our extensive experiments validate these theoretical findings, showing that the trained transformer provides more robust and reliable in-context predictions and coefficient estimates than the 2SLS method, in the presence of endogeneity.

Published as a conference paper at ICLR 2025
TRANSFORMERS HANDLE ENDOGENEITY IN IN- CONTEXT LINEAR REGRESSION
Haodong Liang
Krishnakumar Balasubramanian
Lifeng Lai UC Davis 1 f -fllai0 Fflaioucdavis.edd edu
UC Davis hdl illianguuudaiis.. edu
UC Davis
xbala@ucdavis.. eddu
ABSTRACT
We explore the capability of transformers to address endogeneity in-context lin- ear regression. Our main finding is that transformers inherently possess a mech- anism to handle endogeneity effectively using instrumental variables (IV). First, we demonstrate that the transformer architecture can emulate a gradient-basee based bi-level optimization procedure that converges to the widely used two-stage least squares (2SLS) solution at an exponential rate. Next, propose an in-context pretraining scheme and provide theoretical guarantees showing that the global minimizer of the pre- pre-training loss achieves small excess loss. Our extensive periments validate these theoretical findings, showing that the trained transformer provides robust and reliable in-context predictions and coefficient estimates than the 2SLS method, in the presence of endogeneity.
INTRODUCTION
The transformer architecture /Vaswani et al. 2017) has demonstrated remarkable in-context learning (ICL) capabilities across various domains, such as natural language processing (Devlin et al.] [201 Radford et al., 20199] Brown et al.) (020), computer vision (Dosovitskiy et al. 202 1 Carion et al., 2020), and reinforcement learning (Lee et al.) 2022, Parisotto et al. 2020p. Self-attention mecha- nism, a core component of transformers, allows these models to capture long-range dependencies in data, which is critical for success in these tasks. Despite their impressive performance, the thcoret- ical understanding of transformers remains limited, leaving important questions unanswered about their true capabilities and the underlying mechanisms driving their exceptional results.
Recent efforts to theoretically understand transformers' ICL capabilities have focused on their per- formance in fundamental statistical tasks. Focusing on simple function classes, [Garg et al. (2022) highlighted that transformers, when trained on sufficiently large and diverse data from specific function class, can generalize across most functions of that class without task-specific fine-tuning. Building on this, subsequent work by Bai al. (2023) established that attention layers enable trans- formers to perform gradient descent, implementing algorithms like linear regression, logistic regres- sion, and LASSO; see also Akyürek al. (2023); Von Oswald et al. (2023): Li et al. (2023): ; Fu
et al. (2023) Ahn et al. (2023); Jin et al. (2025). The learning dynamics of transformer was ana- lyzed Huang et al. (2024). Furthermore Zhang al. (2024a b) showed that trained transformers' ICL abilities for linear regression tasks are theoretically robust under certain distributional shifts.
Existing works on analyzing the ICL ability of transformers for linear regression tasks, however, ignore endogeneity and have mainly focused on the exogenous setup where the additive noise is un- correlated with the explanatory variables. Ignoring endogeneity in linear regression leads to biased and inconsistent estimates, resulting from issues like omitted variable bias, simultancity, and mea- surement error, which can distort causal inferences and lead to incorrect policy conclusions (Haus- man| [2001: Wooldridge] 2015 Angrist & Pischke 2009 Greene 2018). Instrumental variable (IV) regression is a widely adopted method to handle endogeneity by utilizing instruments that are corre- lated with the endogenous variables but uncorrelated with the error term (Angrist & Krueger 2001
A naturally intriguing question that therefore arises is:
Can transformers leverage instrumental variables and provide reliable predictions and coefficient estimates, in the presence of endogeneity?

---

Published as a conference paper at ICLR 2025
In this work, aim answer this question and offer new insights in-context linear regression tasks. Our key contributions include:
We demonstrate that looped transformers can address endogeneity in linear regression by lever- aging instrumental variables. Specifically, we show that transformers can implement two-stage least squares (2SLS) regression through a bi-level gradient descent procedure, where each iter- ation executed by a two-layer transformer block. Moreover, the convergence rate to the 2SLS estimator is exponential with respect the number of blocks.
We propose an ICL training scheme for transformers to efficiently handle endogeneity. Under this scheme, show that the global minimizer of the in-context pre-training loss achieves small excess loss compared to the global optimal expected loss.
We evaluate the performance of the trained transformer model through extensive experiments, finding that it not only matches the performance of the 2SLS estimator on standard IV tasks but also generalizes effectively to more complex scenarios, including the challenging cases of weak instruments, non-linear IV, and underdetermined IV problems.
As part of our analysis, we derive the first -onmsyymtootic bound for the 2SLS estimator under random design, providing valuable insights for future theoretical work.
1.1 RELATED WORKS
In-context Learning. Initial works by Garg al. (2022) and Bai et al. (2023) adopted the stan- dard multi-layer transformer architecture to conduct the experiments. Later, Giannou et al. (2023) and Yang et al. (2024) showed that a looped architecture reduces the required depth of transform- ers and exhibits better efficiency in learning algorithms, [ao et al. (2024) illustrated that the looped transformer architecture with extra pre-pre prepprocssiing and post-processing layers can achieve higher ex- pressive power than standard transformer with the same number of parameters. Apart from works concerming the implementability of first-order gradient descent algorithms by transformers, other works have also examined higher-order and non-parametric optimization methods. Specifically, Gi- annou et al. (2024) showed that transformers can emulate Newton's 's method for logistic regression. Cheng et al. (2024) showed that transformers can implement functional gradient descent and hence enable them to learn non-linear functions in-context. Relationship between in-context learning and
Bayesian inference is also studied in Ye al. 2024): Falck et al. 2024))
Nichani et al. 2024) illustrated how the transformers can learn the causal structure by encoding the latent causal graph in the first attention layer. Gool & Bartlett (2024) explored the representational power of transformer for learning linear dynamical systems, Makkuva et al. (2024a]b) Rajaraman al. 2024) Edelman al. (2024) considered ICL. Markov chains with transformers, including both landscape and training dynamics analyses. To the best of our knowledge, we are not aware of prior works on handling endogeniety with transformers.
Instrumental Variable Regression. IV regression has been widely studied in econometrics (An- grist & Krueger 200T: Angrist & Pischke 2009). Recent works in machine learning explored the optimization based approaches for the IV regression problem. [Singh et al. (2019) proposed the kernel IV regression to model non non-linear relationship between variables. Muandet et al. (2020)
proposed that a non-linear IV regression problem can be formulated as convex-concave saddle point problem. Della Vecchia & Basu (2023); Chen al.) (2024); Fonseca et al. (2024) proposed a stochastic optimization algorithm for IV regression.
Notation: Throughout this paper, unless otherwise specified, lower-case letters denote random variables samples, while upper-case letters represent datasets (collections of samples). Bolded letters indicate vectors or matrices, whereas unbolded letters indicate scalars. The notation X., refers to the i-th column, and X, refers to the i-th row of matrix X. Amiu (-) denotes the minimum eigenvalue, and min (-) denotes the minimum singular value matrix. By default, || denotes the Fuclidean norm for vector, or the spectral norm for matrix.
2

---

Published as a conference paper at ICLR 2025
2 ENDOGENEITY AND INSTRUMENTAL VARIABLE REGRESSION
Suppose we are interested in estimating the relationship between response variable E R and predictor variable æ E IR? with endogeneity. Given instruments 12,, we consider the model
y = 3 + and = θ + 62
where 3 E IRP, and θ E RexP are the true model parameters, E R and 62 E IRP are (centered) random noise terms with variance of and covariance matrix Σg. respectively. Further, e2 is an unobserved noise correlated with é1, leading to the correlation between æ and é1, which introduces confounding in the model between and y. Under this setting, the standard ordinary least squares (OLS) estimator is a biased and inconsistent estimator of B (see Wooldridge (2015), Chapter 9). To address this issue, instrumental variable (IV) regression is widely used method to provide consistent estimate for [3.
Definition 2.1 (2SLS estimator). IV regression is a regression model to provide consistent estimate on the causal effect for the endogeneity problem (I). by utilizing the instrument z. Given obser- vational values Z, X, Y) = {(z,. x., 4=)]== the standard approach to estimate the IV regression model 2SLS; see, for example, Wooldridge 2015), Chapter 15.
First stage: Regress X on Z to obtain A
e = =Z
ii. Second stage: Regress Y on Zé to obtain:
ÍSLLS T 2 = y
We introduce the standard assumptions required to show the convergence rate of the above estimator.
Assumption (Instrumental variable). A random variable following conditions:
R is valid IV, if it satisfies the
Fully identification: q 2 p (without loss of generality, assume data Z, X are full rank).
ii. Correlated to æ: Corr(z, x) * O.
ii.. Conditional uncorrelated to y: Corr(z,(1) 0.
In particular, condition (i) above ensures the existence unique solution for 2SULS We refer to
Stock & Watson (2011] Chapter 12) for additional elaborate discussions the above conditions. To derive non-asymptotic convergence rates, we further assume the following regularity conditions.
Assumption (Regularity conditions). Suppose instrument z is a centered random variable. We assume the following conditions hold:
i. Bounded parameters: || S B3, |||| S Be.
ii. Bounded variables: ||||| < B, |<||| S Bz, lel Bey ||E2|l S Bey
ii.. Linear instrument: E [x> =]] (Ok, z).
The boundedness condition in (ii) is required to invoke matrix Bernstein inequalities (Tropp) 2015) in the analysis. We anticipate that this condition may be relaxed to subgaussian or moment conditions by using more sophisticated matrix concentration results.
Theorem 2.1 (MSE of 2SLS estimator). Given Assumptions I and consider clipping operation
if < Bs
clip (5) =
Bg 13
3
if || > B3
When
2 max 462 4e² (a+ (++log log
B62
qe
p² + 1)2K
q
K
9K3
3

---

Published as a conference paper at ICLR 2025
where k= Amin ).. and Ko Amin (exnne mimi ( 6.9
the meun squared error of the 2SLS estimate
213/2
bounded by:
B +
E
||cipp (02sLs) 3|3
SO
4
C
n
K
where C(n) is defined in Equation 42 Σ
2 and c is an absolute constant.
Remark 2.1. We keep the slightly complicated form (3) so that O notation only hides some absolute constant multipliers that are independent of problem-related constants. Note that when Ba
is large enough, we have ((z) + Amin (2.
so ((nz) is also bounded. Thus the error bound
decays with rate (+))
We note that although the consistency of the 2SLS estimator a standard result in econometrics, most existing works focus the asymptotic properties of the estimator. Theorem 2.1 provides the first non-asymptotic bound for estimation error [10esLs - B11 under random design. The detailed proof is provided in Appendix A.1
3 TRANSFORMERS HANDLE ENDOGENIETY
3.1 TRANSFORMER ARCHITECTURE
Denote the input matrix as = [hi..... Nn]] E Dxn, where each column corresponds to one sample vector.
Definition 3.1 (Attention layer). A self-attention layer with M heads is denoted as ATTNa()), with parameters {(Qm: Km; Vm)]me(M] S D Given input H,
H AHT
V H) a(QQ H) (K,, H)) E Dx
(4)
elisent-wise:
M + m=1
h. = [ATTNa H)): h,
Σ
o ((Qmii Kmhz)) Vnhy E RD,
(5)
j=1
where (-) is the ReLU function. Definition 3.2 (MLP layer). An MLP layer is denoted MLPo(.), with parameters (W1, W2) RDD DDXDDD'' Given input H,
MLPa == H (W,,W,
or clement-wise:
h MLP (H)); = h, W20 (,,
Definition 3.3 (Transformer). An L-layer transformer is denoted as TFe(.), with parameters
(0(1-)) (1::2)) Given input H H (0),
H
MLP (ATTN (H--)))) 1= 1=
MIP
ATTN
The output of this transformer's is the final layer output: H (L) TFo ()).
Definition 3.4 (Looped transformer). An I.-looped transformer is special transformer architecture, denoted as LTF (+), with parameters 0 (0(1:b)) o(1:Lo)). Given input H (0))
H
TF H (11)), 1 I....,T.. T..
The output of this looped transformer is the final loop output: H H LTF 0.1 (H
4

---

Published as a conference paper at ICLR 2025
Bai Previous works (c.g., Bai et al. 2023 Zhang al. (2024a)) have shown that transformers can per-
form in-context linear regression by emulating gradient descent (GD) with in-context pretraining. However, these studies have two key limitations. First, their analysis is based single-level op- timization algorithms, which insufficient to demonstrate that transformers can efficiently learn more complex algorithms like 2SLS (Definition 2.1p. Second, most ICI.-related research focuses on the predictive performance of transformers, paying little attention to their ability to provide accu- rate coefficient estimates. We extend the current ICL framework by showing that transformers can implement a bi-level GD procedure (see Section 3.2 with looped transformer architecture (Defini- tion B.4), allowing them efficiently emulate 2SLS and provide coefficient estimates that are at least accurate 2SLS in the presence of endogeneity (as in Equation ①
3.2 GRADIENT DESCENT BASED IV REGRESSION
We first introduce a gradient-based bi-level optimization procedure to obtain the 2SLS estimator in Equation ( Given the dataset (Z, X, Y) = {(z,, th, yi) )=))= the objective function of IV
regression can be formulated as the following bi-level optimization problem:
min C(3) = Σ [ 2; 6,)2,
where
== arg min
Σ [e, 23 0) ==1
(6)
Consider the following gradient updates with learning rates ox TF
(++1)) (*) 1/ (20()) X), g(++1) ( a0 1 z (ZO Y).
(7a)
(7b)
Note that the GD-2SLS updates in Equation 7 are designed to solve Equation (6) that regardless the convergence of ( the GD estimator ((()) will always converge to the 2SLS estimator Equation (2) with exponential rate.
We now show
Theorem 3.1 (Implementing 2SLS with gradient-based method). Given training data
Y, = {(2i, L, Di)) Suppose the learning rates a, 1) satisfy the following conditions:
2
0 0<a<
and 0 イ
0 Paa (Zé)
Z
max
where (-ax(()) denotes the largest singular value of a matrix. Then, the GD updates in Equation 7 converge to the 2SLS estimator at an exponential rate:
(1) B2susll S SO(N'
where, with pl-) denoting the spectral radius of the matrix,
A max(f(a),
= - z - ZT),, <(n) = p(I n/ 2).
(8)
To the best of our knowledge, Theorem 3.1 provides the first theoretical result demonstrating that 2SLS can be efficiently implemented using a gradient-based method, with an exponential conver- gence rate. We provide the proof in Appendix [.I and present simulation results in Appendix C.1 to examine the convergence behavior of the optimization process.
3.3 TRANSFORMERS CAN EFFICIENTLY IMPLEMENT GD-2SLS
The looped transformer architecture (Definition 3,4 as proposed by Giannou et al. (2023), intro- duces an efficient approach to learn iterative algorithms by cascading the same transformer block for multiple times. With the GD updates in Equation 77) we will show that there exists a looped trans- former architecuture that can efficiently learn the 2SLS estimator. We emphasize here that although can implement 2SLS by sequentially attaching two separate GD iterates (each handling OLS for one stage), the overall convergence depends heavily on the convergence of the first stage estimate é. Hence, significantly more number of layers are needed to ensure convergence. In addition, the advantage of looped transformer architecture cannot be fully exploited with this approach.
5

---

Published as a conference paper at ICLR 2025
Theorem 3.2 (Implement step of GD-2SLS with transformer block). Suppose the embedded input matrix takes the form:
2++1
31
nn
Tn +1
Yn
0
H (1))
E R Dx(n+1).
(9)
3 .P
30
3
(I)
(L)
3
2+1
1
1
Given (2),,, there exists a double-layer attention-only transformer block with parameters 021+-12142)) {(Q (2)+)2++2 K (21++1:2(+2
(=12+4122+22) mEMM23-122022 C DDD, where the
ATTN
number of heads M (21+1) 2p, ( 1(p+1) and embedding dimension D ap+3p+
that implements 2SLS gradient update in Equation with any given learning rates
Znt1
Zn
In+1
y1
yn
0
(1+1)
(++1)
(1+1)
2(++) = TF 4::41:2142
(H (2))
E E Px(x1))
(1+1)
1+1)
(1+1)
a (
e
B g
((+1)
(1+1)
(L+1)
n+
I
Our existence proof specifies an attention structure such that one layer updates only the first-stage estimate ( for all samples, followed by another layer to update the parameters ( and (((). Furthermore, as noted in the proof of Theorem 3.2 (ref. Appendix B.2). regardless of the initial values of ()),, (i) and ((),, the structures of the transformer blocks remain the same. This allows exploit the looped transformer architecture to significantly reduce the number of parameters
and improve learning efficiency (Yang et al. 2024).
By cascading the transformer block L times, with Theorem 3.1 one can show that transformers are able to mimic the 2SLS estimator with exponential convergence rate, as described in the following corollary.
1. given learning rates a, 1) and A E (0,1), as defined in Equation {8) there exists a transformer for- milated (:)
Corollary 3.1 (Implementing GD-2SLS with looped transformer). For any 0
TF (LTF 8. o. which consists of an T.-looped transformer LTFa, with L = (E)) and a final atten- tion layer 0' 0' A'ATTN ((Q'., K", V m C DDDxD, where M 2p, M = 2(p+1) 1),
① 0 (1:2)
ATTN [(Q (1:2) (1:2) (1:2) BeelMuay C CRDD M
M' = 2, such that given embedded input H taking the format in Equation g the model output
satisfies:
read,, (TF Besis S Bye,
where (-)d,()) is function that reads the prediction 9a+1 from the output of the transformer.
We emphasize here that construction differs from the implementation of Bai al. 2023 Theo-
rem 4) for OLS in the following aspects:
'This layer updates the prediction fn+1 3 #n+1, which can he constructed with 2 attention heads
using the architecture Bai al. 2023 Theorem 13)

---

Published as a conference paper at ICLR 2025
We apply the square loss defined in Equation simplifies the loss function's 's sum-of-ReLL. representation.
6 to learn the 2SLS estimator, which
ii The dimension of the input embedding is D = QP + g 3, where the extra dimensions store the vectorized parameters ( , ,,))0, and the first stage estimate a(")
iii. We use a two-layer attention-only transformer block to implement 2SLS GD update with the first layer to update the current first-stage estimate á and the second layer to update the parameters (!) and 3()
iv. For each transformer block, the first layer, we equip 2 heads to update each dimension of x (1) E R for all samples. In the second layer, we equip heads to update each column of (() RR** and 3 E RR.
3.4 PRETRAINING AND EXCESS Loss BOUND
With slightly abuse of notations, we denote the (formulated) training prompt
21,k
Zn,k
2n+1,k]
H
1i,6
Fn.k
In+1.
k k= 1= N. E ER(
Vn.k
Note that we denote each training prompt by the subscript 1...., ,. where N the total num- ber of prompts. Each training prompt consists of n labeled training samples {(z,, z, yr) W))>, and one unlabeled query sample (Zu+1, In+1). Our goal is to predict yn++1 given the context provided by the prompt.
We introduce the following ICL. data generating scheme such that endogeneity occurs in the training samples, but does not extend to the query sample. Each training prompt is generated by the in- context distribution P, described by Algorithm [
Algorithm In-Context Distribution
Parameters: Sample size n, clipping thresholds Bz, Br, By. Task parameters O,3, Φb
Σ,, Eu, Ew, a, from meta distribution
2: Output: Training samples {(z,, xi y:) query sample (En+1, #n Bn+1).
3: for -I...,, do
4: Generate: 2; ~ (0, Σ.), 1.),00 u, ~ Ew), wi ~ N(0 (.),,
5: Compute: H; 0 2, $ + wi.
6: Compute: yi = 3 x, + + u; +
7: end for
8: Generate: nn+1 ~ ~N(0, Eg), wnt1 N(0 Ew), cn+1 ~ ~N(,0??)
9: Compute: In+1 = Q zn+1 + Wn+1
10: Compute: Dn+1 = =B xn+1 + En+l-
Clip: 2, = clip,, (z,), H; clpp,, (æx)). yi clip n. (v:) for i 1....,
In Algorithm 1 Rp is the source of endogenous error, R R are the exogenous errors. Note that we have E & u, + E; and 6 4 u, + wi corresponding to the notations in Equation (1. + R.XP, a e R2, E pxpp I are task- parameters following meta distribution π. clip, (-) is clipping operator to bound the norm of input within radius B. We say that the in-context samples [(z), 2., y: w n are drawn from the in-context distribution P. and are sampled from π. One can check that Assumption and Assumption 2 (iii) are directly satisfied with the data generated from the in-context distribution P.
Rº Σ. E R Σ,, Rexp, Σn, RPxP, a,
the task parameters (o, B, 4, Φ. D. E., E.
Following the theoretical framework of Bai et al. 2023 we define the population ICL loss
LIcL (0) = E,, Ep | clip B, (read,, (TF (H
(10)
2 All the clipping operations only for analytical purpose. In practice, the hehavior of the trained trans- former is consistent without the clipping bounds.

---

Published as a conference paper at ICLR 2025
where H is the embedded imput defined in Equation (9) parameter 0 and clipping operation clip applied to each layer output. For simplicity, we denote TFo (H) c, ( (H )))).
TF is the transformer model with
The transformer is trained to minimize the in-context loss in Equation (10) with the following em- pirical loss:
1 N
LIL (0)
U1 TFo (H
(11)
We consider the following constrained optimization problem:
arg min (I)(()),
=
AEO L.M.D',BO
(12)
1 L.M.D' == {0 (0
o(l:L)
max M [E
S M, max D S D', |0|| - Bo), [E[L]
where |0] = max max [E(L] (E(M)]
{IQ m K
M
+ I ||} V + [
m=1
We now establish excess loss bound for the trained transformer model.
Theorem 3.3 (Excess loss bound for in-context pretrained transformer). Suppose condition (i) in
Assumption 2 holds and the meta distribution π satisfies the following conditions:
E [Ø] Σ, + o,]] <ō² and E, [o²] S o²
(13)
Let the in-context distribution P such thut the samples (zi, 2; 1:) are drawn independently from P (ref. Algorithm 7.. With training prompts Hi, k 1..... N, under ICL. loss (TO), the
trained transformer (2) with L 2L - 1, M 2(p++1) + 1), op + 3p 3, D' (attention- only) achieves the following excess loss with probability at least
C:
(0) E,pEEp [1n+11 18,-1,,880 2n+1)) s =00 ((L))
B2 + (2)))²2))))) + B2ō.
B²
n
B3 K
L D2 log(2 + max Ba, R. B, + log log(L/()
B²
q
(m())²²) + 11222 + BB
n
N
where A == min E, [A|H, a, nl 1, and 1h, = E,Ep H. an,n"" ] is close
0..
In practical training, the number of prompts N is usually large enough such that the last term of the above bound is negligible. Thus, given a meta distribution π, the excess loss is dominated by two factors: number of attention layers, and (ii) number of in-context samples. The proof of Theorem 3.3 is provided in Appendix B.3
3.5
EXTRACTING THE REGRESSION COEFFICIENTS
The primary goal of IV regression to estimate the causal effect, i.c. the coefficient / under the stated endogeneity in Equation (T). For 2SLS, the estimated causal effect is given by the coefficients of the endogenous variable in the second stage regression 2 For transformer models, propose a straightforward method to extract these estimated coeffictents by differentiating the output with respect to each dimension of the endogenous variable. The specific approach summarized in Algorithm 2 We observe that the choice of A within reasonable range does not significantly affect the estimation of the coefficients. In practice, usually slightly larger A (for example = 5) can lead to more stable estimation, which is possibly due to the climination of rounding errors during computation.
EXPERIMENTS
4.1 EXPERIMENT SETUP
We conduct a simulation study to evaluate the performance of the ICL-pretrained transformer model in handling endogeneity. We set the maximum input sample size to 51 (n = 50 training samples and
8

---

Published as a conference paper at ICLR 2025
Algorithm 2 Extracting the regression coefficients
Trained transformer model TF g. input matrix H. perturbation A. 2: Output: Estimated coefficient [3.
1: Input:
3: Procedure:
4: Compute the output of the transformer model: y = TF
5: for each dimension à 1, ......? do
6: Copy Ha( = H. Set the k-th dimension of xn+1 to be (In+1)a 1 + A for HA(K).
7: Compute the new output value: YA(K) (Fo(HA(K))
8: Compute the estimated coefficient: 3x Yaa
9:
end for
one query sample), the dimension of endogenous variable = 5, and the dimension of instrument q 10. The training prompts generated using Algorithm [ with task parameters Q B3 t, ob sampled from standard Gaussian distribution, and the covariance matrices Σ,, Σ., E., set to be identity matrices. The noise level a. is set 1. We ignore all the clipping bounds in the experiment (B3, Be. B2, Br, By: Bo, R set to infinity).
The backbone of the transformer block is initialized using GPT-2 settings, with 12 attention heads
(M = 12), 80-dimensional embedding space (D 80) and 2 layers (Lo = 2). following the theoretical guidelines Theorem 3.2 We employ the looped transformer architecture, consisting of 10 identical cascading transformer blocks. The transformer model is trained under the ICL loss
[1] with a batch size of N 64, over a total of 300,000 training steps.
We evaluate the trained transformer model on test prompts that are not included during training. As benchmarks, we compare the transformer's performance against the 2SLS and the OLS estimators, which are obtained by directly fitting the training samples {(2,, Li, M4) J within the text prompts. In contrast, the same trained transformer model is used without any parameter adjustments for each task. We compare the performance of these models from two aspects: the in-context prediction error (ICPE) on the query sample Ya+1, and the mean squared error (MSE) on the coefficient B..
4.2 RESULTS
We first investigate the performance of the trained transformer model over endogeneity tasks with varying training sample sizes from 20 to 50. The results shown in Figure la Under endogeneity, our transformer model achieves similar performance to that of the 2SLS estimator, with only small gaps in ICPE and MSE, both outperforming the OLS estimator.
Next, we examine the performance of the trained transformer model in handling varying levels of IV strength. The strength of an instrument is measured by the correlation between the IV and the endogenous variable. To vary the IV strength, we generate prompts with 2; and æ; following different correlation levels. Specifically, in Algorithm [
we adjust the IV strength by multiplying o by a factor E (0, 2) when generating test prompts. The results are shown in Figure lb
Interestingly, the trained transformer model outperforms the 2SLS estimator in handling weaker IVs (when IV strength < 0.5). This suggests that, beyond merely mimicking 2SLS, the ICL training process may equip the transformer model with more advanced mechanism for handling endogene- ity with weak IVs than the 2SLS estimator. At the same time, when the IV is strong, the transformer model maintains performance comparable to that of the 2SLS estimator.
This finding motivates us to further examine the performance of the trained transformer model in non-standard endogeneity tasks. We consider two scenarios: (a) the IV has a quadratic effect on the endogenous variable, i.e. Ti,A- = 0 2 + error; in Algorithm 1 and (b) the dimension of IV is not sufficient identify the endogenous variable") where we set q (by zeroing out the remaining dimensions of z in test prompts) and 5.
We evaluate the same trained transformer model before, with results presented in Figure 2a and Figure 2b respectively. Once again, the trained transformer model consistently outperforms both
For 2SLS estimate, the actual computation uses pscudoinverse handle rank deficiency.

---

Published as a conference paper at ICLR 2025
Erroe" in-onntext Sample Size
Error vs Iv Strungth
2313
2515 IIEE
0U.5 ICRE
O-S ICPE
USE
Paaaaorrar
Thanalarmer
Emon
625
2.50
3.35
L03
1.35
1.50
1.00
Screnghh
(a)
(b)
Figure 1: The ICL performance of the trained transformer model in endogeneity tasks. We compare in-context prediction error (ICPE) and coefficient MSE versus (a) the number of in-context samples; (b) the IV strength. The curves are averaged over 500 simulations.
2SLS and OLS estimators in handling these non-standard endogeneity tasks. All these results sug- gest that the trained transformer can be generalized effectively to a broader range of endogeneity tasks while still providing reliable in-context predictions and coefficient estimates. To further illus- trate this capability, also examine other cases including multicollinearity, complex non-linear IV, and varying endogeneity strengths, see Appendix CC33
We suspect that, in our pretraining scheme, although the 2SLS estimator already achieves small excess loss, gap remains between the 2SLS estimator and the optimal predictor that the transformer model successfully bridges. Finally, we conclude that through ICL training, the transformer model performs least as well as 2SLS and appears to be a promising tool for handling endogeneity in difficult scenarios.
C.4 C.5
Erron" Context Sample Size (Quadratic
Erron In-Context Sample Size
17-
2615
OL5 ICPE
KCP
OLS MSE
CLS MSE
Thaneboree ICPE lammz MSt
Thinsforme KCE
Samples
Samples
(a)
(b)
Figure 2: The ICL performance of the trained transformer model in non-standard endogeneity tasks:
(a) The IV has quadratic effect on the endogenous variable; (b) The dimension of IV is not sufficient identify the endogenous variable. The curves are averaged over 500 simulations.
5 CONCLUSION
This paper presents novel perspective the transformer model in its ability handle endogeneity in in-context linear regression. We have theoretically shown that the transformer model exists an intrinsic structure that enables to learn the 2SLS algorithm through an efficient GD procedure. We have further provided theoretical guarantee that the trained transformer model can achieve a small excess loss over the optimal loss, under our proposed ICL training scheme. Our simulation study demonstrates that the trained transformer model can achieve comparable performance to the 2SLS estimator handling standard endogeneity tasks. Furthermore, our investigation illustrates that exhibits significantly better performances in handling complex scenarios such as weak instruments, non-linear IV, and underdetermined IV problems, compared to the 2SLS estimator. These results suggest that the ICL pre-preeaat trained transformer model is promising tool for making reliable in-context predictions and coefficient estimates under endogeneity, especially when dealing with non-stande non-standard IV problems.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
The work of H. Liang and L. Lai was supported by National Science Foundation (NSF) under grants CCF-21 12504 and CCF-2232907. The work of K. Balasubramanian was supported by NSF under grants RMS-2413420 and DMS-2053918.
REFERENCES
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn implement preconditioned gradient descent for in-context learning. In Proceedings of Neural Information Processing Systems, New Orleans, LA, USA, December 2023. URL //ap:///axvvv.. .wx.voorg/ org/
0255/2306.000 00297
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm in-context learning? investigations with linear models. In Proceedings of Inter- national Conference on Learning Representations, Kigali, Rwanda, May 2023. URL https: //openzeview.o http://eeeetttoouuoooooooCCoOCC
Joshua D. Angrist and William N. Evans. Children and their parents' labor supply: Evidence from exogenous variation in family size. The American Economic Review, 88(3):450-477, 1998. ISSN 00028282. URL http://tp/. www. jstor. wttgg//aal/e/l//////
Joshua D. Angrist and Alan B. Krueger. Instrumental variables and the search for identification: From supply and demand to natural experiments. Journal of Economic Perspectives, 15(4):69-8)) 85, 2001. URL
https: : www. aewweb. org/ wwt.aarriiccscieecrorrr 120.111...1111111
15.4.69
Joshua D Angrist and Jörn-Steffen Pischke.
Mostly harmless econometries: An empiri- cist's companion. Princeton University Press, 2009. ISBN 9780691120355. URL
https://press.priccccccoo www.eeeettteeeeeeeeeee hwt.t/boos///////ooo http://oooohhaacccoooo http://///.accccee/////011111110
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. In Pro- ceedings of Neural Information Processing Systems, New Orlcans, LA, USA. December 2023. URL http:::/pproceeiiss..eerr. neuripss......uu.........r.....m cc/paper http:/aaprffils/peeeeeeeee tiles/paper/202 File/beeeeeeboeaa333333eee222223 Paper
oepprpaaeeesppppp---o-oereeeeeee Conference.pdl
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray. Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Rad- ford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Pro-
ceedings Neural Information Processing Systems, Virtual, December 2020. URL https:
//arxiv.org/ abtp//aa.//..... 14165
Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In Proceedings of European Conference on Computer Vision, Virtual, August 2020. URL
00055.2873
Xuxing Chen, Abhishek Roy, Yifan Hu, and Krishnakumar Balasubramanian. Stochastic op- timization algorithms for instrumental variable regression with streaming data. In Proceed- ings of Neural Information Processing Systems, Vancouver, Canada, December 2024. URL wtw//bs////oooooooooooo///.aoloa ww..orggaa. abs / 2405. 19463
Xiang Cheng, Yuxin Chen, and Suvrit Sra. Transformers implement functional gradient descent to learn non-linear functions in context. In Proceedings of International Conference Machine Learning. Honolulu, HI. USA. July 2024. URL 0233//cheng22 .tmmmll h
.mw.. pttes...esnmmmm.resse.seessss press/

---

Published as a conference paper at ICLR 2025
Riccardo Della Vecchia and Debabrota Basu. Stochastic online instrumental variable regression: Regrets for endogeneity and bandit feedback. arXiv preprint arXiv:2302.09337, 2023. URL
httts::/arxiv org/
/ttrg//abs///e abs 2302. 09357
Jacob Devlin, Ming-Wai Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep
bidirectional transformers for language understanding. In Proceedings of the North American Chapter of the Association for Computational Linguistics, Minncapolis. MN, USA, June 2019.
URL https: //api
Cwtg//ooppssii... :ttp:///ooppSS::55533599999 52967399
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In Proceedings of International Conference Learning Representations, Virtual, May 2021. URL httpt::////// htttp::////oenntee http:s///oppeeeie.....mom / net wo-.ommiddiiiiiie......o FicbFF
Benjamin I. Edelman, Ezra Edelman, Surbhi Goel, Eran Malach, and Nikolaos Tsilivis. The
evolution of statistical induction heads: In-context learning markov chains. Proceedings of Neural Information Processing Systems, Vancouver, Canada, December 2024. URL https:
Fabian Falck, Ziyu Wang, and Chris Holmes. Is in-context learning in large language models
bayesian? martingale perspective. arXiv preprint arXiv:2406.0079; 2024. URL https:
//arxiv.orq/abait. 02244400007790
Yuri Fonseca, Caio Peixoto, and Yuri Saporito. Nonparametric instrumental variable regression through stochastic approximate gradients. In Proceedings of Neural Information Processing Systems, Vancouver, Canada, December 2024. URL http:::///axxil..og/////////////
05639
Deqing Fu, Tian-Qi Chen, Robin Jia, and Vatsal Sharan. Transformers learn higher-order op- timization methods for in-context learning: A study with linear models. arXiv preprint (tl.:2330.1708.. 2023. URL httthttt///arrt wttpv//.gggab////////
17086
Yihang Gao, Chuanyang Zheng, Enze Xie, Han Shi, Tianyang Hu, Yu Li, Michael K. Ng. Zhenguo Li, and Zhaoqiang Liu. On the expressive power of variant of the looped transformer. arXiv preprint arXiv:2402. 13572, 2024. URtpsttss://axx orq/ abs 22042.113777
Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in- context? a case study of simple function classes. In Proceedings of Neural Information Processing Systems, New Orleans, LA, USA, December 2022. httptttpsp//t.....tttttttott org/ [11666
Angeliki Giannou, Shashank Rajput, Jy- Yong Sohn, Kangwook Lee, Jason D. Lee, and Dimitris Papailiopoolos. Looped transformers as programmable computers. Proceedings of Inter- national Conference Machine Learning, Honolulu, HI, USA, July 2023. URL https:
//proceedings mlr. www.pressssseeec pptpa//202//e/eeepp..peppe2 gitp///annnnnnannaaamam html
Angeliki Giannou, Liu Yang, Tianhao Wang, Dimitris Papailiopoulos, and Jason D. Lee. How well can transformers emulate in-context newton's method? arXiv preprint MXXi:22433003188 2024.
URL htt.hhtsss///xxii wrg.ivvorrrraaae bb0////...22.....// 2311
03183
Gautam Goel and Peter Bartlett. Can transformer represent a kalman filter? In Proceedings of Annual Learning for Dynamics & Control Conference, Oxford, UK, July 2024. URL https:
//proceedings. mlr. www....w................. press/v2422 /tol2//oel2a.a..a.. html
William H. Greene.
Econometric Analysis. Pearson, 8th edition, 2018.
ISBN 978-
0-13-446136-6.
URL https:/ happs://pages.. stern. nyu eeteteedu/aae -aciuu/wwgreee wwweeeeeeeeeeeeeeeeeeeee
Jerry Hausman. Mismeasured variables in econometric analysis: problems from the right and problems from the left. Jourmal of Economic perspectives, 15(4):57-67, 2001. URL https:
7/www. aeaweb. www.orrg//aooo
1.. 45.4.57
12

---

Published as a conference paper at ICLR 2025
Daniel Hsu, Sham M. Kakade, and Tong Zhang. Random design analysis of ridge regression.
Foundations of Computational Mathematics. 14(3):699-600 2014. URL https: //ink.
springer. com/ artiilaeieii 10.00.....0.01 /ttp://///rrraaaa/00.00000000000
Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. In Proceedings
of International Conference on Machine Learning, Honolulu, HI, USA, July 2024. URL https:
//proceedings.oo ml press / v2 w235/ hww.g///aagg224.. html
Yanhao Jin, Krishnakumar Balasubramanian, and Debashis Paul. Meta-learning with generalized
ridge regression: High-dimensional asymptotics, optimality and hyper-covariance estimation.
arXiv preprint wwXiv:2403.19722 2024. htt.sttp::// arxiv. http:oo///aaa.....// 19720
Yanhao Jin, Krishnakumar Balasubramanian, and Lifeng Lai. In-context learning for mix- ture of linear regressions: Existence, generalization and training dynamics. arXiv preprint arXiv:2410. 14183, 2025. URL WRLlhttss://axxiv.ooo abs/ 2410.
Kuang-Hueh Lec, Ofir Nachum, Mengjiao (Sherry) Yang. Lisa Lee, Daniel Freeman, Sergio Guadar- rama, lan Fischer, Winnie Xu, Eric Jang, Henryk Michalewski, and Igor Mordatch. Multi-game decision transformers. In Proceedings of Neural Information Processing Systems, New Orleans, LA, USA, November 2022. URL http::t:: /ttp///aaivv..//.a///a//......aa org abs /bss2211.15136
Yingcong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos. and Samet Oymak. Transform- ers as algorithms: Generalization and stability in in-context learning. In Proceedings of In- ternational Conference on Machine Learning, Honolulu, HI, USA, July 2023. URL https:
mllg..mmlr..oom peress/v/02/1100 /pppoceeiinng........peeee//////
html
Ashok Vardhan Makkuva, Marco Bondaschi, Chanakya Ekhote, Adway Girish, Alliot Nagle, Hyeji Kim, and Michael Gastpar. Local to global: Learning dynamics and effect of initialization for transformers. In ICML 2024 Workshop on Theoretical Foundations of Foundation Models, Vi- enna, Austria, July 2024a. URL / net
Ashok Vardhan Makkuva, Marco Bondaschi, Alliot Nagle, Adway Girish, Hyeji Kim, Martin Jaggi, and Michael Gastpar. Attention with markov: A curious case of single-layer transformers. In ICML 2024 Workshop on Mechanistic Interpretabiliry, Vienna, Austria, July 2024b. URL https: https::///penttt httpse////peereveee...
Krikamol Muandet, Arash Mehrjou, Si Kai Lee, and Anant Raj. Dual instrumental variable regression. In Proceedings of Neural Information Processing Systems. Vancouver. Canada, December 2020. URL https: http::://epppeeeeeeee.aseeeeeeee html
cc//papper/222/haatt/
Eshaan Nichani, Alex Damian, and Jason D. Lee. How transformers learn causal structure with gradient descent. In Proceedings of International Conference on Machine Learning, Honolulu, HI, USA, July 2024. URL arxiv.iivv..//.../ www.wiwww./aaaa.aaaa.a..ao. wwwp//bb//4/..123333
Emilio Parisotto, H. Francis Song, Jack W. Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant M. Jayakumar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, Matthew M. Botvinick, Nicolas Heess, and Raia Hadsell. Stabilizing transformers for reinforcement learn- ing. In Proceedings of International Conference on Machine Learning, Virtual, July 2020. URL html
https::/pproceedings mlr press
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. In OpenAl Technical Report, 2019. https:
semant ccschola org/CorpustD:: 1002255533
Nived Rajaraman, Marco Bondaschi, Kannan Ramchandran, Michael Gastpar, and Ashok Vard- han Makkuva. Transformers on markov data: Constant depth suffices. arXiv preprint
arXiv:2407.17686 2024. URL
abs//405.222.. 17686
Rahul Singh, Maneesh Sahani, and Arthur Gretton. Kernel instrumental variable regression. In Proceedings of Neural Information Processing Systems, Vancou- ver, Canada, December 2019. URL httts://papera: nips cc/paper/
8 8708 1e8keernell inst rumental -ww.nntaallaaaaaaaaaaamaa eeegriissirrreess.on
13

---

Published as a conference paper at ICLR 2025
J.H. Stock and M.W. Watson. Introduction Econometrics. Addison- Wesley, 3rd edi- tion, 2011. ISBN 9780138009007. URL btttp::/ scholars saraaas...a...aaaada... hww.araraaaa............a..... edu/
Joel A. Tropp. An Introduction Matrix Concentration Inequalities. Now Publishers Inc., 2015. ISBN 978-1-60198-838-5 URL http://oo...ggg1........0.00.0.0
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E. ukasz Kaiser, and Illia Polosukhin. Attention all you need. In Pro- ceedings of Neural Information Processing Systems, Long Beach, CA, USA, December 2017. URL https:://proceed neurips. httpp/pprr..lssppppeeeeeeee File/ File/3f5ee2 2tt.:/ffee22554411110011333a33a3
Roman Vershynin. Digeesiinns Probability: An Introduction with Applications in Data Science. Cambridge University Press, 2018. ISBN 978-1-108-415194. URL https:
7/www.
797 19144604237302222233399222222
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In Proceedings of International Conference Machine Learning, Honolulu, HI, USA, 2023. URL UR//ttps://ppreedddaa http://pwccceeiieeoopppeepeeeeee
mlr press pateass/2222//oe/ 22tp:/oooooooaaa22.......... html
Charles F. Westoff and Robert Parke. Demographic and social aspects of population growth. Com- mission Population Growth and the American Future, 1972. URL https://catalo::
Sethitrrust.orr//ecccc///O/0000
J.M. Wooldridge, Introductory Ecoommerre: A Modern Approach. Cengage Learning, 2015. ISBN 9781473754393. URL [RRL[tttt://bookk wowgle cccp:///ockkss1ieeeeeeeaAAAA.
Liu Yang. Kangwook Lee, Robert Nowak, and Dimitris Papailiopoulos. Looped transformers are better at learning learning algorithms. In Proceedings of International Conference on Learning Representations, Vienna, Austria, May 2024. URL https: http:////peeeeeee..o/moom net forum?
[d=HHHRRxDDTXE
Naimeng Ye, Hanming Yang, Andrew Siah, and Hongseok Namkoong. Pre-training and in-context learning is bayesian inference a la de finetti. In ICLR 2024 Workshop on Mathematical and Empirical Understanding of Foundation Models, Vienna, Austria, May 2024. URL https: //openzeview. wwwt///ormmiddtttuuptee
Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in- context. Journal of Machine Learning Research, 25(49):1-55, 2024a. URL /https://www.
Smir.org wwwppppppss/v2555552
1042. .tml
Ruiqi Zhang. Jingfeng Wu, and Peter I. Bartlett. In-context learning of linear transformer block: benefits of the mip component and one-step gd initialization. In Proceedings of Neu- ral Information Processing Systems, New Orleans, LA, USA, December 2024b. URL https:
7/ //araiv..rg//as/22222110
14

---

Published as a conference paper at ICLR 2025
A PROOFS FOR SECTION 2
A. PROOF OF THEOREM
We first introduce the following lemmas that are used in the proof of Theorem 2.1
Lemma A.1 (Bernstein Inequality, from Theorem 6.1.1 in Tropp (015)). Let S ...... Sn be inde- pendent, centered random matrices with common dimension d1 x d2, and assume that each one is almost surely bounded:
E S D(I(( S = =i= 1....... n
With the sum:
S.
and the matrix variance statistic of the sum:
(222) == max [| [
then the following inequality holds:
-E2/2
P 2 $) s (di + d2) exp
for any 2 0.
v(22 + be/3
Lemma A.2 (Inverse Convergence, adapted from Lemma 2.1 in Jin al. (2024)). Suppose have a random invertible matrix Ω and invertible matrix sequence (52()) such that ②2 2g Ω. If there exists constant > 0 such that omin (f2) 2 λ almost surely, then it holds that:
()) P s2-1.
Further, given convergence rate
P (((") 12 2 $) s E(n,
then:
P 2 s2- 2 e < <((n, 12)).
Proof. We have the following decomposition:
- (22
It follows that:
(1(m))----?
< (2)
25 (2)
|2-11
1
12 (2(m0
22
Then
p -
2 22 P
Ω @ 26
12
$ {(n, X²e).
Lemma A.3 (Product Convergence). Let (52 (n)
}, (52 )....
..... {52 (n) be K sequences of matrices
k
such that 2 (m) P, П2 22(m) (m) P Z2.... 22 (m) SK where each || (m) || is almost surely bounded
K
for every l 1, 1,..., K. the dimensions match, then it holds that:
Ω (n) ó
2 (m) P, N S2
15

---

Published as a conference paper at ICLR 2025
Further, given convergence rates:
PP { (2 22
2 s (n,,
P ( 222
2 s (m,,),
P
2 E s Ex(
then it holds that:
P 02 1 ...5%%
22
2x
2 1
Ea 71
(14)
K
A Mr
where M₁ is an upper bound such that ( S M2 almost surely, 4h: 1,...,K. K.
Proof. We begin by showing the case of K 2. By the triangle inequality, we have:
12 12
52
s 12
(n) si("m 22 (m)
+
52 52
( (2(mm - -22 +
( 22
12
(2 22
M
( S2
Using the union bound, we have:
(m) - 2250
2 _
<P
M2
Ω (n) Ω
M1
(2
(n)
-2
_
<P
M2
12
(n)
Ω2
c/2
+P M1
12 (x)
e/2
<61
n,
+ E2
2M1
2M
For any K > 2, suppose the statement
14 holds for k = 2,...,
"....,K K
1. Observe that:
(2 (n) 12 (n) 22 Ωx
52 (1) 2 (2
(n) Uus. ... 2x
12
(n) K
22 ... 2K-12 Ω2Ω 2x
K-1
MMx
2 (m)
Mr Q sxx
K-1
(15)
16

---

Published as a conference paper at ICLR 2025
Then it follows that:
12) (m) (2 S2 10. -2x 26]
K-1
$ P
MK 12 nJ 12 (n) (2)
K-1 2252
2x-1
+
M
/2
-
2xx
2E
k=
P
M 12) m1) 22
(m)
(2)
22 f2
K
K-1
K
K-1
+
IP
MR
52 22x
K
K-1
Σ
n.
KMK I
E
+ EK
n.
K
K- I
e
Me
1=1
K
E
11, K light M
Thus, by induction, the proof is complete.
Remark A.1. In Lemma A.3 consider the special case where Ω2 0. Then the inequality 15 can be simplified as follows:
K
2 m
2
(m) ... 22/%
0
S
M₁
52
k=
And we have the following simplified form:
K
p P(
2
o
2
SPP
MR
(
2 8
k=2
c
S E1 n1
l
M
Lemma A.4 (Deviation Inequality for Minimum Eigenvalue Projected Sample Covariance Ma- trix). Suppose Assumption holds.
11
When n 2 max
the following inequality holds with probability at least 1
K
aKz
3ge Kn
2
Omin (min(() -1 2p(a + 1)B2, log( Kn)
Amain
x P2X 2 a
= Ae,
n
Amin (2z)
where K =
Amtnn
Σx) (
Kn =
Amin 22)o2 (e) Σ, = ==E P2 = 2(2 2)- 1T and Ag is
6.B.?
21,
a lower bound of Amiul ( T 2
Proof. Let
E = PZE2, E = (I - P2)
We have the following decomposition:
= = + E, = Z E E = Z(e + -' + E..
Let V = (Z - -1 Ez. Since P2 E1 0, we have
X / (
(Z
= ( + v Z (++
17

---

Published as a conference paper at ICLR 2025
We can write:
Amin
_ x ExTP2X]
PzX]
Amin (10 Tn ZT z
(6 + v)
(16)
n
Note that in general, for positive semi-definite matrix A, have
u Au
Amin (A) min u/0
u
and
(Bu) ABu
Amin (B" AB)
min uz0
min Amin (A) (Bu)
Bu
uy0
u
Amin (A))min B B).
Thus, from Equation (6). with probability at least 1 - -. we have:
1 x P2X)
Amin (0+ z
Amin
=
( + v)
2 Aedmiu
(00 v) (0 v
It now remains to bound Amiu (o0 T O 1) o min e + V
From
Hsu et al. 2014), for each i E [p] and any given & > 1, with sample size satisfying
Blog +
(17)
Amiu EE;)
we have the following holds with probability at least - 3e-
Ba (q + 2/ + 2t)
B² [q + 2(q + 1)d]
| =11@= -
Note that
W = E. VR = UUA.U
U v. = Σ
1=1
and
9 | = V vv V UU v. Σ
( v
We have
Amin
[E 12 < 1 S. S Amax Σz)l 11
Then
pB2 (a 2(g + 1)6] Amia EnIn
1 12 <
1
(18)
Amin (E,)
A=
k=
Hence, by Weyl's inequality, we have
pB2 [g + 2(q + 1)t]
Umin + V 2 mmin (0)
- > Omin (e)
(19)
Amin Eg)) )
18

---

Published as a conference paper at ICLR 2025
where t is taken to be small enough such that the RHS 0, i.e.
Kon
1cts
(20)
p(q 1) 2(g 1)
where Ko Amin E. Jo ( (
We now rewrite inequality Equation (19) in terms of only. From
2B 5g
condition Equation (77. for any given sample size the range for & is: < S
Kn log g.
(21)
where K = Aman . 6.2
We take
1
K
t= log (Kn) log
log -
2
So that condition Equation 21 is satisfied when ge sufficient condition is:
To satisfy condition Equation (00),
K.
K
Kon pla
log
- n
1))
Note that when 2 ge
we also have:
K
K
K
log
S
q
So sufficient condition to satisfy both Equation
(20) and Equation 2 is:
p² + 1)2K
max
K
qkb
Then the bound Equation (88) can be rewritten as:
K n) 2 pBa q 2(q + log
2p(9 + 1)B?, log [K
S
7
(22)
Amin Eg)m
Amin Egm
Finally, from Equation (16)
2
2p(g + B) log K Amiu (En)m
- [xx P2X 2 >2
Amain
O (0)
-
11
Proof of Theorem
2.1
We denote the observational values (Z,X,Y) (Z,,, W.)l and
E, = {(1,)
1n
(22.)) The 2SLS estimator is given by:
BBSLS
T 26
Tz
[(I 2) 2 x) z Z - x
(I 2) ((
2 x
TTY
(23)
= (X 2T x z=
= - [x] Z(Z 2) _x x z(z Z)-
Define constants A, A > such that the following event A holds with probability at least 1-
I Z
x P
A Xmnin
2 Az, Amiin
2 \ A
(24)
19

---

Published as a conference paper at ICLR 2025
where Amin()) (+) denotes the smallest eigenvalue of matrix, Pz = Z Z denotes the pro-
jection matrix. We will first assume the existence of such X2 Ag. with their values to be determined later.
We first consider the case when event A is true. Let Q:: == E|zz [A], Qz1 == E[zx' |A],
D2 Ei=
2= (z,) Q.x), Σ 2;
Let B B Bux Bze, be some upper upper-ouunds bounds such that |2, -2=0 s B2z
|,,0
Quell f
B Bgsi 2
I BRai || < Bxe, almost surely, for all 1,... n.
The existence of
Bzz, Bew Ba Bze, is guaranteed under Assumption 2 ii..
By Lemma A.1
we have:
Z Z
n i=1
2i2,
P
Q=2
2
A
A
-P-
-Q=
2 8 A
n
PP
22
2=:)
2 ne A
(25)
n²c2/2 /2
< cxp
11 ((12z21A4 (A) + B3 /3
Similarly,
Z
X
Σ 2
P
Qxz
2 E A = P
-Qzz
2 E
n1
n
IP
( - -Qax)
2 nE A
(26)
n²E2/1
S (p + ②) exp
V( (A) + Bs
By Assumption Iii), the instrument z is uncorrelated with the error term 61, which implies |[z = 0. Applying Lemma again, we have:
T E1 2 A
Σ
2,
E
= P
2 2 A
71
= PP
Σ 2,C1,
2 nE A
(27)
i=1
n²²2/2
S (g + exe
(A) + Bze, ne/
With Lemma A.2 and 25 we have:
n n²(22))212
P P[|[n(ZTz)) Z) 2221 2 2 $ 2gexp
"(. (A) + B=sn((REE/33
(28)
X2p==2/2
= eqxxp
( (A) + X2
Note that we have ê -
9 + (ZTZ) E2 = + V. Under event A.
| = | + v < O + | < Bg + By = Bg.
(29)
With Lemma A.3 Reemark
A.I combining
27 29 we have:
n²( (2(22)2/2
x z(z 3, 0 2 {
< (q + 1)eep
1 v (A) + Bze,
no c )/3
Be
n²22/2
B3 B (A) B Bze, ne ne/3
(30)
20

---

Published as a conference paper at ICLR 2025
Additionally, with Lemma A.3
combining
26 28
we have:
{ X Z(Z 2) 2 X
Q
2 8 A
{
220
n²1
A,5
)²12
X1n²
=
1212
S <2(p exp
38.
+ exexp
M((2gn (A) B T1 AgE
\/e3)/3
11S2 (A) - X B Ti 3B 1/3
3B.r
x 12
Xin²E212 12
=2(p+ + exp
+ 2gexp
9B3 v (1:z (A) + Xa A.BrBBuTh B2x
994 v (1;; + B?B2 BrBBBaa B3nnE nE
Applying Lemma A.2
again, we have:
P n (X Z(Z' T
2) 2)-12Tx x) 1Q2x Qzx) ) 2 {
>33E(133)212
Xín
22pp+ + +x)exp
exp
1 ( (A) + A Bu BEZ (AGE)
1)B3yv(0 [A) + +. B
X 2 22/22
X 222/2
-2(p+ + exp
+ exexp
()33,u(Rz.. (A) X₂ Brr B
(B. (A) + 23 B2 BIZnE
(31)
Therefore, have shown that under event A,
n
X Z Z(Z -)-
X
From equation 23 combining
30 and
31
with Lemma
A.3
(Remark A.
have:
2 [A} P { B2SLS 2
=P [[|[( (X Z(Z Z X Z(Z -
ZTE, E o c(AA]
n² n²2222)2/2
I (q + e)exp
B B sc:
A) + BB n(\xE)/33
232Z2212
(q + 1) exp
BIlre A) + Bg nE/3
For the second part of the theorem, let c: (3 LA)
A
we have:
E ||clip Ba (92sLs) [313
E ||clip BR (B2SLs)
311 A [ + E [|clip B, (BeSLs) _113 [
IP [(A))
(32)
E
|1BSSLs
3 A PP (AA) + E |clip n3 ($esLs) B|( A IP {A7
SE
||BSL
31
A]
[[clip (BpsLs) 31 |r] E,
where
||clip, B, (B2sLs) E
-312 < 4B3.
(33)
21

---

Published as a conference paper at ICLR 2025
and
12SLS - 311 A]
L |1Bessu - B|1 2 e c|A] de
% P PIlB2SSS 311 2 vē { de
22
(q exxp
de
B3 Pl ze1 IA)
Bg Bee nv/33
6
1/ exp
E2
de / exp
AgnVE/2 dE
$ (q
2B2 (Sczec (A)
Be.
4B2 v (12 (A)
11exp
(
8B3 (1 IA)
32 B
( (A)
(q 1)
4B2
+
12
9X
8
exp
4 1
261
4B2 (S 261
(A)
8B2
(S IAA)
32B2 B
(1+
+
+
6
1gn2
9X2n2
6
4 + 1)B3
32
B2
12u( [A) +
2
2Bn²
9
(34)
Note that we further have the following bound:
u(S2zr, (A) = maxx
2:61,8
Σ 2501.5
|
[
(E
2501,)
[
y=1
n2 o
[]]]]]
Σ
max
E
2
E
S of
(35)
now remains to determine Az Aa and Bg.
From Theorem 4.6.1 of (Vershynin 2018), when n 2 CB + vB). with probability at least
1 - 2e- + we have:
cB3 + vt) 2 Amiu (E:) 1
1 Z Z n
Xmnin
n
where c is an absolute constant. We rewrite the theorem by taking t = log( K n) (similar to the proof in Lemma A.4). Then we need the following condition to be satisfied:
2
K
1
n 2 B
V + V
log
(
(36)
q
2
We can bound the RHS of Equation
(36 as follows:
2
B V0+ + /
log
K
1
K < 26 B² ( log
(1))
2
$ + 26 2r²
(g+ log
4e B3K
3 2
q
22

---

Published as a conference paper at ICLR 2025
where the second line follows from the inequality log(z) + log(C) 1, with and K
q
C' 4r B So a sufficient condition for Equation 36 to hold is:
4c B.BKK K
3
n 2 4c B
q (g+log log
(37)
2
With condition Equation
37 we have the following bound holds with probability at least
2ge
Kn
2
cB²
+
log
K -
1 =
Amia
Z
Amin (g)
A.
(38)
nn
Furthermore, with Lemma
A.4
when
2 max
qe
p
(q+1) 9k3
we have the following holds
K
with probability at least
Kn
1 [_xT P2X
Amnin
71
Kn "1 2p(a + 1) B3 log
2 A
Tmin (t) -
Amin E2)m
2
cB² /q +
log K n
2p(a + 1)B2 log K
2
9
Amin (2x)
1
Omin (+)
n
Amiin (Ez)
= Aę-
(39)
From Equation
22 and Equation 29 we have:
K 2p(a + B)B2 log 9
Bg Bg +
(40)
Amin
En)n
With E
5ge putting together Equations
(35 38 38
40) into Equation (34). and Equations
33
34) into Equation
32 we have:
E [|clip (B2sLs) 31]]] < E [1,22SLs
A]
4B3c
(g
1)
32BB
20ge B3
12nB2of +
n2
Kn
(41)
B2
+ (()))
0
where
2p(g-1)BB log( (k)) Amiua E.
Be +
B
C(n) :
cB3 (VT (VTt//ugg /lo8( K (n)
2pía+ 1) B2, lool +
Amin Σx)
Omin (
Amiua E. Jn
(42)
thus completing the proof.
23

---

Published as a conference paper at ICLR 2025
B PROOFS FOR SECTION 3
B.1 PROOF OF THEOREM
Lemma B.1. Suppose {12 ..... (2(1), ..... is d d-matrix sequence decaying with exponential rate T, i.c. for some constant <>> 0 and 0 < < L,
[2 S Sor
Then for any O, there exists finite constant:
(1 (1- (e/d)
To
log,
c(1 + (1 (e)ddd
such that
20
II [I + (2(0 (20) 1 5, 1=
and hence
II I + (1+ To
< <vd+e.
Proof. By definition,
(2
D S cr
(k)2
which implies:
(
(k)
S cr
Vi,ik.
Consider the product of any two matrices. By sub-mulippicaatiityy
10
<
120 s cnpttl
which implies:
Scntth Vi,j,k,
is
Similarly, for the product of any number of matrices:
... ..........)) (tt
I thy Vi. ki...., kn
Thus
l2
+ (pns -I|
F I + (2(0))
(I+Q ... (I+ -11p
(43)
(
Σ
() +... + (000 -20
h
Σ
Σ
11 +... + (1
415kst
4
24

---

Published as a conference paper at ICLR 2025
Note that the last inequality can be checked by comparing matrix elements of both sides. For any e > O, we take To [log, (1-r)(e/d) convenience, let
(1++1-1X6/d)) Consider t To and to + Do in 43). For notation
= = Σ To
11 + Σ
11 +
E Toskelem
2ppttt+111
ToSkel
Then
Big
Σ
cr +
Σ
Eyphl++
Σ
e Spbtltm+ +...
To
Toskel
Toskelem
ph To
p 22
k
k2To
k2
CT To
2 C 2T
1
r
1
To
(1 - (1 - cr
8
Thus
d
(I
(20)
+
I
=|
Σ
S
E.
t=bb
F
f-=1
Hence completes the proof.
Proof of Theorem 3.7 In the following proof, we treat Z, X, Y deterministic matrices.
We begin by checking the inner loop
7ap:
+ θ (t-1) θ - T (t((-1)- x)
- (I-n z) (0(1-1)- +nz (X (x-26)
= (I - 1/ "gzzz" 2)²² (0lt 4 +1=² T
(0 (x- z6) + 1) nI 2) z (x (x-20)
[-1
= (I (I-n)) 1/ 7zTZ) 2) (000 - 6) + Σ (I- (I-nZT 1-1- zJ (x -26)
2=
(I 1/ 2)' (0(0) 6) + 1 [ (I - 11 2)4] (nz (1z = ( Xx
z6)
(I 1/2 2)' (0(0) 6 + - (I- 2)0 (2 2) 2 X 6)
2T
(I n 2)' (0(0) ④
With learning rate < 1
- d (2) let n((n) == p (I T7 where p(-) denotes the spectral
radius. Then it follows that ( < k()) < 1. We have:
0 (t) 61 = (T 1 2) -))))())))
s <II(T - nz z)'l (0)) θ
5 x())) ( 6
O(x(n)))
25

---

Published as a conference paper at ICLR 2025
Thus (O(*) converges to @ exponentially with rate k(n).
For the outer loop
7b we have:
3 BESLS - =3((-1) BESLS a0(t-1)T TT (0 (20( - -x)
= -1--00(-11 T 20((-1) (31(-1) - BEsLS) + +00l1111T2 2T (r- (1-)2((-1))
( z z0(1) (3t) - }ESLS)
A13(1)
(I-00 n TT 2
@ T (y-26/1 (ESLS)
A2
(44)
To simplify notations, let
R = ( - Ô = (I 1 2) 2)((01)) (00))
V (I (I-a Z 26
W = RT) z zÔ + θ T 2 ZR + R z ZR
0 < 1 < (2) let 7(a)= = p p(1 p1--00 1=-0@2Tze zé).
With learning rates < a < 02 m (z0)
Then it follows that (a) < We have:
n
<x(n)' -00
(45)
| s sTaa',
and
| = R T T z + èT
z ZR() (1) + -R(() (T) t ZR(
52| zT z R
|
Z||
Ri 2
22x(n)' |0 T 2 /1| (0) 6
+ 2
1 /|| ( - 61
<n(n) (2| T 2
( 66
+ | 7
0
- -61 =6'))
-(x())))
26

---

Published as a conference paper at ICLR 2025
Then from Equation 44 we have:
A,,30
a T z 2
- (B1))-
1=0
I a z 2 a (R z zó + TT ZR + R I ZR )] [3800) - -BSLL))
1=0
I a z z
awio)] (310) (2SLS)
i=0
(1 - -
z 26) z6)
I- a (1 (I II
z6) -1 wo] (B - -Besss
1=0
t-1 -())
(I-a6T (l-a@=zzz6)) zの) -1 w (3 -
I a (I- ab z
i=0
We denote v
I a z z6)
By Lemma B.1
we take 1, co be constant such
that
|w(0|| F S (on( and To = [log.(n) v pcoppt
(1-x(0)
Then we have:
II TToo
[ - )
II To
(1-1- (1- ( wwwh)\ (1
< <vp+1.
(46)
Hence
A A
= V(
w (3(0) Aesis)
t=
II (1 ww |3000 BeSus|
< [
vw
(47)
t=U
To 1
| 11 1) ||300 - -2ass|
< <7(4)'
II
(1- (I- (1-
==0
= =(((())).
Next we consider (2
22
II ( z 20(0)) B ZT (r- 20 (
j=
(I- z z6 - TI- -wi)) (R +0)² T 2 Y 2 Z( (R + b) AOSLS]
(R(+
j=i+1
1-1
Σ
[ (I- z z6)
II j=i+1
I a [ --as T- Tz (0)= w(n)
(RO) + 6) T z Y - Z (R(0) + )) |RSUS]
1-1
1-1
I (I- -www) (n
(R() -6) 2 { Y-22 Z ( 82Sis]
- a V(t-1-i) II
+
j=1+1
For convenience, let A 20) = 22 (0) + B10 where
t-1
t-1
22 == a ((--1-)
(I ww(r))
R ZT Y- 2(R() + +6) Besis]
1=0
27

---

Published as a conference paper at ICLR 2025
t-1 Am231:=>>
t-1
-a
vi----
II (1
=T y-z 2 (((()) + 6) BOSLs
jmit1
Suppose M1, M2 are the upper bounds such that
z Y
R + 6 BESLS
||
$ M,, vi= 0, 0...., 1-11
[-] (1- (
II -=++1
Vi = 0 t- t-. 1. S S2
We know such M1, M2 exist because of the bounds given by (45 and (46). Let NT M, M2. Then
|
M
a (1-1-))) ROOT
S Ma
R
1-1
s Ma
(0) ④
Σ 7(a)~11'x(() 'x())'
[<())
Ma
(0) θ
Σ 7(0)-11
((a))
[aa) - x(1)² -(n)"
o
7(a)
x(T) x(n)
s (max(
and similarly,
t-1
A A22
Σ a / (t-1--1))
1=0
< MM
(1-1-9
s Ma 6
(a)
0(T(a)))
Thus
A2
121 + Am3
< 221 + A20
(48)
< (1mm)(((a),x()))
Therefore, plugging 47 and
488 into 44p. we have:
| 3SSLS
s ((max(7(a))
Hence completes the proof.
28

---

Published as a conference paper at ICLR 2025
B.2 PROOF OF THEOREM 3.2
Proof of Theorem 3.2 For ease of notations, we ignore | in the following proof. Consider the input
matrix taking the form:
21
nn
n
nn
1/1
yn
H
R Dx(n+1),
3
3
(0)
n++1
1
1
1
0
i.e. element-wise,
h (0) 22, yet,, !1))
9 (0) pp
(00))
[= 12 ,.t'", i= i=I....,,+
........
1i n) is the indicator for training sample. We can take any initialization for (0), (0(0) and a(0), To avoid abuse of notations, we omit the superscript of those parameters to be updated in the following proof.
where D qp =p +p 4
Our goal is show that there exists series of attention pa- (1:2) JMEEM C DxD such that 0 (1:2) ATTN updates O 3 on the Do row Do + qp correspond to θ and the updates on row Do + qp + 1 to row Do + p + correspond to B.
Recall the definitions and (5).
(1:2) rameters ATTN {(Q
TT:2) K (T:2)
corresponding rows. i.e, we denote Do == q + ++p-
the updates
1) the first layer, the transformer updates the current first-stage estimate ã.
For m
2k 1, = 1,..., p, define Q!! (1) K!
v-I))
such that:
2:1
+
00
:
(1)
m
h
(0)
2ig
(1)))
m
(0)
(0)
( (0)
m h
e Do+gptptl-
(49)
(0)
gk
T. ik
1
0
0
For m 2k. h: 1,... ,,, define Q K! ((I)
such that:
O 1k
Q (1))
2ig
K (1) (0)
-e-Dotopppte
(50)
gk
a. kk
where e, E D is the standard unit vector with only one at the j-th coordinate. Note that the above are just linear transformations h, hj, hence such matrices ! K( ((1
m must exist.
29

---

Published as a conference paper at ICLR 2025
Then we have:
n+1
h;
(1)
(0)
(0)
K 1
n 11)
T1
n+1
h
(0)
+
0
21 O
Σ
2s1 2 (0)
+ ik
eDttoptpt
1
km
j=1
1m1
n
h
(0)
+
Σ
Σ
211
(0)
Fik
@Dotaptp*k
h
(0)
$
ik
Pik
(0)
eDotoppptk
Thus this layer correctly updates the first-stage prediction values
1 ..... 2 where a!
Ize (.: = Ei=1
2n
1
1,:
We will have:
2n+1
Ø
Un
1n+1
V1
V/n
(0)
H
e
(0)
Up
(
3
3
((1))
4 (1)
a (1)
2) In the second layer, the transformer does the gradient updates on the parameters θ and B.
For m 2k 1, 1 ..... p, define Q m
(2) m
v.(2)
such that:
(0)
2j
Dykty R(1
Q (2)
h
(22
(2)
12
bg)
v nT -(n 1)n Σ
Do
(51)
For m = 2k: 1,..., P. define Q K( v.(
such that:
(0)
2j
.,
Dikty
2
h
(1)
K (2) k22
h
(1)
R(1 - t2)
(2 = -(m + 1)n Σ 1)g+t:
(52)
where
max
((())) 2,11]. Then we have:
I+.
(10
1 24-1
11
(1) (2)
h
(20)
(0) ( i 23 Enntj R(1 - t,))
+
tk Zj k,
(
1ft, = 1}
(0)T
ikk zj Kzsh 4,
30

---

Published as a conference paper at ICLR 2025
and
x2n (2) R. K 1 o (
(0)T .,K 3j + Tzuly - R( 15)
(O)T .
2j + Fzk) 1 1 1}
(0)T
3j + Tzk tj.
So that
2p
(Q 11
(2)
(1)
K m
k (11) v (1)
m
m=1
9 i,k 2j + x33z )
= -nnn+ + 1)t
i,k
(U)
2, - xzbb - (-6) (U)T
Do
k=1
p
= -n + 1)(,7)
Zzl
(0) (0) .,k
2j - -zjk) eDo (k-1)
k=1
1=1
Similarly, for m 2p
1,2p + 2, define Q (2)
K! Vill such that:
7()
# (I)
1
ystz
P
(2) (1) Q
Q2p+1 h
2p+ 1
(2)
R' (1
4
V2p+ hn)) -(n + 1)a
Σ
#
ji
eno
1=1
(53)
-3(00
((1) yjtz
1
P
Q 2ppt2h ())
K.
2p 2p 50 =
R'(1 tz)
V2p-2h) (1) = (n + 1)a
Σ
a(1)
jjl
e
1=1
(54)
where R²=
max
(13(T) Tx:1.. Then
+11
1=,,
20 h
(2)
(1)
K2 2p h nn()))
(B $ n,t, - R'(1 - t,))
(0)) (00 yi 11ls =
(0 2
(1) 8/ ty.
and
a (1o 2 (1)
_ 2002hh?0 (1))
()0 1 + t5bي - R²(1 - t2))
3 (0) T + 8/ 1 =1}
(-1(0) a0 (1) + 91 4/.
So that
2p+2
Σ (10 (@00
11 ()) k Vil2)n ()
1
= =-(n+ + 1)15ax [o (3(O)T ys) (-310) 2000 + w)
Σ
(1)
Ei Da
1=1
= -- + +1),5a Σ
(1))
(30)T&(1 yU eDotppti
31

---

Published as a conference paper at ICLR 2025
Thus the final output. for i ,,..., 1:
2p+2
h
(2)
h
m
j=1
24
Σ
2j
1
00)
I11
Σ 7
== 2/
(0)
P - Izpp
(100
Σ
4 ()) (0) 31
t
144.
(0) 7
Z0 (0)
X
n
Z
(300
a (T)
Z
(Z (()) v)
4
This corresponds to one-step 2SLS GD update of O 3,, according output matrix is:
Therefore, the final
Zn
2++1
In
In+1
V1
Yn
(1)
(2)
g
g()
(
(1)
n+1
Thus the proof is complete. We further note that in construction steps 49 regardless of the initial values (0), same linear transformations on the input vectors. Therefore they are identical across different layers.
50 53
(0),g(()) g00), and (()), the matrices Q
K
V
do the
B.3 PROOF OF THEOREM 3.3
Lemma B.2 (Generalization of pretraining, from Theorem 20 in Bai et al. (2023) Given optimiza- tion problm 112 with probability at least - C. the solution Ô satisfies:
L2DD2 DD') log(2 + max Bo, R, BB)) + log(1/))
(0) < inf (0) + + B² 0.0
N
32

---

Published as a conference paper at ICLR 2025
Proof of Theorem 3.3 We begin by showing the (clipped) 2SLS predictor achieves small excess loss under in-context distribution P:
Bp ( (colipp,, ($2SLs), xn-1) Mn+1) )
Ep (colipns, (B2SLs) - B,+n+1) + (B, In+1) ynot) )
Ep (clip n2 (B2sLs) - xn+2))²
+ 22E [cliip,, (B2SLs) - B, xn+11 ((3, xn+1) - Bb+t))
Ep [(13, 20 yn+1) )²]
Ep (clip (BesLs) B, xn+1) +Ep ((B, - 3
Excess Loss
where Ep[(clip B, (B2sLs) - B, In+1) In+1) xn+1) - yo+1]] = follows from the independence
between (clip (52sLs) - B3 mnt1) and (n+1) - ynt1) with Ep [(3, In+1) yn+1 1)
Ep(en+1l = 0.
To bound the excess loss, we have
T #n+1 (clipn, (B2SLS) - 1/)
Ep (clip (B2sLs) B, an+1) Ep
clip (B2SLs) - .01]
Epp
xx++1 11?
(55)
Ep
||0m~11 11
Ep
[lip (Bosls) - -31²)
B2
+ (n)(( En TEnb + +02) oz)) ))
0
B
K
where the last inequality follows from Next, for the ICL loss, we have
3
(0)
(TF)(H)) -Vn+1) )
Ep
(TFo(H)) (clip (BoSLs), xn+1) + (clips, (B2SLs), xn+1) - vx++1) - 27
)
Ep
(TFo(H)) (clip (BesLs), 2141) + Ep [(clipp, (B2SLs), (n+ Un+1) )²]
E.r -E,, Ep
(TFo (clip (B2sLs),
xn+1)) (oass))xx++4) (clipp (B2sLs), xn+1) - bh+t1]}
+ 2Ep
13 pppp
(TFo(H)) (clip (BesLs), xn+ + Ep (clip B, (BesLs) B, xn+1) Ep (11, vn+1)
2
Ep <clip,,, (BesLs) kn+1) ynt1 |]
+ 2Ep
TFo(H) (clip (B2SLs) x xn+1
(clip (B2SL)) ) I + Ep (clip R3, (B2SLs) B, B, 2
< E.y
Ep
TFe (H)
Ep (1B xn+1) 3m++1)
+ 2Ep
TFo H (clip (B2SLs) xn+1 |] (Ep) [clipp, (B2sLs) B, B.
+ Ep (B, xn+1) - yutz
(56)
33

---

Published as a conference paper at ICLR 2025
know that there exists L 21 1-layer attention-only transformer model 0, with 2(p 1) heads, and embedding dimension D qp + 3p + q 3, such that for any H, given any learning rates a, 7 and A as defined in Equation
From Corollary 3.1
S). the following hold
TF o H) (clip ($2sLs), Tn+1) s (BaA
Denote A = min Ep [A H. a, m)., then under a², , we have:
(Fo H) (clip ($2sLs), Bn+1) so (B((()))))
EpEp
(57)
and
E Ep
(To H) - (clip B3 (B2sLs), ( 0 <B(BAA
(58)
where MX,2 == E Ep
2
H. a", n- is close to O.
With condition [3). from Cauchy-Schward inequality, have:
E Ep (1 yn+il s E
( VEp (ci+1)
=E,lor S o.-
(59)
Also, from 55 we have:
B
+ (m)))²) ))
E Ep (clip,, (B2sis) - B, xn xnn1)) SO
B
q
(60)
K
Further,
EgEE [clip (Gesus) B, xn+1 S
Ep
(clip B3 (B2sLs) - B, an+1)
(61)
B2
0
B
q
C2())62)
72
K
Finally, with (57) 160 61.,
rearranging the terms (56). have:
LIOL (0) Ep [Wn+1
-3,, 11)
(())oe) ) (x
B² +C((n)02
B3 K
0
B²
MA,2 (At)L
q
q
K
B2 K
B2
50
(A+)L
B2
q
C2
+ BZoe
B
q
+ (2()0²
+ +xā2
n
n K
(62)
Thus combining Lemma B.2 with
62) completes the proof.
C ADDITIONAL EXPERIMENTS AND DISCUSSIONS
For all experiments in this section, be consistent with main experiment in Section generate 50 training samples with p 5, q 10, following the data generating process described in Algorithm [ The task parameters O B, Φ. & are sampled from standard Gaussian distribution, the covariance matrices Σ. E., En, are set to be identity matrices, and the noise level o, is set to 1.
*The clipping bound Besus can he matched by adjusting the clipping threshold the last layer of TF
34

---

Published as a conference paper at ICLR 2025
C.I SIMULATIONS VERIFYING THEOREM 3.1
We use the GD-based 2SLS method calculate the following metries:
7 estimate the causal effect B3. For the simulated data, we
2
0.0016,
-0.0212.
6 2
max (ZO)
o2 max (Z)
By Theorem 3.1 the gradient descent converges when E (0.0.0016) and 1 E (0, 0.0212). The overall convergence rate is determined by A == max) where
7(a) = (I- aQT 16TT = 26). <(n) p 11 11 2)
-I-
We first set a 0.0012 and vary 7). The corresponding convergence rates are determined by max(0.87, x(n)). Next, we set 7 0.01 and vary The corresponding convergence curves are determined by max(?(a), 0.82). We compare the estimates (t) with the 2SLS estimate 3SsUS as the iteration proceeds. The convergence results are shown in Figure 3
Convergence IV
(a-0.0012)
Convergence of Estimate with Gradient Method (g=0.01)
u-C 0000000,, A-marto 0.82)=0.94
2,5
g=0 01. Ammax/0.87 0.82)=0.87
u-G 001, /-max(0 89, 0.52)=0 89
q-0 02, Ammax(0.87 0.88)-C.85
a-C 0015, A-maxt0.8B. 0.82)=0.89
g=0.022, A/ 1.0/)11.r
2.0
Amma([] 0.52)=1.53
3.5
1.0
0.0
20
40
50
30
100
0.0
0
20
40
66
80
100
Iteration
Iteration
(a)
Figure 3: The convergence of the GD-based 2SLS method with (a) fixed = 0.0012 and varying and (b) fixed 0.01 and varying a.
The results Figure 3 are consistent with our theoretical analysis in Theorem 3.1 is worth noting that in Figure 3a when 17 is relatively large (or small), the convergence curves exhibit patterns. This is due to the innerloop updates (7a) are converging faster (or slower) than the outer loop updates (7b). However, the overall convergence rate is still determined by A. This pattern doesn't appear in Figure 3h as we set to be moderate value, which ensures that the inner loop and outer doop converge synchronously.
suiggly
Next, we show the bias of the GD estimator. For better convergence, we set a* 22 (Z)) and
17* 1 samples. The results shown in Figure
(2)) We compare the biases of the GD estimator with 50, 100, 150 in-context
35

---

Published as a conference paper at ICLR 2025
Bias of Estimators Iterations
(n-50)
2.00
6D Estimator (n-100)
Estimator (n-150)
OLS Estimator (n=150)
1.25
11800
1.00
0.75
0.50
Iterations
Figure 4: The convergence of the GD-based 2SLS method with a*
a (Zé) and 1/" o (2)
1
1
The biases of 2SLS estimator and OLS estimator at 150 are plotted for comparison.
C.2 DISCUSSIONS ON 2SLS WITH lg NE-REGLLARIZATION
In this section, we briefly discuss generalization of our analysis to the case where the 2SLS estimator is regularized by the penalty (Ridge 2SLS). For this case, the bi-level optimization problem Equation 66) is modified as follows:
min B
((3)
1 n
Σ
y 2, T(3) =
where
aag min
[r, - 0) 1
where X, 2 0 regularization parameters. To solve this problem, the GD updates in Equation 7 is modified as follows:
+ (f+1) = (() 1) z (Z) X) + (
(63a)
((++1) = (()) (T) z (z)
(()) Y) + AB
(63b)
The only difference between Equation 7 and Equation 63 is the additional terms 7 (1) Equa-
tion (63 and (t) in Equation (63b). The convergence analysis of the Ez-regularized GD updates in Equation (63) can be conducted in a similar manner in Theorem 3.1] The only difference is that the convergence rate A is now determined by the spectral radiuses of I 1y I I- z Z + AT), respectively.
rf) and
With the same configuration as in Theorem 3.2 but adding p attention heads in the second layer (i.e. the second layer needs 3p + 3 heads), we can show that transformers are able to implement the fz-regularized GD updates in Equation (63). The proof follows directly from Appendix B.2 with the construction of the new attention heads in the second layer as follows.
For m 2p + 2 + k, ls 1, define Q (2) m
(2) such that:
h
0
h
0
(2)
h -//
(0)
Do+
For m 3p + 3, define
Q
33p43
K(
30+3
(2)
(2)
39.3
such that:
Q
(1)
0
1
3p
(1)
0
3p
t (1) =
- Σ g
e gottl
36

---

Published as a conference paper at ICLR 2025
Then the remaining proof follows exactly the same Appendix B.2 This result indicates that transformers are potentially capable of handling multicollinearit in IV regression problem. We conduct experiments to validate this and the results are shown in Appendix C3
C.3 EXPERIMENTS MULTICOLLINEARITY IV PROBLEM
we examine the case where multicollinearity occurs in the IV re- gression problem. We generate the test prompts in the same way using Algorithm [ but introduce multicollinearity in the endogenous variable æ and instrument
As a supplement to Section 4.2
Specifically, we first generate 4 columns of X, and 9 columns of Z, and then set X.5 N2X ,,4, 100 I,), and 10 10 NN In). The results are shown in Figure 5a As shown in Figure Sa both ordinary OLS and 2SLS estimators fail to estimate the coefficients, while the trained transformer model is still able to provide consistent predictions and coefficient estimates.
Error Contoxt Sample Size (Muttcollinearty)
Error Context
Multicolinearitl
25.5 KIE
Sidge 2515 KPS
25.
Sidge 2515 MSE
OLS ICPE
sidge 0LS ICPE
OLS HSE
Ridge OLS HSE
Th roformer KIE
Transfommr IOPE
Tra roformor MSE
Transfoomer MSE
20
25
25 Number In-context Samples
20 25
Number In-context Samples
(a)
(b)
Figure 5: The ICL performance of the trained transformer model endogeneity tasks with multi- collinearity. (a) 1 collinear column in X, and 1 collinear column in Z. Note that the coefficient MSEs for 2SLS and OLS are both out of range. (b) 2 collinear columns in X, and 5 collinear columns in Z. We compare the performance the fz-regularized 2SLS and OLS estimators. The curves are averaged over 500 simulations.
We further examine a more difficult case where heavy multicollinearity occurs. We first generate 3 column of X, and 5 column of Z, and then set X. -N(2X..- 10-61,) for j 4, 5. and Z.j (. 10-6I,) for j 6, 8,9. 10. For better comparisons, we now compare the performance of the trained transformer model to the fe-regularized 2SLS and OLS estimators (with all regularization parameters set to 1). The results are shown in Figure 5b
These results suggest that the trained transformer model is capable to handle multicollinearity in IV regression problems, even though it has not been specifically trained with multicollinearity tasks.
C.4 EXPERIMENTS ON COMPLEX NON-LINEAR IV PROBLEM
As a supplement to Section 4.2 we examine a more complex scenario where the instrument z has non-linear effect on the endogenous regressor æ. We consider the following data generating process:
y (B.x) +EI, and x g(z)) + E2,
where g RR + RP is a two-layer fully connected neural network with ReLU activation function. Similar to Section 4 the test prompts are generated using Algorithm [ with all task parameters and weights of neural network sampled from standard Gaussian distribution. The results are shown in Figure 6 From this figure, we can see that the trained transformer model still achieves optimal performance in this complex non-linear setting.
C.5 EXPERIMENTS ON VARYING ENDOGENEITY STRENGTH
we examine the performance of the trained transformer model in standard IV tasks with varying endogeneity strengths. The strength of endogeneity is determined
As a supplement to Section 4.2
37

---

Published as a conference paper at ICLR 2025
Error vs In-Context Sample Size (Non-Linear
12
25L5 ICPE
MSE
10
OLS ICPE
MSE
Transformer ICPE
Transformer MSE
Error
Number of In-context Samples
Figure 6: The ICL performance of the trained transformer model in complex non-linear endogeneity tasks where the IV has non-linear effect on the endogenous variable. The curves are averaged over 500 simulations.
by the correlation between and the endogenous Algorithm ! we multiply u by a factor r E (0, 2) when generating test prompts. The results shown Figure 7 which illustrates that the trained transformer model is comparable with the optimal 2SLS estimator in these standard IV tasks, regardless of the endogeneity level.
61- To vary the endogeneity strength, in
Error Endageneity Level
25LS ICPE
MSE
OLS ICPE
MSE
Trnsformer ICPE
Transfomer MSE
Erront
Endogeneity
Level
Figure 7: The ICL performance of the trained transformer model in tasks with varying endogeneity strengths. The curves are averaged over 500 simulations.
C.6 EXPERIMENTS ON REAL-WOOL WORLD DATASET
In this section we provide an example to illustrate how the pretrained transformer model can be applied real-world dataset. We use the dataset from the study of Angrist Evans (1998). This study investigates the effect of childbearing on labor supply. For demonstration purpose, we consider a simplified setup. We focus on subset of the dataset that contains 6421 samples from Alabama. The outcome variable y is mother's labor supply (number of working weeks in a year divided by 52), the endogenous variable x is the number of children (2 2). and the instrument an indicator variable of whether the first and second children are of the same 5
"Research found that parents of same-sex siblings are significantly more likely to go to have an additional child (Westoff & Parke 1972), while it is not directly correlated with mother's labor supply mixture of sex of the first two children considered as randomly assigned.
38

---

Published as a conference paper at ICLR 2025
For each run we randomly select 50 samples from the dataset, and make the boxplot of the estimated over 500 runs As the ground truth effect is unknown, we compare them to the OLS and 2SLS estimates over all samples. The results are shown in Figure g
Boxplot of Estimated Coefficient
0.025
Aois -0.023
Pasus -0.089
0.000
Aur
-0.025
-0.050
Estimated
-0.075
-0.00
Figure 8: The boxplot of transformer's estimates over 500 runs on the labor supply dataset, com- paring to the OLS and 2SLS estimates. 8CPT taken to the median of runs. The gray box represents the interquartile range, where the middle 50% of the estimated values fall. The whiskers of the box indicate the spread of the estimates. Any points falling outside of the whisker can be considered outliers.
The final estimate SGPT 0.091, which suggests that with each increase in the number of children, mother's labor supply is expected to drop 9. 1% (approximately 4.73 weeks per year). This result is closer to the 2SLS estimate 82SLS -0.089 than the OLS estimate foLS = -0.023. This example demonstrates the potential of the pretrained transformer model in handling real-world problems.
C.7 EXPERIMENTAL DETAIL
The training of the transformer in our experiment was conducted on a Windows 11 machine with the following specifications:
GPU: NVIDIA GeForce RTX 4090
CPU: Intel Core 19- 14900KF
Memory: 32 GB DDR5, 5600MHz
The training process took around 10 hours.
'For large enough model that can fit in the entire dataset, this step can be ignored. As shown in the simulation study in Section single estimate is expected perform least good as 2SLS estimator, given the same number of samples.
39


Reviews:
[{"review_id": "WlXH1TIbpN", "reviewer": "Reviewer_feNa", "summary": "The paper extends the theoretical analysis done in previous work that looks at the class of functions that transformers can learn (e.g. simple linear regression) to a more complex set -- those with endogeneity and corresponding instrument variables. The authors show that transformers can learn this function and do as well as the direct solvers in most cases and potentially better in more challenging cases. They supplement this with a theoretical analysis.", "questions": "I want to caveat this review by saying that I have flagged to the AC that this is not my area of expertise.\n\nI wonder how this work helps one to understand what capabilities a transformer should / should not be capable of? Can we say something more concrete / general around what class of functions we can argue that a transformer should be able to solve? How does this relate to larger models ?", "limitations": "", "rating": 6, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I want to caveat this review by saying that I have flagged to the AC that this is not my area of expertise.\n\nThe paper seems interesting: it extends the theoretical analysis done in previous work that looks at the class of functions that transformers can learn (e.g. simple linear regression) to a more complex set -- those with endogeneity and corresponding instrument variables. The authors show that transformers can learn this function and do as well as the direct solvers in most cases and potentially better in more challenging cases. They supplement this with a theoretical analysis.", "weaknesses": "I want to caveat this review by saying that I have flagged to the AC that this is not my area of expertise.\n\nI do not know enough about this area to understand the potential flaws in their statements / etc or more subtle points.\nAt a broad level, their intro / overview makes sense and seems convincing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Xc5PqO3iDZ", "reviewer": "Reviewer_aQCE", "summary": "This paper first analyzes how to use IV regression for endogeneity, where the IV regression is estimated with 2SLS. Then, an error bound is proposed for this process. Furthermore, the paper shows that the transformer can achieve the gradient-based 2SLS for in-context linear regression.", "questions": "Please refer to the weaknesses section.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "I am not familiar with IV regression and endogeneity, so I haven't reviewed the mathematical accuracy of the theorem. My comments are based solely on a basic understanding of the motivation, overall contribution, and presentation. Please consider them with low weights.\n\nOverall, I find the paper well-motivated, with clear writing. The background information and literature review appear thorough.\n\nUnderstanding the mechanism of the Transformer could be valuable for advancing future research in this area.", "weaknesses": "I understand that theoretical analysis requires specific assumptions. I am, however, curious whether it might be possible to extend the theoretical analysis to the non-linear case, as most real-world scenarios tend to be non-linear.\n\nAdditionally, could you provide an example of real-world applications where the proposed analysis could be beneficial? If such examples exist, is it feasible to validate the analysis experimentally?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Z0qS1NvJBw", "reviewer": "Reviewer_XFtG", "summary": "The paper investigates the ability of transformers to address endogeneity in in-context linear regression, proposing that transformers can emulate the two-stage least squares (2SLS) method through gradient-based optimization. Key contributions include demonstrating how transformers can handle endogeneity using instrumental variables, proposing an in-context pretraining scheme with theoretical guarantees for minimal excess loss, and showcasing robust performance in various endogeneity scenarios, including weak and non-linear instrumental variables​.", "questions": "(1) How does the proposed transformer model handle cases with highly correlated instruments? In many practical applications, instrumental variables may exhibit multicollinearity, which could impact the model's stability and coefficient estimates. Could the authors clarify whether the model’s training or architecture includes mechanisms to address such cases?\n\n(2) What led to the choice of hyperparameters in the pretraining scheme? The paper presents a theoretical guarantee of minimal excess loss but does not detail the rationale behind the selection of specific hyperparameters, such as learning rates or number of transformer layers. Could the authors provide insights or guidelines on choosing these parameters for optimal performance?\n\n(3) How sensitive is the model to changes in the strength of endogeneity? While the experiments demonstrate robustness to varying IV strengths, it would be valuable to understand if there is a threshold or particular cases where the transformer’s performance degrades. Could the authors elaborate on the model's sensitivity to different degrees of endogeneity?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "(1) This paper creatively combines transformer architectures with econometric techniques, specifically instrumental variables, to address endogeneity—a novel approach that extends transformers' applicability beyond traditional machine learning domains.\n\n(2) The authors provide rigorous theoretical backing, including a bi-level optimization framework, and offer non-asymptotic error bounds, supporting their claims with comprehensive experiments that validate the model's performance against standard benchmarks like 2SLS.\n\n(3)  By demonstrating that transformers can not only handle endogeneity but also generalize to complex scenarios (e.g., weak and non-linear IV), this work highlights the potential of transformers as a robust tool in econometrics, broadening the scope of in-context learning applications.", "weaknesses": "(1) The paper provides strong theoretical foundations but lacks practical guidance for implementation. More details on the parameter settings, model configurations, and optimization process would enhance reproducibility and help readers better understand how to apply the proposed method.\n\n(2) While the theoretical contributions are thorough, the presentation is complex and could benefit from simplification or visual aids. This would make the bi-level optimization framework and convergence properties more accessible to a broader audience.\n\n(3)  The paper does not address the scalability of the proposed transformer model for handling endogeneity. Including an analysis of computational efficiency or memory requirements would help clarify its suitability for larger datasets or higher-dimensional problems.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)