Paper: How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression? (ICLR-2024)

Reviewer comment (Theory):
The obvious shortcomings of the theoretical results come from the nature of assumptions that deviate from standard practice - the choice of linear attention and a restricting the structure of Q/K/V matrices. Also see Question 1.

Author response:
We acknowledge that our setup deviates from practice by considering linear attention and linear reparameterization. We think the statistical bounds of SGD-trained $\Gamma$ and SGD-trained standard softmax attention model might not be identical. One challenge is that our loss function is convex, but the loss for standard softmax attention is non-convex. So there will be an extra layer of complexity from non-convex optimization in the latter setup. With that being said, our results still provide meaningful insights into the statistical task complexity of ICL. Without a full understanding of the statistical behavior of SGD in the simplified setup, it seems difficult that one can achieve this in more complicated cases.

Reviewer comment (Theory):
While the community is focused on Transformers, I am wondering if the analysis also holds for a different linear parametrization of the function $f$. In the broader context, if similar results hold for another parametrization, then attention would not appear that unique of a function. As an example, imagine a different parametrization that also leads to a dimension-free bound. In the broader context, such a result would then not distinguish what attention brings to the table, when in practice we have ample evidence that other parametrizations don't carry as flexible and generalizable inductive biases.

Author response:
We would like to highlight that attention provides a natural way to parameterize gradient descent (GD), which is a key inductive bias for the model to achieve Bayes-optimal ICL. In equation (1), we perform a linear reparameterization to simplify our analysis, but our simplification still allows equation (1) to realize a GD procedure (in terms of the mean square loss). Other linear parameterization of $f$, while might be easy to pretrain, may not correspond to a meaningful GD procedure, so it might not be able to achieve Bayes-optimal ICL. Therefore, the inductive bias carried by the attention mechanism is not easily replaceable according to our theory in Section 5.

Reviewer comment (Experiments):
Regarding the distributional assumptions of the fixed sized dataset in Assumption 1, it would be great to have experiments where the model is misspecified. Please correct me if I missed, but I don't think I see such an experiment in Appendix A. Misspecification is really important, since for all practical purposes, our models are misspecified, i.e. the data does not really come from the distribution we assume to be. If Transformers can still achieved a good decay of the empirical risk as the number of pretraining tasks increase, it would certainly be a unique characteristic. Also see Question 2.

Author response:
This is a great point. You are correct that neither our current theory nor our current experiments cover misspecified linear regression. Our theory can be extended to deal with non-Gaussian distributions with proper (up to 8-th) moment conditions. However, extra efforts would be required to extend our theory to deal with the correlations between noise and covariates. 

According to your suggestion, we conduct experiments for our attention model in misspecified settings. Please refer to Figure 4 in Appendix A in the revised paper. When we replace the Gaussian noise with an independent and uniformly distributed noise, the attention model can still perform good ICL. However, when the expected response is generated from a nonlinear function (e.g., square and sigmoid functions), the performance of ICL will decrease.

Reviewer comment (Presentation):
I would strongly recommend highlighting the experiments and moving them up to the main text. Perhaps Section 6 could be compressed a little to accommodate.

Author response:
We will move our numerical results to the main body provided we are granted additional pages in the final version.

Reviewer comment (Experiments):
While the theoretical analysis presented is commendable and contributes to our understanding, the paper could be greatly enhanced by including more empirical results. Given that implementing linear regression over a transformer model is feasible and has been achieved in prior work, specific implementations that could improve this paper include: Demonstrating through examples that training models on linear tasks of varying dimensions yields similar performance, supporting the paper's claim that task complexity is independent of the task dimension $d$. Conducting experiments to show that a pre-trained model performs poorly when there is a significant discrepancy between the in-context length during testing and training.

Author response:
**A1-1**. We have provided Figure 2 in Appendix A in the revised paper, which plots curves for risk vs. dimension. The results are consistent with our theory and suggest that the risk does not explicitly depend on the ambient dimension. 

---

**A1-2**. We have provided Figure 3 in Appendix A in the revised paper, plotting curves for risk vs. number of context examples. We observe that when $M$ and $N$ are different (up to $4$ times), the pretrained model exhibits relatively stable performance compared to optimal ridge regression. We suspect that only when $M$ and $N$ are drastically different, the pre-trained model will deviate from optimal ridge regression. We will run additional experiments to verify this.

Reviewer comment (Theory):
The paper’s assertion that task complexity is independent of $d$ is counterintuitive. It is typically expected that the number of tasks required would be on the order of $O(d^2/T)$ since the task distribution or covariance, encompassing at least $d^2$ parameters, needs to be retrieved. The authors are encouraged to provide additional clarification on this aspect of independence. Furthermore, the experiments presented seem to contradict the paper’s statement, as $T$ appears to be exponentially larger than $d^2$, which doesn’t align with the claim that "$T$ could be much smaller than $d^2$".

Author response:
**A2**. We would like to clarify that our effective dimension $D\_{\text{eff}}$ does not explicitly depend on $d$ but it still depends on the eigenvalues of $H$. Though we always have $D\_{\text{eff}} \le d^2$, the equality can be attained in the worst case (e.g., $H=I$ and $T$ is large). Therefore, our bound in the worst case would be on the order of $\tilde{O}(d^2 / T)$, which is consistent with your intuition. 

For Figure 1(a) in the initial version, we would like to point out that $d^2 = 100^2 \approx 2^{13.28}$, so the sample complexity is not exponential in the number of parameters and does not contradict our theory.

Reviewer comment (Writing):
Despite Equation (9) not explicitly showing dependence on $d$, it seems that $d$ is implicitly involved in the matrix operations, which should be addressed for clarity.

Author response:
**A3**. Please refer to our **A2** in the above.

