Transformers Handle Endogeneity in In-Context Linear Regression

Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.

51,719 characters · 14 sections · 37 citation commands

Rendered from LaTeX for readability, not typeset faithfully. Citation keys are highlighted; maths is left as source; figures, tables and equation environments are summarised rather than reproduced; unrecognised commands are greyed out so nothing is silently dropped. Email addresses are removed.

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

Introduction

The transformer architecture vaswani has demonstrated remarkable in-context learning (ICL) capabilities across various domains, such as natural language processing Devlin2019BERTPO,radford2019language,brown2020language, computer vision dosovitskiy2021image, carion2020endtoend, and reinforcement learning lee2022multigame,parisotto2019stabilizing. 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, garg2023transformers 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 statistician established that attention layers enable transformers to perform gradient descent, implementing algorithms like linear regression, logistic regression, and LASSO; see also akyurek2023what,von2023transformers,li2023transformers,fu2023transformers,ahn2024transformers,jin2025incontextlearningmixturelinear. The learning dynamics of transformer trained via gradient descent for in-context learning linear function classes was analyzed in Huang:ICML:24. Furthermore zhang2023trained,zhang2024incontext 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 hausman2001mismeasured, wooldridge2015introductory, angrist2009mostly,greene2003econometric. 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 10.1257/jep.15.4.69. A naturally intriguing question that therefore arises is:

center[center omitted — 163 chars of source]

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

itemize[noitemsep,leftmargin=0.2in] • 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.

Related works

In-context Learning. Initial works by garg2023transformers and statistician adopted the standard multi-layer transformer architecture to conduct the experiments. Later, giannou2023loopedtransformersprogrammablecomputers and yang2024loopedtransformersbetterlearning showed that a looped architecture reduces the required depth of transformers and exhibits better efficiency in learning algorithms. gao2024expressivepowervariantlooped 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, giannou2024transformersemulateincontextnewtons showed that transformers can emulate Newton's method for logistic regression. cheng2024transformersimplementfunctionalgradient 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 ye2024pre,falck2024context.

nichani2024transformerslearncausalstructure illustrated how the transformers can learn the causal structure by encoding the latent causal graph in the first attention layer. goel2024can explored the representational power of transformer for learning linear dynamical systems. makkuva2024local,makkuva2024attention,rajaraman2024transformers,edelman2024evolution 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 10.1257/jep.15.4.69,angrist2009mostly. Recent works in machine learning explored the optimization based approaches for the IV regression problem. singh2020kernelinstrumentalvariableregression proposed the kernel IV regression to model non-linear relationship between variables. muandet2020dualinstrumentalvariableregression proposed that a non-linear IV regression problem can be formulated as a convex-concave saddle point problem. della2023stochastic,chen2024stochasticoptimizationalgorithmsinstrumental,peixoto2024nonparametric proposed a stochastic optimization algorithm for IV regression.

Notation:\quad 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}$. $\lambda_{\min}(\cdot)$ denotes the minimum eigenvalue, and $\sigma_{\min}(\cdot)$ denotes the minimum singular value of a matrix. By default, $\|\cdot\|$ denotes the Euclidean norm for a vector, or the spectral norm for a matrix.

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

align[align omitted — 189 chars of source]

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 wooldridge2015introductory, Chapter 9). To address this issue, instrumental variable (IV) regression is a widely used method to provide a consistent estimate for $\boldsymbol{\beta}$.

defn[2SLS estimator] IV regression is a regression model to provide consistent estimate on the causal effect $\boldsymbol{\beta}$ for the endogeneity problem ((ref)), 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, wooldridge2015introductory, Chapter 15. \begin{itemize} • First stage: Regress $\boldsymbol{X}$ on $\boldsymbol{Z}$ to obtain $\hat{\boldsymbol{\Theta}}$ \begin{align*} \hat{\boldsymbol{\Theta}} = (\boldsymbol{Z}^\top\boldsymbol{Z})^{-1}\boldsymbol{Z}^\top\boldsymbol{X}. \end{align*} • Second stage: Regress $\boldsymbol{Y}$ on $\boldsymbol{Z}\hat{\boldsymbol{\Theta}}$ to obtain: \begin{align} \hat{\boldsymbol{\beta}}_{2SLS} = (\hat{\boldsymbol{\Theta}}^\top\boldsymbol{Z}^\top\boldsymbol{Z}\hat{\boldsymbol{\Theta}})^{-1}\hat{\boldsymbol{\Theta}}^\top\boldsymbol{Z}^\top\boldsymbol{Y}. \end{align} \end{itemize}

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

