Provable In-context Learning for Mixture of Linear Regressions using Transformers

(1): The authors claimed that simulation results showed transformers outperformed the EM algorithm in abstract. However, I did not find those experiments in this paper. Moreover, if the transformers, as the authors claimed, mimic EM algorithm to learn mixtures of linear regression models, then how and why could transformers outperform the EM algorithm?

(2): For the mixture of two linear regression model problems, the authors set $\pi_1=\pi_2=1/2, \beta_1=-\beta_2$. Thus, in Theorem 2.1, the $\beta^{OP}=0$, the optimal prediction of the transformer is $x^T_{n+1}\beta^{OP}=0$. The optimal prediction is always 0 and has no relevance to $x_{n+1}$. I think this is not a good setting for analyzing. Could the author analyze the problem in a more general setting?

(3): The authors only showed the existence of specifically constructed transformers that can mimic EM algorithm to learn mixtures of linear regression models. How transformers in reality can be trained to these parameters is not discussed.

(4): The main component of this paper "transformers can perform multiple gradient ascent steps during each M-step of the EM algorithm" is based on the well-established results in [3], which, to some extent, reduces the contribution and the novelty of this paper.

[3] Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural information processing systems, 36, 2024.

The scope of the paper is a bit limited. It only provides understanding of the in-context learning ability in a quite specific setup. Also, the writing of this paper is not sound enough: it does not provide a clear formulation of the problem setup and the workflow of the EM algorithm. The theorems and lemmas lack intuitions.

1. To further justify why the EM construction is interesting (without analyzing transformer training dynamics), it would be nice if the authors gave some argument (either based on experiments or previous work) suggesting that a transformer does implement an EM-type algorithm.
2. As pointed out in the paper, the current experiments which train the transformer using the pretraining scheme from Section 2.3 suggest that the transformer does not learn to implement EM, as it doesn’t suffer from the same initialization issues as EM.
3. The requirement on the prompt length $n$ being larger than $d$ in Theorem 2.1 seems somewhat limiting, could the authors comment on what might happen in the overparameterized regime where $n \ll d$?
4. Most of the theoretical contributions appear to be somewhat derivative, mostly using prior results from Pathak et al.'24 and Bai et al.'24.

As a reader, I have several doubts and questions - most of them are minor but answers would help my understanding.

1. In the first contribution bullet point, what is meant by oracle predictor? It is not defined in the preceding text.
2. A 2-component mixture implies that for a new datapoints, 2 responses are possible in expectation conditioned on the regressor. However, the transformer output can only provide 1 response via greedy sampling - how are 2 responses collected from the output of the transformer during prediction? I understand what is happening but an algorithm box is important.
3. Can you elaborate on what is meant by "transformers are capable of implementing double loop algorithms such as EM algorithms"? Again I think an algorithm box should answer this.
4. In eq. 3, what is y_i'? Where is the prime coming from? How is y'_{n+1} set?
5. I have a fundamental question - It seems like the authors are not going only after parameter estimation - usually pursued in mixtures literature - but also prediction error. How is it even possible to bound the prediction error unless we know the regressor? Unless something like a min-loss is defined, bounding the MSE in L172-L174 looks fishy. Please decompose the MSE conditioned on the beta's - note that the function f is producing only a single value. If the magnitude of beta increases, then it is impossible to bound MSE in this context.
6. The function read_{\beta} is defined in theorem statement - please provide definition prior to theorem statement. beta is a vector - however in Theorem 2.2, read_y() is used where y is a scalar. This is confusing.
7. Theorem 2.1 does not mention the number of layers in the architecture (along with other details) - How is the FFN layer constructed? Details of the architecture given the dimension of the MoLR problem is critical in understanding.

The main concern is regarding the significance of the results. The main contribution of the paper is in showing that the Transformers can encode the EM algorithm. However, such representability results do not shed light on what algorithmic mechanism Transformers actually learn, which is arguably the more interesting question. Moreover, representability results can be misleading. For example, while numerous works showed that Transformers can encode gradient descent for solving linear regression in-context, recent work shows that the algorithm learned by Transformers demonstrate error convergence rate (as a function of number of layers) similar to second-order iterative methods, which is exponentially faster than first order methods like gradient descent.