Paper: Transformers Learn Higher-Order Optimization Methods for In-Context Learning: A Study with Linear Models (ICLR-2024)

Reviewer comment (Evaluation):
It is unclear why the similarity of errors is a reasonable metric. It would be helpful if the paper could provide more clarity on why the $l_2$ norm is not directly used to measure the distance between predictions instead of the error. In addition, the scale of the error is important to the problem in the linear regression problem. Since normalizing the error is mentioned in the paper, it would be beneficial to provide a more detailed explanation of why this normalization is necessary.

Author response:
**The reviewer asked why we use cosine similarity on errors.** Cosine similarity is better for cases where values are small. Usually errors (also known as residuals) have quite small magnitudes. Directly measuring the prediction will usually yield high similarities. We also ran comparisons using the L2 norm and didn’t find a difference in convergence rates. Our alternative metric (similarity of induced weights) also agrees with our cosine similarity metrics.

Reviewer comment (Theory):
It will be beneficial to discuss whether it is possible to approximate the iterative Newton’s method with causal attention. It will better align the theoretical results and the empirical findings.

Author response:
**The reviewer asked whether we can use Causal Attention for theory.** It can be relaxed if we only care about cases when the number of in-context examples are greater than the data dimension. Under such assumption, in Eqn. 32, $\frac{1}{n} \sum_{j=1}^n x_j x_j^\top$  will become $\frac{1}{t} \sum_{j=1}^t x_j x_j^\top$, when $t \geq d$, this is a valid estimator of the covariance matrix $\hat{\Sigma}$ and the rest of the proof remain similar. However, for $t < d$, this is not a valid estimator any more.

Reviewer comment (Experiments):
The depiction in Figure 3 raises concerns. Given a sufficient sample size and a noise-free data setting, all methods—Newton's method, gradient descent, and transformers—should empirically converge to a correct prediction. Intuitively, this should result in a similarity score of $0.99$ between the final output of the transformer (after the $L$th layer) and the other two methods. However, this is not reflected in Figure 3's right lower corner. This discrepancy prompts further investigation: could additional training iterations, an optimized learning rate, or other modifications enhance the gradient descent results, enabling it to make precise predictions akin to the other methods?

Author response:
**The reviewer pointed out GD should converge to the OLS.** This is indeed the case, however it converges much more slowly. We have thus updated our claims to be about convergence rates. we revised our experiments so that all three methods of interest converge to the OLS solution and the major difference is the rate of convergence.

Reviewer comment (Experiments):
This paper only compared two optimization methods: Newton's method and gradient descent (with a particular step size?). It is very straightforward to see that transformers can implement other first-order methods, including accelerated gradient descent and approximate message passing. It is not clear if there could be another first-order method that is more similar to Newton's method.

Author response:
**The reviewer asked whether another first-order method can be more similar to Newton's method.** In terms of rate of convergence, the lower bound for gradient-based methods is $\Omega(\sqrt{\kappa(S)} \log(1/\epsilon))$ [1]. It is still exponentially slower than Newton’s method. 

[1] A.S. Nemirovski and D.B Yudin. *Problem complexity and method efficiency in optimization.*

Reviewer comment (Novelty):
Construction that transformers can implement Newton's method exists in [1] Section 8.

Author response:
**The reviewer pointed out a missing reference on Looped Transformers that can implement matrix inverse by Iterative Newton’s method.** We thank the reviewer for pointing it out. We believe their construction could be used as a subroutine in an alternative proof of our claim, however one major difference is that our construction explicitly considers regression tasks whereas their construction concerns matrix inversion. Our construction never explicitly calculates the Newton inverse $M_k$ (at step $k$): we always consider $M_k x_j$ for some input token $x_j$. Replacing our construction with looped transformers also requires additional circuits to compute $X^\top y$ and make predictions using $\langle M_k X^\top y, x_\mathrm{test}\rangle$. We will discuss more about this in our next revision.

Reviewer comment (Experiments):
What is the step size for GD in experiments? Is the result robust to the tuning of step size?

Author response:
**Question 1: The reviewer asked about the step size of GD.** We choose the best learning rate for GD via grid search to ensure convergence, and here it’s 0.01. As long as GD converges, it won’t affect our claim on the rate of convergence.

Reviewer comment (Evaluation):
Why, in Figure 3, do Newton's method and Transformer layer 12 have a similarity of 0.994, but GD and Transformer only have a similarity of 0.88? Intuitively, GD and Newton's method should both converge to the minimum-norm OLS solution. In Figure 9, it seems to say that the CosSim of w_{GD} and w_{OLS} is small when the number of in-context examples is less than 20. However, people have proved that GD starting from zero initialization should converge to the min-norm least square solution. Why is there a mismatch between theory and your experiment?