asp[Instrumental variable] A random variable $\boldsymbol{z}\in \mathbb{R}^q$ is a valid IV, if it satisfies the following conditions: \begin{itemize}[noitemsep] • Fully identification: $q\geq p$ (without loss of generality, we assume data $\boldsymbol{Z},\boldsymbol{X}$ are full rank). • Correlated to $\boldsymbol{x}$: $\textsf{Corr}(\boldsymbol{z},\boldsymbol{x})\neq\boldsymbol{0}$. • Conditional uncorrelated to $y$: $\textsf{Corr}(\boldsymbol{z},\epsilon_1)=0$. \end{itemize}

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

asp[Regularity conditions] Suppose instrument $\boldsymbol{z}$ is a centered random variable. We assume the following conditions hold: \begin{itemize} • Bounded parameters: $\|\boldsymbol{\beta}\|\leq B_\beta$, $\|\boldsymbol{\Theta}\|\leq B_\Theta$. • Bounded variables: $\|\boldsymbol{z}\|\leq B_z,\|\boldsymbol{x}\|\leq B_x, |\epsilon_{1}|\leq B_{\epsilon_1}, \|\boldsymbol{\epsilon}_{2}\|\leq B_{\epsilon_2}$. • Linear instrument: $\mathbb{E}\left[x_k\middle| \boldsymbol{z}\right]=\langle\boldsymbol{\Theta}_k,\boldsymbol{z}\rangle$. \end{itemize}

The boundedness condition in (ii) is required to invoke matrix Bernstein inequalities tropp2015introductionmatrixconcentrationinequalities in the analysis. We anticipate that this condition may be relaxed to subgaussian or moment conditions by using more sophisticated matrix concentration results.

thm[MSE of 2SLS estimator] Given Assumptions (ref) and (ref), 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}. $$ When $$ n\geq \max\left\{4c^2B_z^4\left(q+\log\left(\frac{4c^2B_z^4K}{q}\right)-\frac{3}{2}\right), \frac{qe^{\frac{3}{2}}}{K}, \frac{p^2(q+1)^2K}{qK_0^2}\right\}, $$ where $K:=\frac{\lambda_{\min}(\boldsymbol{\Sigma}_z)}{6B_z^2}$ and $K_0:=\frac{\lambda_{\min}(\boldsymbol{\Sigma}_z)\sigma_{\min}^2(\boldsymbol{\Theta})}{2B_{\epsilon_2}^2}$, the mean squared error of the 2SLS estimate is bounded by: \begin{align} \begin{split} \mathbb{E}\left[\|clip_{B_\beta}(\hat{\boldsymbol{\beta}}_{2SLS})-\boldsymbol{\beta}\|^2\right] &\leq\mathcal{O}\left(\frac{q}{n}\left(\frac{B_\beta^2}{K}+{C^2(n)\sigma_1^2}\right)\right), \end{split} \end{align} where $$ C(n):=\frac{\left(B_{\Theta}+\sqrt{\frac{2p(q+1)B_{\epsilon_2}^2\log\left(\frac{K}{q}n\right)}{\lambda_{\min}(\boldsymbol{\Sigma}_z)n}}\right)B_z}{\lambda_{\min}(\boldsymbol{\Sigma}_z)\left(1-\frac{cB_{z}^2\left(\sqrt{q}+\sqrt{\log\left(\frac{K}{q}n\right)-\frac{1}{2}}\right)}{\sqrt{n}}\right)^2\left(\sigma_{\min}(\boldsymbol{\Theta})-\sqrt{\frac{2p(q+1)B_{\epsilon_2}^2\log\left(\frac{K}{q}n\right)}{\lambda_{\min}(\boldsymbol{\Sigma}_z)n}}\right)^2}, $$ $\boldsymbol{\Sigma}_z:=\mathbb{E}[\boldsymbol{zz}^\top],$ and $c$ is an absolute constant.

rmkWe keep the slightly complicated form (ref) so that the $\mathcal{O}$ notation only hides some absolute constant multipliers that are independent of problem-related constants. Note that when $n$ is large enough, we have $C(n)\rightarrow \frac{B_\Theta B_z}{\lambda_{\min}(\boldsymbol{\Sigma}_z)\sigma_{\min}^2(\boldsymbol{\Theta})}$, so $C(n)$ is also bounded. Thus the error bound (ref) decays with rate $\mathcal{O}(\frac{1}{n})$.

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

Transformers Handle Endogeniety

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.

defn[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},$ \begin{align} \tilde{\boldsymbol{H}} = 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}, \end{align} or element-wise: \begin{align} \tilde{\boldsymbol{h}}_i = [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}, \end{align} where $\sigma(\cdot)$ is the ReLU function.

