How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression?

Thank you for your positive feedback!

Q1. 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… Could the authors confirm if the restrictions are significantly restrictive?... I certainly don't discount the fact that the work aims to provide theoretical grounding for similar observations in literature.

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

Q2. “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… 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.”

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

Q3. ... it would be great to have experiments where the model is misspecified.… It would appear that the theorems do not necessarily say anything about misspecification…. the Gaussian assumptions are used for derivations so anything in the Gaussian family would qualify for the bound. Is that the correct assessment?

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

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

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

Q5. In Theorem 4.1, how is the inner product between matrices defined?

A5. For two matrices $A$ and $B$, we define \\langle A, B \\rangle = \\text{tr}(A\^\\top B). We have clarified this (in blue text) before Theorem 3.1 in the revised paper.

Thank you for your positive feedback!

Q1-1. 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$....

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.

Q1-2. 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.

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.

Q2. 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. …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$

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.

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

A3. Please refer to our A2 in the above.

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

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.

Q5. 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$?

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.

Thank you for your strong support!

Q1. ... 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?

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

Q2. … 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?

A2. 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$.

Q3. … It would be interesting to empirically check the training time convergence with different singular value behaviors (as they pick in corollary 4.2) and observe differences in convergence and ICL performance throughout training.

A3. Thank you for your suggestion. We have empirically checked the pretraining convergence under two $H$’s of different spectrums. Please refer to Figure 1 in Appendix A in the revised paper. We can see that under polynomial eigen-decay ($i\^{-4}$) and exponential eigen-decay ($2\^{-i}$), ICL exhibits slightly different convergence behaviors though the trend is the same, which are aligned with our theory.

Thank you for your constructive feedback.

Q1. 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… Could you provide more justification for the pretraining algorithm assumed in equation (6)?

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

Q2. Could you explain the choice of step size in Theorem 4.1?

A2. Equation (7) characterizes a decreasing, piecewise-constant stepsize scheduler. Specifically, the stepsize is initialized from $\\gamma\_0$ and is decreased by a constant factor (we choose 2 in the paper) for every $T/\\log(T)$ steps, where $T$ is the total number of iterates. This is a commonly used stepsize scheduler for SGD in deep learning.

Q3. Why is it reasonable to assume that the initialization $\\Gamma\_0$ commutes with H in Theorem 4.1?

A3. Our assumption is reasonable as it holds when setting $\\Gamma\_0$ to zero or a scalar matrix, that is, $\\Gamma\_0 = c I$ where $c$ is a constant, without prior information. Additionally, it allows setting $\\Gamma\_0$ to any matrix when $H$ is a scalar matrix.

Q4. Could you provide more insight into why $H$ can be assumed to be diagonal WLOG?

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

In this way, the current proof adapts to the eigenspace of $H$.

Q5. Would it be possible to include some numerical results in the main body to support the theoretical findings?

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

Q6. Could you clarify the term "number of contexts"? It reads like the number of in-context samples, but I guess you mean the number of sequences/datasets.

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

Dear Reviewer ozh2, we hope all of your concerns have been addressed. Please let us know if there is anything else that requires our clarification!