Transformers as Unsupervised Learning Algorithms: A study on Gaussian Mixtures
Authors: Zhiheng Chen
Ruofan Wu
Guanhua Fang
Venue: ICLR 2026 Poster (2026)
Keywords: In-context learning
Gaussian Mixture Models
Theory
Primary area: learning theory

Abstract: The transformer architecture has demonstrated remarkable capabilities in modern artificial intelligence, among which the capability of implicitly learning an internal model during inference time is widely believed to play a key role in the understanding of pre-trained large language models. However, most recent works have been focusing on studying supervised learning topics such as in-context learning, leaving the field of unsupervised learning largely unexplored.
    This paper investigates the capabilities of transformers in solving Gaussian Mixture Models (GMMs), a fundamental unsupervised learning problem through the lens of statistical estimation.
    We propose a transformer-based learning framework called Transformer for Gaussian Mixture Models (TGMM) that simultaneously learns to solve multiple GMM tasks using a shared transformer backbone. The learned models are empirically demonstrated to effectively mitigate the limitations of classical methods such as Expectation-Maximization (EM) or spectral algorithms, at the same time exhibit reasonable robustness to distribution shifts.
    Theoretically, we prove that transformers can efficiently approximate both the Expectation-Maximization (EM) algorithm and a core component of spectral methods—namely, cubic tensor power iterations. These results not only improve upon prior work on approximating the EM algorithm,
    but also provide, to our knowledge, the first theoretical guarantee that transformers can approximate high-order tensor operations.
    Our study bridges the gap between practical success and theoretical understanding, positioning transformers as versatile tools for unsupervised learning.

Published as a conference paper at ICLR 2026
TRANSFORMERS AS UNSUPERVISED LEARNING ALGO-
RITHMS:
A STUDY ON GAUSSIAN MIXTURES
Zhiheng Chen't, Ruofan Wu Guanhua Fang21
Shanghai Center for Mathematical Sciences, Fudan University,
2 Department of Statistics and Data Science, School of Management, Fudan University zhchen22@m. fudan, edu en, Nuruofanl 90gmai ccom fanggh& fudan edu.
ABSTRACT
The transformer architecture has demonstrated remarkable capabilities in modern artificial intelligence, among which the capability of implicitly learning internal model during inference time widely believed to play a key role in the understand- ing of pre-trained large language models. However, most recent works have been focusing on studying supervised learning topics such as in-context learning, leav- ing the field of unsupervised learning largely unexplored. This paper investigates the capabilities of transformers in solving Gaussian Mixture Models (GMMs), a fundamental unsupervised learning problem through the lens of statistical esti- mation. We propose transformer- baasoeedr--- learning framework called Transformer for Gaussian Mixture Models (TGMM) that simultaneously learns to solve mul- tiple GMM tasks using shared transformer backbone. The learned models are empirically demonstrated to effectively mitigate the limitations of classical meth- pppeetation-Maximizze
(EM) or spectral algorithms, at the same time exhibit reasonable robustness to distribution shifts. Theoretically, we prove that transformers can efficiently approximate both the Expectation- -axiiitiiionaaaiiiiaaaaiiiooo (EM) algorithm and a core component of spectral methods-namell cubic tensor power iterations. These results not only improve upon prior work on approximat- ing the EM algorithm, but also provide, to our knowledge, the first theoretical guarantee that transformers can approximate high-order tensor operations. Our study bridges the gap between practical success and theoretical understanding, positioning transformers as versatile tools for unsupervised learning.
ods such
INTRODUCTION
Large Language Models (LLMs) have achieved remarkable success across various tasks in recent years. Transformers( (asswani et al., 2017), the dominant architecture in modern LLMs(Brown et al., 2020), outperform many other neural network models in efficiency and scalability. Beyond language tasks, transformers have also demonstrated strong performance in other domains, such as computer vision(Han et al., 2023; Khan al., 2022) and reinforcement learning(L.i al., 2023a). Given their practical success, understanding the mechanisms behind transformers has attracted growing research interest. Existing studies often treat transformers as algorithmic toolboxes, investigating their ability to implement diverse algorithms(Von Von Oswald et al., 2023; Bai al., 2023; Lin et al., 2024; Giannou et al., 2025; Teh et al., 2025)-a perspective linked to Saaa--eaannng(Hoopdaaae
al., 2021).
However, most research has focused supervised learning settings, such regression et al., 2023) and classification Giannou et al., 2025), leaving the unsupervised learning paradigm relatively unexplored. Since transformer models typically trained in supervised manner, unsupervised learning poses inherent challenges for transformers due to the absence of labeled data. Moreover, given the abundance of unlabeled data in real-world scenarios, investigating the mechanisms of
"Equal contribution.
1 'Corresponding to: fangghō Fudan fudan.edu. edu.en
Code available https: / //gi thub. com/ korschach 989/ transformer-

---

Published as a conference paper at ICLR 2026
transformers in unsupervised learning holds significant implications for practical applications. The Gaussian mixture model (GMM) represents one of the most fundamental unsupervised learning tasks in statistics, with rich historical background(DAY, 1969; Aitkin & Wilson, 1980) and ongoing research interest(Zhang et al., 2021; Manduchi et al., 2021; Löffler et al., 2021; Ndaoud, 2022; Gribonval et al., 2021; Yu al., 2021). Two primary algorithmic approaches are existing for solving GMM problems: (1) likelihood- basded methods employing the Expectation Expectation-Maximizaaa Maximizaaiin (EM) algorithm(Dempster et al., 1977; Balakrishnan al., 2017), and (2) moment-based methods utilizing spectral algorithms(Hsu & Kakade, 2013; Anandkumar et al., 2014). However, both algorithms have inherent limitations. The EM algorithm is prone to convergence local optima and is highly sensitive to initialization independent of initialization, requires the number of components to be smaller than the data's dimensionality-ann an assumption that restricts its applicability to problems involving many components in low-dimensional GMMs(Hsu & Kakade, 2013).
201 8; Jin et al., 2016). In contrast, while the spectral method is
In this work, we explore transformers for GMM parameter estimation to address two questions. (i) Can Transformers provably work for GMM in-context? (ii) Can Transformers empirically overcome the drawbacks of both EM algorithm and the spectral method? Our answers are affirmative. We find that meta-trained transformers exhibit strong performance on GMM tasks without the aforementioned limitations. Notably, we construct transformer-based solvers that efficiently solve GMMs with varying component counts simultaneously. The experimental phenomena are further backed up by novel theoretical establishments: We prove that transformers can effectively learn GMMs with different components by approximating both the EM algorithm and a key component of spectral methods on GMM tasks.
Main Contributions.
We propose the TGMM framework that utilizes transformers to solve multiple GMM tasks with varying numbers of components simultaneously during inference time. Through extensive ex- perimentation, the learned TGMM model is demonstrated to achieve competitive and robust performance over synthetic GMM tasks. Notably, TGMM outperforms the popular EM algorithm in terms of estimation quality, and approximately matches the strong performance of spectral methods while enjoying better flexibility.
We establish theoretical foundations by proving that transformers can approximate both the EM algorithm and key component of spectral methods. Our approximation of the EM algorithm fundamentally leverages the weighted averaging property inherent in softmax attention, enabling simultaneous approximation of both the E and M steps. Notably, our approximation results also hold across varying dimensions and mixture components in GMM.
We proved that transformers (with RELU activation) can implement cubic tensor power iterations- a crucial component of spectral algorithms for GMM. The proof is highly dependent the multi-head structure of transformers. To the best of our knowledge, this is the first theoretical demonstration of transformers' capacity for high-order tensor calculations.
Related works. Recent research has explored the mechanisms by which transformers can implement various supervised learning algorithms. For instance, Akyürek et al. (2023), Von Oswald et al. (2023), and Bai et al. (2023) demonstrate that transformers can perform gradient descent for linear regression problems in-context. Lin et al. (2024) shows that transformers are capable of implementing Upper Confidence Bound (UCB) algorithms, as well as other classical algorithms in reinforcement learning tasks. Giannou et al. (2025) reveals that transformers can execute in-context Newton's 's method for logistic regression problems. Teh et al. (2025) illustrates that transformers can approximate Robbins' estimator and solve Naive Bayes problems. Kim et al. (2024) studies the minimax optimality of transformers on nonparametrie regression. Some literature on density estimation using LLMs is discussed in Section A.
Comparison with prior theoretical works in unsupervised learning setting. Several recent studies have investigated the mechanisms of transformer- models in mixture model set et al., 2025a; Jin et al., 2024; He et al., 2025b), Among these, He et al. (2025a) establishes that transformers implement Principal Component Analysis (PCA) and leverages this to GMM clustering. However, their analysis is limited to the two-component case, restricting its broader applicability.
The paper Jin et al. (2024) investigates the in-context learning capabilities of transformers for mixture linear models, a setting that differs from ours. Furthermore, their approximation construction of the transformer is limited to two-component GMMs, leaving the general case unaddressed. While
2

---