Reviewer comment (Theory):
Is it possible to ensure that the first term of Equation (9) will consistently reach zero? Given the exponentially decreasing learning rate $\gamma_t$, this relationship does not seem immediately apparent.

Author response:
**A4**. Yes. Note that the learning rate is piecewise constant and is only decreased by a constant factor after every $\log(T)$ steps. So the constant initial learning rate $\gamma_0$ is used in the first $T/\log(T)$ steps. As a consequence, the first term in equation (9) tends to zero consistently as $T$ tends to infinity.

Reviewer comment (Theory):
In Theorem 4.1, the learning rate $\gamma_0$ is upper bounded. Could the authors elaborate on the reasoning behind this? Is it connected to the implicit convergence rate of gradient descent in linear problems? Additionally, what are the considerations and trade-offs in selecting an appropriate initial learning rate $\gamma_0$?

Author response:
**A5**. If the initial learning rate $\gamma_0$ is too large, the SGD could diverge. The upper bound on $\gamma_0$ is a technical condition to ensure the convergence of SGD. 

For a larger initial learning rate, the first term (i.e., bias error) in equation (9) will be smaller but the second term (i.e., variance error) in equation (9) will be larger. So an optimal initial learning rate requires a balance between the bias and variance errors. Prior knowledge such as the norm of $\Gamma^*_N$ and the scale of $\psi^2 tr(H) + \sigma^2$ could be useful in choosing the initial learning rate.

Reviewer comment (Theory):
I have a question about the theoretical setup. Linear attentions used in practice have a query \( Q \), key \( K \), and value \( V \) matrix. However, the authors use structural modifications in \( Q \), \( K \), and \( V \) to represent the formulation with a single matrix \( \Gamma \) and compute SGD convergence of \( \Gamma \). Without the modification, I believe the optimal solution can be shown to be a 3-matrix factorization of the optimal solution \( \Gamma^* \) given in theorem 1. But what will be the statistical bounds of training these 3 matrices (or any pair among the 3 matrices)? Can the authors discuss whether it is answerable from their theoretical framework and if not, the difficulties one might face to solve?

Author response:
You are correct that the optimal solution can be shown to be a 3-matrix factorization, which can be seen from Appendix B. We think the statistical bounds of SGD-trained $\\Gamma$ and SGD-trained 3 matrices might not be identical. One challenge is that the loss is convex for $\Gamma$ but is non-convex for $Q, K, V$. So there will be an extra layer of complexity from non-convex optimization when SGD is run for the 3 matrices. Our techniques for analyzing SGD in $\Gamma$ might not be enough to deal with this non-convex optimization issue. With that being said, we believe our results still provide meaningful insights into the statistical task complexity of ICL.

Reviewer comment (Theory):
How will the theory change when instead of using training sequences of a single length, we use randomly sampled training sequences with varying lengths? How will the excess risk in theorem 5.3 change then?

Author response:
This is a good question. A model pretrained with a varying context length cannot match the Bayes optimal method for every $M \le N$ (in the setup of Theorem 5.3). This is because to achieve Bayes optimality with $M$ context examples, the model in equation (1) needs to have a parameter $\Gamma$ as a function of $M$. However, the pretrained model in equation (1) has a fixed model parameter (even when the context length is varying during pretraining) that cannot adapt to a varying context length $M$, therefore such a model cannot be Bayes optimal uniformly for $M$.

Reviewer comment (Theory):
The assumption of the pretraining algorithm in equation (6) seems arguable, as the equivalence of the function classes does not necessarily imply identical behavior under different parameterizations in gradient descent.

Author response:
Equation (6) corresponds to the one-pass SGD update for parameter $\Gamma$, which is a natural algorithm. Considering SGD in $\Gamma$ space is meaningful as it provides insights into the statistical hardness of learning the ICL model. We agree with you that the behavior of SGD in $\Gamma$ parameter space is not identical to that of SGD in $Q, K, V$ parameter space. For example, the loss is convex for $\Gamma$ but is non-convex for $Q, K, V$. Nevertheless, considering SGD in $\Gamma$ space simplifies the challenge from non-convex optimization. 

Extending our results to more complicated $Q, K, V$ parameterization is an important direction. However, even for the simplified $\Gamma$ parameterization, there is no such kind of result before our work. Without a comprehensive understanding of the statistical behavior of SGD in the simpler $\Gamma$ space, achieving this in the more complex $Q, K, V$ space appears challenging. As our work is the first one on statistical pretraining task complexity for ICL, we believe this is an important step in this direction and the contribution of our work is significant.

Reviewer comment (Theory):
The assumption that H can be diagonalized without loss of generality (WLOG) is also not sufficiently explained.

Author response:
We assume that $H$ is diagonal to simplify our discussion in “Key idea 1: diagonalization” (and related places in the appendix). This is made without loss of generality. If $H$ is not diagonal, let $H = V \Lambda V^\top$ be the eigen-decomposition of $H$, then we can replace “diagonal matrices” with “matrices of form $V D V^\top$ where $D$ is a diagonal matrix”, and our arguments all go through. Taking equation (30) in Appendix D.4 as an example, we re-define a “diagonalization of an operator” by 
$$
\mathring {\mathcal{O}} (A) = V \text{diag} ( {\mathcal{O}} ( \text{diag}(V^\top A V)  ) ) V^\top.
$$
In this way, the current proof adapts to the eigenspace of $H$.

Reviewer comment (Presentation):
The paper could be enhanced by including some numerical results in the main body.

Author response:
We will move our numerical results (see Appendix A) to the main body provided we are granted additional pages in the final version.

Reviewer comment (Writing):
The terminology used is occasionally confusing, such as the use of "number of contexts".

Author response:
The term “number of contexts” indeed refers to the “number of in-context samples”.