Transformers Handle Endogeneity in In-Context Linear Regression

## Abstract

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.

## 1 Introduction

The transformer architecture (Vaswani et al., ) has demonstrated remarkable in-context learning (ICL) capabilities across various domains, such as natural language processing (Devlin et al., ; Radford et al., ; Brown et al., ), computer vision (Dosovitskiy et al., ; Carion et al., ), and reinforcement learning (Lee et al., ; Parisotto et al., ). Self-attention mechanism, 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 theoretical 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 performance in fundamental statistical tasks. Focusing on simple function classes, Garg et al. () highlighted that transformers, when trained on sufficiently large and diverse data from a specific function class, can generalize across most functions of that class without task-specific fine-tuning. Building on this, subsequent work by Bai et al. () established that attention layers enable transformers to perform gradient descent, implementing algorithms like linear regression, logistic regression, and LASSO; see also Akyürek et al. (); Von Oswald et al. (); Li et al. (); Fu et al. (); Ahn et al. (). Furthermore Zhang et al. (, ) showed that *trained* transformers’ ICL abilities for linear regression tasks are theoretically robust under certain distributional shifts and characterized the corresponding sample complexities.

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 uncorrelated with the explanatory variables. Ignoring *endogeneity* in linear regression leads to biased and inconsistent estimates, resulting from issues like omitted variable bias, simultaneity, and measurement error, which can distort causal inferences and lead to incorrect policy conclusions (Hausman, ; Wooldridge, ; Angrist and Pischke, ; Greene, ). Instrumental variable (IV) regression is a widely adopted method to handle endogeneity by utilizing instruments that are correlated with the endogenous variables but uncorrelated with the error term (Angrist and Krueger, ). 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?*

In this work, we aim to answer this question and offer new insights on in-context linear regression tasks. Our key contributions include:

- •
We demonstrate that looped transformers can address endogeneity in linear regression by leveraging instrumental variables. Specifically, we show that transformers can implement two-stage least squares (2SLS) regression through a bi-level gradient descent procedure, where each iteration is executed by a two-layer transformer block. Moreover, the convergence rate to the 2SLS estimator is exponential with respect to the number of blocks.
- •
We propose an ICL training scheme for transformers to efficiently handle endogeneity. Under this scheme, we show that the global minimizer of the in-context pre-training loss achieves a 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 non-asymptotic 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 et al. () and Bai et al. () adopted the standard multi-layer transformer architecture to conduct the experiments. Later, Giannou et al. () and Yang et al. () showed that a looped architecture reduces the required depth of transformers and exhibits better efficiency in learning algorithms. Gao et al. () illustrated that the looped transformer architecture with extra pre-processing and post-processing layers can achieve higher expressive power than a standard transformer with the same number of parameters. Apart from works concerning the implementability of first-order gradient descent algorithms by transformers, other works have also examined higher-order and non-parametric optimization methods. Specifically, Giannou et al. () showed that transformers can emulate Newton’s method for logistic regression. Cheng et al. () 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 et al. (); Falck et al. ().

Nichani et al. () illustrated how the transformers can learn the causal structure by encoding the latent causal graph in the first attention layer. Goel and Bartlett () explored the representational power of transformer for learning linear dynamical systems. Makkuva et al. (, ); Rajaraman et al. (); Edelman et al. () 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 (Angrist and Krueger, ; Angrist and Pischke, ). Recent works in machine learning explored the optimization based approaches for the IV regression problem. Singh et al. () proposed the kernel IV regression to model non-linear relationship between variables. Muandet et al. () proposed that a non-linear IV regression problem can be formulated as a convex-concave saddle point problem. Della Vecchia and Basu (); Chen et al. (); Peixoto et al. () proposed a stochastic optimization algorithm for IV regression.

Notation: Throughout this paper, unless otherwise specified, lower-case letters denote random variables or samples, while upper-case letters represent datasets (collections of samples). Bolded letters indicate vectors or matrices, whereas unbolded letters indicate scalars. The notation $\boldsymbol{X}_{:,i}$ refers to the $i$-th column, and $\boldsymbol{X}_{i,:}$ refers to the $i$-th row of matrix $\boldsymbol{X}$. By default, $\|\cdot\|$ denotes the Euclidean norm for a vector, or the spectral norm for a matrix.

## 2 Endogeneity and Instrumental Variable Regression

Suppose we are interested in estimating the relationship between response variable $y\in\mathbb{R}$ and predictor variable $\boldsymbol{x}\in\mathbb{R}^{p}$ with endogeneity. Given instruments $\boldsymbol{z}\in\mathbb{R}^{q}$, we consider the model

$$ $\displaystyle y=\boldsymbol{\beta}^{\top}\boldsymbol{x}+\epsilon_{1},\quad\text{and}\quad\boldsymbol{x}=\boldsymbol{\Theta}^{\top}\boldsymbol{z}+\boldsymbol{\epsilon_{2}},$ (1) $$

where $\boldsymbol{\beta}\in\mathbb{R}^{p}$, and $\boldsymbol{\Theta}\in\mathbb{R}^{q\times p}$ are the true model parameters, $\epsilon_{1}\in\mathbb{R}$ and $\boldsymbol{\epsilon}_{2}\in\mathbb{R}^{p}$ are (centered) random noise terms with variance $\sigma_{1}^{2}$ and covariance matrix $\boldsymbol{\Sigma}_{2}$, respectively. Further, $\boldsymbol{\epsilon}_{2}$ is an unobserved noise correlated with $\epsilon_{1}$, leading to the correlation between $\boldsymbol{x}$ and $\epsilon_{1}$, which introduces confounding in the model between $\boldsymbol{x}$ and $y$. Under this setting, the standard ordinary least squares (OLS) estimator is a biased and inconsistent estimator of $\boldsymbol{\beta}$ (see Wooldridge (), Chapter 9). To address this issue, instrumental variable (IV) regression is a widely used method to provide a consistent estimate for $\boldsymbol{\beta}$.

###### Definition 2.1 ( 2SLS estimator) .

IV regression is a regression model to provide consistent estimate on the causal effect $\boldsymbol{\beta}$ for the endogeneity problem ([1](#S2.E1)), by utilizing the instrument $\boldsymbol{z}$.
Given observational values $(\boldsymbol{Z},\boldsymbol{X},\boldsymbol{Y})=\{(\boldsymbol{z}_{i},\boldsymbol{x}_{i},y_{i})\}_{i=1}^{n}$, the standard approach to estimate the IV regression model is 2SLS; see, for example, Wooldridge (), Chapter 15.

- i.
*First stage*: Regress $\boldsymbol{X}$ on $\boldsymbol{Z}$ to obtain $\hat{\boldsymbol{\Theta}}$
$\displaystyle\hat{\boldsymbol{\Theta}}=(\boldsymbol{Z}^{\top}\boldsymbol{Z})^{-1}\boldsymbol{Z}^{\top}\boldsymbol{X}.$
- ii.
*Second stage*: Regress $\boldsymbol{Y}$ on $\boldsymbol{Z}\hat{\boldsymbol{\Theta}}$ to obtain:
$\displaystyle\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}=(\hat{\boldsymbol{\Theta}}^{\top}\boldsymbol{Z}^{\top}\boldsymbol{Z}\hat{\boldsymbol{\Theta}})^{-1}\hat{\boldsymbol{\Theta}}^{\top}\boldsymbol{Z}^{\top}\boldsymbol{Y}.$
(2)

We introduce the standard assumptions required to show the convergence rate of the above estimator.

###### Assumption 1 (Instrumental variable) .

A random variable $\boldsymbol{z}\in\mathbb{R}^{q}$ is a valid IV, if it satisfies the following conditions:

- i.
Fully identification: $q\geq p$ (without loss of generality, we assume data $\boldsymbol{Z},\boldsymbol{X}$ are full rank).
- ii.
Correlated to $\boldsymbol{x}$: $\textsf{Corr}(\boldsymbol{z},\boldsymbol{x})\neq\boldsymbol{0}$.
- iii.
Conditional uncorrelated to $y$: $\textsf{Corr}(\boldsymbol{z},\epsilon_{1})=0$.

In particular, condition (i) above ensures the existence of unique solution for $\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}$. We refer to Stock and Watson (, Chapter 12) for additional elaborate discussions on the above conditions. To derive non-asymptotic convergence rates, we further assume the following regularity conditions.

###### Assumption 2 (Regularity conditions) .

We assume the following conditions hold:

- i.
Bounded parameters: $\|\boldsymbol{\beta}\|\leq B_{\beta}$, $\|\boldsymbol{\Theta}\|\leq B_{\Theta}$.
- ii.
Population condition:
$\|\boldsymbol{z}\|\leq B_{z},\|\boldsymbol{x}\|\leq B_{x},|\epsilon_{1}|\leq B_{\epsilon_{1}}$, almost surely.
- iii.
Sample condition: There exists constants $\lambda_{z},\lambda_{\tilde{x}},\delta_{z\epsilon_{2}}>0$, such that the following event holds with probability at least $1-\xi$:
$\displaystyle\mathcal{A}=\left\{\lambda_{\min}\bigg{(}\frac{\boldsymbol{Z}^{\top}\boldsymbol{Z}}{n}\bigg{)}\geq\lambda_{z},\lambda_{\min}\bigg{(}\frac{\boldsymbol{X}^{\top}\boldsymbol{P}_{Z}\boldsymbol{X}}{n}\bigg{)}\geq\lambda_{\tilde{x}},\left\|\frac{\boldsymbol{Z}^{\top}\boldsymbol{\mathcal{E}}_{2}}{n}\right\|\leq\delta_{z\epsilon_{2}}\right\},$
(3)
where $\lambda_{\min}(\cdot)$ denotes the smallest eigenvalue of a matrix, $\boldsymbol{P}_{Z}:=\boldsymbol{Z}(\boldsymbol{Z}^{\top}\boldsymbol{Z})^{-1}\boldsymbol{Z}^{\top}$ denotes the projection matrix.
- iv.
Without loss of generality, we assume Assumption [1](#Thmasp1) holds under event $\mathcal{A}$, and $\mathbb{E}[\epsilon_{1}^{2}|\mathcal{A}]=\mathcal{O}(\sigma_{1}^{2})$.

The almost sure boundedness condition in (ii) is required to invoke matrix Bernstein inequalities (Tropp, ) in the analysis. We anticipate that this condition may be relaxed to sub-Gaussian or moment conditions by using more sophisticated matrix concentration results. We emphasize that condition (iii) is a mild condition in general, given some common assumptions on the population distributions. See Appendix [C.1](#A3.SS1) for related analysis and empirical justifications. Now we establish the consistency of the 2SLS estimator under random design.

###### Theorem 2.1 (Non-asymptotic error bound of 2SLS estimator) .

Given Assumptions [1](#Thmasp1) and [2](#Thmasp2), for any $\varepsilon>0$, we have:

$$ $\displaystyle\mathbb{P}\big{(}\|\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}-\boldsymbol{\beta}\|\geq\varepsilon\mid\mathcal{A}\big{)}\leq(q+1)\exp\left(-\frac{\frac{\lambda_{\tilde{x}}^{2}n^{2}\varepsilon^{2}}{2}}{(B_{\Theta}+\frac{\delta_{z\epsilon_{2}}}{\lambda_{z}})^{2}\nu(\boldsymbol{\Omega}_{z\epsilon_{1}}|\mathcal{A})+\lambda_{\tilde{x}}(B_{\Theta}+\frac{\delta_{z\epsilon_{2}}}{\lambda_{z}})B_{z}B_{\epsilon_{1}}\frac{n\varepsilon}{3}}\right),$ $$

where $\boldsymbol{\Omega}_{z\epsilon_{1}}:=\sum_{i=1}^{n}\boldsymbol{z}_{i}\epsilon_{1,i}$, and $\nu(\cdot)$ is the matrix variance statistic:

$$ $\displaystyle\nu(\boldsymbol{\Omega}):=\max\{\left\|\mathbb{E}(\boldsymbol{\Omega}\boldsymbol{\Omega}^{\top})\right\|,\left\|\mathbb{E}(\boldsymbol{\Omega}^{\top}\boldsymbol{\Omega})\right\|\}.$ $$

Furthermore, consider clipping operation $\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}):=\begin{cases}\hat{\boldsymbol{\beta}}&\text{if }\|\hat{\boldsymbol{\beta}}\|\leq B_{\beta}\\
\frac{B_{\beta}}{\|\hat{\boldsymbol{\beta}}\|}\hat{\boldsymbol{\beta}}&\text{if }\|\hat{\boldsymbol{\beta}}\|>B_{\beta}\end{cases}$, then the expectation of the estimation error is bounded by:

$$ $\displaystyle\begin{split}\mathbb{E}\left[\|\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}})-\boldsymbol{\beta}\|^{2}\right]&\leq\mathcal{O}\left(\frac{C_{z,x}^{2}q\sigma_{1}^{2}}{n}+\xi B_{\beta}^{2}\right),\end{split}$ (4) $$

where $C_{z,x}:=\frac{(B_{\Theta}+\frac{\delta_{z\epsilon_{2}}}{\lambda_{z}})B_{z}}{\lambda_{\tilde{x}}}$ is a constant depending on the population distributions of $\boldsymbol{z}$ and $\boldsymbol{x}$.

