Transformers can optimally learn regression mixture models
Authors: Reese Pathak
Rajat Sen
Weihao Kong
Abhimanyu Das
Venue: ICLR 2024 poster (2024)
Keywords: transformers
mixture models
linear regression
Primary area: general machine learning (i.e., none of the above)
TL;DR: We show transformers are capable of achieving near-optimal performance on mixtures of regressions data

Abstract: Mixture models arise in many regression problems, but most methods have seen limited adoption partly due to these algorithms' highly-tailored and model-specific nature. On the other hand, transformers are flexible, neural sequence models that present the intriguing possibility of providing general-purpose prediction methods, even in this mixture setting. In this work, we investigate the hypothesis that transformers can learn an optimal predictor for mixtures of regressions. We construct a generative process for a mixture of linear regressions for which the decision-theoretic optimal procedure is given by data-driven exponential weights on a finite set of parameters. We observe that transformers achieve low mean-squared error on data generated via this process. By probing the transformer's output at inference time, we also show that transformers typically make predictions that are close to the optimal predictor. Our experiments also demonstrate that transformers can learn mixtures of regressions in a sample-efficient fashion and are somewhat robust to distribution shifts. We complement our experimental observations by proving constructively that the decision-theoretic optimal procedure is indeed implementable by a transformer.