defn[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' \times D\times D\times D'}$. Given input $\boldsymbol{H},$ \begin{align*} \tilde{\boldsymbol{H}} = MLP_{\boldsymbol{\theta}}(\boldsymbol{H}) := \boldsymbol{H} + \boldsymbol{W}_2\sigma(\boldsymbol{W}_1\boldsymbol{H}), \end{align*} or element-wise: \begin{align*} \tilde{\boldsymbol{h}}_i = [MLP_{\boldsymbol{\theta}}(\boldsymbol{H})]_i := \boldsymbol{h}_i +\boldsymbol{W}_2\sigma(\boldsymbol{W}_1\boldsymbol{h}_i). \end{align*}

defn[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)},$ \begin{align*} \boldsymbol{{H}}^{(l)} = MLP_{\boldsymbol{\theta}_{MLP}^{(l)}}(ATTN_{\boldsymbol{\theta}_{ATTN}^{(l)}}(\boldsymbol{H}^{(l-1)})) , \quad l=1,\ldots,L. \end{align*} The output of this transformer is the final layer output: $\tilde{\boldsymbol{H}}:=\boldsymbol{H}^{(L)}=\textsf{TF}_{\boldsymbol{\theta}}(\boldsymbol{H}^{(0)})$.

defn[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)},$ \begin{align*} \boldsymbol{{H}}^{(l)} = TF_{\bar{\boldsymbol{\theta}}}(\boldsymbol{H}^{(l-1)}), \quad l=1,\ldots,\bar{L}. \end{align*} 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., statistician, zhang2023trained) 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 (ref)). 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 (ref)) with looped transformer architecture (Definition (ref)), 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 (ref)).

Gradient descent based IV regression

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

align[align omitted — 428 chars of source]

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

subequations\begin{align} \boldsymbol{\Theta}^{(t+1)} &= \boldsymbol{\Theta}^{(t)} - \eta \boldsymbol{Z}^\top(\boldsymbol{Z\Theta}^{(t)}-\boldsymbol{X}),\\ \boldsymbol{\beta}^{(t+1)} &= \boldsymbol{\beta}^{(t)} - \alpha \boldsymbol{\Theta}^{(t)\top}\boldsymbol{Z}^\top(\boldsymbol{Z\Theta}^{(t)}\boldsymbol{\beta}^{(t)}-\boldsymbol{Y}). \end{align}

Note that the GD-2SLS updates in (ref) are designed to solve (ref). 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 (ref) with exponential rate.

thm[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: \begin{align*} 0<\alpha<\frac{2}{\sigma_{\max}^2(\boldsymbol{Z}\hat{\boldsymbol{\Theta}})}\quadand\quad0<\eta<\frac{2}{\sigma_{\max}^2(\boldsymbol{Z})}, \end{align*} where $\sigma_{\max}(\cdot)$ denotes the largest singular value of a matrix. Then, the GD updates in (ref) converge to the 2SLS estimator at an exponential rate: \begin{align*} \|\boldsymbol{\beta}^{(t)}-\hat{\boldsymbol{\beta}}_{2SLS}\|\leq\mathcal{O}\left(\Lambda^t\right), \end{align*} where, with $\rho(\cdot)$ denoting the spectral radius of the matrix, \begin{align} \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}). \end{align}

To the best of our knowledge, Theorem (ref) 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 (ref) and present simulation results in Appendix (ref) to examine the convergence behavior of the optimization process.

Transformers Can Efficiently Implement GD-2SLS

The looped transformer architecture (Definition (ref)), as proposed by giannou2023loopedtransformersprogrammablecomputers, introduces an efficient approach to learn iterative algorithms by cascading the same transformer block for multiple times. With the GD updates in (ref), 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.

thm[Implement a step of GD-2SLS with a transformer block] Suppose the embedded input matrix takes the form: \begin{align} \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)} & \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)}. \end{align} 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 (ref) with any given learning rates $\alpha,\eta$: \begin{align*} \boldsymbol{H}^{2(l+1)}=TF_{\boldsymbol{\theta}_{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)}. \end{align*}

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 (ref) (ref. Appendix (ref)), 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 yang2024loopedtransformersbetterlearning.

By cascading the transformer block $\bar{L}$ times, with Theorem $\ref{thm:2SLS GD}$, one can show that transformers are able to mimic the 2SLS estimator with exponential convergence rate, as described in the following corollary.