We note that although the consistency of the 2SLS estimator is a standard result in econometrics, most existing works focus on the asymptotic properties of the estimator. Theorem [2.1](#S2.Thmthm1) provides the first non-asymptotic bound for estimation error $\|\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}-\boldsymbol{\beta}\|^{2}$, under random design. The detailed proof is provided in Appendix [A.1](#A1.SS1).

###### Definition 2.1 ( 2SLS estimator) .

###### Assumption 1 (Instrumental variable) .

###### Assumption 2 (Regularity conditions) .

###### Theorem 2.1 (Non-asymptotic error bound of 2SLS estimator) .

## 3 Transformers Handle Endogeniety

### 3.1 Transformer Architecture

Denote the input matrix as $\boldsymbol{H}=[\boldsymbol{h}_{1},\ldots,\boldsymbol{h}_{n}]\in\mathbb{R}^{D\times n}$, where each column corresponds to one sample vector.

###### Definition 3.1 (Attention layer) .

A self-attention layer with $M$ heads is denoted as $\textsf{ATTN}_{\boldsymbol{\theta}}(\cdot)$, with parameters $\boldsymbol{\theta}=\{(\boldsymbol{Q}_{m},\boldsymbol{K}_{m},\boldsymbol{V}_{m})\}_{m\in[M]}\subseteq\mathbb{R}^{D\times D}$. Given input $\boldsymbol{H},$

$$ $\displaystyle\tilde{\boldsymbol{H}}=\textsf{ATTN}_{\boldsymbol{\theta}}(\boldsymbol{H}):=\boldsymbol{H}+\frac{1}{n}\sum_{m=1}^{M}(\boldsymbol{V}_{m}\boldsymbol{H})\times\sigma((\boldsymbol{Q}_{m}\boldsymbol{H})^{\top}(\boldsymbol{K}_{m}\boldsymbol{H}))\in\mathbb{R}^{D\times n},$ (5) $$

or element-wise:

$$ $\displaystyle\tilde{\boldsymbol{h}}_{i}=[\textsf{ATTN}_{\boldsymbol{\theta}}(\boldsymbol{H})]_{i}:=\boldsymbol{h}_{i}+\sum_{m=1}^{M}\frac{1}{n}\sum_{j=1}^{n}\sigma(\langle\boldsymbol{Q}_{m}\boldsymbol{h}_{i},\boldsymbol{K}_{m}\boldsymbol{h}_{j}\rangle)\cdot\boldsymbol{V}_{m}\boldsymbol{h}_{j}\in\mathbb{R}^{D},$ (6) $$

where $\sigma(\cdot)$ is the ReLU function.

###### Definition 3.2 (MLP layer) .

An MLP layer is denoted as $\textsf{MLP}_{\boldsymbol{\theta}}(\cdot)$, with parameters $\boldsymbol{\theta}=(\boldsymbol{W}_{1},\boldsymbol{W}_{2})\in\mathbb{R}^{D^{\prime}\times D\times D\times D^{\prime}}$. Given input $\boldsymbol{H},$

$$ $\displaystyle\tilde{\boldsymbol{H}}=\textsf{MLP}_{\boldsymbol{\theta}}(\boldsymbol{H}):=\boldsymbol{H}+\boldsymbol{W}_{2}\sigma(\boldsymbol{W}_{1}\boldsymbol{H}),$ $$

or element-wise:

$$ $\displaystyle\tilde{\boldsymbol{h}}_{i}=[\textsf{MLP}_{\boldsymbol{\theta}}(\boldsymbol{H})]_{i}:=\boldsymbol{h}_{i}+\boldsymbol{W}_{2}\sigma(\boldsymbol{W}_{1}\boldsymbol{h}_{i}).$ $$

###### Definition 3.3 (Transformer) .

An L-layer transformer is denoted as $\textsf{TF}_{\boldsymbol{\theta}}(\cdot)$, with parameters $\boldsymbol{\theta}=(\boldsymbol{\theta}_{\textsf{ATTN}}^{(1:L)},\boldsymbol{\theta}_{\textsf{MLP}}^{(1:L)})$. Given input $\boldsymbol{H}=\boldsymbol{H}^{(0)},$

$$ $\displaystyle\boldsymbol{{H}}^{(l)}=\textsf{MLP}_{\boldsymbol{\theta}_{\textsf{MLP}}^{(l)}}(\textsf{ATTN}_{\boldsymbol{\theta}_{\textsf{ATTN}}^{(l)}}(\boldsymbol{H}^{(l-1)})),\quad l=1,\ldots,L.$ $$

The output of this transformer is the final layer output: $\tilde{\boldsymbol{H}}:=\boldsymbol{H}^{(L)}=\textsf{TF}_{\boldsymbol{\theta}}(\boldsymbol{H}^{(0)})$.

###### Definition 3.4 (Looped transformer) .

An $\bar{L}$-looped transformer is a special transformer architecture, denoted as $\textsf{LTF}_{\bar{\boldsymbol{\theta}},\bar{L}}(\cdot)$, with parameters $\bar{\boldsymbol{\theta}}=(\bar{\boldsymbol{\theta}}_{\textsf{ATTN}}^{(1:L_{0})},\bar{\boldsymbol{\theta}}_{\textsf{MLP}}^{(1:L_{0})})$. Given input $\boldsymbol{H}=\boldsymbol{H}^{(0)},$

$$ $\displaystyle\boldsymbol{{H}}^{(l)}=\textsf{TF}_{\bar{\boldsymbol{\theta}}}(\boldsymbol{H}^{(l-1)}),\quad l=1,\ldots,\bar{L}.$ $$

The output of this looped transformer is the final loop output: $\tilde{\boldsymbol{H}}:=\boldsymbol{H}^{(\bar{L})}=\textsf{LTF}_{\bar{\boldsymbol{\theta}},\bar{L}}(\boldsymbol{H}^{(0)})$.

Previous works (e.g., Bai et al. (), Zhang et al. ()) have shown that transformers can perform 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 on single-level optimization algorithms, which is insufficient to demonstrate that transformers can efficiently learn more complex algorithms like 2SLS (Definition [2.1](#S2.Thmdefn1)). Second, most ICL-related research focuses on the predictive performance of transformers, paying little attention to their ability to provide accurate coefficient estimates. We extend the current ICL framework by showing that transformers can implement a bi-level GD procedure (see Section [3.2](#S3.SS2)) with looped transformer architecture (Definition [3.4](#S3.Thmdefn4)), allowing them to efficiently emulate 2SLS and provide coefficient estimates that are at least as accurate as 2SLS in the presence of endogeneity (as in ([1](#S2.E1))).

###### Definition 3.1 (Attention layer) .

###### Definition 3.2 (MLP layer) .

###### Definition 3.3 (Transformer) .

###### Definition 3.4 (Looped transformer) .

### 3.2 Gradient descent based IV regression

We first introduce a gradient-based bi-level optimization procedure to obtain the 2SLS estimator in ([2](#S2.E2)). Given the dataset $(\boldsymbol{Z},\boldsymbol{X},\boldsymbol{Y})=\{(\boldsymbol{z}_{i},\boldsymbol{x}_{i},y_{i})\}_{i=1}^{n}$, the objective funtion of IV regression can be formulated as the following bi-level optimization problem:

$$ $\displaystyle\begin{aligned} \min_{\boldsymbol{\beta}}\quad\mathcal{L}(\boldsymbol{\beta})=\frac{1}{n}\sum_{i=1}^{n}(y_{i}-\boldsymbol{z}_{i}^{\top}\hat{\boldsymbol{\Theta}}\boldsymbol{\beta})^{2},\quad\text{where }\quad\hat{\boldsymbol{\Theta}}:=\underset{\boldsymbol{\Theta}}{\arg\min}\quad\frac{1}{n}\sum_{j=1}^{n}(\boldsymbol{x}_{j}-\boldsymbol{z}_{j}^{\top}\boldsymbol{\Theta})^{2}.\end{aligned}$ (7) $$

Consider the following gradient updates with learning rates $\alpha,\eta$:

$$ $\displaystyle\boldsymbol{\Theta}^{(t+1)}$ $\displaystyle=\boldsymbol{\Theta}^{(t)}-\eta\boldsymbol{Z}^{\top}(\boldsymbol{Z\Theta}^{(t)}-\boldsymbol{X}),$ (8a) $\displaystyle\boldsymbol{\beta}^{(t+1)}$ $\displaystyle=\boldsymbol{\beta}^{(t)}-\alpha\boldsymbol{\Theta}^{(t)\top}\boldsymbol{Z}^{\top}(\boldsymbol{Z\Theta}^{(t)}\boldsymbol{\beta}^{(t)}-\boldsymbol{Y}).$ (8b) $$

Note that the GD-2SLS updates in ([8](#S3.E8)) are designed to solve ([7](#S3.E7)). We now show that regardless the convergence of $\boldsymbol{\Theta}^{(t)}$, the GD estimator $\boldsymbol{\beta}^{(t)}$ will always converge to the 2SLS estimator in ([2](#S2.E2)) with exponential rate.

###### Theorem 3.1 (Implementing 2SLS with gradient-based method) .

Given training data $(\boldsymbol{Z},\boldsymbol{X},\boldsymbol{Y})=\{(\boldsymbol{z}_{i},\boldsymbol{x}_{i},y_{i})\}_{i=1}^{n}$. Suppose the learning rates $\alpha,\eta$ satisfy the following conditions:

$$ $\displaystyle 0<\alpha<\frac{2}{\sigma_{\max}^{2}(\boldsymbol{Z}\hat{\boldsymbol{\Theta}})}\quad\text{and}\quad 0<\eta<\frac{2}{\sigma_{\max}^{2}(\boldsymbol{Z})},$ $$

where $\sigma_{\max}(\cdot)$ denotes the largest singular value of a matrix. Then, the GD updates in ([8](#S3.E8)) converge to the 2SLS estimator at an exponential rate:

$$ $\displaystyle\|\boldsymbol{\beta}^{(t)}-\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}\|\leq\mathcal{O}\left(\Lambda^{t}\right),$ $$

where, with $\rho(\cdot)$ denoting the spectral radius of the matrix,

$$ $\displaystyle\Lambda:=\max\{\gamma(\alpha),\kappa(\eta)\},\quad\quad\gamma(\alpha):=\rho(\boldsymbol{I}-\alpha\hat{\boldsymbol{\Theta}}^{\top}\boldsymbol{Z}^{\top}\boldsymbol{Z}\hat{\boldsymbol{\Theta}}),\quad\kappa(\eta):=\rho(\boldsymbol{I}-\eta\boldsymbol{Z}^{\top}\boldsymbol{Z}).$ (9) $$

To the best of our knowledge, Theorem [3.1](#S3.Thmthm1) provides the first theoretical result demonstrating that 2SLS can be efficiently implemented using a gradient-based method, with an exponential convergence rate. We provide the proof in Appendix [B.1](#A2.SS1) and present simulation results in Appendix [C.2](#A3.SS2) to examine the convergence behavior of the optimization process.

###### Theorem 3.1 (Implementing 2SLS with gradient-based method) .

### 3.3 Transformers Can Efficiently Implement GD- 2SLS

The looped transformer architecture (Definition [3.4](#S3.Thmdefn4)), as proposed by Giannou et al. (), introduces an efficient approach to learn iterative algorithms by cascading the same transformer block for multiple times. With the GD updates in ([8](#S3.E8)), we will show that there exists a looped transformer architecuture that can efficiently learn the 2SLS estimator. We emphasize here that although we 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 $\hat{\boldsymbol{\Theta}}$. 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.

###### Theorem 3.2 (Implement a step of GD- 2SLS with a transformer block) .

Suppose the embedded input matrix takes the form:

$$ $\displaystyle\boldsymbol{H}^{(2l)}=\begin{bmatrix}\boldsymbol{z}_{1}&\cdots&\boldsymbol{z}_{n}&\boldsymbol{z}_{n+1}\\ \boldsymbol{x}_{1}&\cdots&\boldsymbol{x}_{n}&\boldsymbol{x}_{n+1}\\ y_{1}&\cdots&y_{n}&t\\ \boldsymbol{\Theta}_{:,1}^{(l)}&\cdots&\boldsymbol{\Theta}_{:,1}^{(l)}&\boldsymbol{\Theta}_{:,1}^{(l)}\\ \vdots&\vdots&\vdots&\vdots\\ \boldsymbol{\Theta}_{:,p}^{(l)}&\cdots&\boldsymbol{\Theta}_{:,p}^{(l)}&\boldsymbol{\Theta}_{:,p}^{(l)}\\ \boldsymbol{\beta}^{(l)}&\cdots&\boldsymbol{\beta}^{(l)}&\boldsymbol{\beta}^{(l)}\\ \hat{\boldsymbol{x}}_{1}^{(l)}&\cdots&\hat{\boldsymbol{x}}_{n}^{(l)}&\hat{\boldsymbol{x}}_{n+1}^{(l)}\\ 1&\cdots&1&1\\ 1&\cdots&1&0\end{bmatrix}\in\mathbb{R}^{D\times(n+1)}.$ (10) $$

Given $\boldsymbol{H}^{(2l)}$, there exists a double-layer attention-only transformer block with parameters $\boldsymbol{\theta}=\boldsymbol{\theta}_{\textsf{ATTN}}^{(2l+1:2l+2)}=\{(\boldsymbol{Q}_{m}^{(2l+1:2l+2)},\boldsymbol{K}_{m}^{(2l+1:2l+2)},\boldsymbol{V}_{m}^{(2l+1:2l+2)})\}_{m\in[M^{(2l+1:2l+2)}]}\subset\mathbb{R}^{D\times D}$, where the number of heads $M^{(2l+1)}=2p$, $M^{(2l+2)}=2(p+1)$ and embedding dimension $D=qp+3p+q+3$, that implements a 2SLS gradient update in ([8](#S3.E8)) with any given learning rates $\alpha,\eta$:

$$ $\displaystyle\boldsymbol{H}^{2(l+1)}=\textsf{TF}_{\boldsymbol{\theta}_{\textsf{ATTN}}^{(2l+1:2l+2)}}(\boldsymbol{H}^{(2l)})=\begin{bmatrix}\boldsymbol{z}_{1}&\cdots&\boldsymbol{z}_{n}&\boldsymbol{z}_{n+1}\\ \boldsymbol{x}_{1}&\cdots&\boldsymbol{x}_{n}&\boldsymbol{x}_{n+1}\\ y_{1}&\cdots&y_{n}&0\\ \boldsymbol{\Theta}_{:,1}^{(l+1)}&\cdots&\boldsymbol{\Theta}_{:,1}^{(l+1)}&\boldsymbol{\Theta}_{:,1}^{(l+1)}\\ \vdots&\vdots&\vdots&\vdots\\ \boldsymbol{\Theta}_{:,p}^{(l+1)}&\cdots&\boldsymbol{\Theta}_{:,p}^{(l+1)}&\boldsymbol{\Theta}_{:,p}^{(l+1)}\\ \boldsymbol{\beta}^{(l+1)}&\cdots&\boldsymbol{\beta}^{(l+1)}&\boldsymbol{\beta}^{(l+1)}\\ \hat{\boldsymbol{x}}_{1}^{(l+1)}&\cdots&\hat{\boldsymbol{x}}_{n}^{(l+1)}&\hat{\boldsymbol{x}}_{n+1}^{(l+1)}\\ 1&\cdots&1&1\\ 1&\cdots&1&0\end{bmatrix}\in\mathbb{R}^{D\times(n+1)}.$ $$

Our existence proof specifies an attention structure such that one layer updates only the first-stage estimate $\hat{\boldsymbol{x}}_{i}^{(l)}$ for all samples, followed by another layer to update the parameters $\boldsymbol{\Theta}^{(l)}$ and $\boldsymbol{\beta}^{(l)}$. Furthermore, as noted in the proof of Theorem [3.2](#S3.Thmthm2) (ref. Appendix [B.2](#A2.SS2)), regardless of the initial values of $\boldsymbol{\Theta}^{(l)},\boldsymbol{\beta}^{(l)}$ and $\hat{\boldsymbol{x}}^{(l)}$, the structures of the transformer blocks remain the same. This allows us to exploit the looped transformer architecture to significantly reduce the number of parameters and improve learning efficiency (Yang et al., ).

By cascading the transformer block $\bar{L}$ times, with Theorem [3.1](#S3.Thmthm1), one can show that transformers are able to mimic the 2SLS estimator with exponential convergence rate, as described in the following corollary.

###### Corollary 3.1 (Implementing GD- 2SLS with looped transformer) .

For any $0<\varepsilon<1$, given learning rates $\alpha,\eta,$ and $\Lambda\in(0,1)$, as defined in ([9](#S3.E9)), there exists a transformer formulated as $\textsf{TF}_{\boldsymbol{\theta}}(\cdot):=\textsf{TF}_{\boldsymbol{\theta}^{\prime}}(\textsf{LTF}_{\bar{\boldsymbol{\theta}},\bar{L}}(\cdot))$, which consists of an $\bar{L}$-looped transformer $\textsf{LTF}_{\bar{\boldsymbol{\theta}},\bar{L}}$ with $\bar{\boldsymbol{\theta}}=\bar{\boldsymbol{\theta}}_{\textsf{ATTN}}^{(1:2)}=\{(\bar{\boldsymbol{Q}}_{m}^{(1:2)},\bar{\boldsymbol{K}}_{m}^{(1:2)},\bar{\boldsymbol{V}}_{m}^{(1:2)})\}_{m\in[\bar{M}^{(1:2)}]}\subset\mathbb{R}^{D\times D}$, $\bar{L}=\lceil\mathcal{O}(\log_{\Lambda}(\varepsilon))\rceil$, and a final attention layer(^1^11This layer updates the prediction $\hat{y}_{n+1}:=\boldsymbol{\beta}^{(\bar{L})\top}\boldsymbol{x}_{n+1}$, which can be constructed with 2 attention heads using the same architecture as Bai et al. (, Theorem 13)) $\boldsymbol{\theta}^{\prime}=\boldsymbol{\theta}_{\textsf{ATTN}}^{\prime}=\{(\boldsymbol{Q}_{m}^{\prime},\boldsymbol{K}_{m}^{\prime},\boldsymbol{V}_{m}^{\prime})\}_{m\in[M^{\prime}]}\subset\mathbb{R}^{D\times D}$, where $\bar{M}^{(1)}=2p,\bar{M}^{(2)}=2(p+1)$, $M^{\prime}=2$, such that given embedded input $\boldsymbol{H}^{(0)}$ taking the format in ([10](#S3.E10)), the model output satisfies:

$$ $\displaystyle|\textsf{read}_{y}(\textsf{TF}_{\boldsymbol{\theta}}(\boldsymbol{H}^{(0)}))-\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}^{\top}\boldsymbol{x}_{n+1}|\leq B_{x}\varepsilon,$ $$

where $\textsf{read}_{y}(\cdot)$ is a function that reads the prediction $\hat{y}_{n+1}$ from the output of the transformer.

We emphasize here that our construction differs from the implementation of Bai et al. (, Theorem 4) for OLS in the following aspects:

- i.
We apply the square loss as defined in ([7](#S3.E7)) to learn the 2SLS estimator, which simplifies the loss function’s sum-of-ReLU representation.
- ii.
The dimension of the input embedding is $D=qp+3p+q+3$, where the extra dimensions store the vectorized parameters $\boldsymbol{\Theta}^{(l)},\boldsymbol{\beta}^{(l)}$, and the first stage estimate $\hat{\boldsymbol{x}}^{(l)}$.
- iii.
We use a two-layer attention-only transformer block $\bar{\boldsymbol{\theta}}$ to implement a 2SLS GD update ([8](#S3.E8)), with the first layer to update the current first-stage estimate $\hat{\boldsymbol{x}}^{(l)}$, and the second layer to update the parameters $\boldsymbol{\Theta}^{(l)}$ and $\boldsymbol{\beta}^{(l)}$.
- iv.
For each transformer block, in the first layer, we equip $2$ heads to update each dimension of $\hat{\boldsymbol{x}}_{i}^{(l)}\in\mathbb{R}^{p}$ for all samples. In the second layer, we equip 2 heads to update each column of $\boldsymbol{\Theta}^{(l)}\in\mathbb{R}^{q\times p}$ and $\boldsymbol{\beta}^{(l)}\in\mathbb{R}^{p}$.

###### Theorem 3.2 (Implement a step of GD- 2SLS with a transformer block) .

###### Corollary 3.1 (Implementing GD- 2SLS with looped transformer) .

### 3.4 Pretraining and Excess Loss Bound

With slightly abuse of notations, we denote the (formulated) training prompt as:

$$ $\displaystyle\boldsymbol{H}_{k}=\begin{bmatrix}\boldsymbol{z}_{1,k}&\cdots&\boldsymbol{z}_{n,k}&\boldsymbol{z}_{n+1,k}\\ \boldsymbol{x}_{1,k}&\cdots&\boldsymbol{x}_{n,k}&\boldsymbol{x}_{n+1,k}\\ y_{1,k}&\cdots&y_{n,k}&0\end{bmatrix}\in\mathbb{R}^{(p+q+1)\times(n+1)},\quad k=1,\ldots,N.$ $$

Note that we denote each training prompt by the subscript $k=1,\ldots,N$, where $N$ is the total number of prompts. Each training prompt consists of $n$ labeled training samples $\{(\boldsymbol{z}_{i},\boldsymbol{x}_{i},y_{i})\}_{i=1}^{n}$, and one unlabeled query sample $(\boldsymbol{z}_{n+1},\boldsymbol{x}_{n+1})$. Our goal is to predict $y_{n+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 $\boldsymbol{\mathcal{P}}$, described by Algorithm [1](#alg1).

Figure: Algorithm 1 In-Context Distribution $\boldsymbol{\mathcal{P}}$

In Algorithm [1](#alg1), $\boldsymbol{u}\in\mathbb{R}^{p}$ is the source of endogenous error, $\boldsymbol{w}\in\mathbb{R}^{p},{\epsilon}\in\mathbb{R}$ are the exogenous errors. Note that we have $\epsilon_{1,i}=\boldsymbol{\phi}^{\top}\boldsymbol{u}_{i}+\epsilon_{i}$ and $\boldsymbol{\epsilon}_{2,i}=\boldsymbol{\Phi}^{\top}\boldsymbol{u}_{i}+\boldsymbol{\omega}_{i}$, corresponding to the notations in ([1](#S2.E1)). $\boldsymbol{\Theta}\in\mathbb{R}^{q\times p},\boldsymbol{\beta}\in\mathbb{R}^{p},\boldsymbol{\Phi}\in\mathbb{R}^{p\times p},\boldsymbol{\phi}\in\mathbb{R}^{p},\boldsymbol{\Sigma}_{z}\in\mathbb{R}^{q\times q},\boldsymbol{\Sigma}_{u}\in\mathbb{R}^{p\times p},\boldsymbol{\Sigma}_{\omega}\in\mathbb{R}^{p\times p},\sigma_{\epsilon}\in\mathbb{R}$ are task-specific parameters following meta distribution $\boldsymbol{\pi}$. $\textsf{clip}_{B}(\cdot)$ is a clipping operator to bound the norm of input within radius $B$. We say that the in-context samples $\{(\boldsymbol{z}_{i},\boldsymbol{x}_{i},y_{i})\}_{i=1}^{n+1}$ are drawn from the in-context distribution $\boldsymbol{\mathcal{P}}$, and $\boldsymbol{\mathcal{P}}\sim\boldsymbol{\pi}$ if the task parameters $(\boldsymbol{\Theta},\boldsymbol{\beta},\boldsymbol{\Phi},\boldsymbol{\phi},\boldsymbol{\Sigma}_{z},\boldsymbol{\Sigma}_{u},\boldsymbol{\Sigma}_{\omega},\sigma_{\epsilon})$ are sampled from $\boldsymbol{\pi}$.
One can check that Assumption [1](#Thmasp1) and Assumption [2](#Thmasp2)(ii)(iv) are directly satisfied with the data generated from the in-context distribution $\boldsymbol{\mathcal{P}}$.

Following the theoretical framework of (Bai et al., ), we define the population ICL loss(^2^22All the clipping operations are only for analytical purpose. In practice, the behavior of the trained transformer is consistent even without the clipping bounds.):

$$ $\displaystyle L_{\textsf{ICL}}(\boldsymbol{\theta})=\mathbb{E}_{\pi}\mathbb{E}_{\mathcal{P}}[y_{n+1}-\textsf{clip}_{B_{y}}(\textsf{read}_{y}(\textsf{TF}_{\boldsymbol{\theta}}^{R}(\boldsymbol{H}^{(0)})))]^{2},$ (11) $$

where $\boldsymbol{H}^{(0)}$ is the embedded input as defined in ([10](#S3.E10)), $\textsf{TF}_{\boldsymbol{\theta}}^{R}$ is the transformer model with parameter $\boldsymbol{\theta}$ and clipping operation $\textsf{clip}_{R}(\cdot)$ applied to each layer output. For simplicity, we denote $\widetilde{\textsf{TF}}_{\theta}(\boldsymbol{H}):=\textsf{clip}_{B_{y}}(\textsf{read}_{y}(\textsf{TF}_{\boldsymbol{\theta}}^{R}(\boldsymbol{H}^{(0)})))$.

The transformer is trained to minimize the in-context loss in ([11](#S3.E11)) with the following empirical loss:

$$ $\displaystyle\hat{L}_{\textsf{ICL}}(\boldsymbol{\theta})=\frac{1}{N}\sum_{k=1}^{N}(y_{n+1,k}-\widetilde{\textsf{TF}}_{\boldsymbol{\theta}}(\boldsymbol{H}_{k}))^{2}.$ (12) $$

We consider the following constrained optimization problem:

$$ $\displaystyle\begin{split}\hat{\boldsymbol{\theta}}&:=\underset{\boldsymbol{\theta}\in\boldsymbol{\vartheta}_{L,M,D^{\prime},B_{\theta}}}{\arg\min}\hat{L}_{\textsf{ICL}}(\boldsymbol{\theta}),\\ \boldsymbol{\vartheta}_{L,M,D^{\prime},B_{\theta}}:=\{\boldsymbol{\theta}=(\boldsymbol{\theta}_{\textsf{Attn}}^{(1:L)},\boldsymbol{\theta}_{\textsf{MLP}}^{(1:L)}&):\underset{l\in[L]}{\max}\>M^{(l)}\leq M,\>\underset{l\in[L]}{\max}\>D^{(l)}\leq D^{\prime},\>\left|\!\left|\!\left|\boldsymbol{\theta}\right|\!\right|\!\right|\leq B_{\theta}\},\end{split}$ (13) $$

where $\left|\!\left|\!\left|\boldsymbol{\theta}\right|\!\right|\!\right|:=\underset{l\in[L]}{\max}\{\underset{m\in[M]}{\max}\{\|\boldsymbol{Q}_{m}^{(l)}\|,\|\boldsymbol{K}_{m}^{(l)}\|\}+\sum_{m=1}^{M}\|\boldsymbol{V}_{m}^{(l)}\|+\|\boldsymbol{W}_{1}^{(l)}\|+\|\boldsymbol{W}_{2}^{(l)}\|\}$.

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)(iii) in Assumption [2](#Thmasp2) holds and the meta distribution $\boldsymbol{\pi}$ satisfies the following conditions:

$$ $\displaystyle\mathbb{E}_{\pi}\left[\boldsymbol{\phi}^{\top}\boldsymbol{\Sigma}_{u}\boldsymbol{\phi}+\sigma_{\epsilon}^{2}\right]\leq$ $\displaystyle\tilde{\sigma}^{2}\text{ and }\mathbb{E}_{\pi}\left[\sigma_{\epsilon}^{2}\right]\leq\tilde{\sigma}_{\epsilon}^{2}.$ (14) $$

Let the in-context distribution $\boldsymbol{\mathcal{P}}\sim\boldsymbol{\pi}$ such that the samples $(\boldsymbol{z}_{i},\boldsymbol{x}_{i},y_{i})_{i=1}^{n+1}$ are drawn independently from $\boldsymbol{\mathcal{P}}$ (ref. Algorithm [1](#alg1)). With training prompts $\boldsymbol{H}_{k},k=1,\ldots,N$, under ICL loss in ([11](#S3.E11)), the trained transformer in ([13](#S3.E13)) with $L=2\bar{L}+1,M=2(p+1),D=qp+3p+q+3,D^{\prime}=0$ (attention-only) achieves the following excess loss with probability at least $1-\zeta$:

$$ $\displaystyle L_{\textsf{ICL}}(\hat{\boldsymbol{\theta}})$ $\displaystyle-\mathbb{E}_{\pi}\mathbb{E}_{\mathcal{P}}\left[(y_{n+1}-\langle\boldsymbol{\beta},\boldsymbol{x}_{n+1}\rangle)^{2}\right]\leq\mathcal{O}\bigg{(}(\Lambda^{\star})^{\bar{L}}\Big{(}B_{x}^{2}\sqrt{\frac{C_{z,x}^{2}q\tilde{\sigma}^{2}}{n}+\xi B_{\beta}^{2}}+B_{x}\tilde{\sigma}_{\epsilon}\Big{)}$ $\displaystyle+B_{x}^{2}\Big{(}\frac{C_{z,x}^{2}q\tilde{\sigma}^{2}}{n}+\xi B_{\beta}^{2}+\mu_{\Lambda,2}^{\star}\Big{)}+B_{y}^{2}\sqrt{\frac{L^{2}MD^{2}\log(2+\max\{B_{\theta},R,B_{y}\})+\log(1/\zeta)}{N}}\bigg{)},$ $$

where $\Lambda^{\star}:=\underset{\alpha,\eta}{\min}\;\mathbb{E}_{\pi}\mathbb{E}_{\mathcal{P}}[\Lambda|\boldsymbol{H},\alpha,\eta]<1$, and $\mu_{\Lambda,2}^{\star}:=\mathbb{E}_{\pi}\mathbb{E}_{\mathcal{P}}[\Lambda^{2\bar{L}}|\boldsymbol{H},\alpha^{\star},\eta^{\star}]$ is close to 0 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 $\boldsymbol{\pi}$, the excess loss is dominated by two factors: (i) number of attention layers, and (ii) number of in-context samples. The proof of Theorem [3.3](#S3.Thmthm3) is provided in Appendix [B.3](#A2.SS3).

###### Theorem 3.3 (Excess loss bound for in-context pretrained transformer) .

### 3.5 Extracting the regression coefficients

The primary goal of IV regression is to estimate the causal effect, i.e. the coefficient $\boldsymbol{\beta}$ under the stated endogeneity in ([1](#S2.E1)). For 2SLS, the estimated causal effect is given by the coefficients of the endogenous variable in the second stage regression ([2](#S2.E2)). For transformer models, we propose a straightforward method to extract these estimated coefficients by differentiating the output with respect to each dimension of the endogenous variable. The specific approach is summarized in Algorithm [2](#alg2).

Figure: Algorithm 2 Extracting the regression coefficients

We observe that the choice of $\Delta$ within a reasonable range does not significantly affect the estimation of the coefficients. In practice, usually a slightly larger $\Delta$ (for example $\Delta=5$) can lead to a more stable estimation, which is possibly due to the elemination of rounding errors during computation.

## 4 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 one query sample), the dimension of endogenous variable $p=5$, and the dimension of instrument $q=10$. The training prompts are generated using Algorithm [1](#alg1), with task parameters $\boldsymbol{\Theta},\boldsymbol{\beta},\boldsymbol{\Phi},\boldsymbol{\phi}$ sampled from standard Gaussian distribution, and the covariance matrices $\boldsymbol{\Sigma}_{z},\boldsymbol{\Sigma}_{u},\boldsymbol{\Sigma}_{\omega}$ set to be identity matrices. The noise level $\sigma_{\epsilon}$ is set to 1. We ignore all the clipping bounds in the experiment ($B_{\beta},B_{\Theta},B_{z},B_{x},B_{y},B_{\theta},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 $(L_{0}=2)$, following the theoretical guidelines in Theorem [3.2](#S3.Thmthm2). We employ the looped transformer architecture, consisting of 10 identical cascading transformer blocks. The transformer model is trained under the ICL loss ([12](#S3.E12)) 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 $\{(\boldsymbol{z}_{i},\boldsymbol{x}_{i},y_{i})\}_{i=1}^{n}$ 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 $y_{n+1}$, and the mean squared error (MSE) on the coefficient $\boldsymbol{\beta}$.

### 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 are shown in Figure [1(a)](#S4.F1.sf1). 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 $\boldsymbol{z}_{i}$ and $\boldsymbol{x}_{i}$ following different correlation levels. Specifically, in Algorithm [1](#alg1), we adjust the IV strength by multiplying $\boldsymbol{\Theta}$ by a factor $r\in(0,2)$ when generating test prompts. The results are shown in Figure [1(b)](#S4.F1.sf2).

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 a more advanced mechanism for handling endogeneity 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.

Figure: (a)
Refer to caption: /html/2410.01265/assets/Figures/MSE_vs_N_zoom.png

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. $\boldsymbol{x}_{i,k}=\boldsymbol{\Theta}_{k}^{\top}\boldsymbol{z}_{i,k}^{2}+\textsf{error}_{i,k}$ in Algorithm [1](#alg1), and (b) the dimension of IV is not sufficient to identify the endogenous variable(^3^33For 2SLS estimate, the actual computation uses pseudoinverse to handle rank deficiency.), where we set $q=3$ (by zeroing out the remaining demensions of $\boldsymbol{z}$ in test prompts) and $p=5$.

We evaluate the same trained transformer model as before, with results presented in Figure [2(a)](#S4.F2.sf1) and Figure [2(b)](#S4.F2.sf2), respectively. Once again, the trained transformer model consistently outperforms both 2SLS and OLS estimators in handling these non-standard endogeneity tasks. All these results suggest 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. We suspect that, in our pretraining scheme, although the 2SLS estimator already achieves small excess loss, a 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 at least as well as 2SLS and appears to be a promising tool for handling endogeneity in difficult scenarios.

Figure: (a)
Refer to caption: /html/2410.01265/assets/Figures/MSE_vs_N_quadratic_zoom.png

## 5 Conclusion

This paper presents a novel perspective on the transformer model in its ability to handle endogeneity in in-context linear regression. We have theoretically shown that the transformer model exists an intrinsic structure that enables it to learn the 2SLS algorithm through an efficient GD procedure. We have further provided a 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 in handling standard endogeneity tasks. Furthermore, our investigation illustrates that it 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-trained transformer model is a promising tool for making reliable in-context predictions and coefficient estimates under endogeneity, especially when dealing with non-standard IV problems.

## Appendix A Proofs For Section 2

### A.1 Proof of Theorem 2.1

We first introduce the following lemmas that are used in the proof of Theorem [2.1](#S2.Thmthm1).

###### Lemma A.1 (Bernstein Inequality, from Theorem 6.1.1 in Tropp [ 2015 ] ) .

Let $\boldsymbol{S}_{1},\ldots,\boldsymbol{S}_{n}$ be independent, centered random matrices with common dimension $d_{1}\times d_{2}$, and assume that each one is almost surely bounded:

$$ $\displaystyle\mathbb{E}[\boldsymbol{S}_{i}]=\boldsymbol{0},\mathbb{P}(\left\|\boldsymbol{S}_{i}\right\|\leq b)=1,\quad\forall i=1,\ldots,n.$ $$

With the sum:

$$ $\displaystyle\boldsymbol{\Omega}=\sum_{i=1}^{n}\boldsymbol{S}_{i},$ $$

and the matrix variance statistic of the sum:

$$ $\displaystyle\nu(\boldsymbol{\Omega}):=\max\{\left\|\mathbb{E}(\boldsymbol{\Omega}\boldsymbol{\Omega}^{\top})\right\|,\left\|\mathbb{E}(\boldsymbol{\Omega}^{\top}\boldsymbol{\Omega})\right\|\},$ $$

then the following inequality holds:

$$ $\displaystyle\mathbb{P}\{\left\|\boldsymbol{\Omega}\right\|\geq\varepsilon\}\leq(d_{1}+d_{2})\cdot\exp\left(\frac{-\varepsilon^{2}/2}{\nu(\boldsymbol{\Omega})+b\varepsilon/3}\right)\text{ for any }\varepsilon\geq 0.$ $$

###### Lemma A.2 (Inverse Convergence, adapted from Lemma 2.1 in Jin et al. [ 2024 ] ) .

Suppose we have a random invertible matrix $\boldsymbol{\Omega}$ and invertible matrix sequence $\{\hat{\boldsymbol{\Omega}}^{(n)}\}$ such that $\hat{\boldsymbol{\Omega}}^{(n)}\overset{\textsf{p}}{\rightarrow}\boldsymbol{\Omega}$. If there exists a constant $\tilde{\lambda}>0$ such that $\sigma_{\min}(\hat{\boldsymbol{\Omega}})\geq\tilde{\lambda}$ almost surely, then it holds that:

$$ $\displaystyle(\hat{\boldsymbol{\Omega}}^{(n)})^{-1}\overset{\textsf{p}}{\rightarrow}\boldsymbol{\Omega}^{-1}.$ $$

Further, given convergence rate

$$ $\displaystyle\mathbb{P}\left\{\left\|\hat{\boldsymbol{\Omega}}^{(n)}-\boldsymbol{\Omega}\right\|\geq\varepsilon\right\}\leq\xi(n,\varepsilon),$ $$

then:

$$ $\displaystyle\mathbb{P}\left\{\left\|(\hat{\boldsymbol{\Omega}}^{(n)})^{-1}-\boldsymbol{\Omega}^{-1}\right\|\geq\varepsilon\right\}\leq\xi(n,\tilde{\lambda}^{2}\varepsilon).$ $$

###### Proof.

We have the following decomposition:

$$ $\displaystyle(\hat{\boldsymbol{\Omega}}^{(n)})^{-1}-\boldsymbol{\Omega}^{-1}=(\hat{\boldsymbol{\Omega}}^{(n)})^{-1}(\boldsymbol{\Omega}-\hat{\boldsymbol{\Omega}}^{(n)})\boldsymbol{\Omega}^{-1}.$ $$

It follows that:

$$ $\displaystyle\left\|(\hat{\boldsymbol{\Omega}}^{(n)})^{-1}-\boldsymbol{\Omega}^{-1}\right\|$ $\displaystyle\leq\left\|(\hat{\boldsymbol{\Omega}}^{(n)})^{-1}\right\|\left\|\boldsymbol{\Omega}-\hat{\boldsymbol{\Omega}}^{(n)}\right\|\left\|\boldsymbol{\Omega}^{-1}\right\|$ $\displaystyle\leq\frac{1}{\tilde{\lambda}^{2}}\left\|\boldsymbol{\Omega}-\hat{\boldsymbol{\Omega}}^{(n)}\right\|.$ $$

Then

$$ $\displaystyle\begin{split}\mathbb{P}\left\{\left\|(\hat{\boldsymbol{\Omega}}^{(n)})^{-1}-\boldsymbol{\Omega}^{-1}\right\|\geq\varepsilon\right\}&\leq\mathbb{P}\left\{\frac{1}{\tilde{\lambda}^{2}}\left\|\boldsymbol{\Omega}-\hat{\boldsymbol{\Omega}}^{(n)}\right\|\geq\varepsilon\right\}\\ &\leq\xi(n,\tilde{\lambda}^{2}\varepsilon).\end{split}$ $$

∎

###### Lemma A.3 (Product Convergence) .

Let $\{\hat{\boldsymbol{\Omega}}_{1}^{(n)}\},\{\hat{\boldsymbol{\Omega}}_{2}^{(n)}\},\ldots,\{\hat{\boldsymbol{\Omega}}_{K}^{(n)}\}$ be $K$ sequences of matrices such that $\hat{\boldsymbol{\Omega}}_{1}^{(n)}\overset{\textsf{p}}{\rightarrow}\boldsymbol{\Omega}_{1}$, $\hat{\boldsymbol{\Omega}}_{2}^{(n)}\overset{\textsf{p}}{\rightarrow}\boldsymbol{\Omega}_{2},\ldots,\hat{\boldsymbol{\Omega}}_{K}^{(n)}\overset{\textsf{p}}{\rightarrow}\boldsymbol{\Omega}_{K}$, where each $\|\hat{\boldsymbol{\Omega}}_{k}^{(n)}\|$ is almost surely bounded for every $k=1,\ldots,K$. If the dimensions match, then it holds that:

$$ $\displaystyle\hat{\boldsymbol{\Omega}}_{1}^{(n)}\hat{\boldsymbol{\Omega}}_{2}^{(n)}\cdots\hat{\boldsymbol{\Omega}}_{K}^{(n)}\overset{\textsf{p}}{\rightarrow}\boldsymbol{\Omega}_{1}\boldsymbol{\Omega}_{2}\cdots\boldsymbol{\Omega}_{K}.$ $$

Further, given convergence rates:

$$ $\displaystyle\begin{split}\mathbb{P}&\left\{\left\|\hat{\boldsymbol{\Omega}}_{1}^{(n)}-\boldsymbol{\Omega}_{1}\right\|\geq\varepsilon\right\}\leq\xi_{1}(n,\varepsilon),\\ \mathbb{P}&\left\{\left\|\hat{\boldsymbol{\Omega}}_{2}^{(n)}-\boldsymbol{\Omega}_{2}\right\|\geq\varepsilon\right\}\leq\xi_{2}(n,\varepsilon),\\ &\quad\quad\quad\quad\quad\quad\quad\quad\vdots\\ \mathbb{P}&\left\{\left\|\hat{\boldsymbol{\Omega}}_{K}^{(n)}-\boldsymbol{\Omega}_{K}\right\|\geq\varepsilon\right\}\leq\xi_{K}(n,\varepsilon),\end{split}$ $$

then it holds that:

$$ $\displaystyle\mathbb{P}\left\{\left\|\hat{\boldsymbol{\Omega}}_{1}^{(n)}\hat{\boldsymbol{\Omega}}_{2}^{(n)}\cdots\hat{\boldsymbol{\Omega}}_{K}^{(n)}-\boldsymbol{\Omega}_{1}\boldsymbol{\Omega}_{2}\cdots\boldsymbol{\Omega}_{K}\right\|\geq\varepsilon\right\}$ $\displaystyle\leq\sum_{i=1}^{K}\xi_{i}\left(n,\frac{\varepsilon}{K\prod_{k\neq i}^{K}M_{k}}\right),$ (15) $$

where $M_{k}$ is an upper bound such that $\|\hat{\boldsymbol{\Omega}}_{k}^{(n)}\|\leq M_{k}$ almost surely, $\forall k=1,\ldots,K$.

###### Proof.

We begin by showing the case of $K=2$. By the triangle inequality, we have:

$$ $\displaystyle\begin{aligned} \left\|\hat{\boldsymbol{\Omega}}_{1}^{(n)}\hat{\boldsymbol{\Omega}}_{2}^{(n)}-\boldsymbol{\Omega}_{1}\boldsymbol{\Omega}_{2}\right\|&\leq\left\|\hat{\boldsymbol{\Omega}}_{1}^{(n)}\hat{\boldsymbol{\Omega}}_{2}^{(n)}-\boldsymbol{\Omega}_{1}\hat{\boldsymbol{\Omega}}_{2}^{(n)}\right\|+\left\|\boldsymbol{\Omega}_{1}\hat{\boldsymbol{\Omega}}_{2}^{(n)}-\boldsymbol{\Omega}_{1}\boldsymbol{\Omega}_{2}\right\|\\ &\leq\left\|\hat{\boldsymbol{\Omega}}_{2}^{(n)}\right\|\left\|\hat{\boldsymbol{\Omega}}_{1}^{(n)}-\boldsymbol{\Omega}_{1}\right\|+\left\|\hat{\boldsymbol{\Omega}}_{2}^{(n)}-\boldsymbol{\Omega}_{2}\right\|\left\|\boldsymbol{\Omega}_{1}\right\|\\ &\leq M_{2}\left\|\hat{\boldsymbol{\Omega}}_{1}^{(n)}-\boldsymbol{\Omega}_{1}\right\|+M_{1}\left\|\hat{\boldsymbol{\Omega}}_{2}^{(n)}-\boldsymbol{\Omega}_{2}\right\|.\end{aligned}$ $$

Using the union bound, we have:

$$ $\displaystyle\begin{split}&\mathbb{P}\left\{\left\|\hat{\boldsymbol{\Omega}}_{1}^{(n)}\hat{\boldsymbol{\Omega}}_{2}^{(n)}-\boldsymbol{\Omega}_{1}\boldsymbol{\Omega}_{2}\right\|\geq\varepsilon\right\}\\ \leq&\mathbb{P}\left\{M_{2}\left\|\hat{\boldsymbol{\Omega}}_{1}^{(n)}-\boldsymbol{\Omega}_{1}\right\|+M_{1}\left\|\hat{\boldsymbol{\Omega}}_{2}^{(n)}-\boldsymbol{\Omega}_{2}\right\|\geq\varepsilon\right\}\\ \leq&\mathbb{P}\left\{M_{2}\left\|\hat{\boldsymbol{\Omega}}_{1}^{(n)}-\boldsymbol{\Omega}_{1}\right\|\geq\varepsilon/2\right\}+\mathbb{P}\left\{M_{1}\left\|\hat{\boldsymbol{\Omega}}_{2}^{(n)}-\boldsymbol{\Omega}_{2}\right\|\geq\varepsilon/2\right\}\\ \leq&\xi_{1}\left(n,\frac{\varepsilon}{2M_{2}}\right)+\xi_{2}\left(n,\frac{\varepsilon}{2M_{1}}\right).\end{split}$ $$

For any $K>2$, suppose the statement ([15](#A1.E15)) holds for $k=2,\ldots,K-1$. Observe that:

$$ $\displaystyle\begin{split}&\left\|\hat{\boldsymbol{\Omega}}_{1}^{(n)}\hat{\boldsymbol{\Omega}}_{2}^{(n)}\cdots\hat{\boldsymbol{\Omega}}_{K}^{(n)}-\boldsymbol{\Omega}_{1}\boldsymbol{\Omega}_{2}\cdots\boldsymbol{\Omega}_{K}\right\|\\ &\leq\left\|\hat{\boldsymbol{\Omega}}_{1}^{(n)}\hat{\boldsymbol{\Omega}}_{2}^{(n)}\cdots\hat{\boldsymbol{\Omega}}_{K}^{(n)}-\boldsymbol{\Omega}_{1}\boldsymbol{\Omega}_{2}\cdots\boldsymbol{\Omega}_{K-1}\hat{\boldsymbol{\Omega}}_{K}^{(n)}\right\|+\left\|\boldsymbol{\Omega}_{1}\boldsymbol{\Omega}_{2}\cdots\boldsymbol{\Omega}_{K-1}\hat{\boldsymbol{\Omega}}_{K}^{(n)}-\boldsymbol{\Omega}_{1}\boldsymbol{\Omega}_{2}\cdots\boldsymbol{\Omega}_{K}\right\|\\ &\leq M_{K}\left\|\hat{\boldsymbol{\Omega}}_{1}^{(n)}\hat{\boldsymbol{\Omega}}_{2}^{(n)}\cdots\hat{\boldsymbol{\Omega}}_{K-1}^{(n)}-\boldsymbol{\Omega}_{1}\boldsymbol{\Omega}_{2}\cdots\boldsymbol{\Omega}_{K-1}\right\|+\prod_{k=1}^{K-1}M_{k}\left\|\hat{\boldsymbol{\Omega}}_{K}^{(n)}-\boldsymbol{\Omega}_{K}\right\|.\\ \end{split}$ (16) $$

Then it follows that:

$$ $\displaystyle\begin{split}&\mathbb{P}\left\{\left\|\hat{\boldsymbol{\Omega}}_{1}^{(n)}\hat{\boldsymbol{\Omega}}_{2}^{(n)}\cdots\hat{\boldsymbol{\Omega}}_{K}^{(n)}-\boldsymbol{\Omega}_{1}\boldsymbol{\Omega}_{2}\cdots\boldsymbol{\Omega}_{K}\right\|\geq\varepsilon\right\}\\ &\leq\mathbb{P}\left\{M_{K}\left\|\hat{\boldsymbol{\Omega}}_{1}^{(n)}\hat{\boldsymbol{\Omega}}_{2}^{(n)}\cdots\hat{\boldsymbol{\Omega}}_{K-1}^{(n)}-\boldsymbol{\Omega}_{1}\boldsymbol{\Omega}_{2}\cdots\boldsymbol{\Omega}_{K-1}\right\|+\prod_{k=1}^{K-1}M_{k}\left\|\hat{\boldsymbol{\Omega}}_{K}^{(n)}-\boldsymbol{\Omega}_{K}\right\|\geq\varepsilon\right\}\\ &\leq\mathbb{P}\left\{M_{K}\left\|\hat{\boldsymbol{\Omega}}_{1}^{(n)}\hat{\boldsymbol{\Omega}}_{2}^{(n)}\cdots\hat{\boldsymbol{\Omega}}_{K-1}^{(n)}-\boldsymbol{\Omega}_{1}\boldsymbol{\Omega}_{2}\cdots\boldsymbol{\Omega}_{K-1}\right\|\geq\frac{K-1}{K}\varepsilon\right\}\\ &\qquad+\mathbb{P}\left\{\prod_{k=1}^{K-1}M_{k}\left\|\hat{\boldsymbol{\Omega}}_{K}^{(n)}-\boldsymbol{\Omega}_{K}\right\|\geq\frac{1}{K}\varepsilon\right\}\\ &\leq\sum_{i=1}^{K-1}\xi_{i}\left(n,\frac{\varepsilon}{KM_{K}\prod_{k\neq i}^{K-1}M_{k}}\right)+\xi_{K}\left(n,\frac{\varepsilon}{K\prod_{k=1}^{K-1}M_{k}}\right)\\ &=\sum_{i=1}^{K}\xi_{i}\left(n,\frac{\varepsilon}{K\prod_{k\neq i}^{K}M_{k}}\right).\end{split}$ $$

Thus, by induction, the proof is complete.
∎

###### Remark A.1 .

In Lemma [A.3](#A1.Thmlemma3), consider the special case where $\boldsymbol{\Omega}_{1}=\boldsymbol{0}$. Then the inequality ([16](#A1.E16)) can be simplified as follows:

$$ $\displaystyle\begin{split}\left\|\hat{\boldsymbol{\Omega}}_{1}^{(n)}\hat{\boldsymbol{\Omega}}_{2}^{(n)}\cdots\hat{\boldsymbol{\Omega}}_{K}^{(n)}-\boldsymbol{0}\right\|&\leq\prod_{k=2}^{K}M_{k}\left\|\hat{\boldsymbol{\Omega}}_{1}^{(n)}\right\|.\end{split}$ $$

And we have the following simplified form:

$$ $\displaystyle\begin{split}\mathbb{P}\left\{\left\|\hat{\boldsymbol{\Omega}}_{1}^{(n)}\hat{\boldsymbol{\Omega}}_{2}^{(n)}\cdots\hat{\boldsymbol{\Omega}}_{K}^{(n)}-\boldsymbol{0}\right\|\geq\varepsilon\right\}&\leq\mathbb{P}\left\{\prod_{k=2}^{K}M_{k}\left\|\hat{\boldsymbol{\Omega}}_{1}^{(n)}\right\|\geq\varepsilon\right\}\\ &\leq\xi_{1}\left(n,\frac{\varepsilon}{\prod_{k=2}^{K}M_{k}}\right).\end{split}$ $$

###### Proof of Theorem 2.1 .

We denote the observational values $(\boldsymbol{Z},\boldsymbol{X},\boldsymbol{Y})=\{(\boldsymbol{z}_{i},\boldsymbol{x}_{i},y_{i})\}_{i=1}^{n}$, and $\boldsymbol{\mathcal{E}}_{1}=\{\epsilon_{1,i}\}_{i=1}^{n}$. The 2SLS estimator is given by:

$$ $\displaystyle\begin{split}\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}&=\left(\hat{\boldsymbol{\Theta}}^{\top}\boldsymbol{Z}^{\top}\boldsymbol{Z}\hat{\boldsymbol{\Theta}}\right)^{-1}\hat{\boldsymbol{\Theta}}^{\top}\boldsymbol{Z}^{\top}\boldsymbol{Y}\\ &=\left[\left((\boldsymbol{Z}^{\top}\boldsymbol{Z})^{-1}\boldsymbol{Z}^{\top}\boldsymbol{X}\right)^{\top}\boldsymbol{Z}^{\top}\boldsymbol{Z}(\boldsymbol{Z}^{\top}\boldsymbol{Z})^{-1}\boldsymbol{Z}^{\top}\boldsymbol{X}\right]^{-1}\left((\boldsymbol{Z}^{\top}\boldsymbol{Z})^{-1}\boldsymbol{Z}^{\top}\boldsymbol{X}\right)^{\top}\boldsymbol{Z}^{\top}\boldsymbol{Y}\\ &=\left(\boldsymbol{X}^{\top}\boldsymbol{Z}(\boldsymbol{Z}^{\top}\boldsymbol{Z})^{-1}\boldsymbol{Z}^{\top}\boldsymbol{X}\right)^{-1}\boldsymbol{X}^{\top}\boldsymbol{Z}(\boldsymbol{Z}^{\top}\boldsymbol{Z})^{-1}\boldsymbol{Z}^{\top}\boldsymbol{Y}\\ &=\boldsymbol{\beta}+\left(\boldsymbol{X}^{\top}\boldsymbol{Z}(\boldsymbol{Z}^{\top}\boldsymbol{Z})^{-1}\boldsymbol{Z}^{\top}\boldsymbol{X}\right)^{-1}\boldsymbol{X}^{\top}\boldsymbol{Z}(\boldsymbol{Z}^{\top}\boldsymbol{Z})^{-1}\boldsymbol{Z}^{\top}\boldsymbol{\mathcal{E}}_{1}.\end{split}$ (17) $$

We first consider the case when event $\mathcal{A}$ is true. Let $\boldsymbol{Q}_{zz}:=\mathbb{E}[\boldsymbol{z}\boldsymbol{z}^{\top}|\mathcal{A}],\boldsymbol{Q}_{zx}:=\mathbb{E}[\boldsymbol{z}\boldsymbol{x}^{\top}|\mathcal{A}]$, $\bar{\boldsymbol{\Omega}}_{zz}:=\sum_{i=1}^{n}(\boldsymbol{z}_{i}\boldsymbol{z}_{i}^{\top}-\boldsymbol{Q}_{zz}),\bar{\boldsymbol{\Omega}}_{zx}:=\sum_{i=1}^{n}(\boldsymbol{z}_{i}\boldsymbol{x}_{i}^{\top}-\boldsymbol{Q}_{zx}),\boldsymbol{\Omega}_{z\epsilon_{1}}:=\sum_{i=1}^{n}\boldsymbol{z}_{i}\epsilon_{1,i}$.

Let $\bar{B}_{zz},\bar{B}_{zx},B_{zx},B_{z\epsilon_{1}}$ be some upper bounds such that $\left\|\boldsymbol{z}_{i}\boldsymbol{z}_{i}^{\top}-\boldsymbol{Q}_{zz}\right\|\leq\bar{B}_{zz},\left\|\boldsymbol{z}_{i}\boldsymbol{x}_{i}^{\top}-\boldsymbol{Q}_{zx}\right\|\leq\bar{B}_{zx},\left\|\boldsymbol{z}_{i}\boldsymbol{x}_{i}^{\top}\right\|\leq B_{zx},\left\|\boldsymbol{z}_{i}\epsilon_{1,i}\right\|\leq B_{z\epsilon_{1}}$ almost surely, for all $i=1,\ldots,n$. The existence of $\bar{B}_{zz},\bar{B}_{zx},B_{zx},B_{z\epsilon_{1}}$ is guaranteed under Assumption [2](#Thmasp2).

By Lemma [A.1](#A1.Thmlemma1), we have:

$$ $\displaystyle\begin{split}\mathbb{P}\left\{\left\|\frac{\boldsymbol{Z}^{\top}\boldsymbol{Z}}{n}-\boldsymbol{Q}_{zz}\right\|\geq\varepsilon\middle|\mathcal{A}\right\}&=\mathbb{P}\left\{\left\|\frac{\sum_{i=1}^{n}\boldsymbol{z}_{i}\boldsymbol{z}_{i}^{\top}}{n}-\boldsymbol{Q}_{zz}\right\|\geq\varepsilon\middle|\mathcal{A}\right\}\\ &=\mathbb{P}\left\{\left\|\sum_{i=1}^{n}(\boldsymbol{z}_{i}\boldsymbol{z}_{i}^{\top}-\boldsymbol{Q}_{zz})\right\|\geq n\varepsilon\middle|\mathcal{A}\right\}\\ &\leq 2q\exp\left(-\frac{n^{2}\varepsilon^{2}/2}{\nu(\bar{\boldsymbol{\Omega}}_{zz}|\mathcal{A})+\bar{B}_{zz}n\varepsilon/3}\right).\end{split}$ (18) $$

Similarly,

$$ $\displaystyle\begin{split}\mathbb{P}\left\{\left\|\frac{\boldsymbol{Z}^{\top}\boldsymbol{X}}{n}-\boldsymbol{Q}_{zx}\right\|\geq\varepsilon\middle|\mathcal{A}\right\}&=\mathbb{P}\left\{\left\|\frac{\sum_{i=1}^{n}\boldsymbol{z}_{i}\boldsymbol{x}_{i}^{\top}}{n}-\boldsymbol{Q}_{zx}\right\|\geq\varepsilon\middle|\mathcal{A}\right\}\\ &=\mathbb{P}\left\{\left\|\sum_{i=1}^{n}(\boldsymbol{z}_{i}\boldsymbol{x}_{i}^{\top}-\boldsymbol{Q}_{zx})\right\|\geq n\varepsilon\middle|\mathcal{A}\right\}\\ &\leq(p+q)\exp\left(-\frac{n^{2}\varepsilon^{2}/2}{\nu(\bar{\boldsymbol{\Omega}}_{zx}|\mathcal{A})+\bar{B}_{zx}n\varepsilon/3}\right).\end{split}$ (19) $$

By Assumption [1](#Thmasp1)(iii), the instrument $\boldsymbol{z}$ is uncorrelated with the error term $\epsilon_{1}$, which implies $\mathbb{E}[\boldsymbol{z}\epsilon_{1}|\mathcal{A}]=\boldsymbol{0}$. Applying Lemma [A.1](#A1.Thmlemma1) again, we have:

$$ $\displaystyle\begin{split}\mathbb{P}\left\{\left\|\frac{\boldsymbol{Z}^{\top}\boldsymbol{\mathcal{E}}_{1}}{n}\right\|\geq\varepsilon\middle|\mathcal{A}\right\}&=\mathbb{P}\left\{\left\|\frac{\sum_{i=1}^{n}\boldsymbol{z}_{i}{\epsilon}_{1,i}}{n}\right\|\geq\varepsilon\middle|\mathcal{A}\right\}\\ &=\mathbb{P}\left\{\left\|\sum_{i=1}^{n}\boldsymbol{z}_{i}{\epsilon}_{1,i}\right\|\geq n\varepsilon\middle|\mathcal{A}\right\}\\ &\leq(q+1)\exp\left(-\frac{n^{2}\varepsilon^{2}/2}{\nu(\boldsymbol{\Omega}_{z\epsilon_{1}}|\mathcal{A})+B_{z\epsilon_{1}}n\varepsilon/3}\right).\\ \end{split}$ (20) $$

With Lemma [A.2](#A1.Thmlemma2) and ([18](#A1.E18)), we have:

$$ $\displaystyle\begin{split}\mathbb{P}\left\{\left\|n(\boldsymbol{Z}^{\top}\boldsymbol{Z})^{-1}-\boldsymbol{Q}_{ZZ}^{-1}\right\|\geq\varepsilon\middle|\mathcal{A}\right\}&\leq 2q\exp\left(-\frac{n^{2}(\lambda_{z}^{2}\varepsilon)^{2}/2}{\nu(\bar{\boldsymbol{\Omega}}_{zz}|\mathcal{A})+\bar{B}_{zz}n(\lambda_{z}^{2}\varepsilon)/3}\right)\\ &=2q\exp\left(-\frac{\lambda_{z}^{4}n^{2}\varepsilon^{2}/2}{\nu(\bar{\boldsymbol{\Omega}}_{zz}|\mathcal{A})+\lambda_{z}^{2}\bar{B}_{zz}n\varepsilon/3}\right).\end{split}$ (21) $$

Note that under event $\mathcal{A}$,

$$ $\displaystyle\begin{split}\left\|\hat{\boldsymbol{\Theta}}\right\|&=\left\|\left(\boldsymbol{Z}^{\top}\boldsymbol{Z}\right)^{-1}\boldsymbol{Z}^{\top}\boldsymbol{X}\right\|\\ &=\left\|\boldsymbol{\Theta}+\left(\boldsymbol{Z}^{\top}\boldsymbol{Z}\right)^{-1}\boldsymbol{Z}^{\top}\boldsymbol{\mathcal{E}}_{2}\right\|\\ &\leq B_{\Theta}+\left\|\left(\boldsymbol{Z}^{\top}\boldsymbol{Z}\right)^{-1}\boldsymbol{Z}^{\top}\boldsymbol{\mathcal{E}}_{2}\right\|\\ &\leq B_{\Theta}+\frac{\delta_{z\epsilon_{2}}}{\lambda_{z}}:=B_{\hat{\Theta}}.\end{split}$ (22) $$

With Lemma [A.3](#A1.Thmlemma3) (Remark [A.1](#A1.Thmrmk1)), combining ([20](#A1.E20))([22](#A1.E22)), we have:

$$ $\displaystyle\begin{split}\mathbb{P}\left\{\left\|\frac{1}{n}\boldsymbol{X}^{\top}\boldsymbol{Z}(\boldsymbol{Z}^{\top}\boldsymbol{Z})^{-1}\boldsymbol{Z}^{\top}\boldsymbol{\mathcal{E}}_{1}-\boldsymbol{0}\right\|\geq\varepsilon\middle|\mathcal{A}\right\}&\leq(q+1)\exp\left(-\frac{n^{2}(\frac{\varepsilon}{B_{\hat{\Theta}}})^{2}/2}{\nu(\boldsymbol{\Omega}_{z\epsilon_{1}}|\mathcal{A})+B_{z\epsilon_{1}}n(\frac{\varepsilon}{B_{\hat{\Theta}}})/3}\right)\\ &=(q+1)\exp\left(-\frac{n^{2}\varepsilon^{2}/2}{B_{\hat{\Theta}}^{2}\nu(\boldsymbol{\Omega}_{z\epsilon_{1}}|\mathcal{A})+B_{\hat{\Theta}}B_{z\epsilon_{1}}n\varepsilon/3}\right).\end{split}$ (23) $$

Additionally, with Lemma [A.3](#A1.Thmlemma3), combining ([19](#A1.E19))([21](#A1.E21)), we have:

$$ $\displaystyle\begin{split}&\mathbb{P}\left\{\left\|\frac{1}{n}\boldsymbol{X}^{\top}\boldsymbol{Z}(\boldsymbol{Z}^{\top}\boldsymbol{Z})^{-1}\boldsymbol{Z}^{\top}\boldsymbol{X}-\boldsymbol{Q}_{zx}^{\top}\boldsymbol{Q}_{zz}^{-1}\boldsymbol{Q}_{zx}\right\|\geq\varepsilon\middle|\mathcal{A}\right\}\\ &\leq 2(p+q)\exp\left(-\frac{n^{2}(\frac{\lambda_{z}\varepsilon}{3B_{zx}})^{2}/2}{\nu(\bar{\boldsymbol{\Omega}}_{zx}|\mathcal{A})+\bar{B}_{zx}n(\frac{\lambda_{z}\varepsilon}{3B_{zx}})/3}\right)+2q\exp\left(-\frac{\lambda_{z}^{4}n^{2}(\frac{\varepsilon}{3B_{zx}^{2}})^{2}/2}{\nu(\bar{\boldsymbol{\Omega}}_{zz}|\mathcal{A})+\lambda_{z}^{2}\bar{B}_{zz}n(\frac{\varepsilon}{3B_{zx}^{2}})/3}\right)\\ &=2(p+q)\exp\left(-\frac{\lambda_{z}^{2}n^{2}\varepsilon^{2}/2}{9B_{zx}^{2}\nu(\bar{\boldsymbol{\Omega}}_{zx}|\mathcal{A})+\lambda_{z}B_{zx}\bar{B}_{zx}n\varepsilon}\right)+2q\exp\left(-\frac{\lambda_{z}^{4}n^{2}\varepsilon^{2}/2}{9B_{zx}^{4}\nu(\bar{\boldsymbol{\Omega}}_{zz}|\mathcal{A})+\lambda_{z}^{2}B_{zx}^{2}\bar{B}_{zz}n\varepsilon}\right).\end{split}$ $$

Applying Lemma [A.2](#A1.Thmlemma2) again, we have:

$$ $\displaystyle\begin{split}&\mathbb{P}\left\{\left\|n\left(\boldsymbol{X}^{\top}\boldsymbol{Z}(\boldsymbol{Z}^{\top}\boldsymbol{Z})^{-1}\boldsymbol{Z}^{\top}\boldsymbol{X}\right)^{-1}-(\boldsymbol{Q}_{ZX}^{\top}\boldsymbol{Q}_{zz}^{-1}\boldsymbol{Q}_{zx})^{-1}\right\|\geq\varepsilon\middle|\mathcal{A}\right\}\\ &\leq 2(p+q)\exp\left(-\frac{\lambda_{z}^{2}n^{2}(\lambda_{\tilde{x}}^{2}\varepsilon)^{2}/2}{9B_{zx}^{2}\nu(\bar{\boldsymbol{\Omega}}_{zx}|\mathcal{A})+\lambda_{z}B_{zx}\bar{B}_{zx}n(\lambda_{\tilde{x}}^{2}\varepsilon)}\right)+2q\exp\left(-\frac{\lambda_{z}^{4}n^{2}(\lambda_{\tilde{x}}^{2}\varepsilon)^{2}/2}{9B_{zx}^{4}\nu(\bar{\boldsymbol{\Omega}}_{zz}|\mathcal{A})+\lambda_{z}^{2}B_{zx}^{2}\bar{B}_{zz}n(\lambda_{\tilde{x}}^{2}\varepsilon)}\right)\\ &=2(p+q)\exp\left(-\frac{\lambda_{z}^{2}\lambda_{\tilde{x}}^{4}n^{2}\varepsilon^{2}/2}{9B_{zx}^{2}\nu(\bar{\boldsymbol{\Omega}}_{zx}|\mathcal{A})+\lambda_{z}\lambda_{\tilde{x}}^{2}B_{zx}\bar{B}_{zx}n\varepsilon}\right)+2q\exp\left(-\frac{\lambda_{z}^{4}\lambda_{\tilde{x}}^{4}n^{2}\varepsilon^{2}/2}{9B_{zx}^{4}\nu(\bar{\boldsymbol{\Omega}}_{zz}|\mathcal{A})+\lambda_{z}^{2}\lambda_{\tilde{x}}^{2}B_{zx}^{2}\bar{B}_{zz}n\varepsilon}\right).\end{split}$ (24) $$

Therefore, we have shown that under event $\mathcal{A}$,

$$ $\displaystyle n\left(\boldsymbol{X}^{\top}\boldsymbol{Z}(\boldsymbol{Z}^{\top}\boldsymbol{Z})^{-1}\boldsymbol{Z}^{\top}\boldsymbol{X}\right)^{-1}\overset{\textsf{p}}{\rightarrow}(\boldsymbol{Q}_{ZX}^{\top}\boldsymbol{Q}_{zz}^{-1}\boldsymbol{Q}_{zx})^{-1}.$ $$

From equation ([17](#A1.E17)), combining ([23](#A1.E23)) and ([24](#A1.E24)) with Lemma [A.3](#A1.Thmlemma3) (Remark [A.1](#A1.Thmrmk1)), we have:

$$ $\displaystyle\mathbb{P}\left\{\left\|\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}-\boldsymbol{\beta}\right\|\geq\varepsilon\middle|\mathcal{A}\right\}$ $\displaystyle=\mathbb{P}\left\{\left\|\left(\boldsymbol{X}^{\top}\boldsymbol{Z}(\boldsymbol{Z}^{\top}\boldsymbol{Z})^{-1}\boldsymbol{Z}^{\top}\boldsymbol{X}\right)^{-1}\boldsymbol{X}^{\top}\boldsymbol{Z}(\boldsymbol{Z}^{\top}\boldsymbol{Z})^{-1}\boldsymbol{Z}^{\top}\boldsymbol{\mathcal{E}}_{1}-\boldsymbol{0}\right\|\geq\varepsilon\middle|\mathcal{A}\right\}$ $\displaystyle\leq(q+1)\exp\left(-\frac{n^{2}(\lambda_{\tilde{x}}\varepsilon)^{2}/2}{B_{\hat{\Theta}}^{2}\nu(\boldsymbol{\Omega}_{z\epsilon_{1}}|\mathcal{A})+B_{\hat{\Theta}}B_{z\epsilon_{1}}n(\lambda_{\tilde{x}}\varepsilon)/3}\right)$ $\displaystyle=(q+1)\exp\left(-\frac{\lambda_{\tilde{x}}^{2}n^{2}\varepsilon^{2}/2}{B_{\hat{\Theta}}^{2}\nu(\boldsymbol{\Omega}_{z\epsilon_{1}}|\mathcal{A})+\lambda_{\tilde{x}}B_{\hat{\Theta}}B_{z\epsilon_{1}}n\varepsilon/3}\right).$ $$

For the second part of the theorem, let $c:=\left(\frac{3B_{\hat{\Theta}}\nu(\boldsymbol{\Omega}_{z\epsilon_{1}}|\mathcal{A})}{\lambda_{\tilde{x}}B_{z\epsilon_{1}}n}\right)^{2}$, we have:

$$ $\displaystyle\begin{split}&\mathbb{E}\left[\|\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}})-\boldsymbol{\beta}\|^{2}\right]\\ &=\mathbb{E}\left[\|\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}})-\boldsymbol{\beta}\|^{2}\middle|\mathcal{A}\right]\mathbb{P}\left\{\mathcal{A}\right\}+\mathbb{E}\left[\|\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}})-\boldsymbol{\beta}\|^{2}\middle|\mathcal{A}^{c}\right]\mathbb{P}\left\{\mathcal{A}^{c}\right\}\\ &\leq\mathbb{E}\left[\|\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}-\boldsymbol{\beta}\|^{2}\middle|\mathcal{A}\right]\mathbb{P}\left\{\mathcal{A}\right\}+\mathbb{E}\left[\|\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}})-\boldsymbol{\beta}\|^{2}\middle|\mathcal{A}^{c}\right]\mathbb{P}\left\{\mathcal{A}^{c}\right\}\\ &\leq\mathbb{E}\left[\|\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}-\boldsymbol{\beta}\|^{2}\middle|\mathcal{A}\right]+\mathbb{E}\left[\|\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}})-\boldsymbol{\beta}\|^{2}\middle|\mathcal{A}^{c}\right]\cdot\xi,\end{split}$ (25) $$

where

$$ $\displaystyle\mathbb{E}\left[\|\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}})-\boldsymbol{\beta}\|^{2}\middle|\mathcal{A}^{c}\right]\leq 4B_{\beta}^{2},$ (26) $$

and

$$ $\displaystyle\begin{split}&\mathbb{E}\left[\|\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}-\boldsymbol{\beta}\|^{2}\middle|\mathcal{A}\right]\\ &=\int_{0}^{\infty}\mathbb{P}\left\{\|\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}-\boldsymbol{\beta}\|^{2}\geq\varepsilon\middle|\mathcal{A}\right\}d\varepsilon\\ &=\int_{0}^{\infty}\mathbb{P}\left\{\|\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}-\boldsymbol{\beta}\|\geq\sqrt{\varepsilon}\middle|\mathcal{A}\right\}d\varepsilon\\ &\leq\int_{0}^{\infty}(q+1)\exp\left(-\frac{\lambda_{\tilde{x}}^{2}n^{2}\varepsilon/2}{B_{\hat{\Theta}}^{2}\nu(\boldsymbol{\Omega}_{z\epsilon_{1}}|\mathcal{A})+\lambda_{\tilde{x}}B_{\hat{\Theta}}B_{z\epsilon_{1}}n\sqrt{\varepsilon}/3}\right)d\varepsilon\\ &\leq(q+1)\left[\int_{0}^{c}\exp\left(-\frac{\lambda_{\tilde{x}}^{2}n^{2}\varepsilon/2}{2B_{\hat{\Theta}}^{2}\nu(\boldsymbol{\Omega}_{z\epsilon_{1}}|\mathcal{A})}\right)d\varepsilon+\int_{c}^{\infty}\exp\left(-\frac{\lambda_{\tilde{x}}n\sqrt{\varepsilon}/2}{2B_{\hat{\Theta}}B_{z\epsilon_{1}}/3}\right)d\varepsilon\right]\\ &=(q+1)\left[\frac{4B_{\hat{\Theta}}^{2}\nu(\boldsymbol{\Omega}_{z\epsilon_{1}}|\mathcal{A})}{\lambda_{\tilde{x}}^{2}n^{2}}\left(1-\exp\left(-\frac{9\nu(\boldsymbol{\Omega}_{z\epsilon_{1}}|\mathcal{A})}{4B_{z\epsilon_{1}}^{2}}\right)\right)+\left(\frac{8B_{\hat{\Theta}}^{2}\nu(\boldsymbol{\Omega}_{z\epsilon_{1}}|\mathcal{A})}{\lambda_{\tilde{x}}^{2}n^{2}}+\frac{32B_{\hat{\Theta}}^{2}B_{z\epsilon_{1}}^{2}}{9\lambda_{\tilde{x}}^{2}n^{2}}\right)\exp\left(-\frac{9\nu(\boldsymbol{\Omega}_{z\epsilon_{1}}|\mathcal{A})}{4B_{z\epsilon_{1}}^{2}}\right)\right]\\ &\leq(q+1)\left[\frac{4B_{\hat{\Theta}}^{2}\nu(\boldsymbol{\Omega}_{z\epsilon_{1}}|\mathcal{A})}{\lambda_{\tilde{x}}^{2}n^{2}}+\frac{8B_{\hat{\Theta}}^{2}\nu(\boldsymbol{\Omega}_{z\epsilon_{1}}|\mathcal{A})}{\lambda_{\tilde{x}}^{2}n^{2}}+\frac{32B_{\hat{\Theta}}^{2}B_{z\epsilon_{1}}^{2}}{9\lambda_{\tilde{x}}^{2}n^{2}}\right]\\ &=\frac{(q+1)B_{\hat{\Theta}}^{2}}{\lambda_{\tilde{x}}^{2}n^{2}}\left[12\nu(\boldsymbol{\Omega}_{z\epsilon_{1}}|\mathcal{A})+\frac{32B_{z}^{2}B_{\epsilon_{1}}^{2}}{9}\right].\end{split}$ (27) $$

Note that we further have the following bound:

$$ $\displaystyle\begin{split}\nu(\boldsymbol{\Omega}_{z\epsilon_{1}}|\mathcal{A})&=\max\left\{\bigg{\|}\mathbb{E}\Big{[}\Big{(}\sum_{i=1}^{n}\boldsymbol{z}_{i}\epsilon_{1,i}\Big{)}^{\top}\Big{(}\sum_{j=1}^{n}\boldsymbol{z}_{j}\epsilon_{1,j}\Big{)}\Big{|}\mathcal{A}\Big{]}\bigg{\|},\bigg{\|}\mathbb{E}\Big{[}\Big{(}\sum_{i=1}^{n}\boldsymbol{z}_{i}\epsilon_{1,i}\Big{)}\Big{(}\sum_{j=1}^{n}\boldsymbol{z}_{j}\epsilon_{1,j}\Big{)}^{\top}\Big{|}\mathcal{A}\Big{]}\bigg{\|}\right\}\\ &=\max\left\{\bigg{\|}\mathbb{E}\Big{[}\sum_{i=1}^{n}\epsilon_{1,i}^{2}\boldsymbol{z}_{i}^{\top}\boldsymbol{z}_{i}\Big{|}\mathcal{A}\Big{]}\bigg{\|},\bigg{\|}\mathbb{E}\bigg{[}\sum_{i=1}^{n}\epsilon_{1,i}^{2}\boldsymbol{z}_{i}\boldsymbol{z}_{i}^{\top}\Big{|}\mathcal{A}\Big{]}\bigg{\|}\right\}\\ &\leq\mathcal{O}(nB_{z}^{2}\sigma_{1}^{2}).\end{split}$ (28) $$

Plugging ([26](#A1.E26))([27](#A1.E27))([28](#A1.E28)) into ([25](#A1.E25)) completes the proof.
∎

###### Lemma A.1 (Bernstein Inequality, from Theorem 6.1.1 in Tropp [ 2015 ] ) .

###### Lemma A.2 (Inverse Convergence, adapted from Lemma 2.1 in Jin et al. [ 2024 ] ) .

###### Proof.

###### Lemma A.3 (Product Convergence) .

###### Proof.

###### Remark A.1 .

###### Proof of Theorem 2.1 .

## Appendix B Proofs For Section 3

### B.1 Proof of Theorem 3.1

###### Lemma B.1 .

Suppose $\{\boldsymbol{\Omega}^{(1)},\ldots,\boldsymbol{\Omega}^{(t)},\ldots\}$ is a $d\times d$-matrix sequence decaying with exponential rate $r$, i.e. for some constant $c>0$ and $0<r<1$,

$$ $\displaystyle\left\|\boldsymbol{\Omega}^{(t)}\right\|_{F}\leq cr^{t}.$ $$

Then for any $\varepsilon>0$, there exists a finite constant:

$$ $\displaystyle T_{0}=\left\lceil\log_{r}\frac{(1-r)(\varepsilon/d)}{c\left(1+(1-r)(\varepsilon/d)\right)}\right\rceil,$ $$

such that

$$ $\displaystyle\left\|\prod_{t=T_{0}}^{\infty}(\boldsymbol{I}+\boldsymbol{\Omega}^{(t)})-\boldsymbol{I}\right\|_{F}<\varepsilon,$ $$

and hence

$$ $\displaystyle\left\|\prod_{t=T_{0}}^{\infty}(\boldsymbol{I}+\boldsymbol{\Omega}^{(t)})\right\|_{F}<\sqrt{d}+\varepsilon.$ $$

###### Proof.

By definition,

$$ $\displaystyle\left\|\boldsymbol{\Omega}^{(k)}\right\|_{F}=\sqrt{\sum_{i,j=1}^{p}\boldsymbol{\Omega}_{ij}^{(k)2}}\leq cr^{k},$ $$

which implies:

$$ $\displaystyle\left|\boldsymbol{\Omega}_{ij}^{(k)}\right|\leq cr^{k},\quad\forall i,j,k.$ $$

Consider the product of any two matrices. By sub-multiplicativity,

$$ $\displaystyle\left\|\boldsymbol{\Omega}^{(k)}\boldsymbol{\Omega}^{(l)}\right\|_{F}\leq\left\|\boldsymbol{\Omega}^{(k)}\right\|_{F}\left\|\boldsymbol{\Omega}^{(l)}\right\|_{F}\leq c^{2}r^{k+l},$ $$

which implies:

$$ $\displaystyle\left|\left[\boldsymbol{\Omega}^{(k)}\boldsymbol{\Omega}^{(l)}\right]_{ij}\right|\leq c^{2}r^{k+l},\quad\forall i,j,k,l.$ $$

Similarly, for the product of any number of matrices:

$$ $\displaystyle\left|\left[\boldsymbol{\Omega}^{(k_{1})}\boldsymbol{\Omega}^{(k_{2})}\cdots\boldsymbol{\Omega}^{(k_{n})}\right]_{ij}\right|\leq c^{n}r^{k_{1}+k_{2}+\cdots+k_{n}},\quad\forall i,j,k_{1},\ldots,k_{n}.$ $$

Thus

$$ $\displaystyle\begin{split}&\left\|\prod_{t=t_{1}}^{t_{2}}\left(\boldsymbol{I}+\boldsymbol{\Omega}^{(t)}\right)-\boldsymbol{I}\right\|_{F}\\ &=\left\|\left(\boldsymbol{I}+\boldsymbol{\Omega}^{(t_{1})}\right)\left(\boldsymbol{I}+\boldsymbol{\Omega}^{(t_{1}+1)}\right)\cdots\left(\boldsymbol{I}+\boldsymbol{\Omega}^{(t_{2})}\right)-\boldsymbol{I}\right\|_{F}\\ &=\left\|\sum_{t_{1}\leq k\leq t_{2}}\boldsymbol{\Omega}^{(k)}+\sum_{t_{1}\leq k<l\leq t_{2}}\boldsymbol{\Omega}^{(k)}\boldsymbol{\Omega}^{(l)}+\cdots+\boldsymbol{\Omega}^{(t_{1})}\boldsymbol{\Omega}^{(t_{1}+1)}\cdots\boldsymbol{\Omega}^{(t_{2})}\right\|_{F}\\ &\leq\left\|\sum_{t_{1}\leq k\leq t_{2}}cr^{k}\boldsymbol{1}\boldsymbol{1}^{\top}+\sum_{t_{1}\leq k<l\leq t_{2}}c^{2}r^{k+l}\boldsymbol{1}\boldsymbol{1}^{\top}+\cdots+c^{t2-t1+1}r^{t_{1}+\cdots+t_{2}}\boldsymbol{1}\boldsymbol{1}^{\top}\right\|_{F}.\end{split}$ (29) $$

Note that the last inequality can be checked by comparing matrix elements of both sides. For any $\varepsilon>0$, we take $T_{0}=\lceil\log_{r}\frac{(1-r)(\varepsilon/d)}{c(1+(1-r)(\varepsilon/d))}\rceil$.
Consider $t_{1}=T_{0}$ and $t_{2}\rightarrow\infty$ in ([29](#A2.E29)). For notation convenience, let

$$ $\displaystyle\boldsymbol{\Xi}:=\sum_{T_{0}\leq k}cr^{k}\boldsymbol{1}\boldsymbol{1}^{\top}+\sum_{T_{0}\leq k<l}c^{2}r^{k+l}\boldsymbol{1}\boldsymbol{1}^{\top}+\sum_{T_{0}\leq k<l<m}c^{3}r^{k+l+m}\boldsymbol{1}\boldsymbol{1}^{\top}+\cdots.$ $$

Then

$$ $\displaystyle\begin{split}\boldsymbol{\Xi}_{ij}&=\sum_{T_{0}\leq k}cr^{k}+\sum_{T_{0}\leq k<l}c^{2}r^{k+l}+\sum_{T_{0}\leq k<l<m}c^{3}r^{k+l+m}+\cdots\\ &<c\sum_{k\geq T_{0}}r^{k}+c^{2}r^{T_{0}}\sum_{k\geq T_{0}}r^{k}+c^{3}r^{2T_{0}}\sum_{k\geq T_{0}}r^{k}+\cdots\\ &=\frac{cr^{T_{0}}}{1-r}+\frac{c^{2}r^{2T_{0}}}{1-r}+\frac{c^{3}r^{3T_{0}}}{1-r}+\cdots\\ &=\frac{cr^{T_{0}}}{(1-r)(1-cr^{T_{0}})}\\ &\leq\frac{\varepsilon}{d}.\end{split}$ $$

Thus

$$ $\displaystyle\begin{split}\left\|\prod_{t=T_{0}}^{\infty}(\boldsymbol{I}+\boldsymbol{\Omega}^{(t)})-\boldsymbol{I}\right\|_{F}=\left\|\boldsymbol{\Xi}\right\|_{F}=\sqrt{\sum_{i,j=1}^{d}\boldsymbol{\Xi}_{ij}^{2}}\leq\varepsilon.\end{split}$ $$

Hence completes the proof.
∎

###### Proof of Theorem 3.1 .

In the following proof, we treat $\boldsymbol{Z},\boldsymbol{X},\boldsymbol{Y}$ as deterministic matrices.

We begin by checking the inner loop ([8a](#S3.E8.1)):

$$ $\displaystyle\begin{split}\boldsymbol{\Theta}^{(t)}-\hat{\boldsymbol{\Theta}}&=\boldsymbol{\Theta}^{(t-1)}-\hat{\boldsymbol{\Theta}}-\eta\boldsymbol{Z}^{\top}\left(\boldsymbol{Z\Theta}^{(t-1)}-\boldsymbol{X}\right)\\ &=\left(\boldsymbol{I}-\eta\boldsymbol{Z}^{\top}\boldsymbol{Z}\right)\left(\boldsymbol{\Theta}^{(t-1)}-\hat{\boldsymbol{\Theta}}\right)+\eta\boldsymbol{Z}^{\top}\left(\boldsymbol{X}-\boldsymbol{Z\hat{\Theta}}\right)\\ &=\left(\boldsymbol{I}-\eta\boldsymbol{Z}^{\top}\boldsymbol{Z}\right)^{2}\left(\boldsymbol{\Theta}^{(t-2)}-\hat{\boldsymbol{\Theta}}\right)+\eta\boldsymbol{Z}^{\top}\left(\boldsymbol{X}-\boldsymbol{Z\hat{\Theta}}\right)+\eta\left(\boldsymbol{I}-\eta\boldsymbol{Z}^{\top}\boldsymbol{Z}\right)\boldsymbol{Z}^{\top}\left(\boldsymbol{X}-\boldsymbol{Z\hat{\Theta}}\right)\\ &\quad\quad\quad\quad\quad\quad\quad\quad\vdots\\ &=\left(\boldsymbol{I}-\eta\boldsymbol{Z}^{\top}\boldsymbol{Z}\right)^{t}\left(\boldsymbol{\Theta}^{(0)}-\hat{\boldsymbol{\Theta}}\right)+\sum_{i=0}^{t-1}\eta\left(\boldsymbol{I}-\eta\boldsymbol{Z}^{\top}\boldsymbol{Z}\right)^{t-1-i}\boldsymbol{Z}^{\top}\left(\boldsymbol{X}-\boldsymbol{Z\hat{\Theta}}\right)\\ &=\left(\boldsymbol{I}-\eta\boldsymbol{Z}^{\top}\boldsymbol{Z}\right)^{t}\left(\boldsymbol{\Theta}^{(0)}-\hat{\boldsymbol{\Theta}}\right)+\eta\left[\boldsymbol{I}-(\boldsymbol{I}-\eta\boldsymbol{Z}^{\top}\boldsymbol{Z})^{t}\right]\left(\eta\boldsymbol{Z}^{\top}\boldsymbol{Z}\right)^{-1}\boldsymbol{Z}^{\top}\left(\boldsymbol{X}-\boldsymbol{Z\hat{\Theta}}\right)\\ &=\left(\boldsymbol{I}-\eta\boldsymbol{Z}^{\top}\boldsymbol{Z}\right)^{t}\left(\boldsymbol{\Theta}^{(0)}-\hat{\boldsymbol{\Theta}}\right)+\left[\boldsymbol{I}-(\boldsymbol{I}-\eta\boldsymbol{Z}^{\top}\boldsymbol{Z})^{t}\right]\left[(\boldsymbol{Z}^{\top}\boldsymbol{Z})^{-1}\boldsymbol{Z}^{\top}\boldsymbol{X}-\hat{\boldsymbol{\Theta}}\right]\\ &=\left(\boldsymbol{I}-\eta\boldsymbol{Z}^{\top}\boldsymbol{Z}\right)^{t}\left(\boldsymbol{\Theta}^{(0)}-\hat{\boldsymbol{\Theta}}\right).\end{split}$ $$

With learning rate $0<\eta<\frac{2}{\sigma_{\max}^{2}(\boldsymbol{Z})}$, let $\kappa(\eta):=\rho\left(\boldsymbol{I}-\eta\boldsymbol{Z}^{\top}\boldsymbol{Z}\right)$, where $\rho(\cdot)$ denotes the spectral radius. Then it follows that $0<\kappa(\eta)<1$. We have:

$$ $\displaystyle\left\|\boldsymbol{\Theta}^{(t)}-\hat{\boldsymbol{\Theta}}\right\|$ $\displaystyle=\left\|(\boldsymbol{I}-\eta\boldsymbol{Z}^{\top}\boldsymbol{Z})^{t}(\boldsymbol{\Theta}^{(0)}-\hat{\boldsymbol{\Theta}})\right\|$ $\displaystyle\leq\left\|(\boldsymbol{I}-\eta\boldsymbol{Z}^{\top}\boldsymbol{Z})^{t}\right\|\left\|\boldsymbol{\Theta}^{(0)}-\hat{\boldsymbol{\Theta}}\right\|$ $\displaystyle\leq\kappa(\eta)^{t}\left\|\boldsymbol{\Theta}^{(0)}-\hat{\boldsymbol{\Theta}}\right\|$ $\displaystyle=\mathcal{O}(\kappa(\eta)^{t}).$ $$

Thus $\{\boldsymbol{\Theta}^{(t)}\}$ converges to $\hat{\boldsymbol{\Theta}}$ exponentially with rate $\kappa(\eta)$.

For the outer loop ([8b](#S3.E8.2)), we have:

$$ $\displaystyle\begin{split}\boldsymbol{\beta}^{(t)}-\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}&=\boldsymbol{\beta}^{(t-1)}-\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}-\alpha\boldsymbol{\Theta}^{(t-1)\top}\boldsymbol{Z}^{\top}\left(\boldsymbol{Z\Theta}^{(t-1)}\boldsymbol{\beta}^{(t-1)}-\boldsymbol{Y}\right)\\ &=\left(\boldsymbol{I}-\alpha\boldsymbol{\Theta}^{(t-1)\top}\boldsymbol{Z}^{\top}\boldsymbol{Z}\boldsymbol{\Theta}^{(t-1)}\right)\left(\boldsymbol{\beta}^{(t-1)}-\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}\right)+\alpha\boldsymbol{\Theta}^{(t-1)\top}\boldsymbol{Z}^{\top}\left(\boldsymbol{Y}-\boldsymbol{Z\Theta}^{(t-1)}\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}\right)\\ &\quad\quad\quad\quad\quad\quad\quad\quad\vdots\\ &=\underbrace{\prod_{i=0}^{t-1}\left(\boldsymbol{I}-\alpha\boldsymbol{\Theta}^{(i)\top}\boldsymbol{Z}^{\top}\boldsymbol{Z\Theta}^{(i)}\right)\left(\boldsymbol{\beta}^{(0)}-\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}\right)}_{\boldsymbol{\Delta}_{1}\boldsymbol{\beta}^{(t)}}\\ &\quad+\underbrace{\sum_{i=0}^{t-1}\alpha\left[\prod_{j=i+1}^{t-1}\left(\boldsymbol{I}-\alpha\boldsymbol{\Theta}^{(j)\top}\boldsymbol{Z}^{\top}\boldsymbol{Z\Theta}^{(j)}\right)\right]\boldsymbol{\Theta}^{(i)\top}\boldsymbol{Z}^{\top}\left(\boldsymbol{Y}-\boldsymbol{Z\Theta}^{(i)}\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}\right)}_{\boldsymbol{\Delta}_{2}\boldsymbol{\beta}^{(t)}}.\end{split}$ (30) $$

To simplify notations, let

$$ $\displaystyle\boldsymbol{R}^{(t)}$ $\displaystyle\coloneqq\boldsymbol{\Theta}^{(t)}-\hat{\boldsymbol{\Theta}}=\left(\boldsymbol{I}-\eta\boldsymbol{Z}^{\top}\boldsymbol{Z}\right)^{t}\left(\boldsymbol{\Theta}^{(0)}-\hat{\boldsymbol{\Theta}}\right),$ $\displaystyle\boldsymbol{V}^{(t)}$ $\displaystyle\coloneqq\left(\boldsymbol{I}-\alpha\hat{\boldsymbol{\Theta}}^{\top}\boldsymbol{Z}^{\top}\boldsymbol{Z}\hat{\boldsymbol{\Theta}}\right)^{t},$ $\displaystyle\boldsymbol{W}^{(t)}$ $\displaystyle\coloneqq\boldsymbol{R}^{(t)\top}\boldsymbol{Z}^{\top}\boldsymbol{Z}\hat{\boldsymbol{\Theta}}+\hat{\boldsymbol{\Theta}}^{\top}\boldsymbol{Z}^{\top}\boldsymbol{Z}\boldsymbol{R}^{(t)}+\boldsymbol{R}^{(t)\top}\boldsymbol{Z}^{\top}\boldsymbol{Z}\boldsymbol{R}^{(t)}.$ $$

With learning rates $0<\alpha<\frac{2}{\sigma_{\max}^{2}(\boldsymbol{Z}\hat{\boldsymbol{\Theta}})}$, $0<\eta<\frac{2}{\sigma_{\max}^{2}(\boldsymbol{Z})}$, let $\gamma(\alpha):=\rho\left(\boldsymbol{I}-\alpha\hat{\boldsymbol{\Theta}}^{\top}\boldsymbol{Z}^{\top}\boldsymbol{Z}\hat{\boldsymbol{\Theta}}\right)$. Then it follows that $0<\gamma(\alpha)<1$. We have:

$$ $\displaystyle\left\|\boldsymbol{R}^{(t)}\right\|\leq\kappa(\eta)^{t}\left\|\boldsymbol{\Theta}^{(0)}-\hat{\boldsymbol{\Theta}}\right\|,$ (31) $$

$$ $\displaystyle\left\|\boldsymbol{V}^{(t)}\right\|\leq\gamma(\alpha)^{t},$ $$

and

$$ $\displaystyle\begin{split}\left\|\boldsymbol{W}^{(t)}\right\|&=\left\|\boldsymbol{R}^{(t)\top}\boldsymbol{Z}^{\top}\boldsymbol{Z}\hat{\boldsymbol{\Theta}}+\hat{\boldsymbol{\Theta}}^{\top}\boldsymbol{Z}^{\top}\boldsymbol{ZR}^{(t)}+\boldsymbol{R}^{(t)\top}\boldsymbol{Z}^{\top}\boldsymbol{ZR}^{(t)}\right\|\\ &\leq 2\left\|\hat{\boldsymbol{\Theta}}^{\top}\boldsymbol{Z}^{\top}\boldsymbol{Z}\right\|\left\|\boldsymbol{R}^{(t)}\right\|+\left\|\boldsymbol{Z}^{\top}\boldsymbol{Z}\right\|\left\|\boldsymbol{R}^{(t)}\right\|^{2}\\ &\leq 2\kappa(\eta)^{t}\left\|\hat{\boldsymbol{\Theta}}\boldsymbol{Z}^{\top}\boldsymbol{Z}\right\|\left\|\boldsymbol{\Theta}^{(0)}-\hat{\boldsymbol{\Theta}}\right\|+\kappa(\eta)^{2t}\left\|\boldsymbol{Z}^{\top}\boldsymbol{Z}\right\|\\ &\leq\kappa(\eta)^{t}\left(2\left\|\hat{\boldsymbol{\Theta}}\boldsymbol{Z}^{\top}\boldsymbol{Z}\right\|\left\|\boldsymbol{\Theta}^{(0)}-\hat{\boldsymbol{\Theta}}\right\|+\left\|\boldsymbol{Z}^{\top}\boldsymbol{Z}\right\|\left\|\boldsymbol{\Theta}^{(0)}-\hat{\boldsymbol{\Theta}}\right\|^{2}\right)\\ &=\mathcal{O}(\kappa(\eta)^{t}).\end{split}$ $$

Then from equation ([30](#A2.E30)), we have:

$$ $\displaystyle\begin{split}\boldsymbol{\Delta}_{1}\boldsymbol{\beta}^{(t)}&=\prod_{i=0}^{t-1}\left(\boldsymbol{I}-\alpha\boldsymbol{\Theta}^{(i)\top}\boldsymbol{Z}^{\top}\boldsymbol{Z\Theta}^{(i)}\right)\left(\boldsymbol{\beta}^{(0)}-\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}\right)\\ &=\prod_{i=0}^{t-1}\left[\boldsymbol{I}-\alpha\hat{\boldsymbol{\Theta}}^{\top}\boldsymbol{Z}^{\top}\boldsymbol{Z}\hat{\boldsymbol{\Theta}}-\alpha\left(\boldsymbol{R}^{(i)\top}\boldsymbol{Z}^{\top}\boldsymbol{Z}\hat{\boldsymbol{\Theta}}+\hat{\boldsymbol{\Theta}}^{\top}\boldsymbol{Z}^{\top}\boldsymbol{Z}\boldsymbol{R}^{(i)}+\boldsymbol{R}^{(i)\top}\boldsymbol{Z}^{\top}\boldsymbol{Z}\boldsymbol{R}^{(i)}\right)\right]\left(\boldsymbol{\beta}^{(0)}-\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}\right)\\ &=\prod_{i=0}^{t-1}\left[\boldsymbol{I}-\alpha\hat{\boldsymbol{\Theta}}^{\top}\boldsymbol{Z}^{\top}\boldsymbol{Z}\hat{\boldsymbol{\Theta}}-\alpha\boldsymbol{W}^{(i)}\right]\left(\boldsymbol{\beta}^{(0)}-\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}\right)\\ &=\left(\boldsymbol{I}-\alpha\hat{\boldsymbol{\Theta}}^{\top}\boldsymbol{Z}^{\top}\boldsymbol{Z}\hat{\boldsymbol{\Theta}}\right)^{t}\prod_{i=0}^{t-1}\left[\boldsymbol{I}-\alpha\left(\boldsymbol{I}-\alpha\hat{\boldsymbol{\Theta}}^{\top}\boldsymbol{Z}^{\top}\boldsymbol{Z}\hat{\boldsymbol{\Theta}}\right)^{-1}\boldsymbol{W}^{(i)}\right]\left(\boldsymbol{\beta}^{(0)}-\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}\right)\\ &=\boldsymbol{V}^{(t)}\prod_{i=0}^{t-1}\left[\boldsymbol{I}-\alpha\left(\boldsymbol{I}-\alpha\hat{\boldsymbol{\Theta}}^{\top}\boldsymbol{Z}^{\top}\boldsymbol{Z}\hat{\boldsymbol{\Theta}}\right)^{-1}\boldsymbol{W}^{(i)}\right]\left(\boldsymbol{\beta}^{(0)}-\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}\right).\end{split}$ $$

We denote $\boldsymbol{\Psi}:=\alpha\left(\boldsymbol{I}-\alpha\hat{\boldsymbol{\Theta}}^{\top}\boldsymbol{Z}^{\top}\boldsymbol{Z}\hat{\boldsymbol{\Theta}}\right)^{-1}$. By Lemma  [B.1](#A2.Thmlemma1), we take $\varepsilon=1$, $c_{0}$ be a constant such that $\left\|\boldsymbol{W}^{(t)}\right\|_{F}\leq c_{0}\kappa(\eta)^{t}$, and $T_{0}=\lceil\log_{\kappa(\eta)}\frac{(1-\kappa(\eta))}{\|\Psi\|_{F}c_{0}(p+(1-\kappa(\eta)))}\rceil$.

Then we have:

$$ $\displaystyle\left\|\prod_{i=T_{0}}^{t-1}\left(\boldsymbol{I}-\boldsymbol{\Psi W}^{(i)}\right)\right\|\leq\left\|\prod_{i=T_{0}}^{t-1}\left(\boldsymbol{I}-\boldsymbol{\Psi W}^{(i)}\right)\right\|_{F}<\sqrt{p}+1.$ (32) $$

Hence

$$ $\displaystyle\begin{split}\left\|\boldsymbol{\Delta}_{1}\boldsymbol{\beta}^{(t)}\right\|&=\left\|\boldsymbol{V}^{(t)}\prod_{i=0}^{t-1}(\boldsymbol{I}-\boldsymbol{\Psi W}^{(i)})(\boldsymbol{\beta}^{(0)}-\hat{\boldsymbol{\beta}}_{\textsf{2SLS}})\right\|\\ &\leq\left\|\boldsymbol{V}^{(t)}\right\|\left\|\prod_{i=0}^{T_{0}-1}\left(\boldsymbol{I}-\boldsymbol{\Psi W}^{(i)}\right)\right\|\left\|\prod_{i=T_{0}}^{t-1}(\boldsymbol{I}-\boldsymbol{\Psi W}^{(i)})\right\|\left\|\boldsymbol{\beta}^{(0)}-\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}\right\|\\ &<\gamma(\alpha)^{t}\left\|\prod_{i=0}^{T_{0}-1}\left(\boldsymbol{I}-\boldsymbol{\Psi W}^{(i)}\right)\right\|\left(\sqrt{p}+1\right)\left\|\boldsymbol{\beta}^{(0)}-\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}\right\|\\ &=\mathcal{O}(\gamma(\alpha)^{t}).\end{split}$ (33) $$

Next we consider $\boldsymbol{\Delta}_{2}\boldsymbol{\beta}^{(t)}$:

$$ $\displaystyle\begin{split}\boldsymbol{\Delta}_{2}\boldsymbol{\beta}^{(t)}&=\sum_{i=0}^{t-1}\alpha\left[\prod_{j=i+1}^{t-1}\left(\boldsymbol{I}-\alpha\boldsymbol{\Theta}^{(j)\top}\boldsymbol{Z}^{\top}\boldsymbol{Z\Theta}^{(j)}\right)\right]\boldsymbol{\Theta}^{(i)\top}\boldsymbol{Z}^{\top}\left(\boldsymbol{Y}-\boldsymbol{Z\Theta}^{(i)}\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}\right)\\ &=\sum_{i=0}^{t-1}\alpha\left[\prod_{j=i+1}^{t-1}\left(\boldsymbol{I}-\alpha\hat{\boldsymbol{\Theta}}^{\top}\boldsymbol{Z}^{\top}\boldsymbol{Z}\hat{\boldsymbol{\Theta}}-\alpha\boldsymbol{W}^{(j)}\right)\right]\left(\boldsymbol{R}^{(i)}+\hat{\boldsymbol{\Theta}}\right)^{\top}\boldsymbol{Z}^{\top}\left[\boldsymbol{Y}-\boldsymbol{Z}\left(\boldsymbol{R}^{(i)}+\hat{\boldsymbol{\Theta}}\right)\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}\right]\\ &=\sum_{i=0}^{t-1}\alpha\left(\boldsymbol{I}-\alpha\hat{\boldsymbol{\Theta}}^{\top}\boldsymbol{Z}^{\top}\boldsymbol{Z}\hat{\boldsymbol{\Theta}}\right)^{t-1-i}\prod_{j=i+1}^{t-1}\left[\boldsymbol{I}-\alpha\left(\boldsymbol{I}-\alpha\hat{\boldsymbol{\Theta}}^{\top}\boldsymbol{Z}^{\top}\boldsymbol{Z}\hat{\boldsymbol{\Theta}}\right)^{-1}\boldsymbol{W}^{(j)}\right]\\ &\quad\quad\quad\quad\cdot\left(\boldsymbol{R}^{(i)}+\hat{\boldsymbol{\Theta}}\right)^{\top}\boldsymbol{Z}^{\top}\left[\boldsymbol{Y}-\boldsymbol{Z}\left(\boldsymbol{R}^{(i)}+\hat{\boldsymbol{\Theta}}\right)\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}\right]\\ &=\sum_{i=0}^{t-1}\alpha\left[\boldsymbol{V}^{(t-1-i)}\prod_{j=i+1}^{t-1}\left(\boldsymbol{I}-\boldsymbol{\Psi}\boldsymbol{W}^{(j)}\right)\right]\left(\boldsymbol{R}^{(i)}+\hat{\boldsymbol{\Theta}}\right)^{\top}\boldsymbol{Z}^{\top}\left[\boldsymbol{Y}-\boldsymbol{Z}\left(\boldsymbol{R}^{(i)}+\hat{\boldsymbol{\Theta}}\right)\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}\right].\end{split}$ $$

For convenience, let $\boldsymbol{\Delta}_{2}\boldsymbol{\beta}^{(t)}:=\boldsymbol{\Delta}_{21}\boldsymbol{\beta}^{(t)}+\boldsymbol{\Delta}_{22}\boldsymbol{\beta}^{(t)}$, where

$$ $\displaystyle\boldsymbol{\Delta}_{21}\boldsymbol{\beta}^{(t)}:$ $\displaystyle=\sum_{i=0}^{t-1}\alpha\left[\boldsymbol{V}^{(t-1-i)}\prod_{j=i+1}^{t-1}\left(\boldsymbol{I}-\boldsymbol{\Psi}\boldsymbol{W}^{(j)}\right)\right]\boldsymbol{R}^{(i)\top}\boldsymbol{Z}^{\top}\left[\boldsymbol{Y}-\boldsymbol{Z}\left(\boldsymbol{R}^{(i)}+\hat{\boldsymbol{\Theta}}\right)\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}\right],$ $\displaystyle\boldsymbol{\Delta}_{22}\boldsymbol{\beta}^{(t)}:$ $\displaystyle=\sum_{i=0}^{t-1}\alpha\left[\boldsymbol{V}^{(t-1-i)}\prod_{j=i+1}^{t-1}\left(\boldsymbol{I}-\boldsymbol{\Psi}\boldsymbol{W}^{(j)}\right)\right]\hat{\boldsymbol{\Theta}}^{\top}\boldsymbol{Z}^{\top}\left[\boldsymbol{Y}-\boldsymbol{Z}\left(\boldsymbol{R}^{(i)}+\hat{\boldsymbol{\Theta}}\right)\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}\right].$ $$

Suppose $\tilde{M}_{1},\tilde{M}_{2}$ are the upper bounds such that

$$ $\displaystyle\left\|\boldsymbol{Z}^{\top}\left[\boldsymbol{Y}-\boldsymbol{Z}\left(\boldsymbol{R}^{(i)}+\hat{\boldsymbol{\Theta}}\right)\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}\right]\right\|\leq\tilde{M}_{1},\quad\forall i=0,\ldots,t-1,$ $$

$$ $\displaystyle\left\|\prod_{j=i+1}^{t-1}\left(\boldsymbol{I}-\boldsymbol{\Psi}\boldsymbol{W}^{(j)}\right)\right\|\leq\tilde{M}_{2},\quad\forall i=0,\ldots,t-1.$ $$

We know such $\tilde{M}_{1},\tilde{M}_{2}$ exist because of the bounds given by ([31](#A2.E31)) and ([32](#A2.E32)). Let $\tilde{M}=\tilde{M}_{1}\tilde{M}_{2}$. Then

$$ $\displaystyle\begin{split}\left\|\boldsymbol{\Delta}_{21}\boldsymbol{\beta}^{(t)}\right\|&\leq\tilde{M}\left\|\sum_{i=0}^{t-1}\alpha\boldsymbol{V}^{(t-1-i)}\boldsymbol{R}^{(i)\top}\right\|\\ &\leq\tilde{M}\alpha\sum_{i=0}^{t-1}\left\|\boldsymbol{V}^{(t-1-i)}\right\|\left\|\boldsymbol{R}^{(i)}\right\|\\ &\leq\tilde{M}\alpha\left\|\boldsymbol{\Theta}^{(0)}-\hat{\boldsymbol{\Theta}}\right\|\sum_{i=0}^{t-1}\gamma(\alpha)^{t-1-i}\kappa(\eta)^{i}\\ &=\tilde{M}\alpha\left\|\boldsymbol{\Theta}^{(0)}-\hat{\boldsymbol{\Theta}}\right\|\sum_{i=0}^{t-1}\gamma(\alpha)^{t-1}\left(\frac{\kappa(\eta)}{\gamma(\alpha)}\right)^{i}\\ &=\mathcal{O}\left(\frac{\gamma(\alpha)^{t}-\kappa(\eta)^{t}}{\gamma(\alpha)-\kappa(\eta)}\right)\\ &\leq\mathcal{O}(\max\{\gamma(\alpha)^{t},\kappa(\eta)^{t}\}),\end{split}$ $$

and similarly,

$$ $\displaystyle\begin{split}\left\|\boldsymbol{\Delta}_{22}\boldsymbol{\beta}^{(t)}\right\|&\leq\tilde{M}\left\|\sum_{i=0}^{t-1}\alpha\boldsymbol{V}^{(t-1-i)}\hat{\boldsymbol{\Theta}}^{T}\right\|\\ &\leq\tilde{M}\alpha\left\|\hat{\boldsymbol{\Theta}}\right\|\sum_{i=0}^{t-1}\left\|\boldsymbol{V}^{(t-1-i)}\right\|\\ &\leq\tilde{M}\alpha\left\|\hat{\boldsymbol{\Theta}}\right\|\sum_{i=0}^{t-1}\gamma(\alpha)^{t-1-i}\\ &=\mathcal{O}(\gamma(\alpha)^{t}).\end{split}$ $$

Thus

$$ $\displaystyle\begin{split}\left\|\boldsymbol{\Delta}_{2}\boldsymbol{\beta}^{(t)}\right\|&=\left\|\boldsymbol{\Delta}_{21}\boldsymbol{\beta}^{(t)}+\boldsymbol{\Delta}_{22}\boldsymbol{\beta}^{(t)}\right\|\\ &\leq\left\|\boldsymbol{\Delta}_{21}\boldsymbol{\beta}^{(t)}\right\|+\left\|\boldsymbol{\Delta}_{22}\boldsymbol{\beta}^{(t)}\right\|\\ &\leq\mathcal{O}(\max\{\gamma(\alpha)^{t},\kappa(\eta)^{t}\}).\end{split}$ (34) $$

Therefore, plugging ([33](#A2.E33)) and ([34](#A2.E34)) into ([30](#A2.E30)), we have:

$$ $\displaystyle\left\|\boldsymbol{\beta}^{(t)}-\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}\right\|\leq\mathcal{O}(\max\{\gamma(\alpha)^{t},\kappa(\eta)^{t}\}),$ $$

Hence completes the proof.
∎

###### Lemma B.1 .

###### Proof.

###### Proof of Theorem 3.1 .

### B.2 Proof of Theorem 3.2

###### Proof of Theorem 3.2 .

For ease of notations, we ignore $l$ in the following proof. Consider the input matrix taking the form:

$$ $\displaystyle\boldsymbol{H}^{(0)}=\begin{bmatrix}\boldsymbol{z}_{1}&\cdots&\boldsymbol{z}_{n}&\boldsymbol{z}_{n+1}\\ \boldsymbol{x}_{1}&\cdots&\boldsymbol{x}_{n}&\boldsymbol{x}_{n+1}\\ y_{1}&\cdots&y_{n}&0\\ \boldsymbol{\Theta}_{:,1}^{(0)}&\cdots&\boldsymbol{\Theta}_{:,1}^{(0)}&\boldsymbol{\Theta}_{:,1}^{(0)}\\ \vdots&\vdots&\vdots&\vdots\\ \boldsymbol{\Theta}_{:,p}^{(0)}&\cdots&\boldsymbol{\Theta}_{:,p}^{(0)}&\boldsymbol{\Theta}_{:,p}^{(0)}\\ \boldsymbol{\beta}^{(0)}&\cdots&\boldsymbol{\beta}^{(0)}&\boldsymbol{\beta}^{(0)}\\ \hat{\boldsymbol{x}}_{1}^{(0)}&\cdots&\hat{\boldsymbol{x}}_{n}^{(0)}&\hat{\boldsymbol{x}}_{n+1}^{(0)}\\ 1&\cdots&1&1\\ 1&\cdots&1&0\end{bmatrix}\in\mathbb{R}^{D\times(n+1)},$ $$

i.e., element-wise,

$$ $\displaystyle\boldsymbol{h}_{i}^{(0)}=\left(\boldsymbol{z}_{i},\boldsymbol{x}_{i},y_{i}t_{i},\boldsymbol{\Theta}_{:,1}^{(0)},\ldots,\boldsymbol{\Theta}_{:,p}^{(0)},\boldsymbol{\beta}^{(0)},\hat{\boldsymbol{x}}_{i}^{(0)},1,t_{i}\right)^{\top},\quad i=1,\ldots,n+1,$ $$

where $D=qp+3p+q+3$, $t_{i}:=\mathbbm{1}\{i\leq n\}$ is the indicator for training sample. We can take any initialization for $\boldsymbol{\Theta}^{(0)}$, $\boldsymbol{\beta}^{(0)}$ and $\hat{\boldsymbol{x}}^{(0)}$. To avoid abuse of notations, we omit the superscript of those parameters to be updated in the following proof.

Recall the definitions ([5](#S3.E5)) and ([6](#S3.E6)). Our goal is to show that there exists a series of attention parameters $\boldsymbol{\theta}_{\textsf{ATTN}}^{(1:2)}=\{(\boldsymbol{Q}_{m}^{(1:2)},\boldsymbol{K}_{m}^{(1:2)},\boldsymbol{V}_{m}^{(1:2)})\}_{m\in[M]}\subset\mathbb{R}^{D\times D}$ such that $\boldsymbol{\theta}_{\textsf{ATTN}}^{(1:2)}$ updates $\boldsymbol{\Theta},\boldsymbol{\beta}$ on the corresponding rows. i.e, if we denote $D_{0}:=q+p+1$, the updates on row $D_{0}+1$ to row $D_{0}+qp$ correspond to $\boldsymbol{\Theta}$, and the updates on row $D_{0}+qp+1$ to row $D_{0}+qp+p$ correspond to $\boldsymbol{\beta}$.

1) In the first layer, the transformer updates the current first-stage estimate $\hat{\boldsymbol{x}}$.

For $m=2k-1,k=1,\ldots,p$, define $\boldsymbol{Q}_{m}^{(1)},\boldsymbol{K}_{m}^{(1)},\boldsymbol{V}_{m}^{(1)}$ such that:

$$ $\displaystyle\boldsymbol{Q}_{m}^{(1)}\boldsymbol{h}_{i}^{(0)}=\begin{bmatrix}z_{i1}\\ \vdots\\ z_{iq}\\ \hat{x}_{ik}^{(0)}\\ \boldsymbol{0}\end{bmatrix},\boldsymbol{K}_{m}^{(1)}\boldsymbol{h}_{j}^{(0)}=\begin{bmatrix}\boldsymbol{\Theta}_{1k}^{(0)}\\ \vdots\\ \boldsymbol{\Theta}_{qk}^{(0)}\\ -1\\ \boldsymbol{0}\end{bmatrix},\boldsymbol{V}_{m}^{(1)}\boldsymbol{h}_{j}^{(0)}=\boldsymbol{e}_{D_{0}+qp+p+k}.$ (35) $$

For $m=2k,k=1,\ldots,p$, define $\boldsymbol{Q}_{m}^{(1)},\boldsymbol{K}_{m}^{(1)},\boldsymbol{V}_{m}^{(1)}$ such that:

$$ $\displaystyle\boldsymbol{Q}_{m}^{(1)}\boldsymbol{h}_{i}^{(0)}=\begin{bmatrix}-z_{i1}\\ \vdots\\ -z_{iq}\\ \hat{x}_{ik}^{(0)}\\ \boldsymbol{0}\end{bmatrix},\boldsymbol{K}_{m}^{(1)}\boldsymbol{h}_{j}^{(0)}=\begin{bmatrix}\boldsymbol{\Theta}_{1k}^{(0)}\\ \vdots\\ \boldsymbol{\Theta}_{qk}^{(0)}\\ 1\\ \boldsymbol{0}\end{bmatrix},\boldsymbol{V}_{m}^{(1)}\boldsymbol{h}_{j}^{(0)}=-\boldsymbol{e}_{D_{0}+qp+p+k},$ (36) $$

where $\boldsymbol{e}_{j}\in\mathbb{R}^{D}$ is the standard unit vector with only one 1 at the $j$-th coordinate. Note that the above are just linear transformations on $\boldsymbol{h}_{i}$ or $\boldsymbol{h}_{j}$, hence such matrices $\boldsymbol{Q}_{m}^{(1)},\boldsymbol{K}_{m}^{(1)},\boldsymbol{V}_{m}^{(1)}$ must exist. Then we have:

$$ $\displaystyle\begin{split}\boldsymbol{h}_{i}^{(1)}&=\boldsymbol{h}_{i}^{(0)}+\sum_{m=1}^{2p}\frac{1}{n+1}\sum_{j=1}^{n+1}\sigma\left(\langle\boldsymbol{Q}_{m}^{(1)}\boldsymbol{h}_{i}^{(0)},\boldsymbol{K}_{m}^{(1)}\boldsymbol{h}_{j}^{(0)}\rangle\right)\cdot\boldsymbol{V}_{m}^{(1)}\boldsymbol{h}_{j}^{(0)}\\ &=\boldsymbol{h}_{i}^{(0)}+\sum_{k=1}^{p}\frac{1}{n+1}\sum_{j=1}^{n+1}\left[\sigma\left(\sum_{l=1}^{q}z_{il}\boldsymbol{\Theta}_{lk}^{(0)}-\hat{x}_{ik}^{(0)}\right)-\sigma\left(-\sum_{l=1}^{q}z_{il}\boldsymbol{\Theta}_{lk}^{(0)}+\hat{x}_{ik}^{(0)}\right)\right]\cdot\boldsymbol{e}_{D_{0}+qp+p+k}\\ &=\boldsymbol{h}_{i}^{(0)}+\sum_{k=1}^{p}\left[\sum_{l=1}^{q}z_{il}\boldsymbol{\Theta}_{lk}^{(0)}-\hat{x}_{ik}^{(0)}\right]\boldsymbol{e}_{D_{0}+qp+p+k}\\ &=\boldsymbol{h}_{i}^{(0)}+\sum_{k=1}^{p}\left(\hat{x}_{ik}^{(1)}-\hat{x}_{ik}^{(0)}\right)\boldsymbol{e}_{D_{0}+qp+p+k}.\end{split}$ $$

Thus this layer correctly updates the first-stage prediction values $\hat{\boldsymbol{x}}^{(1)}_{1},\ldots,\hat{\boldsymbol{x}}_{n+1}^{(1)}$, where $\hat{\boldsymbol{x}}_{i}^{(1)}:=[\boldsymbol{Z\Theta}^{(0)}]_{i,:}=\sum_{l=1}^{q}z_{il}\boldsymbol{\Theta}_{l,:}^{(0)}$. We will have:

$$ $\displaystyle\boldsymbol{H}^{(1)}=\begin{bmatrix}\boldsymbol{z}_{1}&\cdots&\boldsymbol{z}_{n}&\boldsymbol{z}_{n+1}\\ \boldsymbol{x}_{1}&\cdots&\boldsymbol{x}_{n}&\boldsymbol{x}_{n+1}\\ y_{1}&\cdots&y_{n}&0\\ \boldsymbol{\Theta}_{:,1}^{(0)}&\cdots&\boldsymbol{\Theta}_{:,1}^{(0)}&\boldsymbol{\Theta}_{:,1}^{(0)}\\ \vdots&\vdots&\vdots&\vdots\\ \boldsymbol{\Theta}_{:,p}^{(0)}&\cdots&\boldsymbol{\Theta}_{:,p}^{(0)}&\boldsymbol{\Theta}_{:,p}^{(0)}\\ \boldsymbol{\beta}^{(0)}&\cdots&\boldsymbol{\beta}^{(0)}&\boldsymbol{\beta}^{(0)}\\ \hat{\boldsymbol{x}}_{1}^{(1)}&\cdots&\hat{\boldsymbol{x}}_{n}^{(1)}&\hat{\boldsymbol{x}}_{n+1}^{(1)}\\ 1&\cdots&1&1\\ 1&\cdots&1&0\end{bmatrix}.$ $$

2) In the second layer, the transformer does the gradient updates on the parameters $\boldsymbol{\Theta}$ and $\boldsymbol{\beta}$.

For $m=2k-1,k=1,\ldots,p,$ define $\boldsymbol{Q}_{m}^{(2)},\boldsymbol{K}_{m}^{(2)},\boldsymbol{V}_{m}^{(2)}$ such that:

$$ $\displaystyle\boldsymbol{Q}_{m}^{(2)}\boldsymbol{h}_{i}^{(1)}=\begin{bmatrix}\boldsymbol{\Theta}_{:,k}^{(0)}\\ -1\\ -1\\ \vdots\\ \boldsymbol{0}\end{bmatrix},\boldsymbol{K}_{m}^{(2)}\boldsymbol{h}_{j}^{(1)}=\begin{bmatrix}\boldsymbol{z}_{j}\\ x_{jk}t_{j}\\ R(1-t_{j})\\ \vdots\\ \boldsymbol{0}\end{bmatrix},\boldsymbol{V}_{m}^{(2)}\boldsymbol{h}_{j}^{(1)}=-(n+1)\eta\sum_{l=1}^{q}z_{jl}\boldsymbol{e}_{D_{0}+(k-1)q+l}.$ (37) $$

For $m=2k,k=1,\ldots,p$, define $\boldsymbol{Q}_{m}^{(2)},\boldsymbol{K}_{m}^{(2)},\boldsymbol{V}_{m}^{(2)}$ such that:

$$ $\displaystyle\boldsymbol{Q}_{m}^{(2)}\boldsymbol{h}_{i}^{(1)}=\begin{bmatrix}-\boldsymbol{\Theta}_{:,k}^{(0)}\\ 1\\ -1\\ \vdots\\ \boldsymbol{0}\end{bmatrix},\boldsymbol{K}_{m}^{(2)}\boldsymbol{h}_{j}^{(1)}=\begin{bmatrix}\boldsymbol{z}_{j}\\ x_{jk}t_{j}\\ R(1-t_{j})\\ \vdots\\ \boldsymbol{0}\end{bmatrix},\boldsymbol{V}_{m}^{(2)}\boldsymbol{h}_{j}^{(1)}=(n+1)\eta\sum_{l=1}^{q}z_{jl}\boldsymbol{e}_{D_{0}+(k-1)q+l},$ (38) $$

where $R=\underset{\begin{subarray}{c}i=1,\ldots,n+1\\
t=0,1,\ldots\end{subarray}}{\max}\{||\boldsymbol{\Theta}^{(t)\top}\boldsymbol{z}_{i}||\}$. Then we have:

$$ $\displaystyle\begin{split}\sigma\left(\langle\boldsymbol{Q}_{2k-1}^{(2)}\boldsymbol{h}_{i}^{(1)},\boldsymbol{K}_{2k-1}^{(2)}\boldsymbol{h}_{j}^{(1)}\rangle\right)&=\sigma\left(\boldsymbol{\Theta}_{:,k}^{(0)\top}\boldsymbol{z}_{j}-x_{jk}t_{j}-R(1-t_{j})\right)\\ &=\sigma\left(\boldsymbol{\Theta}_{:,k}^{(0)\top}\boldsymbol{z}_{j}-x_{jk}\right)\mathbbm{1}\{t_{j}=1\}\\ &=\sigma\left(\boldsymbol{\Theta}_{:,k}^{(0)\top}\boldsymbol{z}_{j}-x_{jk}\right)t_{j},\\ \end{split}$ $$

and

$$ $\displaystyle\begin{split}\sigma\left(\langle\boldsymbol{Q}_{2k}^{(2)}\boldsymbol{h}_{i}^{(1)},\boldsymbol{K}_{2k}^{(2)}\boldsymbol{h}_{j}^{(1)}\rangle\right)&=\sigma\left(-\boldsymbol{\Theta}_{:,k}^{(0)\top}\boldsymbol{z}_{j}+x_{jk}t_{j}-R(1-t_{j})\right)\\ &=\sigma\left(-\boldsymbol{\Theta}_{:,k}^{(0)\top}\boldsymbol{z}_{j}+x_{jk}\right)\mathbbm{1}\{t_{j}=1\}\\ &=\sigma\left(-\boldsymbol{\Theta}_{:,k}^{(0)\top}\boldsymbol{z}_{j}+x_{jk}\right)t_{j}.\end{split}$ $$

So that

$$ $\displaystyle\begin{split}&\sum_{m=1}^{2p}\sigma\left(\langle\boldsymbol{Q}_{m}^{(2)}\boldsymbol{h}_{i}^{(1)},\boldsymbol{K}_{m}^{(2)}\boldsymbol{h}_{j}^{(1)}\rangle\right)\boldsymbol{V}_{m}^{(2)}\boldsymbol{h}_{j}^{(1)}\\ &=-(n+1)t_{j}\eta\sum_{k=1}^{p}\left[\sigma\left(\boldsymbol{\Theta}_{:,k}^{(0)\top}\boldsymbol{z}_{j}-x_{jk}\right)-\sigma\left(-\boldsymbol{\Theta}_{:,k}^{(0)\top}\boldsymbol{z}_{j}+x_{jk}\right)\right]\cdot\sum_{l=1}^{q}z_{jl}\boldsymbol{e}_{D_{0}+(k-1)q+l}\\ &=-(n+1)t_{j}\eta\sum_{k=1}^{p}\sum_{l=1}^{q}z_{jl}\left(\boldsymbol{\Theta}_{:,k}^{(0)\top}\boldsymbol{z}_{j}-x_{jk}\right)\boldsymbol{e}_{D_{0}+(k-1)q+l}.\end{split}$ $$

Similarly, for $m=2p+1,2p+2$, define $\boldsymbol{Q}_{m}^{(2)},\boldsymbol{K}_{m}^{(2)},\boldsymbol{V}_{m}^{(2)}$ such that:

$$ $\displaystyle\boldsymbol{Q}_{2p+1}^{(2)}\boldsymbol{h}_{i}^{(1)}=\begin{bmatrix}\boldsymbol{\beta}^{(0)}\\ -1\\ -1\\ \vdots\\ \boldsymbol{0}\end{bmatrix},\boldsymbol{K}_{2p+1}^{(2)}\boldsymbol{h}_{j}^{(1)}=\begin{bmatrix}\hat{\boldsymbol{x}}_{j}^{(1)}\\ y_{j}t_{j}\\ R^{\prime}(1-t_{j})\\ \vdots\\ \boldsymbol{0}\end{bmatrix},\boldsymbol{V}_{2p+1}\boldsymbol{h}_{j}^{(1)}=-(n+1)\alpha\sum_{l=1}^{p}\hat{x}_{jl}^{(1)}\boldsymbol{e}_{D_{0}+qp+l},$ (39) $$

$$ $\displaystyle\boldsymbol{Q}_{2p+2}^{(2)}\boldsymbol{h}_{i}^{(1)}=\begin{bmatrix}-\boldsymbol{\beta}^{(0)}\\ 1\\ -1\\ \vdots\\ \boldsymbol{0}\end{bmatrix},\boldsymbol{K}_{2p+2}^{(2)}\boldsymbol{h}_{j}^{(1)}=\begin{bmatrix}\hat{\boldsymbol{x}}_{j}^{(1)}\\ y_{j}t_{j}\\ R^{\prime}(1-t_{j})\\ \vdots\\ \boldsymbol{0}\end{bmatrix},\boldsymbol{V}_{2p+2}\boldsymbol{h}_{j}^{(1)}=(n+1)\alpha\sum_{l=1}^{p}\hat{x}_{jl}^{(1)}\boldsymbol{e}_{D_{0}+qp+l},$ (40) $$

where $R^{\prime}=\underset{\begin{subarray}{c}i=1,\ldots,n+1\\
t=0,1,\ldots\end{subarray}}{\max}\{|\boldsymbol{\beta}^{(t)\top}\boldsymbol{x}_{i}|\}$. Then

$$ $\displaystyle\begin{split}\sigma\left(\langle\boldsymbol{Q}_{2p+1}^{(2)}\boldsymbol{h}_{i}^{(1)},\boldsymbol{K}_{2p+1}^{(2)}\boldsymbol{h}_{j}^{(1)}\rangle\right)&=\sigma\left(\boldsymbol{\beta}^{(0)\top}\hat{\boldsymbol{x}}_{j}^{(1)}-y_{j}t_{j}-R^{\prime}(1-t_{j})\right)\\ &=\sigma\left(\boldsymbol{\beta}^{(0)\top}\hat{\boldsymbol{x}}_{j}^{(1)}-y_{j}\right)\mathbbm{1}\{t_{j}=1\}\\ &=\sigma\left(\boldsymbol{\beta}^{(0)\top}\hat{\boldsymbol{x}}_{j}^{(1)}-y_{j}\right)t_{j},\end{split}$ $$

and

$$ $\displaystyle\begin{split}\sigma\left(\langle\boldsymbol{Q}_{2p+2}^{(2)}\boldsymbol{h}_{i}^{(1)},\boldsymbol{K}_{2p+2}^{(2)}\boldsymbol{h}_{j}^{(1)}\rangle\right)&=\sigma\left(-\boldsymbol{\beta}^{(0)\top}\hat{\boldsymbol{x}}_{j}^{(1)}+y_{j}t_{j}-R^{\prime}(1-t_{j})\right)\\ &=\sigma\left(-\boldsymbol{\beta}^{(0)\top}\hat{\boldsymbol{x}}_{j}^{(1)}+y_{j}\right)\mathbbm{1}\{t_{j}=1\}\\ &=\sigma\left(-\boldsymbol{\beta}^{(0)\top}\hat{\boldsymbol{x}}_{j}^{(1)}+y_{j}\right)t_{j}.\end{split}$ $$

So that

$$ $\displaystyle\begin{split}&\sum_{m=2p+1}^{2p+2}\sigma\left(\langle\boldsymbol{Q}_{m}^{(2)}\boldsymbol{h}_{i}^{(1)},\boldsymbol{K}_{m}^{(2)}\boldsymbol{h}_{j}^{(1)}\rangle\right)\boldsymbol{V}_{m}^{(2)}\boldsymbol{h}_{j}^{(1)}\\ &=-(n+1)t_{j}\alpha\left[\sigma\left(\boldsymbol{\beta}^{(0)\top}\hat{\boldsymbol{x}}_{j}^{(1)}-y_{j}\right)-\sigma\left(-\boldsymbol{\beta}^{(0)\top}\hat{\boldsymbol{x}}_{j}^{(1)}+y_{j}\right)\right]\cdot\sum_{l=1}^{p}\hat{x}_{jl}^{(1)}\boldsymbol{e}_{D_{0}+qp+l}\\ &=-(n+1)t_{j}\alpha\sum_{l=1}^{p}\hat{x}_{jl}^{(1)}\left(\boldsymbol{\beta}^{(0)\top}\hat{\boldsymbol{x}}_{j}^{(1)}-y_{j}\right)\boldsymbol{e}_{D_{0}+qp+l}.\end{split}$ $$

Thus the final output, for $i=1,\ldots,n+1$:

$$ $\displaystyle\begin{split}\boldsymbol{h}_{i}^{(2)}&=\boldsymbol{h}_{i}^{(1)}+\sum_{m=1}^{2p+2}\frac{1}{n+1}\sum_{j=1}^{n}\sigma\left(\langle\boldsymbol{Q}_{m}^{(2)}\boldsymbol{h}_{i}^{(1)},\boldsymbol{K}_{m}^{(2)}\boldsymbol{h}_{j}^{(1)}\rangle\right)\boldsymbol{V}_{m}^{(2)}\boldsymbol{h}_{j}^{(1)}\\ &=\begin{bmatrix}\boldsymbol{z}_{i}\\ \boldsymbol{x}_{i}\\ y_{i}t_{i}\\ \boldsymbol{\Theta}_{:,1}^{(0)}-\eta\sum_{j=1}^{n}\boldsymbol{z}_{j}\left(\boldsymbol{z}_{j}^{\top}\boldsymbol{\Theta}_{:,1}^{(0)}-x_{j1}\right)\\ \vdots\\ \boldsymbol{\Theta}_{:,p}^{(0)}-\eta\sum_{j=1}^{n}\boldsymbol{z}_{j}\left(\boldsymbol{z}_{j}^{\top}\boldsymbol{\Theta}_{:,p}^{(0)}-x_{jp}\right)\\ \boldsymbol{\beta}^{(0)}-\alpha\sum_{j=1}^{n}\hat{\boldsymbol{x}}_{j}^{(1)}\left(\hat{\boldsymbol{x}}_{j}^{(1)\top}\boldsymbol{\beta}^{(0)}-y_{j}\right)\\ \hat{\boldsymbol{x}}_{i}^{(1)}\\ 1\\ t_{i}\end{bmatrix}\\ &=\begin{bmatrix}\boldsymbol{z}_{i}\\ \boldsymbol{x}_{i}\\ y_{i}t_{i}\\ \boldsymbol{\Theta}_{:,1}^{(0)}-\eta\boldsymbol{Z}^{\top}\left[\boldsymbol{Z\Theta}^{(0)}-\boldsymbol{X}\right]_{:,1}\\ \vdots\\ \boldsymbol{\Theta}_{:,p}^{(0)}-\eta\boldsymbol{Z}^{\top}\left[\boldsymbol{Z\Theta}^{(0)}-\boldsymbol{X}\right]_{:,p}\\ \boldsymbol{\beta}^{(0)}-\alpha\boldsymbol{\Theta}^{(0)\top}\boldsymbol{Z}^{\top}\left(\boldsymbol{Z\Theta}^{(0)}\boldsymbol{\beta}^{(0)}-\boldsymbol{y}\right)\\ \hat{\boldsymbol{x}}_{i}^{(1)}\\ 1\\ t_{i}\end{bmatrix}.\end{split}$ $$

This corresponds to a one-step 2SLS GD update of $\boldsymbol{\Theta},\boldsymbol{\beta}$, according to ([8](#S3.E8)). Therefore, the final output matrix is:

$$ $\displaystyle\boldsymbol{H}^{(2)}=\begin{bmatrix}\boldsymbol{z}_{1}&\cdots&\boldsymbol{z}_{n}&\boldsymbol{z}_{n+1}\\ \boldsymbol{x}_{1}&\cdots&\boldsymbol{x}_{n}&\boldsymbol{x}_{n+1}\\ y_{1}&\cdots&y_{n}&0\\ \boldsymbol{\Theta}_{:,1}^{(1)}&\cdots&\boldsymbol{\Theta}_{:,1}^{(1)}&\boldsymbol{\Theta}_{:,1}^{(1)}\\ \vdots&\vdots&\vdots&\vdots\\ \boldsymbol{\Theta}_{:,p}^{(1)}&\cdots&\boldsymbol{\Theta}_{:,p}^{(1)}&\boldsymbol{\Theta}_{:,p}^{(1)}\\ \boldsymbol{\beta}^{(1)}&\cdots&\boldsymbol{\beta}^{(1)}&\boldsymbol{\beta}^{(1)}\\ \hat{\boldsymbol{x}}_{1}^{(1)}&\cdots&\hat{\boldsymbol{x}}_{n}^{(1)}&\hat{\boldsymbol{x}}_{n+1}^{(1)}\\ 1&\cdots&1&1\\ 1&\cdots&1&0\end{bmatrix}.$ $$

Thus the proof is complete. We further note that in construction steps ([35](#A2.E35))([36](#A2.E36))([37](#A2.E37))([38](#A2.E38))([39](#A2.E39))([40](#A2.E40)), regardless of the initial values of $\boldsymbol{\Theta}^{(0)},\boldsymbol{\beta}^{(0)},$ and $\hat{\boldsymbol{x}}^{(0)}$, the matrices $\boldsymbol{Q}_{m}^{(1:2)},\boldsymbol{K}_{m}^{(1:2)},\boldsymbol{V}_{m}^{(1:2)}$ do the same linear transformations on the input vectors. Therefore they are identical across different layers.
∎

###### Proof of Theorem 3.2 .

### B.3 Proof of Theorem 3.3

###### Lemma B.2 (Generalization of pretraining, from Theorem 20 in Bai et al. [ 2024 ] ) .

Given optimization problm ([13](#S3.E13)), with probability at least $1-\zeta$, the solution $\hat{\boldsymbol{\theta}}$ satisfies:

$$ $\displaystyle L_{\textsf{ICL}}(\hat{\boldsymbol{\theta}})\leq\underset{\boldsymbol{\theta}\in\boldsymbol{\vartheta}}{\inf}L_{\textsf{ICL}}(\boldsymbol{\theta})+\mathcal{O}\left(B_{y}^{2}\sqrt{\frac{L^{2}(MD^{2}+DD^{\prime})\log(2+\max\{B_{\theta},R,B_{y}\})+\log(1/\zeta)}{N}}\right).$ $$

###### Proof of Theorem 3.3 .

We begin by showing the (clipped) 2SLS predictor achieves small excess loss under in-context distribution $\boldsymbol{\mathcal{P}}$:

$$ $\displaystyle\mathbb{E}_{\mathcal{P}}\left[\left(\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}),\boldsymbol{x}_{n+1}\rangle-y_{n+1}\right)^{2}\right]$ $\displaystyle=\mathbb{E}_{\mathcal{P}}\left[\left(\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}})-\boldsymbol{\beta},\boldsymbol{x}_{n+1}\rangle+\langle\boldsymbol{\beta},\boldsymbol{x}_{n+1}\rangle-y_{n+1}\right)^{2}\right]$ $\displaystyle=\mathbb{E}_{\mathcal{P}}\left[\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}})-\boldsymbol{\beta},\boldsymbol{x}_{n+1}\rangle^{2}\right]+2E_{\mathcal{P}}\left[\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}})-\boldsymbol{\beta},\boldsymbol{x}_{n+1}\rangle\big{(}\langle\boldsymbol{\beta},\boldsymbol{x}_{n+1}\rangle-y_{n+1}\big{)}\right]$ $\displaystyle\quad\quad+\mathbb{E}_{\mathcal{P}}\left[\big{(}\langle\boldsymbol{\beta},\boldsymbol{x}_{n+1}\rangle-y_{n+1}\big{)}^{2}\right]$ $\displaystyle=\underbrace{\mathbb{E}_{\mathcal{P}}\left[\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}})-\boldsymbol{\beta},\boldsymbol{x}_{n+1}\rangle^{2}\right]}_{\text{Excess Loss}}+\mathbb{E}_{\mathcal{P}}\left[\big{(}\langle\boldsymbol{\beta},\boldsymbol{x}_{n+1}\rangle-y_{n+1}\big{)}^{2}\right],$ $$

where $\mathbb{E}_{\mathcal{P}}[\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}})-\boldsymbol{\beta},\boldsymbol{x}_{n+1}\rangle\big{(}\langle\boldsymbol{\beta},\boldsymbol{x}_{n+1}\rangle-y_{n+1}\big{)}]=0$ follows from the independence between $\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}})-\boldsymbol{\beta},\boldsymbol{x}_{n+1}\rangle$ and $\big{(}\langle\boldsymbol{\beta},\boldsymbol{x}_{n+1}\rangle-y_{n+1}\big{)}$ with $\mathbb{E}_{\mathcal{P}}\left[\langle\boldsymbol{\beta},\boldsymbol{x}_{n+1}\rangle-y_{n+1}\right]=\mathbb{E}_{\mathcal{P}}[\epsilon_{n+1}]=0$.

To bound the excess loss, we have

$$ $\displaystyle\begin{split}\mathbb{E}_{\mathcal{P}}\left[\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}})-\boldsymbol{\beta},\boldsymbol{x}_{n+1}\rangle^{2}\right]&=\mathbb{E}_{\mathcal{P}}\left[\left\|\boldsymbol{x}_{n+1}^{\top}\left(\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}})-\boldsymbol{\beta}\right)\right\|^{2}\right]\\ &\leq\mathbb{E}_{\mathcal{P}}\left[\left\|\boldsymbol{x}_{n+1}\right\|^{2}\left\|\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}})-\boldsymbol{\beta}\right\|^{2}\right]\\ &=\mathbb{E}_{\mathcal{P}}\left[\left\|\boldsymbol{x}_{n+1}\right\|^{2}\right]\mathbb{E}_{\mathcal{P}}\left[\left\|\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}})-\boldsymbol{\beta}\right\|^{2}\right]\\ &\leq\mathcal{O}\left(B_{x}^{2}\Big{(}\frac{C_{z,x}^{2}q(\boldsymbol{\phi}^{\top}\boldsymbol{\Sigma}_{u}\boldsymbol{\phi}+\sigma_{\epsilon}^{2})}{n}+\xi B_{\beta}^{2}\Big{)}\right),\end{split}$ (41) $$