Published as a conference paper at ICLR 2023
TRANSFORMERS CAN OPTIMALLY LEARN REGRESSION MIXTURE MODELS
Reese Pathak*
Department of Electrical Engineering and Computer Sciences University of California, Berkeley Berkeley, CA 94709, USA
pathakrseecs, berkeeey.ee. edu
Rajat Sen, Weihao Kong, & Abhimanyu Das Google Research Mountain View, CA 94043, USA
ABSTRACT
Mixture models arise in many regression problems, but most methods have seen limited adoption partly due to these algorithms' highly-tailored and model-specific nature. On the other hand, transformers are flexible, neural sequence models that present the intriguing possibility of providing general-purpor prediction meth- ods, even in this mixture setting. In this work, we investigate the hypothesis that transformers can learn optimal predictor for mixtures of regressions. We construct generative process for a mixture of linear regressions for which the -ecisicn-theoreticc optimal procedure is given by data-driven exponential weights on a finite set of parameters. We observe that transformers achieve low mean- squared eπrоr data generated via this process. By probing the transformer's 's output at inference time, we also show that transformers typically make predic- tions that close the optimal predictor. Our experiments also demonstrate that transformers can learn mixtures of regressions in a sample-efficient fashion and somewhat robust distribution shifts. We complement our experimental observations by proving constructively that the decision -theoretc optimal proce- dure indeed implementable by transformer.
INTRODUCTION
In several machine learning applications-fddeet federatted learning (Wang et al. 2021 crowd- sourcing (Steinhardt et al.) (2066) and recommendations systems (Wang et al., (006)) data is col- lected from multiple sources. Each source generally provides small batch of data: for instance in recommendation systems, a user can provide a source of rating data has encountered. Such batches, on their own, are often too small to learn an effective model for the desired application. On the other hand, by pooling many batches together, improvements can typically be made in the quality of the predictors that can be learned.
subset of items that she
An issue with this "pooling" approach is that if it is done carelessly, then the models which are learned may lack personalization (Ting et al. 1999). For instance, in a recommendation system, such approach could yield model that selects similar actions for dissimilar users. A better ap- proach, however, is to model the problem as mixture of distributions: for instance, we can model the sources as arising from m subpopulations, assuming that sources arising from the subpopulation have similar underlying distributions (Kleinberg & Sandler (004). The sources from single sub- population can then be aggregated for the purposes learning. For instance, in the recommendation systems example, users in the same subpopulation might be identified as having similar preferences and tastes for item genres.
Work was completed during internship Google Research

---

Published as a conference paper at ICLR 2023
A supervised learning formulation of the above setup is that the sources arise from subpopulation
Additionally, assume that within each subpop- ulation the input-output pair (a. 3) follows relation of the form y ft (z) 7), where 7) is zero-mean noise, and z R A batch of i.d data from such a source can be represented as {1z,, Hj M5)ر where is the batch size. Given many such batches, each having examples only from one source, the task is to learn the functions {fi well enough to make good predictions on an- other input, sometimes also referred to as a query, 2k+1-- For instance, given the past ratings of an user, we should be able to determine their subpopulation well enough to infer their preferences on an unseen item.
indexed by an integer E [m]] == {1.2, m).
The simplest version of the formulation above additionally imposes the assumption that the functions Ji are linear: (x) (Wi, (,Ai)) This setting has been studied theoretically in (Kong et al.) 2020 Jain et al.] 2023). ). Kong et al. (2020) introduced the problem and designed an algorithm in the setting where there are as many as ((d) batches with size k O(1), and fewer medium- sized batched of size k = O(ym) However, that work imposed strong assumptions on the covariate distribution, which lead to the paper (Jain et al.) 2023). where these assumptions were relaxed. This latter work proposes a different algorithm that even allows covariate distributions to vary among subpopulations. Nonetheless, their algorithm needs to know problem parameters, such as a T,2 T, hypercontractivity, a spectral norm bound the covariance matrix, the noise level, and the number of subpopulations m. It is not clear how these algorithms will fare with model misspecification, or if they could be extended to applications like federated learning where it is unlikely that the correct model is linear, and distributed algorithms are required.
In this work, we ask the question: Is there deep learning architecture that can be trained using standard gradient decent, yet learns mixture models from batched data and can leverage small batches from source make predictions for its appropriate subpopulation? If so, this would allow us to solve these type of mixture problem without needing highly specialized algorithms that could potentially be brittle with respect to knowing the correct form of the mixture model. Moreover, standard (stochastic) gradient descent would naturally extend to distributed training, using standard techniques from federated learning (Wang et al., 2006).
A natural candidate to address this question is the widely used transformer architecture Vaaswan 2017). Motivated by their immense success in NLP (Radford et al., 2019), computer vi- sion HHan et al. 2022) and in context learning abilities demonstrated by large models (Chowdhery jet al. 2022). several recent works have been aimed to analyze whether transformers can learn algo- rithms (Akyürek et al.) 2023 Garg et al.) 2022 Von Oswald et al.) 2023). These papers train decoder only transformers using prompts of the form (21, I f(1)..... Izj f(x, ... W.. (xxx), Tx++1 where the task is to predict f(+++1) after seeing the portion of the prompt preceding it. These papers show empirically that when f is sampled from a class of linear functions, then transformers learn to per- form linear regression in-context. Akyüüek et al. (2023) also show that transformers can represent gradient descent in the case of linear regression provided through a construction.
al.
Another important motivation for studying mixture models in the context of transformers is the question of data mixtures for pretraining large language models (L.LM's). It is well known that for LLM's succeed in various in-context learning tasks the pretraining dataset needs to be diverse, covering various languages, tasks, programming languages and concepts. This is evident from the complex datasheets for training popular LLMs (Chowdhery et al., 2022]. natural question then is to investigate whether transformers can zero in on relevant parts of the pretraining data given new prompt and model select between algorithms (Bai al., 2023). The mixture model studied in this paper can be seen as simple linear or nonlinear analogue of the smae model.
Our contributions:
We demonstrate that transformers can learn mixtures of linear regressions by training on such mixture data and exhibiting near- near-Bayes-o Bayes- error at inference time.
We strengthen this observation by proving constructively that transformers can implement the optimal method for the mixture distribution on which the transformer was trained.
Our experiments show that transformers are sample-efficiens the transformers' perfor- mance is similar (or better) than mpdel-specific methods, when fixing the training set size.
2

---

Published as a conference paper at ICLR 2023
We evaluate certain inference-tine metrics that capture the nearness of predictions made by the transformer versus another predictor. We show these metrics are smallest when taking the comparator to be the decision decision-theoree theoretic optimal method, thereby further corroborating the hypothesis that transformers optimally learn mixtures of regressions.
We also show that the above empirical observations earry over to settings where f are non-linear specifically when they are polynomials or multi-layer perceptron (MIP))s with two hidden layers.
We suggest that transformers tolerate "small" distribution shifts by investigating transform- ers' performance on both covariate and label shifts to the mixture model.
These contributions, taken together, evidence that transformers can optimally, efficiently, and robustly learn mixtures linear regressions.
Related work: The related work can be broadly divided into a thread that studies the theoretical properties of algorithms for estimation or prediction in regression mixture model as well as another thread that studies the empirical and theoretical properties of transformers on learning regression models. Due to space considerations, we present more detailed overview in Appendix A
1.1 GENERATIVE MODEL DATA
Throughout this paper we consider mixture of linear regression, except in Appendix where an ex- tension to nonlinear models is considered. Underlying the mixture of linear regressions, we consider the discrete mixture
(I)
m71
where {w |im e Rd are normalized such that || 112 = v for each [m]] We consider prompts or batches, denoted P = (21, Vi....,... The Vks Fb+1). Here, for noise level O, we have
n, 24 lid. N(O, La), and yı x4 N(KW, 0²)
The goal is then to predict tk+1, the label for the query Tk+1-
1.2 TRANSFORMERS
Transformers are deep neural networks that map sequences to sequences (Vaswani et al., 2017). In this work, we specifically focus on decoder-only, autoregressive transformers. These models are comprised of multiple layers that map an input matrix Hl E R to an output matrix H' E R Here p denotes the hidden dimension, and q is corresponds to the number of input tokens. The output then fed successively to more such layers. Since the computation in each layer is the same (apart from parameters),
describe the computation occurring for the jith column of H, and hý for the jth column H' Additionally, the prefix matrix H is the p x (i 1) submatrix of H obtained by concatenating the first - 1 columns of H
single layer. Write h2 (H₂z) 1isss
layer is parameterized by sequence of weights. Let Wheads denote the number of attention heads and dau denote the hidden dimension for the attention layer and du denote the hidden dimension for the feedforward (i.e., dense) layer. Then, a layer has the following weights:
Ji-1 (W Pheads (WVV Masadda (WK Ahandd R datxxp Ji-1
J=-1
(3)
(WCy Ahendd Rxdaa Win R and Wout Rxd
for each column i E [p], the computation proceeds in the following fashion.
When a = 0, by N(v, 0) we mean the point mass &c. In the 1, the submatrix be interpreted as 0.
3

---

Published as a conference paper at ICLR 2023
Self-attention:
The layer begins with computing the attention vector, a, R", by
Sij softmax (WW H.) hw]t for E [Wheads]: and,
Vhheada W H.ssij
Above, with slight abuse of notation, we define for any integer / >> 0, softmax: R - RR by the formula softmax(()) (e" - Note that, above, Sij R
Feedforward network: The layer then continues by passing the attention vector (along with the original input column h;) through nonlinear dense layer. This is defined by
h = + h; + Wrout ( X(a; + h:))
Above the notation indicates that the map R - R is applied componentwise to its argument. In this work we take the nonlinearity to be the Gaussian error linear unit (GeLU) (Hendrycks & Gimpel 2016) which is defined by
u a(u) = 2
erf ) for any u R.
Above, erf denotes the Gauss error function. The function A: R R2 denotes layer normaliza- tion (Ba et al. 2016), and is given by
A(u) vp 01 T1
1 P
where
G
P
This is form of standardization where is interpreted the mean (averaging the components) and [ Flpl2/p is interpreted as the variance (averaging the squared deviation to the mean).
REPRESENTATION
In this section, we prove that transformers can actually represent the minimum mean squared error procedure. Indeed, let f: P - R, by any procedure which takes prompt P and outputs an estimate ý on the query, and define the mean squared error (MSE) by
MSE() = Ep [US(P) )++1)
Then by standard Bayesian decision theory, under the observational model described in Section I.I it follows that the mean squared error is minimized at the posterior mean f which is given by
Σ w exp
20 Di (KW xx) y;)
SE(P) (1.(P),, xx+1> where ((P)
(4)
E E=1 exp
1. Σ (ww xi) - v;)
2n²
Formally, MSE(F) > MSE(f;), for all (measurable) f. Note above that à does not depend on Ik+1-
Then our main result is that the function f can be computed by transformer.
Theorem 1. There is an autoregressive transformer which implements the function fz in
See Section B for proof of this claim.
For an illustration of the underlying idea behind Theorem 1 see Figure T for an arithmetic circuit that computes the function f in the case m 3, k 2. The objects rij are residuals, defined as
(5) The first layer computes the values (ray). ). the second layer computes the squares of these values, the third layer computes the (scaled) sum of these values over the index i, which runs over the samples in the prompt, excluding the query. The fourth layer, computes the softmax of these sums,
Tú (w x4) fil + 1}, for ie [k], j [m].
When Sig 0.

---

Published as a conference paper at ICLR 2023
1/1
4/2
1/
711
111
221
121
12
112
122
122
113
113
13
23
V31
732
733
11
21
72
22
713
23
731
732
T33
11
1712
22
2g
(ris
+ 23
T31
732
T33
122
P3
732
P1es1 + P2732 + P3733
Figure 1. Illustration of an arithmetic circuit, implementable by transformer, that computes the posterior mean as defined in display 4 Here, Tij are residuals as defined in display (5) and pj are probabilities obtained via softmax operation, defined in display (6). See main text for description of the computation occurring each level.
1 S 20
(w55 x2) - y1)²
exp
Pj :
for [m]]
(6)
-i=1 exp
1 Σ
( x.) uz)²)
And the final layer computes
Pf((+))J
p. 2k+1 = F2()
where the last equation follows from the definitions in display Therefore, the circuit depicted in Figure able to compute the posterior mean f at least for choices 2, m General- izing the circuit to general (k, m) is straightforward; therefore, circuit and demonstrating that each operation: linear transforms in the first and final layers, squaring in the second layer, summation in the third layer, softmax in the fourth layers all implementable by a transformer.
proof amounts exploiting the
EXPERIMENTAL RESULTS
In this section, present results of training transformers on batches described in Section 1.] Our methodology closely follows the training procedure described (Garg et al. 2022). In
notation of Section I.2 our transformer models set the hidden dimension p 256, feedforward Our
network dimension as dg = 4p = 1024, and the number of attention heads as Rheads models have 12 layers. Additional details on the training methodology can be found in Appendix C
We also release our training and simulation code along with this paper.
3.1 TRANSFORMERS CAN LEARN MIXTURES OF LINEAR REGRESSIONS
To begin with, we investigate the performance transformers on mixture models with various numbers of components and varying noise levels. We plot the performance of the transformer when prompted with a prompt P length kt for & & 60. The normalized MSE the mcan-squared emror between the true labels and the estimated labels, divided by the dimension d 20. Above, the
5

---