Author response:
**Question 2: The reviewer pointed out GD should converge to the OLS.** This is indeed the case, however it converges much more slowly. We have thus updated our claims to be about convergence rates. we revised our experiments so that all three methods of interest converge to the OLS solution and the major difference is the rate of convergence.

Reviewer comment (Theory):
The greatest issue is that there are gaps between the theoretical results and experimental results, making it unclear what algorithm TF exactly implements. The author proves that the TF can implement k steps of Newton's iterates using O(k) layers (at least k layers), but the experiments show that TF can approximate 3 Newton steps using a single layer. This suggests that the trained TF may be doing something smarter than Newton's step, raising the question of whether we can discover some new or more efficient algorithms by Transformers.

Author response:
**The reviewer asked whether trained Transformers can do something smarter than Newton.** We agree this is possible. As discussed in Section 5 and Appendix B.3, we showed that Iterative Newton is a specific version of the “sum of moments” method. Trained Transformer is also such a method, but could have better coefficients. Nevertheless, this won’t change the asymptotic rate of convergence (measured as big-O), since Newton’s method matches the lower bound for the problem.

Reviewer comment (Experiments):
The step size tuning for the vanilla GD method and Newton's iterate should be explained, as different step sizes correspond to totally different learning algorithms.

Author response:
**Question 1: The reviewer pointed out GD should converge to the OLS.** This is indeed the case, however it converges much more slowly. We have thus updated our claims to be about convergence rates. In our revised manuscript, we fix this issue: GD, Iterative Newton, and Transformers will all converge to the OLS solution and the major difference is the rate of convergence. GD converges exponentially slower than the other two. The reviewer asked for learning rates used for GD. We choose the best learning rate for GD to ensure convergence, and here it’s 0.01 in well-conditioned cases. As long as GD converges, it won’t affect our claim on the rate of convergence.

Reviewer comment (Evaluation):
Although both Newton's method and vanilla GD should converge to OLS in principle, the maximal similarity between TF and Newton is larger than that between TF and GD in your experiments (e.g., figure 7 shows 0.974 vs 0.907). Since GD and Newton should approximately equal OLS with enough steps, the maximal similarity values should be the same when they converge. The reason for this discrepancy should be explained.

Author response:
**Question 1: The reviewer pointed out GD should converge to the OLS.** This is indeed the case, however it converges much more slowly. We have thus updated our claims to be about convergence rates. In our revised manuscript, we fix this issue: GD, Iterative Newton, and Transformers will all converge to the OLS solution and the major difference is the rate of convergence. GD converges exponentially slower than the other two. The reviewer asked for learning rates used for GD. We choose the best learning rate for GD to ensure convergence, and here it’s 0.01 in well-conditioned cases. As long as GD converges, it won’t affect our claim on the rate of convergence.

Reviewer comment (Evaluation):
There is an issue with Table 1. The claim is 'LSTM is more similar to OGD than to Newton/GD'. To prove this, the three numbers in the right column should be compared: 0.920, 0.808, 0.831. Based on these, LSTM is actually more similar to Newton than to OGD. In the paper, the larger number in each row is highlighted, which actually shows 'OGD is more similar to TF than to LSTM', contradicting the main claim. An explanation is needed.

Author response:
**Question 2: The reviewer asked whether LSTM is also more similar to Newton with 5 steps.** We also updated the OGD experiments by initializing w to zero (previously it was initialized randomly). In the new Fig. 5, LSTM matches well with both OGD and 5-step Newton in terms of in-context performance. However, LSTM and OGD have the same level of forgetting issues but Newton does not. In the updated table, even compared vertically, LSTM matches OGD the most.

Reviewer comment (Theory):
In Theorem 1, it is claimed that 'there is a TF such that for any set of in-context sample...'. This cannot be true, since Newton's iteration requires some condition on α (in the definition of M₀) to converge. For each set of in-context examples, even if the TF implements Newton's iterate, a proper α must be chosen to ensure convergence. This α depends on the eigenvalues of XᵗX, so the Transformer constructed should depend on α, which in turn depends on the dataset {xᵢ, yᵢ}.

Author response:
**Question 3: The reviewer asked how Transformers can decide $\alpha$ which can depend on $S S^\top$ in Eqn. 8.** For Iterative Newton’s method, $\alpha$ can be any value between 0 and  $\frac{2}{||S S^\top||_2}$  to ensure convergence. We pick $\frac{2}{||S S^\top||_2}$ in experiments in order to reduce randomness. Transformers can learn any fixed $\alpha$ as long as it is sufficiently small.