Paper: Transformers can optimally learn regression mixture models (ICLR-2024)

Reviewer comment (Evaluation):
It is unclear why this setting is of interest. While transformers match the performance of the proposed algorithms and could be used as an alternative for mixture of linear regressions, we should consider the cost of training transformers and the time consumed compared to performing any of the other algorithms.

Author response:
We thank you for your detailed comments and review. 

We would like to begin by pointing out that although transformers could be in the worst-case sense slower than other prediction methods, we found that in practice the transformers actually do perform significantly faster than the state of the art algorithms (such as the EM procedure and Subspace Algorithm used as comparisons in our work). This can potentially be explained by the fact that transformers can leverage GPU acceleration and parallelism that some of these other algorithms cannot. 

To answer your question (“do the authors propose the use of transformers for this task?”): indeed, we believe we are the first to propose transformers or mixtures of linear (or as in our updated revision, non-linear) regressions.  Moreover, we want to emphasize the novelty in our Section 3.2: we observe that not only can transformers adapt to unknown mixture models, they can also do so in a sample-efficient way. This is shown by keeping the sample size used to train the transformer as well as other state-of-the-art algorithms fixed and the same, and then showing that the transformer has essentially the same (or even better) performance at inference time. We also want to emphasize that in the non-linear setting, we are not aware of any algorithms which are able to guarantee good prediction of the labels, and surprisingly the transformer seems to be applicable to this setting with no additional changes to the underlying architecture. 
 
Moreover, we would like to point you to Appendix C which describes the answers to your questions about the models and training parameters. Our model is the same as that used in [2], and our training methodology only differs by the parameter choices described in Appendix C.1. If you have questions which are not addressed by that section, please let us know and we can easily update it in the revision. Finally, we have also provided our exact code along with the submission in case there is any concern regarding reproducibility.  

Finally, we thank you for your detailed reading of our original submission; we addressed your concern regarding the size of the identity matrix (it was a typo in the original submission), and added commentary regarding the hidden approximation errors as addressed in [1].

[1] Akyürek, Ekin, et al. "​​ What learning algorithm is in-context learning? Investigations with linear models." The Eleventh International Conference on Learning Representations. 2022.

[2]  Garg, Shivam, et al. "What can transformers learn in-context? a case study of simple function classes." Advances in Neural Information Processing Systems 35 (2022): 30583-30598.

Reviewer comment (Experiments):
More details should be provided for the experimental set-up.

Author response:
Moreover, we would like to point you to Appendix C which describes the answers to your questions about the models and training parameters. Our model is the same as that used in [2], and our training methodology only differs by the parameter choices described in Appendix C.1. If you have questions which are not addressed by that section, please let us know and we can easily update it in the revision. Finally, we have also provided our exact code along with the submission in case there is any concern regarding reproducibility.

Reviewer comment (Reproducibility):
In the appendix of the paper, it is unclear how the proof is implemented. The results of [1] may require the design of specific encodings and entail some error, which was also not analyzed in [1], but the authors should at least show how it is controlled.

Author response:
Finally, we thank you for your detailed reading of our original submission; we addressed your concern regarding the size of the identity matrix (it was a typo in the original submission), and added commentary regarding the hidden approximation errors as addressed in [1].

[1] Akyürek, Ekin, et al. "​​ What learning algorithm is in-context learning? Investigations with linear models." The Eleventh International Conference on Learning Representations. 2022.

Reviewer comment (Experiments):
While empirical results suggest the transformer might be related to the posterior mean, it would be good to have some interpretability results to assess whether this is true in practice, and if your construction in Theorem 1 is practically relevant: is there any evidence that the blocks shown in Figure 1 are actually being learned by the pre-trained transformer? Where are the $w_i^*$ being stored in the weights? Is it necessary to have at least 5 layers in practice, as in your construction? Is it sufficient?

Author response:
While we certainly agree understanding how exactly the transformer stores parameters is an important question, we do want to point out that this is not a claim in the current work. Indeed, we do not claim to show *how* the transformer learns the weights, only that (1) in experiments, it is getting very close to the optimal predictor and (2) it is possible for the transformer to actually implement the optimal predictor since it is within the model class (as shown by Theorem 1). While it is tempting to infer from (2) that the transformer actually does implement the optimal predictor in the manner we suggest, this is by no means our claim. 