Published as a conference paper at ICLR 2023
0.0, components:
variance: components:
1.0
0LS
0LS
Trnsformer
Transormer
0.8
Norma Ized MSE
0.6
Noraa lize
0.4
0.4
0.2
0.2
0.0
0.0
leoght
prompt length
prompt length
1.0, components:
variance: 1.0, companents:
OLS
OLS
Posterior mean
Posterior mean
Argmin
Normalized MSE
Argmin
Normaiized
10n
Ridge
10-1
promptlength
prompt progptiength
Figure 2. Transformer model trained mixture of linear regressions data with and 20 components. Top row mixture data no noise added y (o 0) and bottom row mixture data with noise the optimal penalty taken from either Corollary Dicker (2066) or Section 3.1.1 in Pathak et al. (2023).
added (a 1). We also compare to ridge regression
algorithms that we compare against are
Ordinary least squares (OLS). For a prompt of length k, computes an element à arg win Ejst (w 7 2, "3)² Outputs Wk+1 = w rk+1.
Posterior mean. This is an oracle algorithm. Given prompt P of length k, computes the posterior mean Wc+1 (P) as defined in display (
Argmin This is an oracle algorithm. Given a prompt P of length computes
v = arg min win WE( 1m
ΣΣ
(w (waaz n;)²
(7)
jEkjsk
The prediction is then ūk+1 TT xk+1
Strikingly, we see that the transformer prodictions are as good as-or nearly as good as-the oracle procedures which have knowledge of the true mixture components {wī}. It important to note that OLS is suboptimal in general for mixtures of linear regressions. Nonetheless, the transformer is performing much better than OLS, indicating the trained transformer implements a better predictor which is adapted to the mixtures of linear regressions setting.
In Appendix E we also present result for the extension to non- linear mixtures of regressions. Qualitatively, observe similar behavior where the trained trans- former is competitive with oracle procedures which know the exact mixture model.
Extension to non-linear models:
3.2 COMPARISON OF PERFORMANCE FOR FIXED TRAINING SET SIZE
Next, we investigate whether or not transformers learn mixtures of linear regressions in sample efficient way. To do this, we depart slightly from the training methodology in (Garg et al. 2022)
'oor interpretability of the figures, omit the oracle algorithms above in the noiseless case (o 0) as the is multiple orders of magnitude smaller than the data-driven procedures.

---

Published as a conference paper at ICLR 2023
We first sample fixed training set of size e 115000,30000, 45000, 60000). Then with wiithwtth some hyperparameter tuning to avoid overfitting, as well as a modification to the curriculum training, de-
scrihed in Annendiv C we truin the traneformer ae in that naner We then commare the inference
Training set size: 15000 prompts
Training set size: 30000 prompts
prediction
prediction eror (SA)
Oracle prediction error (EM)
Gracle prediction error (EK)
Posteriar mean IEM)
Posterior mean (EM)
Nomallzed MSE
Posteriar mhan
Normallzzed MSE
Postertor (SA)
Tra reformer
Tra reformer
10-1
10-l
10
20
30
40
50
60
10
20
30
40
50
60
prompt length
prompt length
Training set size: 45000 prompts
Training set size: 60000 prompts
prediction amor
Gracle prediction emor (SA)
Dracle prediction error (EM)
Dracle prediction error (EM)
(EM)
Normalized MSE
Posterior mean (SA)
Normalized MSE
Posterior mean (SA)
Tra neformer
Tra naformer
10-1
10
20
30
40
50
10
20
30
40
60
prompt longth
Figure 3. Comparison of EM, subspace algorithm (SA) in (Jain et al
2023
and Transformers.
The results of our simulation are shown below Figure dures, which have the form of "plug-in" procedures:
We compared against two other proce- S
Posterior mean, EM weights. Here, we first estimate the component means w using batch expectation-Nes expectation-maximaattoo (see Appendix C.2 and Algorithm [ [hereii for details). Then, we form मt the uniform distribution over the estimated weights, and then predict yk441 by (P)..
Posterior mean, SA weights. We follow the approach above, but estimate the weights by using the subspace algorithm (SA), which is Alg. 1 in (Jain et al. 2023).
Note that the 'oracle prediction error' quantities appearing in Figure 3 are essentially the best possi- ble error achievable using the weights estimated by the of weights W estimated by an algorithm. Before normalization by the dimension, it is the noise level plus
min NEW
[w - w
m7
which is easily verified to be the prediction error with oracle knowledge of the nearest element in V the true component mean w under observational model (2). The main take-away from this simulation is that the transformer is able to get very close to the performance of the state-of-the-art model-specific algorithms, even when keeping the sample size the same.
3.3 WHAT IS THE TRANSFORMER ACTUALLY LEARNING?
In this section, we try to understand somewhat better, what algorithm the transformer implement- ing at inference time. To do this, we define the squared distance, for two algorithms f. that map prompt P of length & to a prediction 0k+1 of #x+1
d (S.g) = Ep (((P)-g()))) where & 2 1.

---

Published as a conference paper at ICLR 2023
Number components:
Number compenents:
Argmin 15A)
Argmin (SA)
Argmin (8M)
Argmin [BM)
10-1
Argmin
Argmin (oracle)
Squared distance
Resterior mean (SA)
Squared distance
10-
Posterior moan (SA)
Pesteerion (enank(E
Postericr (EM)
Resterfor mean (oracle)
Posterporsagan.jooe sagan.foracie))
10-2
10-2
50
60
prompt lengtths
prompt lengths
Number of components:
Number components:
10"
Argmin 15A)
Argmin (SA)
Argmin
Argmin (oracle)
Argmin toracle)
Squared distance
Resterior mean (SA)
Squared distance
Pesterior mean (SA)
10-1
Pesterfor
10-1
Posterfor
Posterior mean (aracle)
10-2
prompt longths
prompt longths
Figure 4. Comparing inference- time predictions from transformer versus posterior mean and argmin
procedures with oracle estimated weights. In these simulations, the noise level is set 1.0.
depicts & versus (f. g). taking f to be the transformer, and g to be candidate algorithm listed be below, as l varies between and 60. The compared algorithms are:
Figure 4
Posterior mean, oracle weights. Outputs f=(P) prompt P.
Posterior mean, SA weights. Outputs (P),, with t from the subspace algorithm (SA).
Posterior mean, EM weights. Same as above, but Ω from expectation -aaa--maaiiiionn (EM).
Argmin, oracle weights. Outputs wi P)TT+4 where w P) follows display 7
Argmin, SA weights. Outputs úi TTTT1111 where (P) follows display 7 with wy re-
placed by SA-estimated weights. Argmin, EM weights. Same above, but with EM-estimated weights.
As seen from Figure 4 in all of the simulated settings, the algorithm closest to the transformer at in- ference time is the posterior mean procedure, with the oracle weights. Impressively, this observation holds regardless our choice of the number of mixture components.
3.4 EVALUATION ON COVARIATE SHIFT AND LABEL SHIFT
In this section, evaluate transformers distribution shift settings. The experimental results presented in Figures 55 and 6 The distribution shift settings are described below, where we studied one setting of covariate shift and two settings of label shift.
Covariate scaling: We evaluate the transformer on prompts of length & where the covariates (in- cluding the query) sampled H; N(O, >²Ia) for [k: 1]. This is shift from the training distribution when K # 1. Figure 5 shows the results when taking x E {0.33, 0.5, 1, 2, 3). As we see from the figure, the transformer able handle, to some extent, small shifts, such as K E (0.33, 0.5, 2}. but not shifts much larger than this.
Weight scaling:
We evaluate the transformer on prompts sampled from the mixture distribution
m
scale
Sawz:
where 0.
m
So, the weights w scaled up down by the factor Note that scale n, meaning that no shift. The left panels of Figure T depict results for (0.33, 0.5, 1.0,2,3).
8

---

