Re-examining learning linear functions in context

We thank the reviewer for the comments. We address the following three concerns 1- The models and empirical studies in the paper differ significantly from current large language models, potentially creating a gap between the claims and reality. 2- The findings of the paper have been previously proposed in other works. 3- The paper is missing some key references. As we mentioned we need small models to do training from scratch. We need to train from scratch to avoid “leakage” from uncontrolled pretraining. We did query larger models but we cannot train them from scratch. Thank you for pointing out these two references. They are very helpful and we will certainly cite and discuss them in the final version. Nevertheless with respect to [1], though similar to our work, [1] works with linear attention, whereas we look at attention layers as they actually are used with softmax (thus one could criticize [1] in the way you do in point 1 more than us at least on this score). In addition, [1] uses a new kind of optimization or training with gradients and a special fixed initial point. This means that their architecture and training are quite different from what normally happens with transformers; they are interested in getting a revised transformer-like model to learn linear functions, while we want to find out whether transformers as they actually are learn linear functions or something else. The results for the two architectures are quite different: While [1] says task shift does not affect their models, our task shifts affect the results in an important way, where we take D^{train}_H = N(0,1) (Our D_F) but D^{test}_H = N(0, \sigma) (our D^test_H) for 1 \leq \sigma \leq 10. Figure 1 clearly shows that for transformer models with soft attention, this task shift reduces performance dramatically. We also note unlike [1] that prompts that are too long induce chaotic behavior. In the covariate shift [1] also does something different from what we do. In covariate shift in [1], the distribution in the prompt is shifted but the distribution of the query stays the same. We do something different. When we take a distribution over input points in train D_I and set D^test_I \neq D_I, our shift is not the same; we shift both prompt and query distributions. With covariate shifts we found that the choice of points is important and model performance degrades considerably when the values of the functions on the chosen points lie beyond what we call boundary values. As far as we know we are the first to point out these boundary values and their dependence on model parameters. At least [1] does not do this. Our mathematical formulation of what models do, explains their behavior. We will highlight this in the revised version. With respect to suggested reference [2], they largely follow what many other ICL papers do—offer a proof by construction that under certain assumptions transformers can implement many different algorithms for computing linear functions and other tasks as well. Their empirical experiments show that under suitable training and testing distributions for sampling, transformer models can learn such algorithms. They signal something like our boundary values B, -B and propose to ignore values outside [B, -B] by “clipping” them. We take a very different approach. By ignoring these outside values, we don’t really know what algorithm transformer models have implemented; we demonstrate with out of distribution robustness tests that transformer models don’t use any of the standard algorithms (ridge regression, linear regression,...) but do a different kind of projection and interpolation. We have included a short discussion of these 2 papers in the revised submission; hopefully the revised version will clarify and address your concerns.

We thank the reviewer for the comments. With regards to a more thorough lit review, we will add references to the revised version. After submission to ICLR, we have already added related work not cited in the submitted version : Fu et al. (2023), Xie et al. (2021), Wu et al.(2023), Zhang et al. (2023), Panwar et al. (2023), Bai et al. (2024). We would be glad to have more concrete suggestions about what to add over and above what the other reviewers have suggested. We have posted a new version of the paper with an improved organization and hopefully this will answer your criticism With regards to notation, we will clarify f_{i,\sigma}. What it means is this: f_{i,\sigma} is the ith function sampled from N(0,\sigma] for 1 \leq \sgma \leq 10. We will clarify this notation. We apologize for D^T_F which is a typo, D^t_F is the test distribution for the functions. D^t_I is the distribution for the points x_i in the sequences to which the functions are applied; D_I is the distribution for the training sequences. We will state this more clearly in the final version. Clarify the rationale for studying models of different scales and discuss what insights are gained from these comparisons. We wanted to see which components in actual transformer based models were responsible for ICL and we wanted to see whether ICL improved with scale. This puts us apart from a lot of the literature that has tried to show that under certain assumptions transformers CAN do ICL. We want to know what they DO in practice. This was why we tested a number of models, including Attention Only models and MLP only models. We will improve our description of the results in Table 1 and in section 4.6 where we summarize our findings. The reason we speak of coefficients a, b \in [-1,1] for functions ax +b is that given N(0,1) the model has seen functions with those coefficients many times (~70% of its training data in D_F), and that makes a difference to the overall success of its algorithm. For instance consider the plots in Figure 4 where the coefficients are large This discussion is at the end of the background section. This is regardless of the size of the model though the largest models tested have essentially 0 error and the small ones have an error of 0.1 (see our Table 1).

We thank the reviewer for the remarks and questions. We agree with the reviewer that it is an interesting question why the models don’t do what they are capable of. But we were interested in a different, equally interesting (to us) question: what are the models actually doing in this very simple task? First we showed that they are not learning an algorithm like least squares. In the last section of the paper, we provide a mathematical model of what they are doing; they are treating the sequences not as the graph of a function but as just a sequence. They develop an algorithm using something like Olsson et al’s induction heads to interpolate a next value for the given sequence in the prompt from close-by sequences in their pre-training. Did we try to train with multiple distributions? Yes in section 4 we look at three different distributions for function sampling: N, bimodal and uniform (section 4.2). We do say what the models are actually learning, see section 5 of paper. We will clarify that section in our revised version. We have highlighted it also in the abstract. We will clarify that current training methods are not optimal. We wanted to investigate training and prompting methods already in the literature; we show that the proposed methods don't do what is claimed. We will certainly cite and discuss Giannou, Angeliki, et al. "How Well Can Transformers Emulate In-context Newton's Method?”, an interesting paper, in our revised version. It uses linear self attention, which makes sense for the theoretical results, since they want to show that LLMs can in principle approximate Newton’s method. In this, they are like many other papers we cite, seeking to explain a model behavior by a theoretical construction. Our undertaking is different. We start from actual transformer architectures with soft attention and try to determine what they are actually doing in this task. We in fact show that they don’t do any of the reconstructions that we have seen proposed in the literature. Giannou et al. also only examine differences in sampling the sequences of points in the prompt; i.e. they look at in our notation D_I \neq D^t_I, where D_I is the training distribution of points and D^t_I is the test distribution. In particular they fix a particular function and then input into it points that are in the training distribution and points outside of distribution. In distribution is defined as having been very likely to have been seen in training. However, when we have looked at cases where D_I \neq D^t_I, we saw that sometimes for input values within distribution, we got bad results (see figure 4 for example) when the function’s values on those points were outside what we call boundary values. In addition to examining D_I \neq D^t_I, we also vary the distributions from which we sample functions, D_F and D^t_F–i.e. values of the points chosen. Giannou et al did not do this. Finally, to test Giannou et al use only very small models 4 ah 64 embeddings and go only to 6 layers, we go up to 12L8AH and d_emb = 256. They observe that after four layers, there is no significant improvement in performance by adding more layers. Those results are different from what we find (maybe due to their use of linear attention). We have included some comments on Giannou et al and clarified how we differ from the related work we know about in a revised version of the paper that we have submitted on the site.