where the last inequality follows from ([4](#S2.E4)).

Next, for the ICL loss, we have

$$ $\displaystyle\begin{split}&L_{\textsf{ICL}}(\boldsymbol{\theta})\\ &=\mathbb{E}_{\pi}\mathbb{E}_{\mathcal{P}}\left[\left(\widetilde{\textsf{TF}}_{\boldsymbol{\theta}}(\boldsymbol{H})-y_{n+1}\right)^{2}\right]\\ &=\mathbb{E}_{\pi}\mathbb{E}_{\mathcal{P}}\left[\left(\widetilde{\textsf{TF}}_{\boldsymbol{\theta}}(\boldsymbol{H})-\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}),\boldsymbol{x}_{n+1}\rangle+\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}),\boldsymbol{x}_{n+1}\rangle-y_{n+1}\right)^{2}\right]\\ &=\mathbb{E}_{\pi}\bigg{\{}\mathbb{E}_{\mathcal{P}}\Big{[}\left(\widetilde{\textsf{TF}}_{\boldsymbol{\theta}}(\boldsymbol{H})-\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}),\boldsymbol{x}_{n+1}\rangle\right)^{2}\Big{]}+\mathbb{E}_{\mathcal{P}}\Big{[}\left(\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}),\boldsymbol{x}_{n+1}\rangle-y_{n+1}\right)^{2}\Big{]}\\ &\quad\quad+2\mathbb{E}_{\mathcal{P}}\Big{[}\left(\widetilde{\textsf{TF}}_{\boldsymbol{\theta}}(\boldsymbol{H})-\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}),\boldsymbol{x}_{n+1}\rangle\right)\left(\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}),\boldsymbol{x}_{n+1}\rangle-y_{n+1}\right)\Big{]}\bigg{\}}\\ &\leq\mathbb{E}_{\pi}\bigg{\{}\mathbb{E}_{\mathcal{P}}\Big{[}\left(\widetilde{\textsf{TF}}_{\boldsymbol{\theta}}(\boldsymbol{H})-\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}),\boldsymbol{x}_{n+1}\rangle\right)^{2}\Big{]}+\mathbb{E}_{\mathcal{P}}\Big{[}\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}})-\boldsymbol{\beta},\boldsymbol{x}_{n+1}\rangle^{2}\Big{]}+\mathbb{E}_{\mathcal{P}}\Big{[}\big{(}\langle\boldsymbol{\beta},\boldsymbol{x}_{n+1}\rangle-y_{n+1}\big{)}^{2}\Big{]}\\ &\quad\quad+2\mathbb{E}_{\mathcal{P}}\Big{[}\left|\widetilde{\textsf{TF}}_{\boldsymbol{\theta}}(\boldsymbol{H})-\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}),\boldsymbol{x}_{n+1}\rangle\right|\Big{]}\mathbb{E}_{\mathcal{P}}\Big{[}\left|\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}),\boldsymbol{x}_{n+1}\rangle-y_{n+1}\right|\Big{]}\bigg{\}}\\ &\leq\mathbb{E}_{\pi}\bigg{\{}\mathbb{E}_{\mathcal{P}}\Big{[}\Big{(}\widetilde{\textsf{TF}}_{\boldsymbol{\theta}}(\boldsymbol{H})-\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}),\boldsymbol{x}_{n+1}\rangle\Big{)}^{2}\Big{]}+\mathbb{E}_{\mathcal{P}}\Big{[}\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}})-\boldsymbol{\beta},\boldsymbol{x}_{n+1}\rangle^{2}\Big{]}+\mathbb{E}_{\mathcal{P}}\Big{[}\big{(}\langle\boldsymbol{\beta},\boldsymbol{x}_{n+1}\rangle-y_{n+1}\big{)}^{2}\Big{]}\\ &\quad\quad+2\mathbb{E}_{\mathcal{P}}\Big{[}\left|\widetilde{\textsf{TF}}_{\boldsymbol{\theta}}(\boldsymbol{H})-\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}),\boldsymbol{x}_{n+1}\rangle\right|\Big{]}\left(\mathbb{E}_{\mathcal{P}}\left[\big{|}\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}})-\boldsymbol{\beta},\boldsymbol{x}_{n+1}\rangle\big{|}\right]+\mathbb{E}_{\mathcal{P}}\Big{[}\big{|}\langle\boldsymbol{\beta},\boldsymbol{x}_{n+1}\rangle-y_{n+1}\big{|}\Big{]}\right)\bigg{\}}.\end{split}$ (42) $$