Published as a conference paper at ICLR 2023
1.0, components:
components:
Norma Izzed MSE
0.33
10n
0.33
Norma lized
0.5
0.5
10-1
prompt lengths
prompt lengths
Figure 5. Evaluating transformer performance on covariate shifts.
Weight shift: We evaluate the transformer on weights sampled from the mixture distribution
m d Owy()))
add m
where wy(() w3
Vd 1d.
Thus, andd shifts each component by additive perturbation of norm The right panels of Figure depict the results for E {0, 0.25, 0.5, 0.75, 1.0). Note that & 0 is no shift: add
Noise variance: 1.0, Number components: 5
Noise variance: 1.0, Number of components:
0.25
0.5
Normaiized MSE
Normalizeed MSE
0.75
10n
0.33
10-1
10-1
10
20
30
50
10
20
30
40
60
prompt lengths
prompt lengths
Noise variance: 1.0, Number of components: 20
Noise variance: 1.0, Number companents: 20
101
169
0.25
0.5
Nomalized MSE
Nomalize MSE
0.75
20a
0.33
10
20
30
40
50
60
100 20
30
40
5.0
60
prompt longths
prompt longths
Figure 6. Evaluating transformers on weight scaling (left panels) and weight shift (right panels).
As seen from above, the transformer is fairly sensitive to weight scaling, as seen from the left panels in Figure 6 On the other hand, the transformer can handle small additive weight shifts, such as 0.25, as depicted in the right panels in Figure 6 In Appendix F we also evaluate the minimum and maximum MSE components within the mixture distribution show the sensitivity of the model to different inference-time mixture proportions.
Comparison to posterior mean procedure: In Appendix D we replicate the figures above, with the change that in place of the transformer, evaluate the performance of the posterior mean procedure, ft defined display (4). At a hegh-level, the posterior mean algorithm is less sensitive covariate scaling, but exhibits similar behavior the transformer the two label shift settings.

---

Published as a conference paper at ICLR 2023
4 DISCUSSION
In this work, we studied the behavior of transformers on mixtures of linear regressions, and showed they can learn these mixture models near-optimally, sample-eficiently. and robustly. The fact that general purpose predictors such as transformers can do well in this statistically-compe setting should be practically useful, especially given its flexibility in the settings considered here.
mixture
Additionally, it would be interesting to study the in-context problem as was done in (Garg et al. 2022], but the mixture setting. Here, the mixture distribution would be sampled from distribution over mixture models for each prompt. In general, the decision-theoreee more complicated compute, as implementing the posterior mean would require computing a high- dimensional integral. Nonetheless, possible to approximate the optimal method with trained transformer? We view this as an intriguing direction for future work.
optimal method could be
REFERENCES
Kabir Ahuja, Madhur Panwar, and Navin Goyal. In-context learning through the bayesian prism. arXiv preprini arXiv:2306.0453 2023.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023, 2023.
Lei Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization. CoRR, Abss/607.064500 2016. /tte:://aaarx wwt.://gg//a/aaaaooooo..osssssss 1607. 0645
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Prov- able in-context learning with in-context algorithm selection. arXiv preprint 276.472306.04677, 2023.
Christopher M. Bishop. Pattern recognition and machine learning. Information Science and Statis- tics. Springer, New York, 2006. ISBN 978-0387-31073--2 0-387-31073-8.
Tom Brown, Benjamin Mann, Nick Ryder. Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 202.
Sitan Chen, Jerry Li, and Zhao Song. Learning mixtures of limear regressions in subexponential time via Fourier moments. In STOC.
pdf] 2020.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, al. Palm: Scaling language modeling with pathways. arXiv preprint arXivv2204.023111 2022.
Lee H. Dicker. Ridge regression and asymptotic minimax estimation over spheres of growing dimension. Bernoulli, 22(1):1-37, 2016. ISSN 1350-7265. doi: 10.3150/14-BE1099 URL
https://ddt /oog 3150 BEJ609
Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30595 2022.
Kai Han, Yunhe Wang. Hanting Chen, Xinghao Chen, Jianyuan Guo, Zhenhua Liu, Yehui Tang, An Xiao, Chunjing Xu, Yixing Xu, et al. A survey on vision transformer. IEEE transactions on pattern analysis and machine intelligence, 45(1):87-110, 2022.
Dan Hendrycks and Kevin Gimpel. Bridging nonlinearities and stochastic regularizers with gaussian eтror linear units. CoRR, abs/1606.08415, 2016. URL
08415
Ayush Jain, Rajat Sen, Weihao Kong, Abhimanyu Das, and Alon Orlitsky. Linear regression using heterogeneous data batches. arXiv preprint arXiv:230.001973, 2023.
10

---

Published as a conference paper at ICLR 2023
Jon Kleinberg and Mark Sandler. Using mixture models for collaborative filtering. In Proceedings of the shirty-sixth annual ACM symposium on Theory of computing, pp. 569-578, 2004.
Weihao Kong, Raghav Somani, Zhao Song, Sham Kakade, and Sewoong Oh. Meta-learning for mixed linear regression. In International Conference Machine Learning, pp. 5394-5404. PMLR, 2020.
Yingcong Li, Muhammed Emrullah lidiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning. 2023.
Yuanzhi Li and Yingyu Liang. Learning mixtures of linear regressions with nearly optimal com- plexity. In COLT. arXiv preprint urXiv.1802.07895 2018.
Samuel Müller, Noah Hollmann, Sebastian Pineda Arango, Josif Grabocka, and Frank Hutter. Trans- formers can do bayesian inference. arXiv preprint urXiv:2112.00510 2021.
Reese Pathak, Martin J. Wainwright, and Lin Xiao. Noisy recovery from random linear observations: Sharp minimax rates under elliptical constraints, 2023.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9 2019.
Jacob Steinhardt, Gregory Valiant, and Moses Charikar. Avoiding imposters and delinquents: Ad- versarial crowdsourcing and peer prediction. Advances in Neural Information Processing Systems, 29, 2016.
Kai Ming Ting, Boon Toh Low, and Ian H Witten. Learning from batched data: Model combination versus data combination. Knowledge and Information Systems, 1:83-106, 1999.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 5998 0998-6008, 2017. URL httpppps/ppoccee.ngg...... 3f5ee24354 7
CC/ htte//apprr/111////e Abst html
Santosh Vempala and Grant Wang. A spectral algorithm for learning mixture models. Journal of Computer and System Sciences, 68(4):841-86) 2004.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. International Conference on Machine Learning, pp. 35151-35174. PMLR, 2023.
Jianyu Wang, Zachary Charles, Zheng Xu, Gauri Joshi, H Brendan McMahan, Maruan Al-Shedivat, Galen Andrew, Salman Avestimehr, Katharine Daly, Deepesh Data, et al. A field guide to feder- ated optimization. arXiv preprint arXiv:2107.0663 100-21777.0617. 2021.
Jun Wang, Arjen P De Vries, and Marcel JT Reinders. Unifying user-based and item-based col- laborative filtering approaches by similarity fusion. In Proceedings of the 29th annual interna- tional ACM SIGIR conference on Research and development information retrieval, pp. 501- 508, 2006.
Xinyang Yi, Constantine Caramanis, and Sujay Sanghavi. Alternating minimization for mixed linear regression. In International Conference on Machine Learning, pp. 613-621. PMIR, 2014.
Xinyang Yi, Constantine Caramanis, and Sujay Sanghavi. Solving a mixture of many ran- dom linear equations by tensor decomposition and alternating minimization. arXiv preprint arXiv: 160.:00489 2016.
Ruiqi Zhang, Spencer Frei, and Peter L. Bartlett. Trained transformers leam linear models in-context. arXiv preprint arXiv:2306.099) 2023.
Kai Zhong, Prateek Jain, and Inderjit Dhillon. Mixed linear regression with multiple components.
In Advances in neural information processing systems (NIPS), pp. 2190-2198, 2016.
11

---