Moreover, we want to point out that your question regarding if the "models does resemble the proposed construction" is equally applicable for the paper [1]. In that work, they only show that the transformer *can* implement the algorithms described therein, but they do not claim to show how that happens or even if it is occurring in the trained model. This is for a simple reason: it is very challenging to assess what exact way the transformer learns a given algorithm. Arguably, it is nonetheless extremely relevant for practical problems that the transformer is very close to the optimal predictor, and one should (in our opinion) view the construction as simply a proof of existence of the optimal predictor within the model class of all predictors implementable by an autoregressive (i.e., causal, decoder only) transformer. 

To further clarify this point, we have also added some additional discussion in the introduction of Section B.2 within the appendix. Thank you for your very interesting question again. While we agree it is very important, on the other hand, we also feel that our paper highlights numerous important findings for transformers as applied to the statistical challenging mixtures of linear and nonlinear regressions settings. These findings we believe should influence further theoretical progress and empirical work in this area.

[1] Akyürek, Ekin, et al. " What learning algorithm is in-context learning? Investigations with linear models." The Eleventh International Conference on Learning Representations. 2022.

Reviewer comment (Experiments):
A more extensive empirical analysis would be useful, particularly on how the results vary when changing problem parameters. For instance, how does the performance change as $m$ varies? In particular, it seems difficult to find all the hidden directions $w_j^*$ once $m$ is too large -- does it start resembling OLS at some point? Does increasing the width, number of heads, number of layers change this?

Author response:
Thank you for your suggestion to investigate the relationship between the number of components, m, and the gap between the transformer and linear regression methods. First, we want to point out that it is more natural under our setup (where we have noise and sample weights from a distribution on a norm-constrained set) to consider the convergence of the transformer to ridge regression (with appropriately set penalty), as this is approximately equal to the Bayes estimator for this setup; see the papers [1, Corollary 3] and [2, Section 3.1.1] for additional details to justify this claim. 

Therefore, following your suggestion, we have now revised the paper to compare the ratio of the risk of transformer and ridge regression as a function of the number of components; please see Appendix G and Figure 13. The qualitative takeaway is: as the number of components grow, it is indeed true that the ratio of the risk of ridge regression to transformers is getting smaller, as your question predicted. Note also that in [3] they already show that transformers are close to OLS when sampling the weights from a Gaussian distribution. Note that in their setting, there is no noise, and so the optimal ridge estimator is in fact with penalty taken to be zero, i.e., OLS itself. Therefore, their observations are also consistent with our experiments: in the large m limit OLS and ridge regression are more competitive, but impressively, the transformer is able to adapt to this structure with out any modification to the model architecture.

[1] Dicker, Lee H. “Ridge Regression and Asymptotic Minimax Estimation over Spheres of Growing Dimension.” Bernoulli, vol. 22, no. 1, Feb. 2016.

[2] Reese Pathak, Martin J. Wainwright, Lin Xiao. “Noisy recovery from random linear observations: Sharp minimax rates under elliptical constraints.” arxiv, 2023

[3] Garg, Shivam, et al. "What can transformers learn in-context? a case study of simple function classes." Advances in Neural Information Processing Systems 35 (2022): 30583-30598.

Reviewer comment (Novelty):
I'm not sure that the new work really addresses limitations in the existing literature. The main motivation is that existing methods are potentially brittle and that their theoretical guarantees do not extend to the model misspecification setting. It isn't clear that this work really demonstrates much of an improvement in that regard. As a result, it isn't clear to me what this approach really offers (other than perhaps simplicity?). It too does not come with any guarantees more generally -- or maybe I have misunderstood?

Author response:
We would like to begin by addressing your question regarding the limitations that this work addresses. Indeed, we would like to point out that in many settings, such as federated learning, one has to select the choice of the number of components in the mixture model by hand. For instance see papers [1, Section 2] and [2, Section 4.2]. This is quite tedious, and it is indeed of interest to have general purpose prediction methods which do not need to deal with the unknown parameter of the number of components in a mixture model. 