Published as a conference paper at ICLR 2026
they assume ReLU the activation function contrary to the conventional choice of softmax- their theoretical proofs rely on a key lemma from prior work Pathak et al. (2024) that assumes softmax activation, thereby introducing an inconsistency in their assumptions. The paper He et al. (2025b) studies the performance of transformers on multi-class GMM clustering, a setting closely related ours. However, our work focuses parameter estimation rather than clustering. We give discussion of our theoretical improvements over their work in detail in the following paragraph. From an empirical perspective, their experiments are conducted small-scale transformer, which fails to validate their theoretical claims.
function-contrary
Sharpness of our results. Our theoretical analysis fully leverages key architectural components of Transformers: the query-key-value mechanism, multi-head attention, and the properties of the activation function. It is worth pointing out that our result improves the prior work for EM ap- proximation in several points: First, Our analysis shows that Transformers can approximate L-step EM algorithms with just O(L) layers, significant improvement over prior work (He al., 2025b) which requires O(KL) layers (dependent on the number of components K). Second, unlike He et al. (2025b), which needs number attention heads M - +oo to get valid bounds, our results hold with M (1), aligning better with real-world designs. Third, our approximation bounds scale polynomially in dimension d, unlike He et al. (2025b)'s 's exponential dependence- crucial improvement for high-dimensional settings. We believe our results and proofs can offer profound insights for subsequent theoretical research on transformers.
Organization. The of paper is organized follows. In Section 2, some background knowledge is introduced. In Section 3, we present the experimental details and findings. The theoretical results proposed in Section 4, and some discussions are given in Section 5. The proofs and additional experimental results are given in the appendix.
Notations. We introduce the following notations. Let [n] == 21,2, .n]. All vectors are represented as column vectors unless otherwise specified. For vector E Rd we denote ||v|| as Euclidean norm. For two sequences a, and 6, indexed by n, we denote a, = (b,,) if there exists universal constant C' such that a, I Cb, for sufficiently large n.
2
METHODOLOGY
2.1 PRELIMINARIES
The Gaussian mixture model (GMM) is cornerstone of unsupervised learning in statistics, with deep historical roots and enduring relevance modern research. Since its early formalizations(DAYY, 1969; Aitkin & Wilson, 1980), GMM has remained a fundamental tool for clustering and density estimation, widely applied across diverse domains. Recent advances have further explored the theoretical foundations of Gaussian Mixture Models (GMMs)LLöfllrr et al., 2021; Ndaoud, 2022; Gribonval et al., 2021), extended their applications in incomplete data settings(Zhang et al., 2021), and integrated them with deep learning Trameworks(Mandddh et al., 2021; Yu et al., 2021). Due to their versatility and interpretability. GMMs remain indispensable in unsupervised learning, effectively bridging classical statistical principles with modern machine learning paradigms. We consider the (unit-variance) isotropic Gaussian Mixture Model with K components, with its probability density function as
p(x = Σ T&M(Z;;
(1)
where ((z; is the standard Gaussian kernel, Lee o(x; 14) (2mja/2 exp (-2 - T) -1).
The parameter 0 is defined O U μ., where == T2, nK), Th E R and μ {mi, (pi,22," M2, [K] i (4k E Rd, E [K]. We take samples X = (Xthe(N) from model (1).
(XiheeM can be also rewritten
X; My. Zi,
where (UAheeN) are i..dd discrete random variables with P(y = k) = T for E [K] and {Z.)ee are i.idd. standard Gaussian random vector in Idd.
The EM algorithm(Dempstee et al., 1977) remains the most widely used approach for GMM parameter estimation. Due space constraints, propose the algorithm in Section B. Alternatively, the
3

---

Published as a conference paper at ICLR 2026
spectral algorithm(Hsu & Kakade, 2013) offers efficient moment- nomenttbased approach that estimates parameters through low-order observable moments. A key component of this method is cubic tensor decomposition( (Anaaiionnnnnnnnaaaaaa et al., 2014). For brevity, we defer the algorithmic details to Section B.
Next, we give a rigorous definition of the transformer model. To maintain consistency with existing literature, we adopt the notational conventions presented in Bai et al. (2023), with modifications tailored to our specific context. We consider a sequence of N input vectors { C RD,, which can be compactly represented as an input matrix H [hi....,Al hy] E RDN,, wheee each h2 corresponds to column of H (also referred to token).
Here we introduce several useful definitions and their full notations are given in Appendix C. Definition (Attention layer). A (self- (self-Jattention layer with M heads is denoted as Atine ()) with parameters attn = {{Vm m. Qm, Km) (mmm)mm[M)
ImE[M]] C. DxD,
Definition 2 (MLP layer). (token-wise) MLP layer with hidden dimension D' is denoted as MLP L (-) with parameters mip (W1, W2) E D'xD xRDxD',
Definition 3 (Transformer). An L-layer transformer, denoted as TF ()) is composition of L self-attention layers each followed by an MLP layer:
T (H)- = MLP (1) (Attn 2
(G) att
(...MLP
(1)
Attn (1) arEE
2
(H) )))
2.2 THE TGMM ARCHITECTURE
recent line of work(Xie et al., 2021; Garg et al., 2022; Bai et al., 2023; Akyürek et al., 2023; Li al., 2023b) has been studying the capability of transformer that functions as data-driven algorithm under the context of in-context learning (ICL). However, in contrast to the setups therein where inputs consist of both features and labels, under the unsupervised GMM setup, there is no explicitly provided label information. Therefore, we formulate the learning problem as learning an estimation algorithm instead of learning prediction algorithm in the case of ICL. A notable property of GMM is that the structure of the estimand depends on an unknown parameter K, which is often treated as a syyper-aramette in GMM estimation(Titeen (Ttterrngton al., 1985; McLachlan & Peel, 2000). For clarity of representation, we define an isotropic Gaussian mixture task as T = (0, X, K), where X is i.i.d. sample generated according to ground truth 0 according to the isotropic GMM law and K is the configuration used during estimation which we assume to be the same the number of components of the ground truth 0. The GMM task is solved via applying some algorithm A that takes X and K as inputs and outputs an estimate of the ground truth 0 = A(X; K).
In this paper, propose transformer- bassed architecture, transformers-for- fow--oosi-aasaan---- mixttrees (TGMM), as a GMM task solver that allows flexibility in its outputs, while at the same time being parameter-efficient,, as illustrated in Figure 1: A TGMM model supports solving different GMM tasks with K E IC == (KI..., K.}. Given inputs N data points X E RdxN and a structure configuration of the estimand K. TGMM first augments the inputs with auxiliary configurations about K via concatenating it with a task embedding P = embed( K),, i.e., H = [X]]P] and use a linear Readin layer to project the augmented inputs onto a shared hidden representation space for several estimand structures (Ki,..., K2), which is then manipulated by a shared transformer backbone that produces task-aware hidden representations. The TGMM estimates are then decoded by task-specific Readout modules. More precisely, with target decoding parameters of K components, the Readout module first performs an attentive-pooling operation(Lce et al., 2019):
O = -V..H)SoftMaa ((,,)) Q,)) E (d(++))*
where Vo, K. E R(d+K)xD Q E [dd+K)xK The estimates for mixture probability are then extracted by a row-wise mean-pooling of the first K rows of O, and the estimates for mean vectors are the last d rows of O. We wrap the above procedure as { (ikiiejk] Readoute... (H). TGMM is parameter-efficienn in the sense that it only introduces extra parameter complexities of the order Ossd) in addition the backbone. We give more detailed explanation of the parameter efficiency of TGMM in appendix Section D. We wrap the TGMM model into the following form:
TGMMe X: K) = Readoute... (TF o (Readine,
[[eaeine__(X(]lembeeee (XX]embed
Above, the parameter θ = (0rr, O p consists of the parameters in the transformer O and the parameters in the Readin and the Readout, functions O... O
4

---

Published as a conference paper at ICLR 2026
GMM inputs
Transformer backbone
Readouts
Outputs
K K₁
Readouts
(fi Mi Heml
K' = K2
Readin
TF
Readout
Ka (T Rit Jkml
K = K
Readout,
K. (Tai fir Jamm
Figure Illustration of the proposed TGMM architecture: TGMM utilizes a shared transformer backbone that supports solving different kind of GMM tasks via task-specific Readout strategies.
Algorithm TaskSampler
Algorithm 2 (Meta) Training procedure for
Require: sampling
distributions
TGMM
Pu.PP. PK-
1: Sample the type of task (i.e., number of mix-
Require: task dimension d task types IC (Ki....; K2), number of tasks n per step, number of steps T.
ture components) K PK-
2: Sample GMM task according to the type of task
1: Initialize a TGMM model TGMM
2: for t 1 do
(1u, T); 1 Pus pm
3: Sample n tasks {Ti)eejn) independently
using the TaskSampler from Algorithm 1.
4: Compute the training objective
where
tu pKk.
{T k
T.n
O (Q
as in (2).
3: Sample the size of inputs N PN-
Update (t-1) into ((+) using any gra- dient based training algorithm like Adam
4: Sample the data points X
(X1...., .N) Lid. p(-10).
6: end for
5: return An (isotropic) GMM task T
7: return Trained model TGMM (T).
(X, K.).
2.3 META TRAINING PROCEDURE
We adopt the meta teta-training framework as in Garg et al. (2022); Bai et al. (2023) and utilize diverse synthetic tasks to learn the TGMM model. In particular, during each step the learning process, we first use a TaskSampler routine (described in Algorithm 1) to generate batch of n tasks, with each task having probably distinct sample size. The TGMM model outputs estimates for each task, i.e., {ea, tx [K] TGMMe (X; K). Define मे == [#*]ee[K] and ğú == { ( For a batch of tasks (T.) (X,, 0,, K, )e)ee()) denote by 0, UTUUT, and 0, -ū TGMMe (X,; K.),
i [n]. Then the learning objective is thus:
1 (e) = 7I M
(p m.) + + (T, T.).
where f,. and E, are loss functions for estimation of and ग, respectively. We will by default use square loss for f,, and cross entropy loss for e, Note that the task sampling procedure relies on several sampling distributions Pa: PN. PK. which are themselves dependent upon some global configurations such as the dimension d as well as the task types K. We will omit those dependencies on global configurations when they are clear from context. The (meta) training procedure is detailed in Algorithm 2.
EXPERIMENTS
In this section, we empirically investigate TGMM's capability learning to solve GMMs. We focus on the following research questions (RQ): RQ1 Effectiveness: How well do TGMM solve GMM problems, compared to classical algorithms?
RQ2 Robustness: How well does TGMM perform over test tasks unseen during training?
RQ3 Flexibility: Can we extend the current formulation by adopting alternative backbone architec-
tures relaxing the isotropic setting more sophisticated models like anisotropic GMM?
5

---

Published as a conference paper at ICLR 2026
3.1 EXPERIMENTAL SETUP
We pick the default backbone of TGMM similar to that in Garg et al. (2022); Bai et al. (2023), with a GPT-2 type transformer encoder(Radford et al., 2019) of 12 layers, 4 heads, and 128- do--dimessional hidden state size. The task embedding dimension is fixed at 128. Across all the experiments, we use AdamW(L W(Loscchilov et al., 2017) as the optimizer and use both learning rate and weight decay coefficient set to 10-1 without further tuning. During each meta-training step, we fix the batch size to be 64 and train 10 steps. For the construction of TaskSampler, the sampling distributions are defined follows: For рk. We sample K uniformly from {2, 3, 4,53; For Pμ.. given dimension d and number of components K, sample each component uniformly from [-5, 5] Additionally, to prevent collapsed component means(Ndaoud. 2022). we filter the generated mean vectors with a maximum pairwise cosine similarity threshold of 0.8. For Ps. given K, we sample each иTą uniformly from [0.2, 0.8] and normalize them to be a probability vector; For Pv. Given a maximum sample size No. we sample N uniformly from [No/2 No)- The default choice of No is 128. During evaluation, we separately evaluate 4 tasks with 2, 3, 4, 5 components, respectively. With a sample size of 128 and averaging over 1280 randomly sampled tasks.
Metrics. We use f2 errro as evaluation metrics in the experiments. We denote the output of the TGMM 0 == P1, R2, P2, ... ........K [x, The rigorous definition is
1
(2117800 -mll (Fa()) -m:)² m))) KE[K]
Σ
K
where ē is the permutation such that ā argmin, EkeiK |I --il² We obtain the per-
mutation via solving linear assignment program using the Jonker- 2016). We also report all the experimental results under two alternative metrics: cluster-laasifficttio accuracy and log- lideeoooooodd in Section H.2.
Jonker-Volgenan Volgenant algorithm(Crouse
3.2 RESULTS AND FINDINGS
RQI: Effectiveness
We compare the performance of a learned TGMM with the classical EM algorithm and spectral algorithm under 4 scenarios where the problem dimension ranges over {2, 8, 32, 128). The results reported in Figure 2. We observe that all three algorithms perform competitively (reaching almost zero estimation error) when K = 2. However, as the estimation problem gets more challenging as K increases, the EM algorithm gets trapped in local minima and underperforms both spectral and TGMM. Moreover, while the spectral algorithm performs comparably with TGMM, it cannot handle cases when K > d, which is effectively mitigated by TGMM, with corresponding performances surpassing those of the EM algorithm. This demonstrates the effectiveness of TGMM for learning an estimation algorithm that efficiently solves GMM problems.
X
EM
BM
Special
Sentra
Saccrll
IGMI
32
1.5
125
135
Figure 2: Performance comparison between TGMM and two classical algorithms, reported in E2-error.
RQ2: Robustness To assess the robustness of the learned TGMM, we consider two types of test-time distribution shifts:
1. Shifts in sample size N Under this scenario, evaluate the learned TGMM model on tasks with
sample size ieest that are unseen during training.
2. Shifts in sampling distributions Under this scenario, test the learned TGMM model tasks that are sampled from different sampling distributions that are used during training. Specifically, the same training sampling configuration as stated in Section 3.1 and test the following

---

Published as a conference paper at ICLR 2026
perturbed sampling scheme, with jig = μe + toEE where pa sid. Unif (1-5, Ek k E [K] and {=*1 independent with {Jes (E())
aid. NO.
126
12 12
D4-228
120
325
125
Figure 3: Assessments of TGMM under test-time task distribution shifts I: A line with Na + Nest draws the performance of a TGMM model trained over tasks with sample size randomly sampled in |graan /2, Ng and evaluated over tasks with sample size Aiest We can view the configuration 128 + 128 as an in-distribution test and the rest as out-of-distribution tests.
In Figure 3, report the assessments regarding shifts in sample size, where we set Niest to be 128 and vary the training configuration No to range over {32, 64, 128), respectively. The results demonstrate graceful performance degradation out-of-domain testing performance in comparison to the in-domain performance. To measure performance over shifted test-time sampling distributions, we vary the perturbation scale Op E {0, 1,..,, 10) with problem dimension fixed results are illustrated in Figure 4 along with comparisons to EM and spectral baselines. As shown in the results, with the increase of the perturbation scale, the estimation problem gets much harder. Nevertheless, the learned TGMM can still outperform the EM algorithm when K 2. Both pieces of evidence suggest that our meta-traine meta-training procedure indeed learns an algorithm instead of overfitting to some training distribution.
The
Sanctral
Specisua
UUN
23
25
Figure 4: Assessments of TGMM under test-time task distribution shifts II: eg-error of estimation when the test-time tasks Tte are sampled using mean vector sampling distribution plea different from the one used during training.
RQ3: Flexibility Finally, we initiate two studies that extend both the TGMM framework and the (meta) learning problem of solving isotropic GMMs. In our first study, investigated alternative architectures for the TGMM backbone. Motivated by previous studies( et al., 2024) that demon- strate the in-context learning capability of linear attention models such Mamba series(Gu & Dao, 2023; Dao & Gu, 2024). We test replacing the backbone of TGMM with a Mamba2(Dao & Gu, 2024) model with its detailed specifications and experimental setups listed in Section H.I. The results are reported in Figure 5, suggesting that while utilizing Mamba2 as the TGMM backbone still yields non-trivial estimation efficacy, is in general inferior to transformer backbone under comparable model complexity.
In our second study, we adapted TGMM to be compatible with more sophisticated GMM tasks via relaxing the isotropic assumption. Specifically, we construct anisotropic GMM tasks via equipping it with another scale sampling mechanism Pa. where for each task we sample a ~ softplus(ō) with ā being sampled uniformly from [-1, 1]d. We adjust the output structure of TGMM accordingly so that its outputs can be decoded into both estimates of both mean vectors, mixture probabilities, and scales, which are detailed in Section H.1. Note that the spectral algorithm does not directly apply to anisotropic setups, limiting its flexibility. Consequently, we compare TGMM with the EM approach and plot results in Figure with the fz-error metric accommodating errors from scale estimation. The
7

---

Published as a conference paper at ICLR 2026
results demonstrate similar trend as in evaluations in the isotropic case, showcasing TGMM as versatile tool in GMM learning problems.
Additional experiments We postpone some further evaluations to Section H. where we present complete report consisting of more metrics and conduct several ablations on the effects of backbone scales and sample sizes.
R 3
R -444
vansbera
Yonsbrrer
traefore rumbaz
treaforee rumbaz
22
33
125
125
Figure 5: Performance comparisons between TGMM using transformer and Mamba2 as backbone, reported in f fr-orr
K
194M
TO/M
TGMM
11.0
300
2.5
3.0
1.0
60
0.
0.0
Figure 6: Performance comparison between TGMM and the EM algorithm on anisotropic GMM tasks, reported in f2-error.
Remark 1. One might be concerned with the fairness of comparisons between TGMM pre-training and EM/spectral method. We would like to point out that the only additional information that TGMM receives during veta-trainnng is the (implicitly provided) distributional information. The empirical results show that TGMM can generalize beyond the meta taattttaiiing distribution.
4 THEORETICAL UNDERSTANDINGS
In this section, we provide some theoretical understandings for the experiments.
4.1 UNDERSTANDING TGMM
We investigate the expressive power of transformers-for Gaussian-mixturese in Section 3. Our analysis presents two key findings that elucidate the transformer's effectiveness for GMM estimation: 1. Transformer can approximate the EM algorithm; 2. Transformer can approximate the power iteration of cubic tensor.
as demonstrated
Transformer can approximate the EM algorithm. We show that transformer can efficiently approximate the EM algorithm (Algorithm B. see Section B) and estimate the parameters of GMM. Moreover, we show that transformer with one backbone can handle tasks with different dimensions and components simultaneously. The formal statement appears in Section F due to space limitations.
Theorem 1 (Informal). There exists 2T.llayer transformer TF such that for any d S do, K < Ko and task T = (X,0, K) satisfying some regular conditions, given suitable embeddings, TF approximates EM algorithm L steps and estimates O efficiently.
Transformer can approximate power iteration of cubic tensor. Since directly implementing the spectral algorithm with transformers proves prohibitively complex, we instead demonstrate that transformers can effectively approximate its core computational step-the power iteration for cubic tensors (Algorithm in Anandkumar et al. (2014); single-layer transformer can approximate the iteration step:
Section B). Specifically, we prove that
((+ = T 00) o0) j JeN,
where denotes the identity matrix and T represents the given cubic tensor. For technical tractability, we assume the attention layer employs ReLU activation function. The formal statement appears in Section G due to space limitations.
8

---

Published as a conference paper at ICLR 2026
Theorem 2 (Informal), There exists 2L-layer transformer TFe with ReLU activation such that for any d I dy. T E Raxddd and vl0) E IR,, given suitable embeddings, TFe implements steps of (3) exactly.
We give some discussion of the theorems in the following remarks.
Remark 2. (1) Theorem demonstrates that a transformer archisecture can approximate the EM algorithm for GMM tasks with varying numbers of components using a single shared set of param- eters (i.e., one backbone (). This finding supports the empirical effectiveness of TGMM (RQ1 in Section 3.2). Additionally, Theorem 2 establishes that transformers can approximate power iterations for third-order tensors across different dimensions, further corroborating the model's ability to generalize ucross GMMs with varying component counts.
(2) Theorem holds uniformly over sample sizes N and sampling distributions under mild regularity conditions, aligning with the observed robustness of TGMM (RQ2 in Section 3.2).
Remark 3. Different "readout" " functions are also required to extract task cask-specific parameters in our theoretical analysis, aligning with the architectural design described in Section 2.2. For further discussion, refer to Remark F.3 in Section F:2.
4.2 PROOF IDEAS
Proof Idea of Theorem We present a brief overview of the proof strategy for Theorem 1. Our approach combines three key components: (1) the convergence properties of the population-EN, algorithm( ( won Caramanis, 2020), (2) concentration bounds between population and sample quantities (established via classical empirical process theory), and (3) a novel transformer architecture construction. The transformer design is specifically motivated by the weighting properties of the softmax activation function, which naturally aligns with the EM algorithm'' update structure. For intuitive understanding, Figure 7 provides a graphical illustration of this construction. The full proof in Section F.
Compate [mr]( (x](x via softrmax actention
Approodmate log
Compute (TT, Pa)s via softrmux attention Attention Layer
Approximate log x and x" clean some termu
clean some terms
Attention Layer
MLP Layor
MLP Layer
Eatep
Matep
Figure 7: (Informal version) Transformer Construction for Approximating EM Algorithm Iterations. The word "clean" means setting all positions of the corresponding vector to zero.
To approximate (3), perform two-dimensional computation within a single-layer transformer. The key idea is to leverage the number of attention heads M to handle one dimension while utilizing the Q, K, V structure in the attention layer. Specifically, let T (Tig.m)is..eld and vن = (v! )ee]4] Then, (3) can be rewritten as g(++1) = Lj.meid UyUtT..j.ms where T.j.m (T. ke|4 Rdd This operation where each head processes a dimension of size d (Figure 8). The complete construction and proof are provided in Section G.
Proof Idea of Theorem 2.
be implemented using d attention heads,
R h
h
K.,hh
V.nh
T
Figure& 8: Illustration of implementing (3) via multi-head attention structure, where colored boxes denote corresponding implementation components. Here denotes the ReLU function.
CONCLUSION AND DISCUSSIONS
In this paper, we investigate the capabilities of transformers in GMM tasks from both theoretical and empirical perspectives. Our work is among the earliest studies to investigate the mechanism of transformers in unsupervised learning settings. Our results establish fundamental theoretical guaran- tees that Transformers can efficiently implement classical algorithms-suee as the EM algorithm and spectral methods. This is consistent with our empirical finding that the performance of our meta-training algorithm can interpolate between EM and the spectral method. It also opens room for future improvement of attention- based meta-training algorithms in a broader class of unsupervised learning problems. We discuss the limitations and potential future research directions in Section E.

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
We appreciate the constructive comments from the anonymous reviewers. Guanhua Fang is partly supported by Shanghai Educational Development Foundation (23CGA02) and National Natural Science Foundation of China (nos. 12301376).
REFERENCES
Murray Aitkin and Granville Tunnicliffe Wilson. Mixture models, outliers, and the em algorithm. Technometrics, 22(3):325-33) 1980.
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, 2023. URL https:///pperrviiie net forum?? id=0g0X4H8yN4I.
Animashree Anandkumar, Rong Ge, Daniel Hsu, Sham M. Kakade, and Matus Telgarsky. Tensor decompositions for learning latent variable models. Journal of Machine Learning Research, 15 (80):2773-2832, 2014. URL http: / htllmmmrrooo .rgg wwtggg/aprs////e/aa//ar/ge/sseee /v15. anandkumar 4b html.
Yu Bai, Fan Chen, Huan Wang. Caiming Xiong, and Song Mei. Transformers as statis- ticians: Provable in-context learning with in-context algorithm selection. In Advances in Neural Information Processing Systems, volume 36, pp. 57125-57211. Curran Asso- ciates, Inc., 2023. URL htw.hhttss//prrcceeeiiiiee wwwrips. http:nnns..rrrrrrrreeeeeeeeeeeee 2023 2002:f/2222222eee313330222222222 cctpppprrfffilllspppppeeeee pdff.
Sivaraman Balakrishnan, Martin J. Wainwright, and Bin Yu. Statistical guarantees for the EM algorithm: From population sample-based analysis. The Amnals of Statistics, 45(1):77 120,
2017. doi: 10.1214/16-ASS URL https: http:///ooo.......1111.211 1a0.://116A0113..11111110
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Adítya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS '20, Red Hook, NY, USA, 2020. Curran Associates Inc. ISBN 9781713829546.
David F Crouse. On implementing 2d rectangular assignment algorithms. IEEE Transactions on Aerospace and Electronic Systems, 5274): 2016.
Tri Dao and Albert Gu. Transformers are ssms: generalized models and efficient algorithms through structured state space duality. In Proceedings of the 41st International Conference on Machine Learning, pp. 10041-10071, 2024.
N. E. DAY. Estimating the components of a mixture of normal distributions. Biometrika, 56(3): 463-474, 12 1969. ISSN 0006-3444, doi: 10. 10.1093/tiomme/56.3.433 URL https: /doii org/ 10.1093/bicmet.
3. 463.
A. P. Dempster, N. M. Laird, and D. B. Rubin. Maximum likelihood from incomplete data via the algorithm. Journal of the Royal Statistical Society. Series B (Methodological), 39(1):1-38, 1977. ISSN 00359246. URL http: /www. jstor. org/ sttbst/eee888855555
Rasool Fakoor, Pratik Chaudhari, Jonas Mueller, and Alexander J Smola. Trade: Transformers for density estimation. arXiv preprint urXiv:2004.00 0224004400244,,. 2020.
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30504 2022.
10

---

Published as a conference paper at ICLR 2026
Angeliki Giannou, Liu Yang, Tianhao Wang, Dimitris Papailiopoulos, and Jason D. Lee. How well can transformers emulate in-context newton's 's method? In The 28th International Conference on Artificial Intelligence and Statistics, 2025. URL https: https///pppeeeeeeeeemmmm net wtw.://wwttetttumem///// worrmoriiiooo
ej51229Wol
Rémi Gribonval, Gilles Blanchard, Nicolas Keriven, and Yann Traonmilin. Statistical learning guarantees for compressive clustering and compressive mixture modeling. Mathematical Statistics and Learning, 3(2):165-257. 2021.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint ar arXiv:2312.00755 2023.
Kai Han, Yunhe Wang. Hanting Chen, Xinghao Chen, Jianyuan Guo, Zhenhua Liu, Yehui Tang, An Xiao, Chunjing Xu, Yixing Xu, Zhaohui Yang, Yiman Zhang, and Dacheng Tao. A survey on vision transformer. JEEE Transactions on Pattern Analysis and Machine Intelligence, 45(1): 87-110, 2023. doi: 10A.109//AM..222...11.5
Yihan He, Yuan Cao, Hong-Yo Yu Chen, Dennis Wu, Jianqing Fan, and Han Liu. Learning spectral methods by transformers, arXiv preprint urXiv:2501.01312 2025a.
Yihan He, Hong- Yu Chen, Yuan Cao, Jianqing Fan, and Han Liu. Transformers versus the algorithm in multi-class clustering. arXiv preprint arXiv:2502.06007, 2025b.
Timothy Hospedales, Antreas Antoniou, Paul Micaelli, and Amos Storkey. Meta-learning in neural networks: A survey. IEEE transactions on pattern analysis and machine intelligence, 44(9): 5149-5169, 2021.
Daniel Hsu and Sham M. Kakade. Learning mixtures of spherical gaussians: moment methods and spectral decompositions. Proceedings of the 4th Conference on Innovations in Theoretical Computer Science, ITCS 13, pp. 11-20, New York, NY. USA, 2013. Association for Computing Machinery, ISBN 9781450318594, doi: 10.1145/2422436222244589 URL https: /ttp////oiii./ /oo.orgg////oooooooooo..oo/oo 10. 10.1145/242248 2ee.24422232222222....2 2422439.
Chi Jin, Yuchen Zhang, Sivaraman Balakrishnan, Martin J. Wainwright, and Michael I. Jordan. Local maxima in the likelihood of gaussian mixture models: structural results and algorithmic consequences. In Proceedings of the 30th International Conference on Neural Information Processing Systems. NIPS' 16, pp. 4123 4131, Red Hook, NY, USA, 2016. Curran Associates Inc. ISBN 9781510838811
Yanhao Jin, Krishnakumar Balasubramanian, and Lifeng Lai. Provable in-context learning for mixture of linear regressions using transformers. arXiv preprint arXiv:2410.14188 2024.
Salman Khan, Muzammal Naseer, Munawar Hayat, Syed Waqas Zamir, Fahad Shahbaz Khan, and Mubarak Shah. Transformers vision: survey. ACM Comput. Surv,, 54(10s), September 2022. ISSN 0360-0300. doi: 10. 10.1145/3505244 URL https: //ooi. org/ 10. 111./11511122
Juno Kim, Tai Nakamaki, and Taiji Suzuki. Transformers are minimax optimal nonparametric in-context learners. Advances in Neural Information Processing Systems, 37: 37:106667-1067133, 2024.
Jeongyeol Kwon and Constantine Caramanis. The em algorithm gives sample-optimalig for learning mixtures of well-separated gaussians. In Proceedings of Thirty Third Conference Learning Theory, volume 125 of Proceedings of Machine Learning Research, pp. Jul 2020.
PMLR, 09-12
Juho Lee, Yoonho Lee, Jungtack Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set trans- former: A framework for attention based permutation inna-riiiinntt neural networks. In International conference on machine learning, pp. 3754-3753. PMIR, 2019.
Wenzhe Li, Hao Luo, Zichuan Lin, Chongjie Zhang, Zongqing Lu, and Deheng Ye. survey on transformers in reinforcement learning. Transactions on Machine Learning Research, 2023a. ISSN 2835-8856. URL https: http/////p//.w..////////// net www.ee/e/oooodooeeeeeeeeeeeeeeee Certification.
Survey
11

---

Published as a conference paper at ICLR 2026
Yingcong Li, Muhammed Emrullah Iidiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning. In International conference on machine learning, pp. 19565-19594 PMLR, 2023b.
Licong Lin, Yu Bai, and Song Mei. Transformers as decision makers: Provable in-context reinforee- ment learning via supervised pretraining. In The Twelfth International Conference on Learning Representations, 2024. URL https: httpa:://ppeeevieeeeeee net/ ttt/://oummmiide// 2i 7s533.
Toni Liu, Nicolas Boulle, Raphaël Sarfati, and Christopher Earls. Density estimation with llms: geometric investigation of in-context learning trajectories. In Y. Yue, A. Garg, N. Peng, F. Sha, and R. Yu (eds.), International Conference on Representation Learning, volume 2025, pp. 22163- 22197, 2025. URL https:: httpps:///proooeeiigggggsgs/ iic.rnnnnssiii.lii..m wwtp://tcc..eeeeeee.eeeee/e wwtpapprrcppppppppeeeeeee pdff pdf.
2025/
Matthias Löffler, Anderson Y Zhang, and Harrison H Zhou. Optimality of spectral clustering in the gaussian mixture model. The Annals of Statistics, 49(5):2506-2530) 2021.
Ilya Loshchilov, Frank Hutter, et al. Fixing weight decay regularization in adam. arXiv preprint 17Xiv:1711.05100 5:5, 2017.
Laura Manduchi, Kieran Chin-Cheong, Holger Michel, Sven Wellmann, and Julia E. Vogt. Deep conditional gaussian mixture model for constrained clustering. In Proceedings the 35th Interna- tional Conference on Neural Information Processing Systems, NIPS '21, Red Hook, NY, USA, 2021. Curran Associates Inc. ISBN 9781713845393.
Geoffrey McLachlan and David Peel. Finite mixture models. John Wiley Sons, 2000.
Song Mei. U-nets as belief propagation: Efficient classification, denoising, and difflusion in generative hierarchical models. ArXiv, abs/2404. 18444, 2024. URL htt.tp/s///xxiv..ggaaaa///////// wttp://.gggab////0/// 18444.
Ankur Moitra, Algorithmic aspects of machine learning. Cambridge University Press, 2018.
Mohamed Ndaoud. Sharp optimal recovery in the two component gaussian mixture model. The Annals of Statistics, 50(4):2096 2222,,, 2022.
Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jaewoong Cho, Samet Oymak, Kang- wook Lee, and Dimitris Papailiopoulos. Can mamba learn how to learn? a comparative study in-context learning tasks. In Proceedings of the 4sst International Conference Machine Learning, pp. 19793-39812 2024.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Z. Yang. Zachary Vito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high- performance deep learning library. In NeurIPS, pp. 8024 0044--005, 2019.
Reese Pathak, Rajat Sen, Weihao Kong, and Abhimanyu Das. Transformers can optimally learn regression mixture models. In The Twelfth International Conference on Learning Representations, 2024. URL https: http///pppeeeee.eeee/ee/ee net wwwumoorriioooooo stk
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.
Rylan Schaeffer, Mikail Khona, and Sanmi Koyejo. In-context learning of energy functions. In ICML 2024 Workshop In- Context Learning, 2024. URL htttps://openreview... wttp://wwwwww.w..wtteteee fformm? [d=99II3E2aaS:
Nimrod Segol and Boaz Nadler. Improved convergence guarantees for learning Gaussian mixture models by EM and gradient EMM. Electronic Journal of Statistics, 15(2):4510- 4544, 2021.
Kai Shen, Junliang Guo, Xu Tan, Siliang Tang, Rui Wang, and Jiang Bian. A study on relu and softmax in transformer. arXiv preprint 202.022302.06461. 2023.
12

---

Published as a conference paper at ICLR 2026
Anzo Tch, Mark Jabbour, and Yury Polyanskiy. Solving empirical bayes via transformers. arXiv preprint arXiv:2502.09844, 2025.
D.M. Titterington, A.F.M. Smith, and U.E. Makov. Statistical Analysis of Finite Mixture Distributions. Wiley, New York, 1985.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Proceedings of the 31 International Conference on Neural Information Processing Systems, NIPS' 17, pp. 6000-6010, Red Hook, NY, USA, 2017. Curran Associates Inc. ISBN 9781510860964.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, Joao Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 35151-35174, PMLR, 23-29 Jul 2023. URL https: / httpp/ppooeenn...e./o/// mlx press / v2222/22///////// www.://22oooooaaaa2222mamaaa html.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite. Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-the-ar art natural language processing. In Proceedings of the 2020 Conference Empirical Methods in Natural Language Pro- cessing: System Demonstrations, pp. 38-45, 45, Online, October 2020. Association for Computational Linguistics. URL hww.hhtttt/s/t....
ac web. org/ htt.nnnooooooooooooooooooo..oooo hoy 2020 0. cmnnp-ddemee 6.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.00 1..0000 2021.
Bin Yu, Chen Chen, Ren Qi, Ruiqing Zheng, Patrick J Skillman-Lawrence,, Xiaolin Wang, Anjun Ma, and Haiming Gu. segmai: a gaussian mixture model for clustering single-cell rna-seq data based on deep autoencoder. Briefings in bioinformatics, 22(4):bbaa316, 2021.
Yi Zhang, Miaomiao Li, Siwei Wang, Sisi Dai, Lei Luo, En Zhu, Huiying Xu, Xinzhong Zhu, Chaoyun Yao, and Haoran Zhou. Gaussian mixture model clustering with incomplete data. ACM Trans. Multimedia Compul. Commun, Appl., 17(1s), March 2021. ISSN 1551-6857. doi: 10.1 1451340033888 8. URL https:tt////. http:://di..ooggg/// 1145/333313133388111..11
Ruofei Zhao, Yuanzhi Li, and Yuekai Sun. Statistical convergence of the EM algorithm on Gaussian mixture models. Electronic Journal of Statistics, 14(1):632 660, 2020. doi: 10.
13

---

Published as a conference paper at ICLR 2026
Appendix
Table of Contents
A Literature Density Estimation using LLMss
14
B Algorithm Details
15
C Full Notation of Network Architecture
15
D On the parameter efficiency of TGMM
16
E
Limitations and future work directions
F Formal statement of Theorem I and proofs
17
E.I
Formal statement of Theorem 1 Construction of transformer architecture and formal version of Figure 7
17
F2
19
F3 Convergence results for EM algorithm
23
F4 Proof of Theorem F.1
30
G Formal statement Theorem 2 and proofs
31
G.1 Formal statement of Theorem 2
31
G.2 Proof of Theorem G.8
32
H More on empirical studies
34
H.l More on experimental setups
34
H.2 A complete report regarding different evaluation metrics
..e.............................
35
On the impact inference-tine sample size
35
H.4 On the impact of backbone scale
35
Disclosure of LLM usage. We used LLMs solely for grammatical correction, language polishing, and drafting preliminary code snippets.
Organization the Appendix. The appendix is organized as follows. Section A provides additional literature review. Section B formally presents the GMM algorithms referenced in Section 2, and Section details the complete notation for the network architecture. We analyze the parameter efficiency of TGMM in Section D. Section E discusses limitations and outlines directions for future work. Rigorous statements and proofs of Theorem and Theorem 2 are provided in Section F and Section G, respectively. Additional experimental details are included in Section H.
Additional notations in the Appendix. The maximum between two scalars a, is denoted as a V b. For vector 0 E Rd. let ||u] == max |oy] be its infinity norm. We use O to denote the zero vector and e, E IRd denote the i-th standard unit vector in Rd For matrix A IRds xda we denote ||A]] l := supp|=/1,=1 ||/l as operator norm. We use Ó(-) to denote O(-) with hidden log factors. For clarify, we denote the ground-trun truth parameters of GMM with a superscript ", i.c. { M k throughout this appendix.
A LITERATURE ON DENSITY ESTIMATION USING LLMS
Recent studies have explored the capabilities of large language models (LLMs) for in-context probability density estimation. For instance, Liu et al. (2025) interprets LLM learning as an adaptive form of Kernel Density Estimation, revealing divergent learning trajectories compared to traditional methods. Schaeffer et al. (2024) introduces more general framework for in-context learning by modeling unconstrained energy functions, enabling effective learning even when input and output
14

---

Published as a conference paper at ICLR 2026
spaces are mismatched. Meanwhile, Fakoor et al. (2020) leverages seff-attention mechanisms to perform empirical density estimation across heterogeneous data types. Whereas these efforts prioritize paper focuses the theoretical expressive
empirical performance in distribution estimation, power of transformers, specifically in the context of GMM estimation.
B ALGORITHM DETAILS
We state the classical algorithms of GMM mention in Section 2 in this section.
Algorithm B.1 EM algorithm for GMM
Require: iX e [N]], (()) {mi 11 TK MK
(0) (0)
(0) (0)
1: j 4-0 0
2: while not converge do
3: E-step: we (i+1) ((+t)(())) (X:) = x MoX
1 [N], A [K]
DiENy nie o(Xein (5))
4: M-step: (f+1)
Eiei w/ (x,) (1-1)
EiIN w6 + (x.xx,
A E [K]
14
EieM wh (x.)
5: j j+J+l +1
6: end while
Algorithm B.2 Spectral Algorithm for GMM
Require: {X,, E [N]} 1: Compute the empirical moments M2 and M3 by
M2 N [E[N]
Σ X, & X, - Lda
M3
Σ X [E(N]
Σ
+ 0z X,)
N
2: Do first K-th singular value decomposition((VVD) for Me: M2 R U and let W 1/2
B = !D2
3: Do first K-th robust tensor decomposition (Algorithm
Anandkumar al. (2014), see
Algorithm B.3) for M3 (WW,WW)
M3
A 03
KE
return म4 -==2 fix = Ar E [K].
C FULL NOTATION OF NETWORK ARCHITECTURE
Definition 4 (Attention layer). A (self-)attention layer with M heads is denoted as Attno with parameters @attn = ((V, m, Qm, Km) ) ME(M]] C DxD, On any input sequence BXX
M = Attne HH) H + DM/11 Vm softmax
(Q. RXN (
In vector form,
((1amtt)) ((Km)))) N ) Vnhy- )
h [Attn
(H), h; AM Σ
softmax
Here
softmax is the activation exp(va) for u E Rd a exp(v:.
function defined by softmax (u)
exp(e1)
Ei exp(e:)
15

---

Published as a conference paper at ICLR 2026
Algorithm B.3 Robust Tensor Power Method
Require: symmetric tensor dx number of iterations L, N.
Ensure: the estimated eigenvector/eigenvalue pair; the deflated tensor.
for ==
do
2: Draw vg
(T) uniformly random from the unit sphere in IR
3: for = 1 to N do
4:
Compute power iteration update:
T( (/ IT( (3 10
5:
6: end for
7: end for
8: Let r" == arg [ (T "N " )]-
9: Do N power iteration updates (line 5) starting from (T") "N to obtain é.
10: Set Å = T( 6, 6).
return the estimated cigenvector/cigeevalue pair (0, A); the deflated tensor T A6
The Multilayer Perceptron(MLP) layer is defined as follows. Definition 5 (MLP layer). A (token-wise) MLP layer with hidden dimension D' is denoted as MLP (:) with parameters Om1 = (W1, W2) E DD'xD XD x DxD' On any input sequence H E DXN,
= MLMP (H) = H (2
where a R + R is the ReLU function. In vector form, we have h, = h, + W
Then
the above definitions define the transformer model.
Definition 6 (Transformer). An I.-laer transformer, denoted as TF (.), is a composition of
self-attention layers each followed by an MLP layer:
TF (H) = MLP (2)
(Attn MLP
(
(ALLI (H)
Here the parameter θ (0) ((V1) m Q k! mEsas(O)
o
(1:L)
A alp
(1:2)) consists of the attention layers
+
(
attn
attn
R DxD, the MLP layers O (W(") w
DD x
D(D
D ON THE PARAMETER EFFICIENCY OF TGMM
Aside from its backbone, the extra parameters in TGMM comprises the following:
Parameters in the task embedding module This part has parameter count of d
Parameters in the Readin layer This part has a parameter count of O((dusk + d) x D).
Parameters the Readout layer This part has a parameter count of O(sDD), which comprises of parameters from distinct attention mechanisms.
As dans is typically of the order O(D), we conclude that the total extra parameter complexity of the order (s)D), which in practice is often way smaller than the parameter complexity of the backbone, i... of the order O(LD?) Meanwhile, a naive implementation of adapting transformer architecture to solve distinct GMM tasks require different transformer backbone. As the complexity of backbone often dominate those of extra components, the TGMM implementation can reduce parameter complexity by (approximate) factor of 1/s in practice.
E LIMITATIONS AND FUTURE WORK DIRECTIONS
First, while our theoretical analysis focuses on the approximation ability transformers, the opti- mization dynamics remain unexplored. This is common theoretical challenge in ICL literature;
16

---

Published as a conference paper at ICLR 2026
see Bai et al. (2023); Lin et al. (2024); Giannou al. (2025). Second, approximating the full spectral algorithm (Algorithm B.2; see Section B) presents a significant challenge, which we leave for future work. Third, our study is limited to the expressivity of transformers classical GMM tasks; exploring their performance on other unsupervised learning tasks is an interesting direction that warrants further investigation.
F FORMAL STATEMENT OF THEOREM AND PROOFS
For analytical tractability, we implement Readin as an identity transformation and define Readout to extract targeted matrix elements hence they are both fixed functions. Actually, we also need "Readout" functions to get the estimated parameters for different tasks, see Remark F.3. To theoretical convenience, we use the following norm of transformers, which differs slightly from the definition in Bai et al. (2023).
|||| = max EEIL]
max mE|M
{
Q
K!
V
[,] + |//
1ell|++/ww/ +1w33
Then the transformer class can be defined as
D M, Be) = (Te. ||||| S Be, D S D', M s (.)
F.1 FORMAL STATEMENT OF THEOREM 1
First, introduce some notations. We define Tiin min, n.. Pr. max; /// min; m.. We use Pa = |l44 - اill to denote the pairwise distance between components and Rmin = miniys Raj. Rmaz (maxişjj R2,) V (max;e] |Ix Without the loss of generality, we assume that Rmax 2 1. For dimension and components adaptation, we assume d I dg and K S Kg- Since in practice the sample size N is much larger than the number of components K, we assume that N is divisible by K, i.e. N/K E N. Otherwise, we only consider the first K [N/K] samples and drop the others. We encode X
{X:) into an input sequence H the following:
H
X1 P1
X2
X pv
RxN
P: =
0: ri
(4)
P2
Tlog
Op
0, =
TiLK
E -ooo-Koo-1
r
E ERD (2d0033k))
(5)
C%
[e%
O3Ka
where X; = [x,,00 oTo-ol' 1lg [ log OT, -k--kl
In K. 0Z--dl C E RR and
0 E Ko denotes the i% standard unit vector. To match the dimension, = - (2dg-
5Ko + 2). We choose D Odo + Ko) to get the encoding above. For the initialization, we choose
(0) Tlog = log Togr(()) (0) p = H; c = ||m (0) I2.
To guarantec convergence of the EM algorithm, adopt the following assumption for the initializa- tion parameters, consistent with the approach in Kwon & Caramanis (2020).
(A1) Suppose the GMM has parameters {(F} m5) i E [K]] such that
Rmin 2 C K))
and suppose the mean initialization 11 ..... Hk satisfies
(0) (0)
Vi E [K] |4000 μنر
Romin
16
Also, suppose the mixing weights are initialized such that
Vi E [K], 7 न.
S /
17

---

Published as a conference paper at ICLR 2026
We denote the output of the transformer T TTH == =x] 11 IT2 matched indices. Define
122
WK. (K
TF and assume
D = max KE[K]
[114 - will (1 क. /m,)}
Now we propose the theorem that transformer can efficient approximate the EM Algorithm (Algo- rithm B.I), which is the formal version of Theorem 1.
Theorem F.1. Fix < 6, 3 < 1 and 1/2 < < 1. Suppose there exists a sufficienlly large universal constant C 128 for which assumption (AI) holds. If N is suffcient large and 1/(100Kg) sufficient small such that
č1
Ryax [mmax V d) log (4k)
(1 a)) miin
N
+ 02
Rmaxx d 1
log( L.V.
5
1 a 2-22 Ne < 2
d
2)
and
Kdlog( CN N
C3
c(N,,,,,0) = (1 a) Tmin
4.V
+ E4
+N
Rmax + d + 2d
Amin
8
< <a(1- - B),
hold, where 01- are universal constants, Č 28822882222 (vdd vdd+ R + - 2(L + 1y-layer transformer TF such that
Then there exists
DT S as4 + 1- (N,,N,
(6)
holds with probability at least satisfying:
8. Moreover, TFe falls within the class with parameters
Odo + Ko), D' Rm Penax + do)e M 0(1,, log Be S Bma (Bmax + da)).
Notably, (6) holds for all tasks satisfying d S do and K S. Ko, where the parameters of transformer + remains fixed across different tasks T.
Remark F.1. From Theorem F.1, if we take -NN-3//2--17 and L O(log N), then have
D Ô
which matches the canonical parametric error rate. Remark F.2. We give some explanations for the notations in Theorem F.1. Define
DREM
(1 -Mall -
[t -7, -m./m:)}
-aax &E|K]
where (p()) *(()) liE(k] are the parameters obtained at the j-th iteration of the population-Ei algorithm (see Section F.3 for details). In the convergence analysis of the population- algorithm (Kwon & Caramanis, 2020), it is shown that after the first iteration, the parameters lie in small neighborhood of the true parameters with high probability (i.e., DPEM S a for some 1/2 S 1). Furthermore, the authors prove that the algorithm achieves linear convergence (i.e., DPEM < BDDEM holds. Following their
for
for j Nr and some 0 3 1) with high probability if D2EM notations, here a represents the radius of the neighborhood after the first iteration, while the linear convergence parameter. Finally, controls the approximation error of the transformer.
18

---

Published as a conference paper at ICLR 2026
F.2 CONSTRUCTION OF TRANSFORMER ARCHITECTURE AND FORMAL VERSION OF FIGURE
In this section, we give the transformer architecture construction in Theorem F.l. We denote wij wy i [V]. k [K] in this subsection for simplicity. Recall that we have assumed that d S do. K & Ko and N is divisible by K(N/K E N). We first restate the encoding formulas in (4):
X1
X²
XN
H
0
02
0 N
E D
P1
PN
log
P4%K
0 52o6+
0,
G%K
E R
1oo+1
: :
E R 322--22-33Ktt)
W;
e;%K
Wilog
π
where X, [X]_ 0%
Ky=Ki K. T 14% [M 02 T. W;
Tog
log
[w. 00 Ko xl W, llg [w log
x]
[
00, Kg xl G%L E IR and 09% E R Ko de-
Kg
notes the 1%Kttt -tt standard unit vector. For the initialization, we choose Tlog = log (0) #4 μ4
C; || 1. and w; Wo = 0k, E [K]. Finally, take H H which is defined in ??.
Then in ststep, we consider the following attention structures: we define matrices Q K(1), V(), such that
Odo
X; Ing
-yyak 05%K
Kg
Q
h
(0)
K
k =
v(2)//)) (0) =
Odg+1
1
2 C/S6K
ej%
0
0
0 (2do +
and use the standard softmax attention, thus
6.00
Attn
(
(0)
N
N (((0D))) (0) T (
)
h
softmax
v(nn
Ø=1
(0)
exp -X. M1%K + 22 5(
5%
h
B
exp (-X,T Me + tuu
==
h
vann
B ==1
1%
h
Σ
h
(0)
Odo
Oko
h (0)
Odo+1
ie [N)
w:
0 - +2Ko
where W
-(1)
((0())
10
10 5
(1)
(1)
0, ... 0) T
Rko
19

---

Published as a conference paper at ICLR 2026
Then two-layer MLP to approximate log z and clean all Rlg Ba% and CI,K, which is 0 Kg
X;
Ogo
Od++1
h
MLP !!!! [
(
, [NT,
W log
0p - 5ko
8;%K
where w log = logw; Notice that although log x is not defined at O, the MLP approximation is well defined with some value which we do not care because will not use it in the M-step. Similarly, for any 8%2 = 1, E 4. we have
X; 0 Ko
((+1)//2
h MLP (300))
(Attum (19)2)
(H (H(-)))
(C+1)/22
iE E [N],
mip
log
0ko
e,%.K
((e+1)/2) where Welog
((++))/2
logw;
In M-step, we consider the following attention structures: we similarly define matrices Q K(2) (
2 such that
O
0 ОKo
X;
h
(1)
Cj%K
K 22) k
Wilog
v (2)
h
0
0 Ko
Оgo
0 D +3oo
and
0
Ko
0do+1
(2) ph-0= 0,
K h (2) 0. v(2n
0
O
w
0 10- (2do 4oo
Then we get
h 2) Attn (2) (H))
( ((())1)))
T
hh
ΣΣ Σ
softmax
V@h (1)
(1)
= -h
(1)
1
(2)
h
(1)
Σ
1
v! n
N
20
N
11%K
f=1
20

---

Published as a conference paper at ICLR 2026
Odo
Odo
Ogo
=(1)
BygK
h +
+
O Ko
Oko
Ko
π (1)
Oko
10 -k-(22(33Ko+33
0 D- 1Ko
Odo
Oko
=()) Fi%K 0
hg +
j E [N].
Ok.
0Ko
#0)
Similarly, two-layer MLP approximate log æ, x² and clean all wi Wslog and ,. which is
X
# log
(1)
%
h,² = MLP (2)
h
Ko
XK
.J E [M]
Og
0 (2-
9,%K
where 1lo =
(1)
log
(1)
(1)
9%A
= =1
Similarly, for any 8%2
0, l E N.. we have
X X/2
(2/2)
P 1
(7
GAK
h1 MLP (1922)
(1362)
(H
0Ko
. E [N].
OKa
0
Op (2do+55x++)
1
05%K
where 1lg
(1/2)
loga
(1/2)
(1/2)
10 12
(2/2)
Thus, get # and ( j [K] after 28 layers of transformer constructed above. (The last- layer MLP block retains π as an output parameter without cleaning it.) Our transformer construction summarized in Figure which is the formal version of Figure 7 in Section 4.2.
Remark F.3. The output of transformer (21)) is a large matrix containing lots of elements. To get the estimated parameters, need to extract specific elements. In details, ((L)
21

---

Published as a conference paper at ICLR 2026
h
(2L)
by
(2L)
R
DxN, where
X
G
h
(2L)
0g
E [N].
Ka
0p (2dg
Cj%K
We use the following linear attentive pooling to get the parameters:
(V
QQ,) E R(((+K)XK,
N
where Qu =[qo1," 9okk
R(d+K)xK
K. Vo R(4KK)NN satisfying
Ke,
K.h
(2L)
e,xk
Voh
(2L)
(2L)
(())
qoi =
Od
Thus by N/K E N, we have
O: N
1
Σ
qod (o,
=
K N [L)
(2])
E (d+
, [K].
N K
As
JE[N]
Finally, we get
[L]
(L))
A (4)
[q01."
9ox]
f
px
step:
M step:
The Pisu cuea
Op
uppeosiimnte Jogz 22, cleam W. WT MLP
Bisk
Embedding stop
Owe Fesu:
0p
X
Ryan
27
Qhi
Kly
WF
csp (Ine Fymee X. kyna baleeteme
Vhy
Attention
Q.
quity
Og Moan
by you
Vilu
Valu
Attention
approsimate in Miss Heur MLP
Dii Ouy
Og
Embedding of M step
Figure 9: Transformer Construction for Approximating EM Algorithm Iterations. The pink box represents the state of tokens, while the blue box represents the structure of different parts of the network. The term "clean" means setting all positions of the corresponding vector to zero.
22

---

Published as a conference paper at ICLR 2026
F.3 CONVERGENCE RESULTS FOR EM ALGORITHM
F.3.1 CONVERGENCE RESULTS FOR POPULATION-EM ALGORITHM
First, we review some notations. Recall that Tiin min; n7. Px max; o// min; n.. Raj
||g] -15 Rmin mingşs Rij and Rmax (max+e) Rej) V (max;e [k] || Without the loss of generality, we assume that Rmar 2 For clarity, we restate assumption ((A1),, which is consistent with Kwon & Caramanis (2020).
(A1) Suppose the GMM has parameters {(rj. M5) j E [K]]
such that
Rain 2 C Vlog(p*K),
(7)
and suppose the mean initialization [1] HkK satisfies
(0) (0)
Vi E [K], (0) p
Rmin
(8)
16
Also, suppose the mixing weights are initialized such that
Vi E [K],
न.. S 12/
For population-EM, the algorithm can be presented as
w, (X) exp - || - 12il112)
(E-step):
ΣØ= Tyy exp( (1 - J4y 112/2)
(M-step):
n E E[ws] n E E[w,x|/E[ur]
The following results gives linear convergenve guarantees of population -MM, which comes from Kwon & Caramanis (2020).
Theorem F.2 (Kwon & Caramanis (2020), Theorem I, part i). Let C 2 64 be a universal constant for which assumption ((Al)) holds. Then, after one-step population EM update, we have
Vi E [K], |- x म. /2, | 14 S 1/2.
(10)
Now define
Dm = max (1l - 1 v | - .;|/ (E[K]
and
D max (Il4) - HIll V | - x;ç1/
<E[K]
The linear convergence of population-EM is stated by the following theorem.
Theorem F.3 (Kwon& Caramanis (2020), Theorem 1, part ii). Let > 128 be a large enough universal constant. Fix 0 < < 1. Suppose the separation condition (7) holds and suppose the initialization parameter satisfies Dm S a, then D S BD Dm for some 0 < 3
Remark F44 Here the contraction parameter B is only dependent with C and a. In other words, if we fix E (0,1), then for any 3 E (0, 1), there exists a large enough C such that Theorem F.3 holds. For details, see Appendix E in Kwon & Caramanis (2020).
Combing Theorem F.2 and Theorem F.3, we can get the linear convergence of population-EM algorithm.
F.3.2 CONVERGENCE RESULTS FOR EMPIRICAL-EM ALGORITHM
Now we consider the empirical-EM. EM,, i.e., Algorithm B.I. For convenience, the algorithm can be presented as
expl - Xr 1:ll2/2)
(E-step):
wx())) Wii
Σ TTy expl -1 11 112/2)
Σ w (Xe)
(M-step):
wi (xe))), u1t
wi (X.)
Σn wi(Xe)
1=
23

---

Published as a conference paper at ICLR 2026
Similarly, can define Dm, and D in empirical sense.
For the linear convergence of empirical-EM. we have the following theorem.
Theorem F.4. Fix 6, < and
1. Let 2 128 be a large enough universal constant.
Suppose the separation condition (7) holds and suppose the initialization parameter satisfies Dw. S If n is suffcient large such that
Ca Kdlog funif = (1 a) min
< a(1
8)
n
where C -= =7223(Vd max
I-a )2 and is universal constant. Then
D S 3Dm + Cunif S
uniformly holds with probability at least
Proof. First, have
(Xe) m
x
1
ws( (Xe) E (us (X)
|E (x0] -xtl
(T) (TT).
By Theorem F3, we get
1 π.
(II)
[x,]] オł1 s 3Dm
And by Lemma F.2, we have
Kd Kd Cja
(1)
We (Xt) [x,(
Tmin
n
where Č 18K2 ( + 2Rmar + ) and ē is a suitable universal constant. Thus, by taking
3 = 3. 61 8/2 and suitable 2 m ni m/m I 3m + Eunif a, Vi E [K].
For the second term, we have
(X)) 433
11 11l
wi (Xe) 143) Blw;(X)(X mt)1
11 (X)(X m2)Ill
(III) (TV),
By Theorem F.3 and Remark F.4 we get,
(IV)
E [[(x((X)EE - m3)1l
x
[[[w (X)x m111l
(1 a) a)mr
E Elus X))
E3 Iu (X)X)
14
a)m??
E E[w (X)]
+ a 3Dm
<
24

---

Published as a conference paper at ICLR 2026
where (1) follows from
1/1/ S And by Lemma F.3, we have
(
= _
1
Σ
we (X:) 14) (X); 14511
1
(Xi) (5)) E[w, m2)i)
(1
a)
Kd logo Sza)
02
<
(1 min
where Č 18K2 (v vd+ RRmaz
T )² and ē is suitable universal constant. Thus, by taking 3= (1-a) a)3, 82 =8/2 and suitable è, I1 - mill - 3Dm + Eunif S a, Vi [K].
In conclusion, if we take 3 = (1--aa =11-aa/(1+ a)3, č = E V V C C(3,a) > 128 large enough such that Theorem F.3 holds, and take 81 = 82 = 6/2 and use union bound argument, then we get D S 3D + Eunif S
We need the following technical lemma. Lemma F1 (Segol & Nadler (2021), Lemma B.2.). Fix < 8 1. Let B Bk C. Idd be Euclidean balls of radii ......, T... Define B & Bg C kk and r [ Tk. Let X be
random vector in R& and W RRd B IR* where k S d. Assume the following hold:
There exists a constant T, 2 1 such that for any E B,e > 0. and OF E B which satisfies then Ex [sup || (X, W 0 || S Le.
max |10, 10.-051
2. There exists constant R such that for any 0 B. (W(X,0)ll S R.
Let XI...., Xn be i.i.d. random vectors with the same distribution as X. Then there exists a universal constant ē such that with probability at least 1 8.
Kdlog (1+ 12nfr
sup BEB
w - (X,0) S RV
-
(11)
t=1
Remark F.S. There is one difference between Lemma F.I and LemmaB.2. in Segol & Nadler (2021),
in Lemma F.I, we use 1 12nL" replace 1871r, thus we avoid the condition T1, .
1.
Hence we can get the uniform convergence of (x,(X,) 0) and w. (X,0) - n&) E [K]. Our proof similar to Segol & Nadler (2021), except that we consider the variation of both and μ. From now on, we denote 0; = {T,, p. , 0 = {0
Lemma F.2. Fix 0 < 8 < 1 and 0 < < 1. Consider the parameter region D. (Dm S a). Let X1, Xn ~ GMM(T", ,²') then with probability at least 6,
idd
Kolog( En Vi (K),
sup BED
we (Xe,
(12)
t=1
where Ò 18K2 2 (vd 2Rmaz + + and ē is a suitable universal constant.
Proof. The proof similar to the proof of Lemma 5.1 in Segol Nadler (2021). For simplicity, we omit it.
L
Lemma F.3 Fix 0 < 8 < and 0 < a < 1. Consider the parameter region D. == {Dm I a}. Ler XI, Xx, itd GMM(T", .p²) with Rmin satisfying (7), then with probability at least 5,
On Kdlog(
sup OED
we (Xe, (Xe (Xt-Xaaa - 14) - El,,((X - -47)] <
Vi [K], (13)
n
1=1
where
Č 36K2 (Vd 2R } and is suitable universal constant.
25

---

Published as a conference paper at ICLR 2026
Proof. The proof is similar the proof of Lemma 5.4 in Segol & Nadler (2021)(Notice that the condition (36) in Segol & Nadler (2021) is trivial in our case). For simplicity, we omit it.
For the first step empirical-EM, we have the following results.
Theorem FS. Fix < and 1/2 < a < 1. Let 2 128 be a large enough universal constant for which assumption ((AI)) holds. n is suffcient large such that
Rmaz Rmaaz V d) log
6K 5
1
Estep1
(1 a Tinn
n
where č is universal constant. Then
D S 2
Estepl S
holds with probability at least 1-6.
Proof. Notice that only need simple concentration uniform concentration in this theorem. We use the same definition term (7), (IT) as in the proof of Theorem F.4. First, by Theorem F.2, we have (II) S 1/2. Since 0 I w;(X) I 1, by a standard concentration of bounded variables, we can get
č1 Tymón
log(
K 61 Vi E [K]
(1)
(14)
n
where č₁ is universal constant. Taking > č₁ and d1 8/2, we have
n
m?
1
2K log
a, Vi E [K].
2
Tmin
n
For the second term, have
|
1
wi
(X))xe ps
not
B(u,(Xx)XX (X2) Xz)
E Elec,(Xe (Xe)l 11
(XeXxx
E [Xe)
no
1=1
(V) + (VT).
By Theorem F.2, we have (VT) < 1/2. For (V), by triangle inequality,
1 ग
(Xe)
(V)
w (X))xx
Elw;(Xx)Xx
Elw:((e)
Blw;(Xe))
...
| [ Xe]l
w: (X))))
E w. (X,)
|m
Blucc( (x)]l (15)
E (Xe)l
Using Lemma B.I and Lemma B.2 Zhao et al. (2020), we can get (X,) S ||Xellva S og Rymaa Vi [K]. Hence by Lemma B. in Segol & Nadler (2021), with probability at least 82.
Rmaad log
3K dz
wi (X,)x, Elu, (X,) S sa
Vi E [K],
n
where ēN is an universal constant. And by Theorem F.2, have
(xx)xxX))))l Rmna < 2R 2Rnaa Vi [K].
w. (Xe)
26

---

Published as a conference paper at ICLR 2026
Finally, by (14),
log( K 61
Elux((Xe)]l S E1
n
Combining all terms together and taking
d2 8/2 we can bound (15) by
3K Rmand log
K log(
(V) v
21
82
26j Rmaz
5
7
n
č6
Rmar Rmar V d) log (6K) (+)
S (1
0) miin
n
Taking ē 226 Co we get
1
ē
6K Rmaz Rmax V d) log
1 1-55ll S 2 (1 a) mmin
+
S
n
F.3.3 CONVERGENCE RESULTS FOR TRANSFORMER-BASER EM SECTION F.2
We first state some useful approximation lemmas.
Lemma F44 (Lemma 9 in Mei (2024)). For any A > 0, 0, take M [2log A/07 +1 +1 E N. Then there exists {(az, wy. by JEjr] with
sup layl S 2A, sup |wyl < 1, sup |by1 < A,
(16)
such that defining logs : R + R by
M
logs(r)) = Σ aj ReLUUWwx + bz),
j=1
have logg is non-decreasing on [1/A, A), and
sup (ELAAAAA
[lg((=) log, (x)) 8.
Remark F66. There is small improvement M =[2log A/0 + compured to M = [2A/0]] + Mei (2024). Further more, easy to check that logs(x) S log for x [0A]
Lemma E., For any A {la wy b,) with
0, 8 >
take M [2A4/0T] +1 N. Then there exists
sup |a,1 2A,
sup [uyl I 1, sup bb, < A.
(17)
)
such that defining Φs RR + R by
M
<s(x) = Σe a, ReLU(Wwr + b,),
1=1
have dss is non-decreasing on [-A, A), and
(0s(x)) - x²| 1
sup xE =E[-A,A]
Proof. Similar Lemma F.4. Omitted.
27

---

Published as a conference paper at ICLR 2026
Lemma F.6 (Lemma A.1 in Bai et al. (2023)), Let IN).. Then we have
IP (118112 d(1 + 8)2) S Se-0612
Lemma E.7 (Lemma 18 in Lin et al. (2024)). For any u, R", have
ea
c²
log
log
2 vll
||e|]
|E>l|
Corollary F.1. For any u, E
Rd,
have
e"
exp vllo))
|ealla
ley
1
Proof. This follows directly from Lemma F.7 and simple calculations.
Now we propose the results for vansformer-based EM. Similar to Section F.2, we use notations with superscript to represent the output of the transformer-based EM.
Theorem F.6. Fix 0 < 8 and 1/2 L. Let C Z 128 be large enough universal constant for which assumption ((A1)) holds. If n is sufficient large and & I 1/100 sufficient small such that
Rmaz (Rmaz V d) log 12K
a) Mmin
n
( (1+ log( 2n)
+ 02
Rmax d 1+
ne <
where E₁ universal constants. Then there exists 2-layer transformer TF such that D,, holds with probability at least - 5. Moreover, TF falls within the class F with parameters satisfying:
= Odo + Ko), D' s Ō (K Rmaxe (Rmax + do)e 1)) M = 0(1),
log Be O(KK Rmax RRmax du)).
Proof. Recall that Dt = max;e[K] (11 14ll V | - x x///x²)) Thus
Dt < max (lM 1451 v | - n&7 l/T/// + max (11 - - V | - x/1/=²)
[E[R]
kE[K]
= D + max
(IlA) 1 |म l/-1/
sE[K]
We first claim that with probability least 8/2,
log 2n
max (1AA - V V [ -: 1/
S 02 Rmax
1+
ne. (18)
[E[K]
d
Then by Theorem F.5, with probability at least - 8, we have
Dt Dt + max (IlT
p | m, 1/
EE[K]
1
ē
Rmoa Rmmax V d) log
12K
2 log 22
+
E2 Riaa + d1+1
1
ne
2 (1 a. Tmin
a.
28

---

Published as a conference paper at ICLR 2026
Now only need prove (18). By the construction in Section F.2, that wy; in first step
can be well calculated, thus | - x; = 0 and the error comes only from the calculation of { Recall that p
Di We.e and 2i=1 wes
Ee= exp (log(wss) X
A
E exp (log(wx))
Recall that
m. exp( ( - /211
We = Σラ= my exp( -1 - 13513/2) 1 1+E jii
T1 exp (mj ui) X 11//13/22 1lax11/22)
By the initial condition (9) and (8), we have
1aa - Hill S Rmax + 2 *
16 Rmin = 0(0m(mma)) || O O(R
Since Xe Lid. GMM( n²), using Lemma F.6, with probability at least 8/2, we have
|Xell Rmaa + d (1+1
2 log( (g
sup
O(max d).
rE|n]
Combine all things together, we get that with probability at least 5/2,
W S exp (Ö(Ko Pemax [Rmax + do))): vl E [n]] and E [K].
Thus taking A exp (Ö(Ko (Rmax + do)) and and 5 in Lemma F.4, get
log -logl 1/A,A| < E. Then by Lemma F.7, we have
E== exp log(we))
Xe
Ei exp (log we) Em exp (log was)
11 M
n exp log(wes))
exp (weg
Xe
exp (log wei
Dim exp [og((w))
Ex exp (log wis)
{=1
exp (log(ue))
exp (log wei) ===1 exp (log wes)
< sup ||Xell (E(M)
Σ
El=1 exp
(log(wes))
< n Rmax + d v
12n /2
2 (Tog(wus)) - (log(w>)))
8
exp
2n
< 4n Rina d + 22d
Vi E [K].
Thus (18) is proved. The parameter bounds can be directly computed by the construction in Section F.2 and Lemma E.4.
Theorem F.7. Fix 0 < 6,3 < and 1/2 < a 1. Let C 2 128 be a large enough universal constant. Suppose the separation condition (7) holds and suppose the initialization parameter input to transformer satisfies Dm S If n is suffcient large and KoE I 11/100 sufficient small such that
Kdlog( Cp)
&
22n
(1 - a) Tmin < a1 B). ((n,
82
+ n
Rimx + d 2d
E
n
Tmin
29

---

Published as a conference paper at ICLR 2026
where E₁, 02 universal constants, Č 144K3 1-a).. transformer TF such that
Then there exists 2-layer
t S 3Dm e(n,, a
uniformly holds with probability least Moreover TFe falls within the class F with parameters satisfying:
D Odo + Ko), D' s Omax (Rmax + do)e
M 0(1),
log Be O (Rmax do)).
Proof. Similar to the proof of Theorem F6, using Theorem F.4, we only need to prove that with probability least 1 0/2,
2n
max (11A - M V | - n. 1/ S ê2 iE[K]
+ n Rmax + d / 2d
(19)
Tiin
8
Define ut (ue,1 ... [e, T. ūe ( ke,k) where we, log ग. + u4 X 1/2 |i,l?
and ūæ loga; n4. X 1/2||/i|12 By the construction in Section F:2 and Corollary F.1, we have
1/ Welloo =
11| Te
etee
e
S exp (2]l uelloo) (0lloo)) 1, Vl E [n].
l
Now taking 5 = &, A = ((1 a)Tmin) in Lemma F.4 and 5 E/K, A (Bmax + a)² in Lemma F.5, have |lie uelloo S 32/2, hence
|i wellx S exp (2/ ûe ue (lo)) S. exp(3e) S 62, Ve [n].
Then by the construction in Section F.2, have
Ge, Vi [K].
(20)
For the term | 14 I]. we can calculate it similar to the proof of Theorem F6. First, we recall that with probability at least 8/2,
w < exp (Ō(KgRmaxl (Rmax + do))). Al E
[r]] and E [K].
Similarly, for We we can also get(just calculate again) that with probability at least 1 - 8/2,
/ < exp (Ö(Ko (Rmax do))): Al [m]] and i [K].
Then following the same argument in Theorem F.6, taking A exp (Ö (m(KoBmax (Bmax + do)) and and 8 e in Lemma F.4, we have also
2n
lAt - 1t || < 4n Rmax + + 2d
E, E [K].
(21)
Combining (20) and (21), (19) is proved. The parameter bounds can be directly computed by the construction in Section F.2, Lemma F.4, Lemma F.5 and the parameter A, 5 taken in the proof.
F.4 PROOF OF THEOREM F.
First, by Theorem F.6 and the first condition Theorem F.I. such that
there exist 2-layer transformer TF e.
D <
(22)
30

---

Published as a conference paper at ICLR 2026
holds with probability least 1 - 5/2. Then using Theorem F.3, (22) and the second condition in Theorem F.I, there 2-layer transformer e, such that
D S 3De, + + ec(n,,,6/2,0) S a,
Thus, for any N, uniformly holds with probability at least 8/2. Denote as 0 UeElL) by reduction, we have
De De S SB (1+ +1+ +....
(1 (t,)) 5/2,,a))
uniformly holds with probability at least 1 5/2. Combine all things together, have, for any Le N.
D Dy ton D + (1+ 1++B+ 3 ................................ ...................... +
S 2/ a 1
(n, E,, 6,82,a)
holds with probability at least 1 - 8 (Note that the definitions of c(-) ) in Theorem F.7 and Theorem F.1 differ slightly). The parameter bounds can be directly computed by Theorem F.6 and Theorem F.7. The theorem is proved.
G FORMAL STATEMENT OF THEOREM 2 AND PROOFS
Following Section F, we implement Readin as an identity transformation and define Readout to extract targeted matrix elements hence they are both fixed functions.
G.1 FORMAL STATEMENT OF THEOREM 2
In this section, we give the formal statement of Theorem First, we need to introduce the embeddings of the transformer. Let T be the matrix representation of the cubic tensor T, which is
T_111
T22,1
T.d,1
T,,,2
T,2,2
T = [t1, t2, ... tal =
E IR &Rde
:
T.1,d
T,2,d
T.d.al
where T.i.j = (Tis., Teaj: Ta..s) E R i,j E [d]. For dimension adaptation, we assume d 5 dg. The augment version of T is defined
T T.1.1 T,,22
T:21
T dod
T,22
T. do.22
T= E2,
Ta,
Edol
=
(23)
T 1,1dd
?22,d
T :dd
where Tiis E Rdd Ifi s and S d, T T.ks the following input sequence:
[ [Ja-d]
T Else Tiks
Otg We construct
T( C;
tl
E2
Ea
H
Dxd
p
1
(24)
pi
p2
Pa
d
0,5
where E, E & is defined as (23), )(() unit vector Re We choose D (dh)) and D D df 2dg- to get the encoding above. Then we give rigorous definition of ReLU-activvted transformer following Bai et al. (2023).
O)T 0 do e; E Io denotes the i-th standard
Definition 7 (Rel.U-attention layer). jattenlattennoon layer activated by ReLU function with M
heads is denoted as Attne.... (.) with parameters θ Oatta {(V... m. Qm. Km me[M)] C RR DxD, On
any input sequence E I8 D
R H = Attu
(H) = H+ DM H) H) (QmH)) E DDNN
N
31

---

Published as a conference paper at ICLR 2026
In vector formor
h: Altu (H)]; h, + Σm=1 N Eタ= o ((Q (KKm,)) Vnhy,
Here (xx) x V 0 denotes the ReLU function.
The MLP layer is the same as Definition The ReLU-activated transformer defined as follows.
Definition (ReLU-activated transformer). An L-layer transformer, denoted as TF is a compo- sition of L
layers each followed by MLP layer:
TF = ML.
(Attn R ((L)
(..
(Attn (H)
(
Above, the parameter C
O
(1:2)
consists of the atlention layers + attn
atta
mlp
{(V m Q k! mepario
R DxD and the MLP layers O (f)
mip (w()) (1) (
R DD
C
Dx
Similar to Section 4.1, We consider the following function class of transformer.
F F(L,
= [T ||||| s Be. D D', s SM,LE
Now we can give the formal statement of Theorem 2.
Theorem G.8 (Formal version of Theorem 2). There exists transformer TF with ReLU activation such that for any d S do, T E Rddddxd and (()) E TRd., given the encoding (24), TF implements steps of (3) exactly. Moreover, TF falls within the class F with parameters satisfying:
D'
O(da), Be S O(1).
Remark G.1. In fact, Theorem G.8 is also hold for attention-only transformers since the MLP layer do not use in the proof. do that, we only need add another head in every odd attention layer to clean the terms (dv;). For details, see the proof.
Remark G.2. Readers might question why the normalization step is omitted our theorem. The key challenge is that we have absolutely no knowledge of a lower bound for [[ ,l)),, wl Without this bound, approximating the normalization step becomes infeasible.
Remark G.3. The use of the ReLU activation function here is primarily for technical convenience and does not alter the fundamental nature of the attention mechanism. Several studies have demon- strated that transformers with ReLU-based attention perform comparably to those using softmax attention( Shen et al., 2023; Bai et al., 2023; He et al., 2025a).
G.2 PROOF OF THEOREM G.8
Proof. For simplicity, we only proof the case that (xx) = x in the attention layer. For ReLU activated transformer, the result can be similarly proved by ReLU(x) ReLU(-z) and the (xx) case. Hence omit the notation ø in the following proof. We take H H. the first attention layer, consider the following attention structures:
O da
Q
e;
K
h
TT))]
vOn
(0)
0
0
0
d
O
32

---

Published as a conference paper at ICLR 2026
After the attention operation, have
h
(1)
R Attn
2 1
(H))
h
(0)
+
d
1
Σ
((Qw)!) T (k(N))) vaing
j=1
0
de
E
O
h
+
d
i E [d].
(0)
du
(0)
0
Then we use a two-layer MLP to implement identity operation, which is
t
T())
h
MLP
d [d].
E
kip
dr!
(0)
Now attention layer with do + heads implement the power iteration step of the cubic tensor. Consider the following attention structure:
Ogg
(0) Vmm
dd
(0)
K ((2) m
n
v m
h
T T...m
m E [do].
m
0
and
0 di (
1 0
d
K
(2)
h
(2)
(2)
h
do
do+1
0
do+1
After the attention operation, we have
(2) = Attn
[H())]
(2)
dg
d
h
mm (n
(K(')
vieln (
m=
j=l
+
1do+1 (1
(2)
K(2) Kott n1)
(2)
do h.
j=1
dg
d
O2g
db
h
do
(0)
m (0))
T f.j.
+ d
d
- 0)
m=
J=1
j=1
0
h
7(1
+
0
E
(1)
E [d],
dv! (0)
where (1) [T(1T OT
and (1)
Ey.mela um (00) T j.m-
(0)
33

---

Published as a conference paper at ICLR 2026
Then two-layer MLP clean the term dv;
which is
t
h
MLP
(2) olp
(1.2)
h
e [d].
0
Similarly, for any E N4. we have
E,
(22)) ML.P e
att cita
(MIP
(Attn
(1) cita
(2
1
i
[d].
d
0
The parameter bounds can be directly computed by the construction above. The theorem is proved.
H
MORE ON EMPIRICAL STUDIES
H.1 MORE EXPERIMENTAL SETUPS
Anisotropic adjustments We consider anisotropic Gaussian mixtures that takes the following form: A K K-component anisotropic Gaussian mixture distribution is defined with parameters O
μ σ. where π { - T o E E, μ Mg Pic 2 R k [K] and σ {oi 02,-- o E Rd. k [K]. sample X; from the aforementioned anisotropic
GMM is expressed as:
X; = My + oy. Zz
(25)
where {y: (yile are iid discrete random variables with Pyy k) गे for k [K] and {Z, [Z.)ee[M
are iid standard Gaussian random vector in IRd, Analogous to that in the isotropic case and overload some notations, we define an anisotropic GMM task to be T (X, 0, K).
To adapt the TGMM framework to be compatible to anisotropic problems, we expand the output dimension of the attentive pooling module from (d + K) x K (d+ 2K) x K, with the additional K rows reserved for the estimate o of o, with the corresponding estimation loss function augmented with a scale estimation part:
1 Σ + + (o.
(
(26)
where the loss function &a is chosen as the mean-square loss. During the experiments, we inherit configurations from those of isotropic counterparts, except for the calculation of the e-reroo metric, where we additionally considered contributions from the estimation eтrоr of scales.
Configurations related Mamba2 architecture We adopt a Mamba2 Dao & Gu (2024) model comprising 12-layers and 128-dimensional hidden states, with the rest to approximately match the number of 12-layer transformer with 128-dimensional hidden states. As the Mamba series of models are essentially recurrent neural networks (RNNs), we tested two different kinds of Readout design with either (i). the attentive pooling module as used in the case
chosen as
of transformer backbone and (ii). more natural choice of using simply the last hidden state to decode all the estimates, as RNNs compress input information an ordered fashion. We observe with
from empirical investigations that using attentive pooling yields better performance Mamba2 backbone. The other training configurations are cloned from those in TGMM experiments with transformer backbones.
34

---

Published as a conference paper at ICLR 2026
Software and hardware infrastructures Our framework is built upon PyTorch Paszke al. (2019) and transforme Wolf et al. (2020) libraries, which are open-source software released under BSD-style and Apache license" ?. The code implementations will be open-sourced after the reviewing process of this paper. All the experiments are conducted using 8 NVIDIA A100 GPUs with 80 GB memory each.
H.2 A COMPLETE REPORT REGARDING DIFFERENT EVALUATION METRICS
In this section, we present complete reports of empirical performance regarding the evaluation problems mentioned section Aside from the E2-error metric that was reported in section 3.2, we additionally calculated all the experimental performance under the following metries:
Clustering accuracy We compare estimated cluster membership with the true component assign- ment, after adjusting for permutation invariance mentioned in section 2.3.
Log-likelihood We compute average log- likelihood as standard metric in unsupervised statistical estimation.
The results are reported figure 10, 11, 12, 13 and 14, respectively. According to the evaluations, the learned TGMM models show comparable clustering accuracy against the spectral algorithm and outperform EM algorithm when > 2 across all comparisons. Regarding the ligliikel metric, TGMM demonstrates comparable performance with the other two classical algorithms in comparatively lower dimensional cases. i.e., E {2, 8), but underperforms both baselines in larger dimensional problems. We conjecture that is might be due to the fact that EM algorithm is essentially a maximum-likelihen algorithm Dempster et al. (1977), while the TGMM estimation objective (2) not explicitly related to lkkelihood-basee training.
H.3 ON THE IMPACT INFERENCE-TING SAMPLE SIZE N
Motivated by the classical statistical phenomenon that estimation quality tends to improve with sample size, we test whether TGMM's estimation performance increases as N goes up. We run corresponding experiments by varying the sample size be N E {32, 64, 128} during both train and inference, while controlling other experimental configurations same as those in section 3.1. The results are reported in fz-error, clustering accuracy as ligeliill and summarized in figure 15. The results exhibit a clear trend that aligns with our hypothesis, justifying the TGMM learning process as learning statistically meaningful algorithm for solving GMMs.
H.4 ON THE IMPACT BACKBONE SCALE
The scaling phenomenon is among the mostly discussed topics in modern AI, as choosing a suitable scale is often critical to the performance of transformer-based architectures like LLMs. In this section we investigate the scaling properties of TGMM via comparing performances produced by varying sizes of backbones that differ either in per-layer width (i.e., the dimension of attention embeddings) or in the total number of layers L. With the rest hyper-parameters controlled to be the same as those in section 3.1. The results are reported in three metrics and summarized in figure 16 and figure 17, respectively. According to these investigations, while in general larger-sized backbone yields slightly better performance as compared to smaller ones. The performance gaps remain mild especially for tasks with relative lower complexity, i.e., K 2. Consequently, even 3-layer transformer backbone is able to achieve non-trivial learning performance for solving isotropic GMMs, a phenomenon that was also observed in recent work He et al. (2025b).
35

---

Published as a conference paper at ICLR 2026
Specto
TCMM
TOMM
EM
TO
TOMM
175
22t
136
125
1M6
(a) E--error
K 2
x
X 5
1.0
1.0
1.0
EM
Conte
Spectral TGMM
Spectral
0.4
123
0.2
(b) clustering accuracy
K
K
-25
-2.
-225
-2.
bg bood
Bood
-6.0
-75
-75
-7.5
-100
-100
-10L
LM
LM
LM
LM
Speriial TOMM
Spennal
Soecral
Spectal
-1511
-150
TGMM
-15.
TGMM
-15.
TCNM
(c) log-likelihood
Figure 10: Performance comparison between TGMM and two classical algorithms, reported in three metrics.
36

---

Published as a conference paper at ICLR 2026
32-112
12-1126
32-1126
64-122
64-1128
-1--122
176
178-1171
170-1176
2-
32-1221
64-1178
128 1128
E2-error
1.0
66-1128
128 1225
0/6
0.0
12-
12-1128
648
64-122
66--128
128 1128
128 +128
120 -128
0.2
27
136
0.2
176
0.2
35
176
(b) clustering accuracy
2=2
K =
K
K 5
0.09
0..0
0.0
-3.5
-2.5
Hig bakkoc
-3.0
-7.5
2.56
7.5
-7.6
122
128
12-
32--128
12.6
12.
12.8
12.6
16.0
(c) log-likelihood
Figure 11: Assessments TGMM under test-time task distribution shifts line with Mi test draws the performance of TGMM model trained over tasks with sample size randomly sampled N /2, Ni and evaluated over tasks with sample size A We can view the configuration 128 + 128 as in in-iitrrbutton test and rest out-of-distributte tests.
37

---

Published as a conference paper at ICLR 2026
EM
EM
EM
Sawrr
Spectrol
125
Spereo TOMM
125
Spertal
TOMM
TGNK
110
D30
TCMM
73
23
20
22-error
K
X
X
1.0
1.0
1.0
Spectral
EM
CARDE
0/6
0.6
Spectall
GGMM
Spectral
Spectral TGMM
82
02
0.2
10
10
10
(b) clustering accuracy
K
1ig Bardb
log ber Rood
12
12
12
(c) log-likelihood
Figure 12: Assessments of TGMM under test-time task distribution shifts II: ey-error of estimation when the test-time tasks T are sampled using mean vector sampling distribution p)as different from the used during training.
38

---

Published as a conference paper at ICLR 2026
2
traneforrer rrombta?
tannformer
nonts??
tondooeer
toaedoree ner
numbee
numbel
222
135
8
175
37
133
(a) E--error
K 2
K
K
R =
1.0
1.0
1.0
1.0
Calte
06
96
123
0.2
(b) clustering accuracy
K
K 5
-25
-2.
-225
-2.
-6.0
-75
-0011
-100
-100
-11L
tranfwwer numbea2
tondlrrer maamia
-25
mondeerrer
trarofrm mardu2
-15
-150
-150
mambaa
1150
(c) log- likeeiklihoodd
Figure 13: Performance comparisons between TGMM using transformer and Mamba2 as backbone, reported in three metrics.
39

---

Published as a conference paper at ICLR 2026
18.5
19.5
195
10.
TCWM
TOWM
TOMM
TCMM
180
100
10.3
10.0
Fo4ter
7.5
125
12
126
125
(a) f2-error
K -
R 5
LM
1.01
TGMM
TOMM
0.8
Clide
0-1
0-1
0.4
0.4
(b) clustering accuracy
1ng bolbice
bollood
EM
EM TENM
EM TGNM
TOMM
TIQUM
(c) log-likelihood
Figure 14: Performance comparison between TGMM and the EM algorithm on anisotropic GMM tasks, reported in three metrics
40

---

Published as a conference paper at ICLR 2026
128
122
-18
-64
22-error
X
X
3.0
3.0
1.0
FUIDE DEC
C ADC
CLIDE
0-6
1170
122
82
02
126
0.2
125
(b) clustering accuracy
3
N
X
23
20
lig Rorooo
Dichoo
Dacooo
-0
32232
-91
-40
-10
-10
(c) log-likelihood
Figure 15: Performance comparison between TGMM models trained under varying configurations of sample-size. For example, N = 64 means that the model is trained over GMM tasks with (randomly chosen) sample sizes within the range [32, 64] and tested tasks with sample size 64.
41

---

Published as a conference paper at ICLR 2026
011
022
01 111
71 77:250
--2
-11
0 ---
D--256
22-error
X
X
K
3.0
1.0
FUIDE DEC
ASC
CARDE
0.6
0-6
0.6
-25
-256
122
0:
126
0.2
33
125
(b) clustering accuracy
3
5
Harboo
Rackooo
Bacooo
-20
-3
200
20
128
120
-221
(c) log-likelihood
Figure 16: Performance comparison between TGMM under backbones of varying scales I: We fix embedding size at d = 128 and tested over different number of transformer layers L E {3, 6, 12). Results are reported in three metries.
42

---

Published as a conference paper at ICLR 2026
2-112
22-error
K
X
X
X K
AEC
D6
86
0.6
0.5
05
05
65
122
-25
126
33
125
(b) clustering accuracy
2
K
K
10.0
-15
-2.5
Bood
-75
-75
-75
-7.5
-120
-125
-12
-12
-150
-15.
-15.
(c) log-likelinn logliiced
Figure 17: Performance comparison between TGMM under backbones of varying scales II: We fix the number of transformer layers at L = 12 and tested over different number of hidden states d {128, 256, 512). Results reported in three metries.
43


Reviews:
[{"review_id": "LkLTUORnra", "reviewer": "Reviewer_4R4H", "summary": "The paper proposes TGMM, a transformer-based learning algorithm that estimates parameters of Gaussian Mixture Models across varying dimensions d and numbers of components K. Empirically, TGMM competes with or exceeds EM and matches spectral methods while handling cases ($K>d$) where spectral fails, and shows robustness to distribution and sample-size shifts. Theoretically, the authors provide interpretations of their algorithm by proving transformers can approximate EM updates and tensor power iterations used in spectral methods.", "questions": "1. What is the data used for training TGMM? How does the model training cost compare in quantity with the computational cost of EM/spectral methods?   \n2. How does the model perform in small separation regimes close to the information theoretical lower bound (see \\[1\\])? \n\nReferences:\n\n1. Regev, Oded, and Aravindan Vijayaraghavan. \"On learning mixtures of well-separated gaussians.\" 2017 IEEE 58th Annual Symposium on Foundations of Computer Science (FOCS). IEEE, 2017\\.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. This paper proposes a new algorithm that uses transformer architectures for GMM learning.  \n2. The authors empirically validate the effectiveness, robustness and flexibility of TGMM, and provide theoretical justifications by proving approximability results of their method.", "weaknesses": "1. While the authors try to demonstrate the effectiveness of TGMM via experiments, it actually performs worse than spectral methods in Figure 2 and Figure 4\\. These figures also fail to decouple statistical errors and algorithmic errors. Spectral methods, as theoretically guaranteed, will have errors converging towards 0 with enough samples. The authors should at least show the same (empirically or theoretically) for TGMM. To demonstrate effectiveness, in my opinion, they should also compare sample complexity of TGMM to baselines and show improvements.  \n2. While spectral methods fail in the degenerate case of $K\\>d$, there are sum-of-squares methods that overcome this issue (see \\[1\\]). Similarly, while EM easily gets trapped in local minima, people have proposed pruning based variants to overcome this (e.g., see \\[2\\]). The authors should conduct more comprehensive baseline comparisons before claiming performance superiority., \n\nReferences.\n\n1. Liu, Allen, and Jerry Li. \"Clustering mixtures with almost optimal separation in polynomial time.\" Proceedings of the 54th Annual ACM SIGACT Symposium on Theory of Computing. 2022\\.  \n2. Dasgupta, Sanjoy, and Leonard Schulman. \"A two-round variant of em for gaussian mixtures.\" arXiv preprint arXiv:1301.3850 (2013).", "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": "K3IfybUV2a", "reviewer": "Reviewer_TJ7f", "summary": "The paper considers the use of transformers to solving Gaussian Mixture Models (GMMs) by learning to directly map from the input samples to the associated GMM parameters considering a simple GMM with identity matrix as covariance whereas mean and cluster proportions are learned. The resulting TGMM is trained to solve multiple GMM tasks using the same transformer backbone. The work is related especially to He et al. (2025b) but the current work focuses on estimating the parameters of the GMM as opposed to the cluster labels. It is demonstrated that L conventional EM steps can be approximated using \\mathcal{O}(L) transformer layers with bounds relying by a constant on the attention heads with polynomial scaling in dimension, consequently improving upon the scaling laws of the work of He et al. It is further discussed how power iterations in spectral methods can similarly be approximated. Notably, the proposed TGMM can also learn GMMs with different number of clusters as the transformer structure used an encoder that takes as input an embedding representation of K concatenated the inputs and a decoder structure that is specific to the number of components to be produced in the output. The TGMM is demonstrated to perform well when compared to EM and spectral GMM estimation whereas a Mamba2 backbone is further considered and found to have inferior performance to the transformer backbone.", "questions": "The loss function optimized in equation (2) contains an implicit weighting between accuracy of mean estimates by least squares and accuracy of cluster sizes by cross-entropy loss. How should these losses be balanced in general? And why not use the standard log-likelihood as loss?\n\nEM and spectral methods are widely used, but it would be relevant also to compare against a simple SGD procedure based on Adam or similar optimization minimizing the GMM likelihood for comparison. How would such simple procedure compare?\n\nWhat is the reason for not including a diagonal variance in the modeling which would make the approach more useful and seems to be straight forward to implement?\n\nWhat is the practical relevance of the procedure – it would strengthen the paper to consider real datasets and the quality of the learned clusters. This would however likely also require expanding the modeling procedure to at least have cluster centers with diagonal variance parameters. What is hampering the procedure to be expanded to such setting and evaluated on real data?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "•\tThe approach is sound and improves upon the related work especially of He et al 2025b.\n\n•\tThe results are compelling and the procedure appears promising when compared to EM and spectral estimation of GMMs.\n\n•\tThe design to accommodate different number of components seems logical and useful.\n\nOriginality: The approach expands upon current efforts utilizing Transformers in the context of unsupervised learning in GMM especially leveraging the recent works of He et al 2025a,b with substantial improvements as outlined in the summary. Theoretical properties wrt. the Transformer surrogating EM and spectral inference procedures is further investigated. The contribution in all provides an overall sufficient and interesting contribution.\n\nClarity: The paper is very clearly written and well presented.\n\nQuality: The paper well presents the developed methodology as well as theoretical insights and the experimentation is generally well executed but limited to synthetic data.\n\nSignificance: The approach expands upon previous recent work and make significant contributions including better scaling properties and insights to Transformers capabilities mimicking unsupervised learning procedures in the context of EM and Spectral based inference for the GMM. The experimentations are solid but limited to synthetic data whereas the GMM assuming isotropic variance makes the practical use very limited reducing the impact of the work.", "weaknesses": "The considered GMM formalism is quite limited and it would strengthen the contribution to consider more realistic settings such as GMMs with diagonal covariance clusters. It is unclear why this would form a major challenge in the presented framework as it technically just requires expanding the readout function to have parameters for the variances of similar dimensionality as the produced means in the readout function and loss functions based on an additional squared error loss term for the diagonal variances.\n\nThe analysis are purely considering synthetic data and it is unclear why the approach is not also tested on real data. This would substantially expand and verify the practical utility of the developed framework. Currently, the limited variance formulation makes it likely impractical for real data analyses.\n\nThe presented performance metric is in terms of least squares error to ground truth cluster centers as well as mixing proportions. I appreciate the L2 metric in the main paper but think the log-likelihood metric considered in the supplementary should also be part of the main-paper as the EM-based procedure is based on optimizing the log-likelihood whereas the TGMM is minimizing the LS error of the parameters. This will provide a more comprehensive evaluation in the main paper in terms of recovering the underlying parameters and how well the model accounts for the underlying density. Furthermore, it would be interesting to display the mean and mixture proportion errors separately at least in the supplementary to understand if there are discrepancies in these errors. It would also be natural to quantify the mixture proportion errors using cross-entropy as optimized.", "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": "HGMmWYntAa", "reviewer": "Reviewer_9Fgx", "summary": "The classical Gaussian mixture model problem (GMM) states as follows.\nSuppose we have a set of $N$ samples, $X_1,\\dots,X_N$, drawn from a mixture of $K$ (isotropic) Gaussians $\\sum_{i=1}^K \\pi_i \\phi(\\cdot\\;; \\mu_i)$ where $\\phi$ is the $d$-dimensional Gaussian with covariance being the identity matrix, $\\pi_i > 0$, $\\sum_{i=1}^K \\pi_i = 1$ and $\\mu_i\\in \\mathbb{R}^d$ for all $i = 1,\\dots, K$.\nOut task is to design an algorithm that takes $X_1,\\dots,X_N$ as input and returns $\\hat \\pi_1, \\dots, \\hat \\pi_K, \\hat \\mu_1,\\dots, \\hat\\mu_K$ such that there exists a permutation $\\sigma$ on $\\{1,\\dots, K\\}$ that $\\hat \\pi_{\\sigma(i)}$ (resp. $\\hat \\mu_{\\sigma(i)}$) is close to $\\pi_i$ (resp. $\\mu_i$) for all $i = 1,\\dots, K$.\n\nThis paper studies the approach of using transformers to solve GMM.\nMore precisely, the authors propose an architecture that first encodes the samples $X_1,\\dots,X_N$, parses them into a multi-layer transformer and then decodes the output.\nThe authors use the $\\ell_2$-error to measure the quality of the output, i.e. $\\frac{1}{K}\\sum_{k=1}^K (\\frac{1}{d}\\lVert \\hat \\mu_{\\sigma(i)} - \\mu_i\\rVert^2 + (\\hat\\pi_{\\sigma(i)} - \\pi_i)^2)$.\nThe authors show that the aforementioned architecture can approximate the Expectation-Maximization (EM) algorithm and the spectral algorithms which are the classical algorithms for solving GMM.", "questions": "- Abstract: It may be helpful to spell out the full name of TGMM for the first time.\n\n- Line 154: ``... i.i.d. ...''", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- The paper provides an interesting connection between the modern approach and the classical approach in machine learning. \nI believe this connection can open up a new pathway to explaining the fundamental structures of different models.\n\n- The presentation is clear.\nReaders of different levels of expertise should be able to grasp the central idea of this paper.", "weaknesses": "- The paper is generally good.", "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": "RbbcOsf76G", "reviewer": "Reviewer_YKHK", "summary": "The paper studies whether transformers can simulate unsupervised learning algorithms with a focus on Gaussian mixtures. The paper provides both empirical and theoretical evidence showing that there exists a transformer that approximates the EM algorithm well.\nNext, they propose a learning framework named TGMM for transformers to learn Gaussian mixtures.", "questions": "Please refer to weaknesses.", "limitations": "", "rating": 2, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The figures are well presented\n- The empirical evidence well demonstrated the proposed TGMM framework is effective.", "weaknesses": "- The theoretical result seems to be weaker than [1], where [1] shows the softmax version and this paper mainly focuses on ReLU transformers.\n- I think the presentation of this paper can be further improved. To me, this paper is more of a theoretical focus than empirical based on its settings. However, most of the main text are used to describe empirical results. \n- The proposed TGMM framework seems to be trivial. Its setup is very similar to the pretraining setup in (Bai et al.2023), where it samples tasks from a meta distribution $\\pi$.\n- It is unclear that how the proposed TGMM framework benefits training from a theoretical perspective.\n\n\n\n[1] Transformers versus the EM Algorithm in Multi-class Clustering (He et al, 2025)", "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)