From Corollary [3.1](#S3.Thmcoro1), we know that there exists a $L=2\bar{L}+1$-layer attention-only transformer model $\boldsymbol{\theta}$, with $M=2(p+1)$ heads, and embedding dimension $D=qp+3p+q+3$, such that for any $\boldsymbol{H}$, given any learning rates $\alpha,\eta$ and $\Lambda$ as defined in ([9](#S3.E9)), the following holds(^4^44The clipping bound on $\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}$ can be matched by adjusting the clipping threshold on the last layer of $\widetilde{\textsf{TF}}_{\boldsymbol{\theta}}$.):

$$ $\displaystyle\left|\widetilde{\textsf{TF}}_{\boldsymbol{\theta}}(\boldsymbol{H})-\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}),\boldsymbol{x}_{n+1}\rangle\right|\leq\mathcal{O}\left(B_{x}\Lambda^{\bar{L}}\right).$ $$

Denote $\Lambda^{\star}:=\underset{\alpha,\eta}{\min}\;\mathbb{E}_{\pi}\mathbb{E}_{\mathcal{P}}\left[\Lambda|\boldsymbol{H},\alpha,\eta\right]$, then under $\alpha^{\star},\eta^{\star}$, we have:

$$ $\displaystyle\mathbb{E}_{\pi}\mathbb{E}_{\mathcal{P}}\left[\left|\widetilde{\textsf{TF}}_{\boldsymbol{\theta}}(\boldsymbol{H})-\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}),\boldsymbol{x}_{n+1}\rangle\right|\right]\leq\mathcal{O}\left(B_{x}(\Lambda^{\star})^{\bar{L}}\right),$ (43) $$