Regarding additional benefits of our work, we believe that transformers is potentially desirable due to its ability to adapt to the unknown mixture structure. In our original experiments, this was shown by employing the exact same transformer architecture on models with different numbers of components (5 or 20) and different noise levels (0 or 1). In our revision we also show transformers can extend to nonlinear models (see more discussion below). In all of these cases, we show that our trained transformer can achieve the same error as the posterior mean and argmin algorithms which are the decision-theoretically optimal methods and are very strong comparisons since they are oracle methods that have knowledge of the model. This, we believe, is quite striking: transformers have absolutely no model knowledge, yet nonetheless are performing as well as methods which have exact knowledge of the precise model. 

Moreover, we believe the sample complexity comparison done in Section 3.2 is very novel. Most other transformer papers do not consider the sample complexity of achieving a transformer which achieves low generalization error. Our experiment fixes the training sample size and shows that indeed transformers are nearly as sample-efficient as recently proposed, quite complicated methods which are model-specific. This is quite relevant for practice, where one may not have the luxury of seeking unlimited data and must cope with a finite, and perhaps only moderate-in-size sample. We view these results as quite surprising: transformers have absolutely no explicit knowledge of the model structure and yet nonetheless are achieving error on par with state-of-the-art methods, and with exactly the same sample complexity.

Reviewer comment (Novelty):
It also feels like this result is a bit preliminary and could potentially encompass a wider range of mixture models and statistical settings.

Author response:
Finally, we have further strengthened our results to address your concern about the limited scope of our original set of statistical settings. For instance, in appendix E, we have done extensive additional experiments, where we present new simulation results when the conditional mean of the label given the covariate x is either a multivariate polynomial or a function which itself is expressible via a multilayer perceptron. In both cases, our previous observations from linear mixture models do in fact carry over: the trained transformer is able to achieve nearly the Bayes-optimal risk (that is, the risk achieved by the oracle which knows the exact model used to generate the data). Please see Appendix E for further discussion and details. We believe that these experiments do support the claim that transformers can adapt to unknown (even nonlinear) mixture models, which should be of considerable interest to practitioners who work with data that is reasonably modeled as coming from a mixture distribution.

Reviewer comment (Writing):
Minor typo/suggestion — "gradient descent would naturally extends".

Author response:
Finally, we addressed the typos you pointed out. Thank you for your very close reading of our initial submission.

Reviewer comment (Writing):
Minor typo/suggestion — "better predictor to adapted to the mixtures of linear regressions setting".

Author response:
Finally, we addressed the typos you pointed out. Thank you for your very close reading of our initial submission.

Reviewer comment (Novelty):
I am not convinced about the significance of this work. To clarify that, I would like the authors to address the following questions, how and when can one use the observation in the paper? In the introduction, federated learning is mentioned as a possible application. But, federated learning systems do not use linear models and if they do the number of components is known in advance (which is arguably the biggest advantage of the proposed viewpoint in the paper). I acknowledge though that this work can be a stepping stone towards a mixture of non-linear models, which potentially can have more impact.

Author response:
Thank you for the detailed review. 

We would like to begin by addressing your comments on federated learning. We want to point out that in federated learning applications, it is indeed not always the case that “the number of components is known in advance.” For instance, please see the papers [1, Section 3] and [2, Section 4.2]. In these works, the number of clusters is either learned or treated as a hyperparameter; it is not fixed ahead of time. In fact, as the latter paper says, it is often advisable to “find clusters for several values of [the number of clusters] and use the best one for each client separately using a hold-out set of samples.” 

Regarding motivation, we also wanted to point out that yet another motivation for this work can be found in pretraining LLMs, where one wishes to use diverse data (such as mixture data) with the hope that the transformer can zero-in on the relevant aspect (such as the correct mixture component). We study the simplest linear analogue of this. We also have revised the introduction to include this motivation as well as reference recent literature sharing in this motivation (see the blue text in the introduction of the revision).

Reviewer comment (Evaluation):
I do not have experience with mixture models. But, to me, it seems that a clear advantage of using mixture models is the access to the underline model and the mixture components. This is somewhat lost here. Having an approximation for the posterior mean only is nice, but I assume that in many cases one wants to evaluate each mixture separately in order to make a choice.

Author response:
Regarding other comparisons, please also note that throughout our paper we already compare the transformer to the optimal method which is the running posterior mean: see for instance our Figure 3. Note that the posterior mean is an oracle method since it knows the exact mixture distribution; surprisingly, the transformer is nonetheless able to achieve nearly the same performance as the posterior mean. This should be interpreted as: the transformer is doing essentially as best as possible for the settings described in our paper.

