Transformers Handle Endogeneity in In-Context Linear Regression

Thank you for your valuable comments on our paper and the positive evaluation. We appreciate the opportunity to address your concerns.

[W3] Scalability Analysis

Thank you for your suggestion. As a reference, the training of the transformer in our experiment was conducted on a Windows 11 machine with the following specifications:

• GPU: NVIDIA GeForce RTX 4090
• CPU: Intel Core i9-14900KF
• Memory: 32 GB DDR5, 5600MHz

The training process took around 10 hours. These are mentioned in our supplementary "readme" file. We also added this information to Appendix C.8.

We also emphasize that our current work is focused on the existence of transformer implementing 2SLS. In practice, the parameters of transformer are learned by gradient descent-type algorithm. Understanding the training dynamics, including scalability, computational complexity and memory efficiency is certainly an important question that we plan to work on in the future. However, it is clearly beyond the scope of the current paper.

[Q1] How does the proposed transformer model handle cases with highly correlated instruments?

The issue of multicollinearity can be resolved by ridge regression. Although not directly specified in our paper, previous works have studied that transformers have intrinsic mechanism to implement ridge regression in-context (for example, see Bai et al. (2024), Section 3.1). This mechanism can be easily extended to 2SLS case by considering each stage as a constrained optimization problem. We provide a dicussion in Appendix C.3, and we add some experiments to examine the capability of transformer in handling IV regression with multicollinearity in Appendix C.4.

[Q2] What led to the choice of hyperparameters in the pretraining scheme?

Our hyperparameter settings of number of attention heads, embedding space dimension, and number of transformer layer in each transformer block follow the theoretical results of Theorem 3.2. For the number of transformer blocks needed, one may refer to Appendix C.2. The empirical results suggest that 20 iterations (blocks) are sufficient for the algorithm to converge. We set the number of blocks to be 10 in our experiments which is mainly out of the consideration of training time, and it turns out that the pretrained model is powerful enough. While we expect the model performance to be further improved with a larger model size, our experiment focuses on the minimum sufficient case directed by our theoretical results.

Note that the learning rates $\alpha, \eta$ in GD-2SLS are not hyperparameters of our transformer model. Instead, for any values of $\alpha$ and $\eta$, there exists a corresponding transformer architecture to implement the GD-2SLS algorithm, so we claim that there exists a transformer architecture that can implement GD-2SLS with global optimum $\alpha$ and $\eta$, determined by the training data distributions. For specific training of the transformer model, as a common practice, we use the Adam optimizer with default initial learning rate 0.001.

[Q3] How sensitive is the model to changes in the strength of endogeneity?

Thank you for this good question. Based on your question, we provide an experiment in Appendix C.6 to examine the sensitivity of the transformer model to the strength of endogeneity. The results show that the transformer model is robust to different endogeneity levels, and the performance is comparable with the optimal 2SLS estimator.

We sincerely hope that we have addressed many of your questions, in which case, we would greatly appreciate if you could increase your score as you see fit.

Thank you for your valuable comments on our paper and the positive evaluation. We appreciate the opportunity to address your concerns.

[W1] Can the theoretical analysis be extended to nonlinear case?

We agree that most of the real-world problems are non-linear in nature. However, we also emphasize that for various applications in econometrics, it is common to assume linearity between y and x, for the interpretability reasons. Regarding the relation between x and z, our experiments in the original draft show that the pre-trained transformer model provides consistent coefficient estimates when the relationship between x and z is quadratic. In the revised draft, we additionally examine a more complicated non-linear case in Appendix C.5.

While it is an interesting question to extend our theoretical analysis to the case of non-linear relationships (both between y and x, and x and z), some of the underlying assumptions need to be revisited, and we expect there will be structural changes in the proof. We leave this as a future work.

[W2] Real-world applications and experiments.

Thank you for your suggestion. To illustrate the performance of pre-trained transformers on real-world data, we now include an example in Appendix C.7. In this example, we analyze the effect of childbearing on mother's labor supply, with instrument being the indicator variable of whether the first two siblings are of the same sex. The transformer gives an estimate that with each increase in the number of children, the mother's labor supply is expected to drop 4.73 weeks per year. This result is statistically close to the 2SLS estimate, a widely agreed upon benchmark for this dataset.

We sincerely hope that we have addressed many of your questions, in which case, we would greatly appreciate if you could increase your score as you see fit.

Thank you for your valuable comments on our paper and the positive evaluation. We appreciate the opportunity to address your concerns.

[Q1] I wonder how this work helps one to understand what capabilities a transformer should / should not be capable of?

We emphasize that our work provides a theoretical analysis of the transformer model's ability to perform in-context instrumental variable (IV) regression and deliver reliable coefficient estimates under endogeneity. Specifically, we demonstrate the existence of an intrinsic attention structure that enables transformers to execute GD-based 2SLS. We then propose a pretraining scheme and prove that under this scheme, the pretrained transformer achieves a small excess loss. Furthermore, our experiments reveal that the pretrained transformer model not only effectively addresses standard IV tasks but also generalizes well to more challenging scenarios, providing superior estimates compared to standard 2SLS and OLS estimators.

At a general level, your question is quite open-ended. Various studies, such as Wei et al. (2022) and Sanford et al. (2024), have explored the general algorithmic representational capabilities of transformers. However, the results in these works are necessarily coarse, as they operate in a broad and general framework. In contrast, our work is motivated by a recent line of research, including Bai et al. (2024), Akyürek et al. (2023), Von Oswald et al. (2023), Li et al. (2023), Fu et al. (2023), and Ahn et al. (2024), which focus on specific statistical problems and provide more refined, fine-grained results. However these prior works do not address the challenge of handling endogeneity. Our work fills this gap by establishing precise bounds on the in-context learning capabilities of transformers for addressing endogeneity in linear models through instrumental variable regression.

[Q2] Can we say something more concrete / general around what class of functions we can argue that a transformer should be able to solve?

In contrast to several works demonstrating that transformers can implement gradient descent and Newton's method for standard single-level optimization problems (e.g., argmin-type tasks), our analysis fundamentally argues that transformers are capable of executing a single step of bi-level gradient descent using two attention layers. Based on our findings, we believe that the proposed framework can be extended to other tasks that can be effectively formulated as bi-level optimization problems. These include various statistical learning challenges, such as hyperparameter tuning, adversarial machine learning, and fairness-related problems.

[Q3] How does this relate to larger models?

It is somewhat unclear to us what you mean by "larger models." If you are referring to larger transformer architectures, works such as Wei et al. (2022) and Sanford et al. (2024) offer insights in that direction. We would greatly appreciate it if you could clarify your question by providing more specific details.

Wei, Colin, Yining Chen, and Tengyu Ma. "Statistically meaningful approximation: a case study on approximating turing machines with transformers." Advances in Neural Information Processing Systems 35 (2022): 12071-12083.

Sanford, Clayton, Daniel J. Hsu, and Matus Telgarsky. "Representational strengths and limitations of transformers." Advances in Neural Information Processing Systems 36 (2024).

We sincerely hope that we have addressed many of your questions, in which case, we would greatly appreciate if you could increase your score as you see fit.