coro[Implementing GD-2SLS with looped transformer] For any $0<\varepsilon<1$, given learning rates $\alpha,\eta,$ and $\Lambda\in(0,1)$, as defined in (ref), there exists a transformer formulated as $\textsf{TF}_{\boldsymbol{\theta}}(\cdot):=\textsf{TF}_{\boldsymbol{\theta}'}(\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\footnote{This 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 statistician} $\boldsymbol{\theta}'=\boldsymbol{\theta}_{\textsf{ATTN}}'= \{(\boldsymbol{Q}_m', \boldsymbol{K}_m', \boldsymbol{V}_m')\}_{m\in[M']}\subset \mathbb{R}^{D\times D}$, where $\bar{M}^{(1)}=2p, \bar{M}^{(2)}=2(p+1)$, $M'=2$, such that given embedded input $\boldsymbol{H}^{(0)}$ taking the format in (ref), the model output satisfies: \begin{align*} |read_y(TF_{\boldsymbol{\theta}}(\boldsymbol{H}^{(0)}))-\hat{\boldsymbol{\beta}}_{2SLS}^\top\boldsymbol{x}_{n+1}|\leq B_x\varepsilon, \end{align*} 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 statistician for OLS in the following aspects:

itemize• We apply the square loss as defined in (ref) to learn the 2SLS estimator, which simplifies the loss function's sum-of-ReLU representation. • 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)}$. • We use a two-layer attention-only transformer block $\bar{\boldsymbol{\theta}}$ to implement a 2SLS GD update (ref), 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)}$. • 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$.

Pretraining and Excess Loss Bound

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

align*[align* omitted — 332 chars of source]

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

algorithm[algorithm omitted — 2,028 chars of source]

In Algorithm (ref), $\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 (ref). $\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 (ref) and Assumption (ref)(ii)(iii) are directly satisfied with the data generated from the in-context distribution $\boldsymbol{\mathcal{P}}$.

Following the theoretical framework of statistician, we define the population ICL loss\footnote{All the clipping operations are only for analytical purpose. In practice, the behavior of the trained transformer is consistent even without the clipping bounds.}:

align[align omitted — 241 chars of source]

where $\boldsymbol{H}^{(0)}$ is the embedded input as defined in (ref), $\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 (ref) with the following empirical loss:

align[align omitted — 195 chars of source]

We consider the following constrained optimization problem:

align[align omitted — 570 chars of source]

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.

thm[Excess loss bound for in-context pretrained transformer] Suppose condition (i) in Assumption (ref) holds and the meta distribution $\boldsymbol{\pi}$ satisfies the following conditions: \begin{align} \mathbb{E}_\pi\left[\boldsymbol{\phi}^\top\boldsymbol{\Sigma}_u\boldsymbol{\phi}+\sigma_\epsilon^2\right]\leq& \tilde{\sigma}^2 and \mathbb{E}_\pi\left[\sigma_\epsilon^2\right]\leq\tilde{\sigma}_\epsilon^2. \end{align} 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 (ref)). With training prompts $\boldsymbol{H}_k, k=1,\ldots,N$, under ICL loss in (ref), the trained transformer in (ref) with $L=2\bar{L}+1, M=2(p+1), D=qp+3p+q+3, D'=0$ (attention-only) achieves the following excess loss with probability at least $1-\zeta$: \begin{align*} L_{ICL}(\hat{\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}\bigg((\Lambda^\star)^{\bar{L}}\left(B_x^2\sqrt{\frac{q}{n}\left(\frac{B_\beta^2}{K}+C^2(n)\tilde{\sigma}^2\right)}+B_x\tilde{\sigma}_\epsilon\right)\nonumber\\ &+B_x^2\left(\frac{q}{n}\left(\frac{B_\beta^2}{K}+C^2(n)\tilde{\sigma}^2\right)+\mu_{\Lambda,2}^\star\right)+B_y^2\sqrt{\frac{L^2MD^2\log(2+\max\{B_\theta,R,B_y\})+\log(1/\zeta)}{N}}\bigg), \end{align*} 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$.

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 (ref) is provided in Appendix (ref).

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 (ref). For 2SLS, the estimated causal effect is given by the coefficients of the endogenous variable in the second stage regression ((ref)). 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 (ref).

algorithm[algorithm omitted — 1,049 chars of source]

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 elimination of rounding errors during computation.

Experiments

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 (ref), 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 (ref). We employ the looped transformer architecture, consisting of 10 identical cascading transformer blocks. The transformer model is trained under the ICL loss ((ref)) 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}$.

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 (ref). 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 (ref), 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 (ref).

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[figure omitted — 674 chars of source]

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 (ref), and (b) the dimension of IV is not sufficient to identify the endogenous variable\footnote{For 2SLS estimate, the actual computation uses pseudoinverse to handle rank deficiency.}, where we set $q=3$ (by zeroing out the remaining dimensions of $\boldsymbol{z}$ in test prompts) and $p=5$.

We evaluate the same trained transformer model as before, with results presented in Figure (ref) and Figure (ref), 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. To further illustrate this capability, we also examine other cases including multicollinearity, complex non-linear IV, and varying endogeneity strengths, see Appendix (ref),(ref),(ref). 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[figure omitted — 733 chars of source]

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.

Acknowledgements

The work of H. Liang and L. Lai was supported by National Science Foundation (NSF) under grants CCF-2112504 and CCF-2232907. The work of K. Balasubramanian was supported by NSF under grants DMS-2413426 and DMS-2053918.