Transformers Handle Endogeneity in In-Context Linear Regression

TL;DR

This paper demonstrates that transformers can inherently handle endogeneity in linear regression by emulating instrumental variable methods, providing a new approach that improves robustness and reliability over traditional techniques.

Contribution

It introduces a transformer-based method that emulates 2SLS for endogeneity correction and offers theoretical guarantees for its effectiveness.

Findings

Transformers can emulate 2SLS at exponential convergence.

The proposed pretraining scheme achieves small excess loss.

Trained transformers outperform 2SLS in robustness and reliability.

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 $(\textsf{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 $\textsf{2SLS}$ method, in the presence of endogeneity.