and

$$ $\displaystyle\mathbb{E}_{\pi}\mathbb{E}_{\mathcal{P}}\Big{[}\big{(}\widetilde{\textsf{TF}}_{\boldsymbol{\theta}}(\boldsymbol{H})-\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}),\boldsymbol{x}_{n+1}\rangle\big{)}^{2}\Big{]}\leq\mathcal{O}\left(B_{x}^{2}\mu_{\Lambda,2}^{\star}\right),$ (44) $$

where $\mu_{\Lambda,2}^{\star}:=\mathbb{E}_{\pi}\mathbb{E}_{\mathcal{P}}\left[\Lambda^{2\bar{L}}|\boldsymbol{H},\alpha^{\star},\eta^{\star}\right]$ is close to 0.

With conidition ([14](#S3.E14)), from Cauchy-Schwarz inequality, we have:

$$ $\displaystyle\mathbb{E}_{\pi}\mathbb{E}_{\mathcal{P}}\Big{[}\big{|}\langle\boldsymbol{\beta},\boldsymbol{x}_{n+1}\rangle-y_{n+1}\big{|}\Big{]}$ $\displaystyle\leq\mathbb{E}_{\pi}\left[\sqrt{\mathbb{E}_{\mathcal{P}}\left(\epsilon_{n+1}^{2}\right)}\right]=\mathbb{E}_{\pi}[\sigma_{\epsilon}]\leq\tilde{\sigma}_{\epsilon}.$ (45) $$

Also, from ([41](#A2.E41)), we have:

$$ $\displaystyle\mathbb{E}_{\pi}\mathbb{E}_{\mathcal{P}}\left[\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}})-\boldsymbol{\beta},\boldsymbol{x}_{n+1}\rangle^{2}\right]\leq\mathcal{O}\left(B_{x}^{2}\Big{(}\frac{C_{z,x}^{2}q\tilde{\sigma}^{2}}{n}+\xi B_{\beta}^{2}\Big{)}\right).$ (46) $$

Further,

$$ $\displaystyle\begin{split}\mathbb{E}_{\pi}\mathbb{E}_{\mathcal{P}}\left[\big{|}\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}})-\boldsymbol{\beta},\boldsymbol{x}_{n+1}\rangle\big{|}\right]&\leq\sqrt{\mathbb{E}_{\pi}\mathbb{E}_{\mathcal{P}}\left[\langle\textsf{clip}_{B_{\beta}}(\hat{\boldsymbol{\beta}}_{\textsf{2SLS}})-\boldsymbol{\beta},\boldsymbol{x}_{n+1}\rangle^{2}\right]}\\ &\leq\mathcal{O}\left(B_{x}\sqrt{\frac{C_{z,x}^{2}q\tilde{\sigma}^{2}}{n}+\xi B_{\beta}^{2}}\right).\end{split}$ (47) $$