Published as a conference paper at ICLR 2023
A RELATED WORK
The related work can be broadly divided into two categories: (i) theoretical works on learning mix- ture models and (ii) analyzing theoretically and empirically the learning abilities of transformers.
In the context of (i), there numerous works that study the well known mixed linear regression problem with a batch size of 1 (Vempala & Wang 2004, Yi et al. [2014; 2016 Chen et al. |020 Li
& Liang 2018 Zhong et al. 2016). In general the problem NP-Hard as shown (Yi al. 2016).
Therefore most of the above works with the exception of (Li & Liang 2018) makes the assumption
that the covariates of all the mixture components are isotropic Gausstans. However, even with this strong assumption the time-complexity of all these algorithms are at least super-polynomical in m
rendering them impractical.
Kong et al. (2020) pioneered the study of the problem in batch setting where they were motivated by meta-learning multiple tasks. They showed that they can recover well separated mixture of linear models from batched data with polynomial dependence on d, m and the inverse of the fraction of the smallest mixture component. However, this work still had the isotropic covariate assumption. Recent work @Jain al. (023) removed this assumption and further improved the sample complexity and the length of the medium size batches that is required for learning. We compare the training sample complexity of learning using transformers with that of the latter, as well as the popular EM method (Zhong et al. 2016), modified to work with batched data.
In the context of (ii), following the emergence of several hundred billion parameter large language models (LLM)'s ' like (Radford et al. 2019;; Chowdhery et al., (022). it has been observed that such models can learn from few examples supplied in prompt during inference (Brown et al. 2020).
This ability to learn in-context has been studied in simpler settings in many recent works (Garg
al. 2022; Von Oswald al. 2023 Akyürek et al., 2023 Zhang et al. 2023). (Garg et al. 2022)
showed empirically that transformers can learn to perform linear regression in context. (Akyürek
jet al., 2023) then showed that transformers can represent gradient decent for linear regression in
context. A similar result was shown in (Von Oswald et al. (2033) but using linear self attention.
Zhang et al. (2023) go one-step further by showing that gradient flow in linear self-attention based transformers learn to do population gradient decent for linear regression. More general algorithm learning behavior has been demonstrated in (Li et al. 2023) and they also provide stability bounds for in-context learning.
Note that none of these prior works imply ability of transformers to learn mixture models from batch non batch setting. Müäler et al. 2): (Ahuja al. (2023) look at in context learning from a Bayesian perspective. Müller et al. (2021) show that transformers fitted on the respective prior can emulate Gaussian processes. (Ahuja et al.] (223) has section on learning multiple function classes in-context where they empirically study gaussian mixture models with two mixture components. However, they do not study the representation learning problem and training sample complexity is not investigated in depth.
The work in (Bai et al., 2023) is also extremely relevant to this work. They generalize the results in (Garg et al. [202] Akyürek et al. 2023) Von Oswald et al., (023) to show that transformers can represent a general version of in-context gradient descent which lets them show constructions for several in-context algorithms like linear, ridge regression and LASSO, gradient descent in 2-layer neural networks, in-context learning of GL.Ms. Perhaps most related is the part of the paper that deals with model selection. In particular the paper shows that tranformers can choose between two models classes either by a Post-ICL train-test split within context or Pre-ICL distribution testing in-context. Note that these results do not generalize our setting as (i) they do not yield construction of the Bayes Opt in the mixture setting in Section (ii) their experiments include learning with a mixture of different noise settings or distinguishing between a mixture of classification and regression, which is very different from having different 1
B PROOF OF THEOREM 1
In this section, present the proof of Theorem I We begin, in Section B. by stating some pre- liminaries, such as the necessary operators we need to show that the transformer can implement. We then present the proof, assuming that these operators are transformerr-eprresenaabe in Section B.2
12

---

Published as a conference paper at ICLR 2023
Finally, the proof of the representation capacity of these operators by transformers is provided in Section B.3
B.I OPERATORS THAT A TRANSFORMER CAN IMPLEMENT
We now list some operators, for matrix H R that output matrix H' e R The following list includes all the operators we need.
For columns with index E I, outputs Il' where Hige Hgst and the remaining entries unchanged. Here, l' = = (e- k) and K' > k, so that entries are copied "down" within columns i I. Note, we assume > and that k & that the operator is well-defined.
copy_down( H; k, k', E,)):
copy [ver k W, (. For columns with index E I, outputs H' with H Ke H The remaining entries stay the same. Here entries from column i are copied "over" column i.
mul(H; R. I):
For columns with index i E I, outputs H' where
Hirtes = Hktt, Hett,, for I E {0. (0....,e-. k... for i E [k. f. The remaining entries stay the same.
aff H; k K, k", ,,,, e, W. W bII): For columns with index I, outputs H' where Ha WHee,, + W W'eee, b.
Note that t" k" + 8" where W R² W' R and t= = k 8. We
assume 6, 8,8 5 > 0. The remaining entries of H are copied over to H1' unchanged.
scaled [ H: a, b,e,W,i,,),
Outputs matrix H' with entries
Hw44. Σ
Hb4t. tor (
jeI
The set I is causal, so that Ic [i 1J. unchanged. soft H; H:K,e,K) For the final column outputs a matrix with entries
The remaining entries of H are copied over to H',
Hutra
for LE 1, 1.. kj.
Et=0 Heres
The remaining entries of H are copied over to H', unchanged.
The important property of the above list of operators is that all be implemented in a single layer of a autoregressive transformer.
Proposition 1. Each of the operators copy down, copy over, mul, aff, scaled_ _agg, and soft, be implemented by a single layer of an autoregressive transformer.
See Section B.3 for proof of this claim.
B.2 PROOF OF THEOREM 1
In this section, we present the proof of Theorem assuming Proposition T We need to introduce a bit of notation:
Ieven (k) = {2j ie [&]] and Ik) = (2j- -ije [k + 1]].
Additionally, we define W E R to have rows wt Rd, which as we recall from (10.. are the true mixture weights.
We begin by assuming that the input prompt P is provided H R 42 This matrix is such that the only nonzero entries are ((()) 1id:2:11 = xy E Rd for each j [k 1].
Additionally, 10) = yj for each j E [k] Then, by leveraging the operators described above, (9)
we can see that (P) 2d 4.2kk++++1 where the matrix (8) is constructed by the following process:
13

---

Published as a conference paper at ICLR 2023
H copy [oy__down(HH0010 1, (oda
H ( 1, 1, 2d, (k))
H copy_down| 1, (,
H aff ( 1,2d 2, 1,
m 1, W 1ax1, 0, (k)
{2k + 11}
11 mul (1 2,20 2d m m 1,I even (k))
( scaled Booodd_gg(HH) : 25 m 2,20 2m 1, 1,2d m + 2,
1, L,Teven (k))
h soft ( 2d + m + 2, 2d 2m + 1,2d + 2m +
(()) mul( + 2, 2d + 2m 2,2d + 3m 2, 2d m + 1(2k + 1])
H aff 8 2d 3m +, 1, 2d 2, 4m 2d d+1, 4m m, 0, 0, (2k+
1)))
The process above illustrated Section B.2.1 By Proposition I each operation above is imple- mentable by layer of autoregressive transformer. Therefore, this completes the proof.
Discussion the construction: Note that in the operation aff used above, we do assume that the transformer stores the weights W" While at first glance this may seem like a strong assumption, we note that in our setting it is actually sistently as the number of prompts tends to infinity while keeping the prompt length size fixed.(For instance, see the paper Jain et al. (2023) and references in section therein for more discussion.) We view it as interesting line of future work to investigate to what extent the transformer stores approximate versions of the true mixture parameters w explicitly. The main goal the above con- struction is simply to show that it is indeed possible to implement the optimal method viz-a-viz transformer.
cermmtion-thhoreticall possible learn the weights con-
B.2.1 ILLUSTRATION OF PROOF OF THEOREM
We illustrate the steps taken by the transformer to implement the softmax operation. To begin with, the matrix input to the transformer is modelled as below, the case where k Below, (y1, .,.... Rd. Throughout we only show the nonzero entries (i.e., missing rows and columns are always assumed 0). Then, our input
0)
(0) [
32
0/
33
After the copy_down operation, we have
H
271
422
1/2
233
27
253
After the copy_over operation, we have
H
Ø1
//
422
1/2
253
422
After another cowy_down operation, we have
F1
01
02
'/2
C3
H
422
32
0
0
0
After the aff operation, we have,
Σ1
V1
Ø2
V2
Ø3
(4)
F1
F1
02
T2
0
1/1
0
1/2
0
0
0
T3
14

---

Published as a conference paper at ICLR 2023
Note that r; W/ x; y: for i k + 1 and otherwise rk+11 we obtain
WM #k+1- After the mul operation,
y1
22
D2
23
21
22
33
21
y2
71
0
72
73
12
0
Above the square should be interpreted element wise the vectors Ty. operation, we obtain
Then, after the scaled _agg
y1
22
12
253
11
22
222
23
(6)
y1
1/2
0
0
71
0
12
73
ri
0
72
2
1
(rf
+
r2)
Then, after the softmax operation, we obtain
/1
1/2
03
I
22
22
23
(T))
0
y1
0
1/2
0
r1
12
13
0
r7
0
12
(ri
rÿ)
0
0
0
p
Here, p softmax( (rî + r2)). Finally, after yet another mul, obtain
Ø1
Ø1
22
1/
23
21
22
22
23
y1
1/2
0
(8)
0
11
0
12
13
0
2
0
72
20
1
(rí
+ +2)
p
0
0
0
0
por3
Above, denotes elementwise multiplication.
Finally, after an aff operation, obtain
y1
0/2
Ø3
F1
21
T2
.22
23
0
y1
0
V/2
0
(9)
T1
T2
F3
0
ri
0
(rT
+ +2)
20
0
0
0
0
P
poro 13
0
0
0
0
[H)
Note that the bottom-right entry of (9) contains the output f fP)). In other words, the desired
result if H (9)
4m222242 Note that H E RPxx where p 2d + 4m 2, and q 2k 1.
B.3 PROOF PROPOSITION
To begin with, we recall a few definitions, introduced in recent work (Akyürek et al., 2023). Definition (RAW operator). The Gradd-Arihmetic--Wite (RAW) operators are maps on matrices, RPx4 RPxy
RAW. (H: Oz Oy, 0x, H', where {8,+)
Here is causal set-valued map, with (i) = [ - 1]. The operators 8,0 denote elementwise multiplication and addition, respectively. The entries H' are given by
Hk, = 0x AgHHs
max{
Σ Hz:)). and
(8a)
ven())
Hkr H
(8b)
15

---

Published as a conference paper at ICLR 2023
for each i [g]. Note that above 1 [p]] \ K, and for some positive integer O₂ R G, R and Oc R R|
In AnJAkkrrekco et al. (022), they show that the RAW operator can be implemented in one autoregressive transformer layer They also argue that (with a slight change in parameterization) that the mul and Therefore, we simply need to argue that the operators soft, copy_down, copy_ over, and scaled_ agg are all implementable by a transformer.
aff operators are wransfomer-implementtae
To begin with, note that, by inspection, we have, with 8 e- k.
copy_over( H; k, K, [, = RAW a H:; [k, { Ø, [K', K + 6]. 18+1. O, 0,14+1,
(9a)
copy_down( H: k. K M, e,,,,, E,T) RAW H: [k 4 Ø, [K, M' + +5], T6 0 Is+1, 77))
(9b)
scaled_agg( H; at, k f, M, T) = RAW @ H: [k, 4 Ø, [K' k' ++ 5]. To+1, 0,841, mi.)
(9c)
Above, note that the intervals [a, b] are just the integers between and (inclusive) and that we have defined
ii - 1}
i>?
{i)
iE and T²,(1)
j=i
(1)
+0
Ø
otherwise
otherwise
Ø
otherwise
Therefore, the displays 9
establish the following result.
Lemma 1. The operators copy_over, sopy_down and scaled agg all implementable via the RAW operator in single autoregressive transformer layer.
Finally, in Section B we demonstrate the following result.
Lemma 2. The softmax operation implementable by an autoregressive transformer.
This completes the proof of Proposition
B.3.1 PROOF OF LEMMA
In order to implement the softmax operation, we need to introduce few other operations:
div(H:jj K,,, E) For columns with index E I, outputs H' where Hitts
Hu4t,i/ H3 for all t kk.... The remaining entries of H are copies as is into H'.
mov(H: H: k. K,,,K,,,,, I)
For columns with index i E I, outputs H' where H Hatt. k). The remaining entries of H are copies as is into H'
for all / E (O...., ...,..... ..,
sigmoid H: k. R²) In the final column q. outputs H' with H remaining entries of Il copies is into Hl'
14e HE The
The operations div, mov are special cases of the same operations as introduced in the paper Akyürek (et al. (2023). Thus, we only need to demonstrate that sigmoid is cansformeriimlleennaaee As- suming this for the moment, note that the softmax operation softmax is then implementable by the following sequence of operations. Let H denote the input to the softmax layer, and let s = Hggt.q- Using the affine operation (aff) together with the softmax operation (soft) ues 1/(1 + eft). Using the affine operation (aff) together with the div operation, we can invert these values compute ex Finally, we can compute the sum of these values S Σ, e with an affine operation (aff) and we can divide by this sum using another div operation. The result values are e*/ Σ, e*s, which is the softmax of the vector A move operation (mov) then can move these values into the correct locations. He:(k+e-)).qr
compute the val-
required.
Thus, to complete the proof, we need to show how to implement the sigmoid operation. For this, we can begin by using the affine operation to insert a value of in the final column, and another affine operation to insert a 2 x 2 identity matrix in the first 2 columns of H. Then by selecting WK and Qº select the identity matrix and to select (Hr.. 1), respectively, we can ensure that
'Note that in that paper. they more or less ignore the aggregate approximation error. Following this ap- proach, also ignore the lower order terms arising from the approximation of the underlying operations such multiplication, done in Akyürek al. (023)
16

---

Published as a conference paper at ICLR 2023
W H 12x2 and h = (Hk,i. 1), and the corresponding softmax values in the self-attention layer are 8; (o Ha,)). 1/
1)). By selecting W to select the identity matrix in His and WC to select the first value of and place it in position k' + t, we can ensure that a = =1/e"" He, + I)e where e, denotes the jth standard basis vector. This value precisely the sigmoid, as needed. To place this value in the correct location, we simply set the feedforward network matrices Win Wrout O. Then, to preserve the output, we need to delete the value 1, and identity matrices placed into H at the beginning; this can clearly be done by two affine operations.
C ADDITIONAL DETAILS ON TRAINING METHODOLOGY
Our training approach closely follows that of Garg et al. (2022) and Akyürek et al. (2023). After some hyperparameter optimization, we settled on the choice hidden dimension 256, attention heads, and 12 layers. We trained our transformers using Adam, with a constant step size of 0.1. We used curriculum training, as Garg et al. (2022), with the exception of Figure where the sample size was fixed. Our curriculum phases were 2000 steps each, with batch size of 64. The final stage of training had 250000 steps with 64 batches. All of our figures presented mean squared errors computed over batch sizes of 256. The dimension of the original covariates was d 20 throughout this paper.
C.I DETAILS ON FIXED SAMPLE-SIZE TRAINING
In this setting, we used hyperparameter tuning over the dropout parameter, p (0, 0.05, 0.13, and found the following choices be best, for Figure]
for 15000, we took p 0.0.
for 30000, we took 0.1.
for = 45000, we took = 0.0.
for 60000, we took = 0.1.
We also used curriculum training in this setup, but obtained the samples by subsampling the fixed dataset. This was done by first randomly sampling batch from the full dataset, and then randomly dropping and shuffling the prefix of each prompt so as to obtain a prompt of the shorter, desired length. Otherwise, the entire procedure was the same as the other figures, as described above.
C.2 BATCH EXPECTATION MAXIMIZATION (EM) ALGORITHM
is variant of the standard expectation- catiiiiinmmiimizziiiiio method (see, for instance, Section 14.5.1 in (ishhop(2006) For completeness, describe the algorithm formally here. Note that 6 denotes the standard univariate Gaussian pdf below. For notation, we also denote the prompts
Batch expectation-maxiiizziee
p(i) = (a
(6)
. Fx
(8)
(8)
Uk
x41))
for i e [n]]
The algorithm is then stated below Algorithm
17

---

Published as a conference paper at ICLR 2023
Algorithm 1 Batch expectation-maaxiiizzttoo for discrete mixture of linear regressions with Gaus-
sian noise
Require: Length & prompts {P(() noise variance o >> 0, number of components m > 0.
Initialize π (0)) E [0, 10" drawn uniformly on the probability simplex.
Initialize w (0) R drawn uniformly on the sphere of radius vd for j [m]]
Initialize Taj
(0)
for all [m].j E [m].
while have not converged do update prompt-component assignment probabilities,
(6) (z) (T
(1) 11 @6
(++1) Tig
for all ie [n],je e [m].
Di (4) 1o)
- I
(1) k
update the marginal component probabilities by the formula
1 "
(1+1)
(++1)
for all j [m]
76
T1
update the parameter estimates by solving.
W (t+1) arg min WERd
Σ
n
(++1)
w wxxx1)) ) (ui")
for all e [m].
1=1--1
update the iteration counter, t - t+ 1. end while
return final set of component centers, { 13-1
In our implementation we stop (or declare the algorithm converged) if tmaxs or if
max min 100 w ((-1) $ E.
In our experiments took fmax 20000 and 0.001.
D COMPARISON TO DISTRIBUTION SHIFT WITH THE POSTERIOR MEAN ESTIMATOR
In this section, we replicate the figures presented Section 3.4 except we evaluate the distribution shift settings on the posterior mean procedure, f defined in display (4).
18

---

Published as a conference paper at ICLR 2023
1.0, components:
components:
0.33
0.33
109
Norma Izzed MSE
Norma lized
10-1
prompt lengthh
prompp lengths
varlance: 10.0, components:
Noise varlance: 10.0, components:
0.33
0.33
10
x 10"
10
prompe longths
prompt lengths
Figure 7. Posterior mean algorithm on covariate scaling distribution shift setting.
1.0, components:
components:
10
Normalized MSE
100
Normalized
100
0.33
0.33
0.5
0.5
10"4
prompt longths
prompt longths
104
Noise varlance: 10.0, Number components:
Noise varlance: 10.0, components:
0.33
Nomalized MSE
Normalized
100
10a
prompt lengths
prompt lengths
Figure Posterior mean algorithm on weight scaling distribution shift setting.
19

---

Published as a conference paper at ICLR 2023
1.0, compenents:
components:
0.5
0.75
Norma Izzed MSE
0.75
1.0
Norma lieed
1014
10-l
prompt lengiths
prompp lengths
Noise varlance: 10.0, components:
Noise varlance: 10.0, components:
169
0.5
0.75
Normalized MSE
0.75
Normalized
10-1
prompt lengths
prompt lengths
Figure 9, Posterior mean algorithm on weight additive shift setting.
E EXTENSIONS TO NON-LINEAR MIXTURES OF REGRESSIONS
In this section, we consider nonlinear extensions to the model described in Section We consider discrete mixtures of the form
m
(10)
where { are functions lying a nonlinear function class F, mapping Rd to R. We consider prompts or batches, denoted P (æ1 ,1,...,, Mes T&+1). Here, for noise level 0, we have
f~ k..d. N(O. and 2, N(f(zz),
(11)
The goal is then to predict yk+1, the label for the query Tk+1-
We train the transformer models here as described in the main text, with the exception that the data is generated according the distribution above. We also compare, inference time, to the following generalizations of the posterior mean algorithm, denoted fPMA. and the argmin procedure, denoted fAM. These are maps from the observed prompt to function J, given by
E exp
exp
202 (-
1 Σm (1 yi)²)
fPMA
and,
E exp
> (f7 yi)²)
2n2
fAM(P) = arg min jin
Σ (f(xi) - vi)
Thus, in order to specify the setting fully, we need only define our particular choice of function class, F. We do this in each of the subsequent sections.
20

---

Published as a conference paper at ICLR 2023
E.1 MULTIVARIATE POLYNOMIALS
In this section, we consider degree-2 polynomials in R given by
F= f()) = Σ anjz,z ayy E R 1-=1
To compute these efficiently, we select azy w,w, where the weights wy R are such that (Ei- w²)² = d, which ensure that E(y²) = 3d + 1. When plotting the mean squared error (MSE) in the figures below, we divide by 3d + 1. The interpretation is that when the normalized MSE is equal to 1, we are doing no better than predicting 0 for each æ; MSE significantly lower (on the normalized scale) indicates substantial improvement.
In addition to the comparison with the posterior mean and argmin procedures, we also compare against polynomial regression, where y regressed on the degree-2 monomials of the form Ty where j E [d]. We additionally present results with a ridge penalty at the noise level.
Noise variance: 1.0, Number of components: 5
Noise variance: 1.0, Number of components: 20
Rry reg.
Rry reg
Rey neg. wirdige pernlly Tranctimar
Ray neg W rdie pensity Tranfomer
10a
Pocietto mean
Rodtetor mean Argman
Argmin
Nomokized MSE
Nomakzed
10
20
20
40
60
20
50
prompt lengih
prompslangth
Transformer model trained mixture of polynomial regressions data with and 20 components, presented with added noise(a Co.paapaii of trained transformer posterior mean and argmin algorithms. We also compare against degree-2 polynomial regression with and without ridge penalty at the noise level.
Figure 10.
E.2 LAYER FUNLY-CONNECTED NEURAL NETWORKS
variance: Number components:
6 -x10
2 layer MN
Transformer
Posterfor mean
4 x10
Argmin
x
10 1
20
50
prompt longth
Figure 11. Transformer model trained on mixture of layer regression data with components, presented with added noise (a 0.1). Comparison of trained transformer to posterior mean and argmin algorithms. We also compare against layer ML. regression.
We consider the case where f.'ss are hidden layer MLP's similar. Note that (Garg et al. 2022 also had non-linear experiments with one hidden layer MLP's but not in the mixture setting. In particular our functions take the form,
(x( (Wao(Uee + ((9) + 1
(13)
21

---

Published as a conference paper at ICLR 2023
There are such fixed MLP's making the mixture components for generating the data. Each prompt is generated from one component after adding additive noise with a fixed variance. In Fig- ure [ show the performance of our method against posterior mean and argmin procedure. We also include the result from regressing a single MI.P (of the same size as data generating functions) per prompt called layer NN in the figure. The f used in this experiment have 10 neurons in both hidden layers. We set m = 5 and d = 10. The noise variance is set to 0.1 which is roughly 10 percent of the signal for the networks chosen. We can clearly see that the trained transformer behaves very similar to the posterior mean algorithm which is optimal in this setting. Moreover, as expected the MLP regression does not work well since trains one model per prompt.
F SENSITIVITY TO DIFFERENT MIXTURE PROPORTIONS
In this section, we extend our simulation results by considering what happens if the inference-time mixture propotions change. Consider mixture distribution of the form
Σ
P:O
where the weights wy are the same as in model algorithm under is given by the average
10 Note that the mean squared error any
MSE(a) - E: Bu: Eu;
Σ
(f(P) 9
where the expectation under u indicates the distribution of prompt P under the generative model with the prompt weight fixed to be wi.
Evidently, by considering the maximum and minimum choice of mixture distribution n, we have
min Eng [U(PP) 1)+1) < MSE(T)
& max Eug [U((P) -1k+1)
(m(m))
[e[m]
This holds for any choice of mixture distribution क where the weights are we but the mixture pro- portions p; arbitrary.
Therefore, motivated by this incquality, we plot the maximum and minimum MSE for each algorithm considered in the main text, for each choice of prompt length. The results are presented in Figure 12 The main takeaway that the models are fairly insensitive to changing the mixture proportions, indicated by the minimum and maximum of the MSE for each model being quite close to each other.
components:
varisnce:
components:
01 (mma)
mmss)
31S
31S
Ridge reg. (max)
Ridge reg (max)
Rodge reg. gnni
Risge reg (nm)
Transformer (max)
Transformer (max)
Transformer
Transformer
Festertor meen (mai)
Nomollaed MSE
Postwreer mean (mac)
Hormakzed
Testener meen (min))
Posterer mean (min)
Argmin (mas)
Argmin (mas)
Argminn (min]
Argmnin (min))
10
20-2
pramptlength
prompslength
Figure 12. Transformer model trained on mixture of linear regressions data with 5 and 20 components. We also compare to ridge regression with the optimal penalty taken from either Corollary 3 in Dicker (2016) or Section 3.1.1 in Pathak et al. (2023). Each algorithm has two lines to indicate the component with the maximum and minimum MSE, for each prompt length.
22

---

Published as a conference paper at ICLR 2023
G COMPARISON TO LINEAR REGRESSION METHODS AS GROWS LARGE
When the number of components m grows large, the optimal method, when o > O, is a ridge regression estimator (asymptotically, the prompt length grows large); this is justified by Corollary 3 in Dicker (2116) and the result in Section 3.1.1 Pathak et al. (2023).
Thus, investigate what extent this is true in our experiments, with transformers. In Figure T3 we plot the mean squared enrors of ridge regression divided by the mean squared error of our trained transformer. The training of the transformer follows the main text of the paper, and the ridge gression estimates were computed for each k and m on batch of 2048 samples. The qualitative takeaway that as m grows large, this ratio of risks is getting smaller. Since when is large the optimal estimator the optimal estimator should be closer to ridge regression, this can be interpreted as an indication that the same architecture transformers used in our paper is able to be adapted to any structure of mixture distribution with essentially modifications.
4.5
13
2.50
Ratio of MSE: Rodge Transforree
4.0
12
12
10.00
2.45
9.75
11.0
2.40
9.50
3.0
10
10.5
2335
233
2.0
10.0
8.75
1.5
8.50
20
30
components
Figure 13. Comparison the ratio of MSEs of ridge regression divided by the MSE of transformer for various choices of prompt length & and number of mixture components m.
23


Reviews:
[{"review_id": "oOcSAjPv5K", "reviewer": "Reviewer_Pvi3", "summary": "The paper shows that transformer architectures can implement a mixture of linear regressions. Namely, it can implement the optimal solution that uses the true underline model parameters. The authors showcase their claim via a sequence of experiments.", "questions": ".", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The idea of the paper is original and novel.\n* The paper is written well. The illustration of the proof adds to the clarity of it and gives a good intuition.\n* In general, the experiment section is good (although I think it misses some things; but more on that in the next section).\n* Code was provided and the results seem to be reproducible.", "weaknesses": "* 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.\n* 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.\n* 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. \n* 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.\n* 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "d2HNcVEnfZ", "reviewer": "Reviewer_wAvt", "summary": "The authors argue that transformers provide a simple mechanism to efficiently and accurately learn mixtures of linear regression models.  They prove that the optimal solution is representable by such transformers and then experimentally verify that sample complexity/accuracy is on par with existing methods for this task.", "questions": "Questions are in the weakness section above. \n\nMinor typos/suggestions:\n- \"gradient descent would naturally extends\"\n- \"definitions in display (4).\" -> \"definitions in (4).\"\n- \"better predictor to adapted to the mixtures of linear regressions setting\"", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "It is a simply presented, well-articulated problem and solution.  The presentation is clear and more or less self-contained.", "weaknesses": "I'm not sure that the new work really address 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?\n\nAlso, it also feels like this result is a bit preliminary and could potentially encompass a wider range of mixture models and statistical settings.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BmnXfdxZuC", "reviewer": "Reviewer_1RLx", "summary": "The paper studies the performance of transformers on in-context learning problems consisting of mixtures of linear regression models, where the prompt consists of input-output pairs coming from a linear model with one of $m$ different target weight vectors.\nThe authors show that the posterior mean in this problem may be implemented with a well-chosen transformer, and that empirically the predictions behave similarly to this algorithm, in particular outperforming OLS and EM approaches.", "questions": "see weaknesses", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper makes an interesting contribution in the recently popular literature on using transformers for in-context learning regression models.\n\nThe posterior mean in eq.(4) is particularly interesting as a desired goal, as it requires a more complex transformer architecture than related papers, which combines both in-context algorithmic operations, and some \"knowledge\" from the data distribution, in the form of the $w_i^*$ vectors.\n\nThe experiments seem to give promising evidence that this target model may indeed resemble the one learned by transformers.", "weaknesses": "Two points would significantly strengthen the paper:\n\n* 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?\n\n* 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?\n\n* related work: is your setting covered by the general setting in [this paper](https://arxiv.org/abs/2306.04637)?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iSA7dKZ6Lg", "reviewer": "Reviewer_KcGJ", "summary": "The authors explore whether transformer architectures can be trained to learn mixtures of linear models from batched data. They find that Transformers are surprisingly effective at this empirically, and they support their empirical results with a constructive proof that the optimal solution is representable in a transformer architecture.", "questions": "I am most curious about whether you have experimented with any of the questions raised in your next steps. For example,\n\n> in practice, the regression function within each component could potentially be nonlinear. To what extent do transformers perform well in these settings? \n\n> In general, the decision-theoretic optimal method could be more complicated to compute, as implementing the posterior mean would require computing a high-dimensional integral. Nonetheless, is it possible to approximate the optimal method with a trained transformer? \n\nI would be very happy to increase my score if you could show some experiments that show whether Transformers easily generalize beyond the linear Gaussian case (or not - a negative result could also be interesting if it is explained).", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I should note that this paper is out of area for me, and while I followed all the details, I may have missed some of the broader context.\n\nI thought this a was well executed paper:\n - They give a nice constructive argument that shows that transformers can implement mixtures of linear regressions. \n - The experiments are very interesting in that they not only show that mixtures of linear regressions can be learned (this is perhaps not surprising given Garg et al's recent results showing this for linear regression), but that they are also competitive in terms of sample complexity with state of the art algorithms for this problem. I would have expected that you pay a larger cost for generality.", "weaknesses": "When I got to the end of the experimental section, I felt that there was a missed opportunity to look at whether Transformers allow one to easily go beyond the linear mixtures setting. While it is very interesting that Transformers are competitive with recent specialist algorithms for the linear mixture setting, I think the key advantage of a black-box method like a transformer is the ability to directly apply it to settings whether the linear mixture assumptions fail. Investigating how performance degrades (if at all) would have been interesting and would have potentially allowed you to show where Transformers outperform existing approaches.\n\nThis is brought up in the future work section of the discussion, but I think it would have been better to include it in this paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0fuNvldMHA", "reviewer": "Reviewer_sqwx", "summary": "This paper studies the ability of transformer models to perform mixture of linear regressions. Specifically, the authors train transformer models in the task of mixture of linear regressions, in which $k$ possible weight vectors are sampled with equal probability. They show that  there exists a decoder-based transformer that can implement the Bayes optimal for this task. Furthermore, they compare the performance of the transformer models with previously proposed algorithms.", "questions": "1. Did the authors perform experiments in which the weight vectors $w_i^*$ are not sampled with equal probability?\n2. What models exactly the authors used? GPT2? \n3. How many samples did they use to train the models and how many to test them? Which is also the sequence length for in-context learning? Is the sequence length during training the same as the one during inference? \n4. How in the proof of lemma 2, the authors simply select $W_K H:,i = I_{2x2}$ since this matrix is required to have the second dimension equal to the sequence length.\n5. Do the authors propose the use of transformers for this task ? \n6. In the discussion section it is mentioned that \"The fact that transformers... quite useful for practical problems\". Could the authors mention some of those problems in which transformers would be useful ?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "This paper performs extensive experimental comparisons on the transformers' performance with other algorithms. They also extend the setting of previous work from linear regression to mixture of linear regressions.", "weaknesses": "It is unclear why this setting is of interest. I understand that transformers match the performance of proposed algorithms and that they could be used as an alternative for mixture of linear regressions. However, we should think of the cost training transformers and the time consumed, compared to performing any of the other algorithms.  The current models are used to perform language tasks and we are unaware on how optimization/numerical tasks could be merged with language tasks. \n\nFurthermore, I think that more details should be provided for the experimental set-up.\nLooking at the appendix of the paper, it is unclear to me how this proof is implemented. I think that the results of [1] 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. I find in general the proof to be very high level and I had trouble verifying that it is correct. \n\n[1]: Akyürek, Ekin, et al. \"What learning algorithm is in-context learning? investigations with linear models.\" arXiv preprint arXiv:2211.15661 (2022).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)