Reviewer comment (Experiments):
Unless I didn't understand something, a potential issue for taking the viewpoint of this paper is that training transformers can be demanding in computation and data. I would expect that in most cases one would like to use these types of models in exactly the opposite cases, i.e., small training sets with limited computation.

Author response:
Regarding the practicality of our setting, we would also like to point out that in the revision, we have attempted to take yet a further step towards addressing more practical and nonlinear models. For instance, in Appendix E, we have done extensive additional experiments, where we present new simulation results when the conditional mean of the label (given the covariate x) is either a multivariate polynomial or a function which itself is expressible via a multilayer perceptron. In both cases, our previous observations from linear mixture models do in fact carry over: the trained transformer is able to achieve nearly the Bayes-optimal risk (that is, the risk achieved by the oracle which knows the exact model used to generate the data). Please see Appendix E for further discussion and details. We thank you for pointing out that nonlinear data “can have more impact” and hope the strong performance exhibited in our experiments can address this aspect of your previous concerns.

Reviewer comment (Experiments):
In section 3.1, under the noisy case, it is not surprising that the OLS model does not work as well since it doesn't model the noise. In my opinion, a more appropriate comparison would be an OLS model with a hyper-parameter for the noise variable which is chosen based on a validation set. Especially in light of the fact that a grid search was done for the dropout rate of the proposed approach. Conversely, and perhaps even more appropriate, is to compare to a Bayesian model and either optimize the noise via the marginal likelihood (or ELBO) or give it a full Bayesian treatment.

Author response:
Regarding other predictors, we disagree that the reason the OLS model does “not work as well” is due to a lack of “model[ing] the noise.” In contrast, OLS is actually known to be decision-theoretically optimal in the scenario where the learner only has knowledge of the noise level (see for instance paper [3, Theorem 1] below). Instead, it is due to the norm constraint we employ on the linear regression model that improvements can be had (along with, of course the Bayesian structure of the generative process: a discrete mixture). Thus, we do agree with you that a better procedure would use the structure of the parameter (although, notably this is a somewhat unfair comparison, since the transformer does NOT have this knowledge). 

Nonetheless, following your suggestion we implemented the optimal ridge regression model under the assumption that the parameter is norm constrained, which consists of regularizing the parameter at the level of the noise level (see for instance the papers [4, Corollary 3] and [5, Section 3.1.1] below for justification of this claim). Moreover, as you seemed to suggest, this can also be seen as a MAP estimator under the Gaussian prior which concentrates on the scaled sphere in R^d. These results are now included in the revised version of Figure 2. Importantly, transformers still significantly outperform ridge regression. 
Note that the only change in the comparison qualitatively is that ridge regression is better than OLS when the prompt length is similar to the ambient dimension.

Reviewer comment (Experiments):
To complement the question of "What is the transformer actually learning?" in section 3.3., I believe that some form of evaluation on out-of-distribution data be done. I suspect that the transformer learns an approximation for the posterior mean only in regions of in-distribution, but outside of it, it will behave in an arbitrary fashion. On the other hand, we know exactly how the posterior mean solution will behave in every region. Perhaps the authors can verify that using a similar experiment to the one in section 3.3 or on a simple 2D problem. If that is indeed the case, then how can you guarantee that indeed the solution found by the transformer matches that of the posterior mean?

Author response:
Regarding out of distribution data, we would like to point out that figures 5 and 6, when compared to figures 8 and 9 in the appendix actually already comprises the comparison you asked for. We compare the performance of the transformer on various out of distribution settings and show that it behaves quite similarly to the posterior mean algorithm on these settings. We also have some discussion of this point in the paragraph immediately before the “Discussion” heading in the main text. To clarify one point, we also do agree that the “transformer learns an approximation for the posterior mean only in regions of in-distribution.” Nonetheless, the similarity to the performance of the posterior mean algorithm on covariate scaling experiments shows that it does gracefully degenerate in the same manner as the optimal in-distribution algorithm when evaluated in OOD settings. We would also like to point out that our methodology to justify the similarity of the transformer and the posterior mean is precisely the same as the paper [6].