Finally, with ([43](#A2.E43))([44](#A2.E44))([45](#A2.E45))([46](#A2.E46))([47](#A2.E47)), rearranging the terms in ([42](#A2.E42)), we have:

$$ $\displaystyle\begin{split}&L_{\textsf{ICL}}(\boldsymbol{\theta})-\mathbb{E}_{\pi}\mathbb{E}_{\mathcal{P}}\left[(y_{n+1}-\langle\boldsymbol{\beta},\boldsymbol{x}_{n+1}\rangle)^{2}\right]\\ &\leq\mathcal{O}\left(B_{x}^{2}\Big{(}\mu_{\Lambda,2}^{\star}+(\Lambda^{\star})^{\bar{L}}\sqrt{\frac{C_{z,x}^{2}q\tilde{\sigma}^{2}}{n}+\xi B_{\beta}^{2}}+\frac{C_{z,x}^{2}q\tilde{\sigma}^{2}}{n}+\xi B_{\beta}^{2}\Big{)}+B_{x}(\Lambda^{\star})^{\bar{L}}\tilde{\sigma}_{\epsilon}\right).\\ &\leq\mathcal{O}\left((\Lambda^{\star})^{\bar{L}}\Big{(}B_{x}^{2}\sqrt{\frac{C_{z,x}^{2}q\tilde{\sigma}^{2}}{n}+\xi B_{\beta}^{2}}+B_{x}\tilde{\sigma}_{\epsilon}\Big{)}+B_{x}^{2}\Big{(}\frac{C_{z,x}^{2}q\tilde{\sigma}^{2}}{n}+\xi B_{\beta}^{2}+\mu_{\Lambda,2}^{\star}\Big{)}\right).\end{split}$ (48) $$

Thus combining Lemma [B.2](#A2.Thmlemma2) with ([48](#A2.E48)) completes the proof.
∎

###### Lemma B.2 (Generalization of pretraining, from Theorem 20 in Bai et al. [ 2024 ] ) .

###### Proof of Theorem 3.3 .

## Appendix C Additional Experiments

For all experiments in this section, to be consistent with our main experiment in Section [4](#S4), we generate $n=50$ training samples with $p=5,q=10$, following the data generating process described in Algorithm [1](#alg1). The task parameters $\boldsymbol{\Theta},\boldsymbol{\beta},\boldsymbol{\Phi},\boldsymbol{\phi}$ are sampled from standard Gaussian distribution, the covariance matrices $\boldsymbol{\Sigma}_{z},\boldsymbol{\Sigma}_{u},\boldsymbol{\Sigma}_{\omega}$ are set to be identity matrices, and the noise level $\sigma_{\epsilon}$ is set to 1.

### C.1 Empirical Justification for Assumption 2

In this section, we aim to provide a better understanding of the values of $\lambda_{z},\lambda_{\tilde{x}}$ and $\delta_{z\epsilon_{2}}$ required to satisfy condition ([3](#S2.E3)).

We first investigate $\lambda_{z}$. According to Rudelson and
Vershynin [], given that random matrix $\boldsymbol{Z}$ has independent sub-Gaussian entries, one can show that

$$ $\displaystyle\mathbb{P}\left(\lambda_{\min}\left(\frac{\boldsymbol{Z}^{\top}\boldsymbol{Z}}{n}\right)\leq\left(1-\sqrt{\frac{q}{n}}-\frac{n}{n-q}\varepsilon\right)^{2}\right)\leq\frac{C}{1-\sqrt{q/n}}\exp\left(-c^{\prime}q\varepsilon^{3/2}\right).$ $$

This inequality implies that setting $\lambda_{z}$ slightly lower than $(1-\sqrt{\frac{q}{n}})^{2}$ is sufficient to ensure a small tail probability.

For $\delta_{z\epsilon_{2}}$, consider the case when $\boldsymbol{z}$ and $\boldsymbol{\epsilon}_{2}$ are uncorrelated (i.e, $\mathbb{E}[\boldsymbol{z}\boldsymbol{\epsilon}_{2}^{\top}]=\boldsymbol{0}$), and both $\boldsymbol{z}$ and $\boldsymbol{\epsilon}_{2}$ are almost surely bounded by $B_{z}$ and $B_{\epsilon_{2}}$, respectively. Applying the matrix Bernstein inequality (see Lemma [A.1](#A1.Thmlemma1)) directly yields:

$$ $\displaystyle\begin{split}\mathbb{P}\left(\left\|\frac{\boldsymbol{Z}^{\top}\boldsymbol{\mathcal{E}}_{2}}{n}\right\|\geq\varepsilon\right)&\leq(q+p)\exp\left(-\frac{n^{2}\varepsilon^{2}/2}{\nu(\boldsymbol{\Omega}_{z\epsilon_{2}})+B_{z}B_{\epsilon_{2}}n\varepsilon/3}\right)\\ \end{split}$ $$

where $\nu(\boldsymbol{\Omega}_{z\epsilon_{2}})$ increases on the order of $\mathcal{O}(n)$. Thus, the tail probability decreases sharply at a rate of $\mathcal{O}\left(\exp(-n)\right)$. For the case where $\mathbb{E}[\boldsymbol{z}\boldsymbol{\epsilon}_{2}^{\top}]\neq\boldsymbol{0}$, a similar conclusion can be reached by applying the Bernstein inequality on the cetered random matrix $\frac{\boldsymbol{Z}^{\top}\boldsymbol{\mathcal{E}}_{2}}{n}-\mathbb{E}[\boldsymbol{z}\boldsymbol{\epsilon}_{2}^{\top}]$.

To validate $\lambda_{\tilde{x}}$, we note that the entries of $\frac{\boldsymbol{X}^{\top}\boldsymbol{P}_{Z}\boldsymbol{X}}{n}$ are not independent, making it challenging to derive a non-asymptotic bound. However, it is known that $\lambda_{\min}\left(\frac{\boldsymbol{X}^{\top}\boldsymbol{X}}{n}\right)$ can be lower bounded with high probability. Since $\boldsymbol{z}$ and $\boldsymbol{x}$ are correlated and $\textsf{rank}(\boldsymbol{Z})>\textsf{rank}(\boldsymbol{X})$, projecting $\boldsymbol{X}$ onto the column space of $\boldsymbol{Z}$ should preserve the eigenproperties of the sample covariance matrix (i.e. the minimal spread along any dimension should not decrease significantly towards 0).

To further verify our statements, we conduct 10000 experiments and record the values of $\lambda_{\min}(\frac{\boldsymbol{Z}^{\top}\boldsymbol{Z}}{n}),\lambda_{\min}(\frac{\boldsymbol{X}^{\top}\boldsymbol{P}_{Z}\boldsymbol{X}}{n})$, and $\|\frac{\boldsymbol{Z}^{\top}\boldsymbol{\mathcal{E}}_{2}}{n}\|$. The histograms shown in Figure [3](#A3.F3) indicate that mild choices on $\lambda_{z},\lambda_{\tilde{x}}$ and $\delta_{z\epsilon_{2}}$ are sufficient to satisfy the sample condition ([3](#S2.E3)). Specifically, we have $\xi=0.01,\lambda_{z}=0.2137,\lambda_{\tilde{x}}=0.2011,$ and $\delta_{z\epsilon_{2}}=3.3617$ as empirical guidelines.

Figure: (a)
Refer to caption: /html/2410.01265/assets/Figures/sampcond1.png

### C.2 Simulations Verifying Theorem 3.1

We use the GD-based 2SLS method ([8](#S3.E8)) to estimate the causal effect $\boldsymbol{\beta}$. For the simulated data, we calculate the following metrics:

$$ $\displaystyle\frac{2}{\sigma_{\max}^{2}(\boldsymbol{Z}\hat{\boldsymbol{\Theta}})}=0.0016,\frac{2}{\sigma_{\max}^{2}(\boldsymbol{Z})}=0.0212.$ $$

By Theorem [3.1](#S3.Thmthm1), the gradient descent converges when $\alpha\in(0,0.0016)$ and $\eta\in(0,0.0212)$. The overall convergence rate is determined by $\Lambda:=\max\{\gamma(\alpha),\kappa(\eta)\}$, where

$$ $\displaystyle\gamma(\alpha)$ $\displaystyle:=\rho\left(\boldsymbol{I}-\alpha\hat{\boldsymbol{\Theta}}^{\top}\boldsymbol{Z}^{\top}\boldsymbol{Z\hat{\Theta}}\right),$ $\displaystyle\kappa(\eta)$ $\displaystyle:=\rho\left(\boldsymbol{I}-\eta\boldsymbol{Z}^{\top}\boldsymbol{Z}\right).$ $$

We first set $\alpha=0.0012$ and vary $\eta$. The corresponding convergence rates are determined by $\Lambda=\max(0.87,\kappa(\eta))$. Next, we set $\eta=0.01$ and vary $\alpha$. The corresponding convergence curves are determined by $\Lambda=\max(\gamma(\alpha),0.82)$. We compare the estimates $\hat{\boldsymbol{\beta}}^{(t)}$ with the 2SLS estimate $\hat{\boldsymbol{\beta}}_{\textsf{2SLS}}$ as the iteration proceeds. The convergence results are shown in Figure [4](#A3.F4).

Figure: (a)
Refer to caption: /html/2410.01265/assets/Figures/gradconv1.png

The results in Figure [4](#A3.F4) are consistent with our theoretical analysis in Theorem [3.1](#S3.Thmthm1). It is worth noting that in Figure [4(a)](#A3.F4.sf1), when $\eta$ is relatively large (or small), the convergence curves exhibit some suiggly patterns. This is due to the innerloop updates ([8a](#S3.E8.1)) are converging faster (or slower) than the outer loop updates ([8b](#S3.E8.2)). However, the overall convergence rate is still determined by $\Lambda$. This pattern doesn’t appear in Figure [4(b)](#A3.F4.sf2) as we set $\eta$ to be a moderate value, which ensures that the inner loop and outer loop converge synchronously.

Next, we show the bias of the GD estimator. For better convergence, we set $\alpha^{\star}=\frac{1}{\sigma_{\max}^{2}(\boldsymbol{Z\hat{\Theta}})}$ and $\eta^{\star}=\frac{1}{\sigma_{\max}^{2}(\boldsymbol{Z})}$. We compare the biases of the GD estimator with $n=50,100,150$ in-context samples. The results are shown in Figure [5](#A3.F5).

Figure: Figure 5: The convergence of the GD-based 2SLS method with $\alpha^{\star}=\frac{1}{\sigma_{\max}^{2}(\boldsymbol{Z\hat{\Theta}})}$ and $\eta^{\star}=\frac{1}{\sigma_{\max}^{2}(\boldsymbol{Z})}$. The biases of 2SLS estimator and OLS estimator at $n=150$ are plotted for comparison.
Refer to caption: /html/2410.01265/assets/Figures/bias_vs_iteration.png

## Peer Reviews

confidence: 1
contribution: 3
initial_rating: 6
presentation: 3
questions: I want to caveat this review by saying that I have flagged to the AC that this is not my area of expertise.

I 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 ?
soundness: 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.

The 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.
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.
weaknesses: I want to caveat this review by saying that I have flagged to the AC that this is not my area of expertise.

I do not know enough about this area to understand the potential flaws in their statements / etc or more subtle points.
At a broad level, their intro / overview makes sense and seems convincing.

---

confidence: 2
contribution: 2
initial_rating: 6
presentation: 3
questions: Please refer to the weaknesses section.
soundness: 3
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.

Overall, I find the paper well-motivated, with clear writing. The background information and literature review appear thorough.

Understanding the mechanism of the Transformer could be valuable for advancing future research in this area.
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.
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.

Additionally, 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?

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
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?

(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?

(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?
soundness: 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.

(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.

(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.
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​.
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.

(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.

(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.

Meta-review: 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. All reviewers gave positive feedback while the confidence scores are relatively low. The paper's quality is high and authors engaged with rebuttals, so I see no clear reason to reject this work. I recommended an acceptance.

Decision: Accept (Poster)