What Should Embeddings Embed? Autoregressive Models Represent Latent Generating Distributions
Authors: Liyi Zhang
Michael Y. Li
R. Thomas McCoy
Theodore Sumers
Jian-Qiao Zhu
Thomas L. Griffiths
Venue: Accepted by TMLR ()

Abstract: Autoregressive language models have demonstrated a remarkable ability to extract latent structure from text. The embeddings from large language models have been shown to capture aspects of the syntax and semantics of language. But what should embeddings represent? We show that the embeddings from autoregressive models correspond to predictive sufficient statistics. By identifying settings where the predictive sufficient statistics are interpretable distributions over latent variables, including exchangeable models and latent state models, we show that embeddings of autoregressive models encode these explainable quantities of interest. We conduct empirical probing studies to extract information from transformers about latent generating distributions. Furthermore, we show that these embeddings generalize to out-of-distribution cases, do not exhibit token memorization, and that the information we identify is more easily recovered than other related measures. Next, we extend our analysis of exchangeable models to more realistic scenarios where the predictive sufficient statistic is difficult to identify by focusing on an interpretable subcomponent of language, topics. We show that large language models encode topic mixtures inferred by latent Dirichlet allocation (LDA) in both synthetic datasets and natural corpora.

Published in Transactions on Machine Learning Research (07/2025)
What Should Embeddings Embed? Autoregressive Models Represent Latent Generating Distributions
Liyi Zhang Deportment of Computer Science Princeton University
shong. boang.iyf@pprnneeto..d
Michael Y.
michackyjk@sstnnfordd edu
Department of Computer Science Stanford University
R. Thomas McCoy Department of Linguistics and Wu Tsai Institute
tom.mccoy@yole.coo edu
Yale University
Theodore R. Sumers Anthropic
ted@anthropic.com com
Jian-Qiao Zhu Deportment of Computer Science
je5204@princetorv edu
Princeton University
Thomas L. Griffiths
tomg@ppinnetonn.co edu
Departments of Psychology and Computer Science Princeton University
Reviewed on OpenReview: https: // openreview. net/ forum? id= YyMACp98K2
Abstract
Autoregressive language models have demonstrated remarkable ability to extract latent structure from lext. The embeddings from large language models have been shown to capture aspects of the syntax and semantics of language. But what should embeddings represent? We show that the embeddings from autoregressive modcls correspond to predictive sufficient statistics. By identifying settings where the predictive sufficient statistics are interpretable distributions over latent variables, including exchangeable models and latent state models, we show that embeddings of autoregressive models encode these explainable quantities of interest. We conduct empirical probing studies to extract information from transformers about latent generating distributions. Furthermore, we show that these embeddings generalize to out-of fit-o-f-istrrbuution cases, do not exhibit token memorization, and that the information we identify is more casily recovered than other related mcasures. Next, we extend our analysis of exchangeable models to more realistic scenarios where the predictive sufficient statistic is difficult to identify by focusing on an interpretable subcomponent language, topics. We show that large language models encode topic mixtures inferred by latent Dirichlet allocation (LDA) in both synthetic datasets and natural corpora.
1 Introduction
Autorogressive language models (LMs) are trained predict the next token in sequence (o.g., Bengio et al.
2000). Many large language models (LLMs) use the autoregressive objective for pretraining (e.g., Radford
al. 2019) and their document-level embeddings have been shown to capture elements of latent structure

---

Published in Transactions on Machine Learning Research (07/202) 2025)
X1
X
Xa
Xati
X
Xgt
Xprt
Three data generation processes where prediction of the next token In+11 is independent from previous tokens 2mm given predictive sufficient statistic. The left corresponds to exchangeable data, the middle to latent state models, and the right to discrete hypotheses. The relevant predictive sufficient statistics are the sufficient statistic for 0, 2n+1, h (or p(0[zim)) ((2++1(21m)) and p(h|r respectively). We show the embeddings learned by autoregressive transformers represent this information.
Figure 1:
that appear in text, such as agent properties (Andreas) 2022) and syntax (Hewitt & Manning 2019 While it is intuitive that representing this information assists in next token prediction, we lack formal understanding of which aspects of text should be represented and why. Given that these embeddings are often used for downstream tasks such as sequence classification and sentiment analysis, understanding what they represent and why is key to interpreting and building on these models.
Analyzing the representations formed by LLMs is challenging partly due to polysemanticity, where a neuron may activate for several distinct concepts (Cunningham et al.! (023). Previous work has developed mothods for probing LLM internal representations for specilic concepts (Li al. 2023 Meng et al.) 2022 Zheng el al.,
2023 Tenney et al., 2019). These efforts can be guided more effectively by developing a general theory of what aspects of the data embeddings should represent.
In this work, we investigate several cases where the representations of autoregressive LMs can be formally connected with those of a Bayes-optimal agent. Because of their expressive architecture and their objective, LM embeddings should encode latent structure such that the next word xx+1 is independent from previous words Fiss when conditioned on that structure. This corresponds to the notion of a predictive sufficient stotistic. Using this connection, we show that optimal content of embeddings can be identified in two major cases: 1) independent identically distributed data, where the embodding should capture the sufficient statistics of the data; 2) latent state models, where the embedding should encode the posterior distribution over the next latent state given the data. An application of case that we explore is topic models based on latent Dirichlet allocation (LDA), where the LM embedding encodes the topic mixtures of a given document. Wo use probing to confirm that the relevant information can be decoded from LM embeddings, and that content. that is not expected to be captured via predictive sufficient statistics is more challenging recover.
Next, extend our analysis of exchangeable models to more realistic scenarios by generating synthetic data where only fraction of words are generated from LDA. We confirm that language models represent topic distributions when a sufficiently large portion of the text serves this semantic component. We finally show that the encodings of pretrained LLMs (GPT-2, LLAMA 2, and BERT) contain information analogous to that extracted by LDA from two natural corpora, supporting the hypothesis that LLMs implicitly perform Bayesian inference in the naturalistic setting.
Our analysis provides concrete examples where we can analytically identify what embeddings should represent, and empirical confirmation that this information is indeed captured in models trained with an autoregressive objective. By linking LM representations to Bayesian inference, our approach supports recent arguments that the behavior of LMs can be interpreted through comparison with Baycs-optimal agents (Mikulik et al.) 2020). Bayesian inference encourages the agent to summarize data by constructing generative process. Thus, understanding this connection can help LLM researchers hypothesize what features are encoded in LLM internals as done in mechanistic interpretability (Cunningham et al.) 2023); by recovering latent generating distributions from LM, we show that LMs not only learn concepts, but also learn their uncertainty representation, complementing recent studies at the behavioral level (Gruver al. 2024
2

---

Published in Transactions on Machine Learning Research (07/ (07/2025)
2 Related work
Embedding analysis in language models
The embeddings produced by language models have been
investigated in detail (uupta et al. 2015; Köhn
2015 Ettinger ct al. 2016 Adi et al. 2017 Hupkes et al.)
2018; for reviews, see Rogers et al. (2020) and Belinkov (2022). They have been shown to capture aspects
of the latent structure of text, including part of speech (Shi et al. 2016 Belinkov et al. 2017). sentence
structure (Tenney et al.) 2019 Hewitt & Maannig 2019 Liu et al. 2019 Lin al., 2019), sentiment (Radford
et al. 2017), semantic roles (Ettinger et al., 2016 Tenney et al.. 2019), and agent properties (Andreas 2022).
However, our work is motivated from fundamentally different perspective. Instead of focusing on what is captured in the emboddings of these models, we provide general theory of language models as Bayes-optimal agents to explain why these kinds of structure might. be represented due to statistical propert ies of the training data. Characterizing the ideal representations of a Bayes-optimal agent complements recent efforts in mechanistic interpretability (Nanda et al., 2023 Cunningham et al.] 2023 Cammarata et al. 2020) by identifying optimal representations of data; this can inform the search for interpretable features that are used in LLMs' circuit-level computations.
Implicit Bayesian inference Several previous papers have analyzed LLMs by making a connection to Bayesian inference. Of these, Xie et al. (2021), McCoy et al. (2023), and Wang et al. (2024b) analyze the in-context learning behavior of LLMs. However, we study what models should encode based on the autoregressive objective that is typically used to train LLMs. [Zheng et al. (2023) focus on topic models embedded in LSTMs, while we extend the connection to more general cases. [Wang et al.) (2024a) provides a novel method to extract from language models concepts that come in the form of mixed membership model (including the topic model). Our focus is more general and explanatory we analyze why concepts that come in the form of sufficient stats should be embedded by autoregressive language models.
Belief state inference (BST) on latent state models (Shai et al. 2024) decodes the distributions of the current latent state from the transformer embeddings on observed sequences. Our more general theory shows that this ability comes from the embedding encoding predictive sufficient statistics for the sequence, which more directly relates to the distribution of the next latent state given observations. Concurrently, Ye al (2024) also used de Finetti's theorem to support probabilistic reasoning in transformers, but focused on in-contest learning in the context of Bayesian linear regression. We analyze transformer internal representations and identify optimal embeddings in data that violate the exchangeability assumption.
Theory Tishby & Zaslavsky 2013 gave a theoretical analysis of optimal embeddings in deep networks from an information-theoree perspective. We take a complementary Bayesian approach, explicitly connecting optimal embeddings to prodictive sufficient statistics and presenting extensive exporiments confirming that the predicted content of embeddings is tracked by transformer-based LMs. Metalearned RNNs have also been shown to encode information equivalent to Bayesian posterior distribution (Mikulik et al. 2020
Furthermore, recent work has also demonstrated that transformers behave like the Bayes-optimal predictor in linear regression settings (Panwar et al., 2024 Garg et al., 2022 Akyürek et al. 2023 and can approximate the posterior predictive distributions of probabilistic models such as Gaussian processes and Bayesian neural networks (Müller et al. (022). We extend this analysis to general autoregressive language models and consider more general generative processes and what posterior distributions they should capture in these cases.
3 Optimal embeddings
Assume we have a sequence Fian and an autoregressive language model that predicts the next item in the sequence, p(zn+1 [#1..). We denote the LM embedding for sequence Zim as é, = f(1m)) The distribution p(#at1l#1sd is function g(Øx) of this embedding, with that function implemented by the final layer of the neural network instantiating the LM such that the probability of the next element 4x41 only depends on dn- This establishes our question: what should on represent. in order to accurately predict xx+1?
The key idca behind our approach that we can identify situations where Φn contains all of the information from Fis required to predict 2n+1 by using the notion of sufficient statistic (Gelman et al., 2004). Given distribution p(x) with parameters 0, statistic s(x) is sufficient for 0 if the conditional distribution of x
3

---

Published in Transactions on Machine Learning Research (07 (07/2025)
given does not depend on A. In other words, if we only know 8, estimate @ just well as if we know 2. In the autoregressive setting, we care about predictive sufficiency (Bernardo & Smith 2000). A statistic s(11m) is predictive sufficient for the sequence Fim if
p(ra+1 Fim) = p p( (((()))))
(1)
Crucially, if a model performs autoregressive modeling perfectly, it learns a predictive sufficient statistic.
In the remainder of this section we describe two cases where predictive sufficient statistics can be easily identified and could plausibly be represented by a neural network (Figure 1]. These cases correspond to common applications of machine learning models as well as classic models for text such as hidden Markov models (Jurafsky & Martin 2008). First, when 21. are independently sampled conditioned on an unknown parameter, On needs only represent the sufficient statistic of this sequence. Second, when Tim are generated by state space model (in the discrete case, hidden Markov modcl), Pn needs only represent the posterior distribution over the next latent slate given Ti.n- In each case explain how p( [m)m)) factorizes to make it possible for Iim to be summarized by some on and identify the form of the corresponding g(o,).
3.1 Case 1: Exchangeable models
Predictive sufficiency is straightforward to establish in ezchangeable models, where the probability of a sequence remains the same under permutation of the order of its elements. That is, a sequence is exchangeable if p(11.N) = p(=+(1:N)) for some permutation æ. Any exchangeable model can be re-expressed in terms of the 2; being sampled independently and identically distributed according a latent distribution p(x|0) parameterized by 0, with p(æ1.N) II, ((2,)))p(()))))))) ())( d0 (Gelman al. 2004). This idea leads to the following proposition:
Proposition. Given an exchangeable sequence øn,N where each x; is of dimension day and given functions f Rads + Rd g Rd + Rd such that, for each 1 < s N, g f(r1n) = (1n4t1/11n VIn+1, f(rim) is a sufficient statistic for Tim-
In other words, if we have a perfect autoregressive predictor that is composable into o f, the output of f is sufficient statistic for its sequence input.
Proof. The result follows from the fact that for exchangeable sequences, general sufficiency is equivalent to predictive sufficiency (Bernardo & Smith 2000). Because p(zn+1|z1m) = ((f(#im)) Vn, f(11m) is a predictive sufficient statistic for the sequence, and it is also a sufficient statistic.
The resulting sufficient statistic also fully specifies the postcrior on the parameters of the generating distribution, p(0 (.))). Sufficient statistics are easily identified for a wide range of distributions, including all exponential family distributions (Bernardo & Smith 2000), and are easy represent. Since the LM's predictive distribution decomposes into the form above, this result gives us strong predictions about the contents of embeddings for models trained exchangeable data.
3.1.1 Case 1.1: Discrete hypothesis spaces
In more specific version of Case 1, assume each 2; is generated independently from some unknown generative model, and let H. denote the discrete set of hypotheses h about the identity of this model. In this case, Zin are exchangeable, but any sufficient statistics might be difficult to identify. The posterior predictive distribution can be written as
p(Zn+1l21m)
Σ (( hEW
By an argument similar to Case 2, p(h Fin) predictive sufficient statistic in this model and an embedding thus need only capture this posterior distribution.
4

---

Published in Transactions on Machine Learning Research (07 (07/2025)
3.1.2 Case 1.2: Topic Models
Latent Dirichlet Allocation (LDA; Blei al. (2001)) is exchangeable generative model that is widely used for modelling the topic structure of documents. A document is generated from mixture of K topics. Each topic is distribution over the vocabulary; e.g., topic corresponding to geology might assign high probability lo words such as mineral or sedimentary. We detail the full generative process in Appendix A.2
After LDA is trained on a corpus, the inferred quantities can be used to explore the corpus. We denote by 0 as the latent variable that stands for each document's underlying topic mixture. Because topic mixture i and word distribution form the sufficient statistic, our theory suggests that autoregressive LMs should implicitly encode the topic structure of document.
3.2 Case 2: Latent state models
In latent state model, each T3 is generated based on a latent variable 2;- These 2 are interdependent,, with 2; being generated from a distribution conditioned on 74-1- Common latent state models include Kalman filters (where zj and 2; are continuous and the emission and transition functions, p(zz(2;) and p(z; (zz]zz4-), 1), are linear- Gaussian (Kalman 1960) and hidden Markov models (where the 2; are discrete) (Baum & Petric 1966). In latent state model, the posterior predictive distribution is
p( | p( (20tn-11 (p(===1 ( dznz1-
(2)
In this case, the posterior over the next latent state p(=n+1 captures all of the information in Zimn relevant to predicting In+i, rendering In+1 independent of Fim when conditioned on this statistic. If we were to drop conditional independence assumptions in latent state model, one would need ((2n+1/24421 (im)) instead of p( 1 for the equality. Thus, (( is predictive sufficient statistic in this model. However, there is another predictive sufficient statistic. We can write
p(Znt1l11m)) = / p(sn+alle))(za (1 den,
(3)
showing p(zn|z1mm) to also be predictive sufficient statistic. Which of these will be favored will depend on how easily the relevant integrals can be approxinated.
More formally, we consider an embedding a representation on such that a fixed operator g(dn) can produce p(rn+1 J. The distributions p(2m+1/21.0) and p(2n (im) satisfy this characterization, with Equations 2 and 3 showing that the relevant operator is the integral of p(#=+ 1 over 1n+1 or p(*n+1|zz) over 2, and 2n+1- That operator can be easily approximated linearly and hence by a single layer of neural network. The embedding on thus needs only represent p(zn+1 (mi:n) or (21n), depending on the relative case of approximating the relevant integrals in specific autoregressive predictor.
3.3 Probing embeddings to recover predictive sufficient statistics
These cases specify information that should be encoded in neural network embeddings. This sets up the second element of our approach, which is building probes to check this hypothesis. We focus on transformers (Vaswani et al. 2017), since they widely used as language models. We denote the decoding target for each sequence Timh by vector tn- Given trained transformer, we decode t, by training second model (a probe) to predict the target tn from the embedding on of the corresponding document Zim- The embedding is defined to be the last layer embedding, which is what researchers and practitioners typically use as a document representation. The probe g maps from the sequence embedding to the target. To ensure that the relevant statistical information contained in the LM, not in the probe, we keep the probe simple by In each case, the probe
defining it as a linear layer with softmax activations: g(Φ.)) = Doftmaxx(innea((( uses the layer, last-token embedding its input, unless othree specified.
5

---

Published in Transactions on Machine Learning Research (07/2025)
Table 1: Probing results on discrete hypothesis spaces. In general, the probe achieves strong performance recovering 784-length vector, and performance increases as task difficulty decreases. Average and standard error across 10 random seeds are reported.
Equal Wiath
Unequal Width
Sample Size
Accuracy f
Squared Loes
Accuracy ?
Squared Loss
0.005
0.3%
0.004
50
99.5 0.1%%
0.008 0.001
88.5 0.3%
0.159 0.003
4 Empirical analysis
We have identified cases where a predictive sufficient statistic is expected to be encoded by an autoregressive model. In this section, we use probing analyses on transformers to empirically validate this hypothesis. Ablation studies memorization, parsimonious properties of embeddings, and decoding alternative quantities further support this idea. We use the Adam optimizer (Kingma & Ba. 2014) in all experiments. Code is available al github.com//haaae liyi/lmm--nbeeddnng.
4.1
Simple exchangeable models
4.1.1
Generative model
Bayesian conjugate models We start with three data generating distributions: Gaussian- Gaussian-Gamma, Beta-Bernoulli and Gamma-Expome Exponenaxppnentiall models.
For M Gaussian-Gamma sequences (x((k Jisasar as with N tokens each, each sequence i.i.d. generated by mean and precision parameter sampled from prior. For sequence 2((h, the generative process is,
GAMMA(a0,88)
Th
M4 1 Tk ~ N(Ho: (1)T))=
(K)
14,, Th ~ N(HR, ( for E (1,...NN,
where an. 90 Ho, Ao are fixed hyperparameters. The full generative processes for the other two models are given in Appendix A.2
The Gaussian Gamma Gamssian-Gamma model generates from a Gaussian distribution with two unknown parameters, mean μ and precision т. The predictive distribution for the next token xn-1 1 is
p(an+ ( f plznn+ (oriill,T,a T (/om)) 1m
The optimal Bayesian agent uses the same prior distributions as the Gaussian and Gamma that generate p and т. It will analytically infer the ground truth posterior p(u, p(1,T( T(m)s) for any stream of data it sees, and use this posterior for predicting next tokens. A Bayesian agent also use other suitable priors and converge to the optimal posterior. To be consistent with other exchangeable conjugate models, we denote 0 (p, r) to indicate latent variables whose posterior distribution is predictive sufficient.
Discrete hypothesis spaces Each sequence in this datasct consists of two-dimensional points uniformly sampled from a rectangular region in 2D space (cf. Tenenbaum, 1998). The hypothesis space H. is the set of all rectangles whose corner points are pairs of integers in (0, 1, ..., 7). The generative process uniformly samples rectangle brect from the set of all rectangles H. Then, each token in a sequence is sampled uniformly from the region defined by hreet-
4.1.2 Probing experiments
Implementation details For each dataset, the out-of-distributio (OOD) hyperparameters are chosen such that OOD data are centered on a different mean and has different spread. The exact choices are detailed in Appendix A.3 along with hyperparameter sweep in Appendix A.5

---

Published in Transactions on Machine Learning Research (07 (07/202) 2025)
(a) Gaussian-Gamma.
(b) Beta- Bernoulli.
(c) Gamma Gamma-Exp.
Figure 2: Probe recovery of transforner-learreed sulficient isate (blue) and ground truth sulficient statistic (red) on the y-axis, across 1000 test datapoints on the x-axis. In the plot above, the datapoints are sorted based on thoir ground truth sufficient statistic. The first row shows paramcters probed in the non-OOD case (from left to right: Gaussian mean μ, Gaussian precision т, Bernoulli mean, and Exponential mean). The second row shows the corresponding information the OOD case.
3.4
(a) hypothesis space with equal-width rectan- gles (left), and unequal widths (right).
(b) Colored by mean along y- axis of the generating rectangle.
(c) Colored by spread along y-axis of the generating rectangle.
Figure 3: (a): Two discrete hypothesis spaces H used in experiments. Any continuous rectangle contained within the axes (c.g., the red or the orange rectangle) is valid hypothesis h E H. The data consist of a sequence of points sampled uniformly from the target reclangle. (b) and (c): Two-dimensional representation of embeddings of all validation datapoints (the setup is unequal width and sample size = 50). The two subfigures show the same emboddings, colored by propertics of the true generating rectangles.
In discrete hypothesis spaces, rectangles with seven unit blocks each side, resulting in a size-784 hypothesis space, i.c., 784 possible rectangles from which sequences are drawn. The unit lengths are either 1, or 0.4 and 0.16 alternating (Figure 3al.
Experimental results We hypothesize that a transformer trained on these datasets should come represent the sufficient statistics of the corresponding distribution.
Sufficient statístic. Results for the three Bayesian conjugate models are shown in Figure where the probe successfully decodes the sufficient statistics 0. For hypothesis spaces, we probe the distribution p(h)aim)) which is a length-784 simplex vector. Results are shown in Table I In general, the true hypothesis can be found with high accuracy, even though the number of classes is high.
Figure g visualizes how embeddings represent cach hypothesis through Principal Components Analysis (PCA) that reduces the 128-dim embeddings 2D, and the points are colored by properties of each rectangle. Embeddings clustered into different. regions based the geometry of the true generating rectangle (Figure

---

Published in Transactions on Machine Learning Research (07/202) 2025)
d
Figure 4: Probing over the first 10 tokens themselves using the 10th token embedding of the transformer. Aside from perfectly encoding the 10th token, this embedding does not show memorization over the other 9 tokens as suggested by the noise in probe recovery.
3b and Bcb. where, 3b for instance, the top ycllow clusters correspond the upper-most roctangles, and dark blue clusters correspond to the lower-most rectangles. Additionally, the embeddings encode the distance between corners of the generating rectangle (Figure Bc). We also find the same color pattern (but rotated 90 degrees) when the horizontal instead of vertical axis is used to color the embeddings.
Moments of the posterior distribution. A consequence of knowing the sufficient statistic is finding the true posterior p(0[ I1:N), and we show that the moments of this posterior can be decoded from the transformer embedding (see Figure S and 9 in the Appendix). The moments are functions of the mean and variance of the stream of observed data. In these conjugate models, might be unsurprising for a transformer to encode the mean of the sequence it processes, because the optimal strategy for its loss function is always predict the mcan of the sequence that it sees far. However, encoding the variance would not be directly related to this strategy and would support the argument that it infers sufficient statistics.
Out-of distribution simulations. The analytical nature of the Bayesian predictor means that it is robust to datasets generated far from the prior - it would simply update its postcrior based on the data. Thus, we probe the transformer on out-of-distrrbuttion (OOD) datasets that are generated from a distribution of the same form but with distinct hyperparameters. Figure 2 as well as Sa in the Appendix, suggest. that the recovery of sufficient statistic and moments of the posterior distribution is generalizable to OOD cases, just Bayes-optimal agent would generalize.
Memorizing the context or storing sufficient statistics of the context? A possible confound behind the probe is that, instead of having a probe that uses the ceuusforer-emmeeeee transformer may memorize a set of tokens, and the probe infers quantities of interest from these memorized inputs. We begin by making sure that the sufficient statistics over the first ten tokens can be decoded from the probe on the tenth token in the Gaussian-Gammi case (Figure 11 in Appendix A.5). Then, we probed the token values themselves to look for memorization. Figure 4 suggests that memorization is generally absent. The 10th token embedding recovers the 10th token perfectly, but cannot recover the other 9 tokens. For the other 9 tokens, correlation exists between probe results and true token values, but some correlation is expected because even single token can revcal information about the generating distribution. However, the noise suggests that the model is finding sufficient statistics rather than memorizing.
sufficient statistics, the
Embeddings are parsimonious Our analysis so far has focused on whether embeddings contain infor- mation that allows decoding sufficient statistics. However, our hypothesis is stronger than this: since the sufficient statistics are enough to encode the relevant information from the data, a model need only represent that information. To explore whether transformers construct. such parsimonious representations, we fixed the models trained on different distributions above and used multi-layer perceptron (MLP) to directly predict their last-layer, last-token embeddings using the sufficient statistics of the training sequences as input.
Results show that in most cases the sufficient statistics capture well over 50% of variance in the embeddings (Table X in Appendix [A.5). The results show that sufficient statistics capture a substantial amount of variance in transformer embeddings, establishing that the relationship between embeddings and sufficient statistics runs in both directions: we can decode sufficient statistics from embeddings, and we can predict embeddings from sufficient statistics.
8

---

For an HMM generating data with M sequences with tokens cach, we formulate the generative process as Published in Transactions on Machine Learning Research (07 (07/2025)
Table 2: Probing target quantities in HMM dataset using different transformer token embeddings. The target (2)m) is a simplex vector found by running the forward-back backward algorithm, and the target 2 is a scalar standing for the most likely latent class found by Viterbi algorithm. Average and standard error across 10 random seeds are reported.
= 0.5
1
Target Quantity
Embedding
Accuracy
Squared Loss 4
Accuracy f
Squared Loss J
p( ia (
#im
90.8 1.7%
0.004
1.9%
0.011 + 0.003
plan (4))))))
86.5 1.3%
0.066 1 0.011
82.2 1.8%
0.067 1 0.011
planea (Amm)) P Mennea (Pime)) 2nss
Zime1
66.6 2: 6.6%
0.072 = 0.013
65.9 6.8%
0.058 + 0.011
Fim
53.2 + 3.5%
0.356 + 0.014
50.8 3.6%
0.278 + 0.012
Fir
59.8 3.8%
61:1 4.7%
Pyll
80.8
77.5 2.8%
4.2
Hidden Markov model
4.2.1 Generative model
A. ~ DIRICHLETC()) for C {I, 0}
B. DIRICHLETY (8) for c E {l,..., C}
20 (AIRGORICAL(T)
CATEGORICAL.(::.)
24+1 CATEGORICAL(A.))
where C, a are initialized and denote, respectively, the number of classes, vocabulary size, and a list of probabilities on the number of classes to initialize the first latent state. 7,6 are scalar hyperparameters that are also initialized and fixed, and they represent the evenness of the samples from the Dirichlet distributions. A, B, and 2, result, represent the transition matrix, the emission matrix, and the latent states, respectively.
We use the forward-backward algorithm [Rabiner 1989 to compute the posterior p(zatal Fr:a), and explore whether this distribution can be decoded from the transformer embedding on Fi:n-
4.2.2 Probing experiments
Implementation details We choose C 4, V 64, 0.5, and set π to be uniform in our experiments. We also vary to control the level of difficulty: how distinct is one class from another.
Results Our theoretical treatment suggests that the transformer should encode the predictive sufficient statistic p(z,+1 [xI.)) or p(=n) The latter is the target used in belief state inference with transformers (Shai et al. 2024). However, there are other natural decoding targets that could be used. A simple target is the onc-hot vector corresponding to the most likely hidden state ên+1- This is not a predictive sufficient statistic and does not encode the full information about the sequence relevant. to future prediction. Our analysis thus suggests that will thus provide a poorer match to the information contained in the embedding.
Table 2 suggests that the transformer encodes the predictive sufficient statistics. Furthermore, performance is better for hypothesized measure (first and second row) than for related quantities (all other rows). Additionally, it is more difficult to recover the mcasure used in belicf state inference (sccond row), although it is intuitively easier to learn: the larget probability is over the same latent variable, but the emission of this latent variable is observed, unlike in our hypothesized measure where is unobserved. This may be because if the embedding were to directly encode p(2, 111),, the transformer output layer after the embedding would additionally need to encode the transition matrix and additional integration computations, potentially making the task more challenging than encoding p(z,+1 (1.))
9

---

Published in Transactions on Machine Learning Research (07/2025)
Table Topic prediction by the autoregressive transformer (AT), BERT, LDA, and word-embedder (WE) on LDA-generated synthetic datasets, along with standard deviations across random seeds. Hyperparameter a defines the dataset generation process, where a higher a means more difficult task with underlying topics being more evenly distributed. AT and BERT have similar performance in the casiest setting, but AT performs well in harder settings where BERT performances worsen. End-to-end WE achieves stronger performance than language models, and LDA matches expectations by providing an upper bound in performance.
Method
Accuracy
loss
Tor. var. loss
AT
82.8% + 0.5%
0.041 = 0.001
0.141 + 0,001
0.5
BERT
83.6% + 1%
0.036 = 0.003
0.131 + 0,005
LDA
87% +: 0.6%
0.029 1
0.117 0.001
WE
85.8% 1.3%
0.03 0.001
0.119 0.002
AT
75.5% 1 0.8%
0.044 + 0.001
0.144 + 0.001
0.8
BERT
51.5% + 1.7%
0.111 + 0.005
0.233 + 0.011
LDA
82.6% + 0.5%
0.036 + 0.001
0.133 + 0.004
WE
80.9% + 0.5%
0,029 =
0.116 + 0.001
70.5%
1.6%
0.045
0.146 0.003
BERT
46.
3.3%
0.1 0.004
0.222 0.006
LDA
79.6% 1 1.4%
0.045 - 0.004
0.147 0.006
WE
79.4% + 1%
0.027 +
0.113 + 0.001
0.5
Dataset usee to train BENT
y
train
1 12
0.6
0.4
02
0:0
-135 -152 -135
Dalasel the probe
Dalasel usec the probe
AT probe accuracies.
(b) BERT probe accuracies.
(c) Accuracy neg, perplexity
Figure 5: Figure 5a and 5b control experiments showing AT (left) and BERT'S (middle) probe validation performance synthetic data. For each AT and BERT, five models are trained and validated on five datasets, with each dataset generated by distinct topic model. Colors show probe accuracy. A cell on row i and column j corresponds to model i on dataset j, so the diagonal corresponds to a model on its own datasct. For AT, performance is only strong on the dataset with the same generating topic model, suggesting that the underlying statistical model, not the probe taking different word embeddings, is responsible for performance a relationship that is also present for BERT but a lesser degree. Figure 5cl 20NG probe classification performance (accuracy) negative perplexity measured at 100 different tokens. The dots are colored by the position percentile. Probe performance increases with lower perplenity.
4.3
Topic Models
We aim at recovering from language models the topic mixture 8 that draws words in document i.
4.3.1
Probing experiments
Experiment setup The dataset is bags- oo-dords generated by LDA. We set the vocabulary size V 103. number of lopics K = 5, and generated N 10a documents that are each 100 words long.
Models We trained four models: an autoregressive transformer decoder (AT), BERT, LDA, and an end-to- end word embedder (WE). BERT uses masked objective instcad of the antoregressive one, and implemented by a small version called BERT-TINY (uure al.) 2019). LDA is used to establish an upper-bound for model performance. We also include model intended to provide upper bound for embedding performance:
10

---

Published in Transactions on Machine Learning Research (07/2025)
IDA ditepcitt
LDA datapoin:: 100
LBA: deteschtt
Satarsint
0.0
06-
6:3
64-
0:39
62
8:55
3.2
00
Jotape
Secaporn
0.5
114
048
03-
033
34
64-
0.5
622
61
Figure 6: Distribution of synthotic data topics predicted by LDA (row 1) and AT classifier (row 2) for different validation datapoints. Predictions are stick plots, and ground truth is bar chart in the background. They both exhibit learning of topic mixtures by trying to match the distribution, in addition to top-1 agreement.
word embedder which is matrix that maps from the vocabulary space to the AT BERT embedding dimension and is end-to-end trained with the probe that predicts topics.
For each model, the hidden sizes and final-layer embedding sizes are 128. AT has 4 decoder layers to match the size of BERT-TINY. More hyperparamctors are detailed in Appendix A.3
Metrics We use accuracy, L2 loss, and total variation loss to measure both classification performance and recovery of the topic mixture distribution. Accuracy is defined as how often the top topic predicted by the classifier's mixture agrees with the top topic from ground truth. The remaining loss measures apply to the whole topic mixture.
Results Figure G shows five examples of ground truth, mansformet-prediccee and LDA- topic mixtures. Table 3 shows that all models demonstrate success at recovering latent topics on at least the easiest. setting (i.e., a = 0.5), being able to return both top-topic accuracy and the topic distribution spread. Between AT and BERT, the probe on AT able to infer latent topic structures more difficult tasks (i.c., a = 0.8, 1) whereas the probe on BERT shows deteriorating performance. LDA outperforms both AT and BERT as expected because However, the strong WE performance suggests that the probe can predict topics mainly from stand-alone words. This raises the question whether AT and BERT are learning underlying statistical model, or are simply uniquely embedding each word and making the probe mainly responsible for topic recovery.
specified exactly to learn a dataset generated by the other manually initialized L.A.
4.3.2 Controlling for probe performance
In this section, we conduct control experiments that suggest that language models learn an underlying statistical model, making them-rather than the topic probes taking word emboddings -mainly responsible for successful topic recovery. If AT BERT performs well just because it gives each word unique embedding from which trained probe suffices recover the topic mixture, then probe top of the language model should additionally predict topic mixtures from different underlying topic model than the one that generated the language model's training data.
To assess this, we generate five datasets using five distinct topic models under the setting of a = 0.5. By definition, these topic models initialized five different sets of topics, or distributions over words (B....,Bk]. Jy.j E f1,2,3,4 3, 4.. The transformer trained on set j should encode the mixture of topics 0 that corresponds to topics (fI,-.,,kk) 8x))
11

---

Published in Transactions on Machine Learning Research (07 (07/2025)
LPPA
Car
LDDA
Car
Car
EPA
L.DA
Cat
LDA
ACCURTT
V=O0.C
VH005
LDA
24
Figure 7: Left: Illustration of two sequences (upper: p R 0.5; lower: p 2 0.25). Each box is a word, coming from either the -ggeeerrete semantic component (blue), or one of the three categorical distributions that indicate three syntactic classes (grey). Right: Probing accuracy across semantic proportion p on HMM-LDA generated datasets. Lower 7 corresponds to more sparse transition matrices. The language model successfully encodes semantics if it occupies a considerable enough proportion.
Table 4: 20NG topic prediction performance based on different LLMs. Trained LLMs substantially outperform the wull GPT-2 model, supporting the hypothesis that the training process encourages LLMs to implicilly develop topic models. Autoregressive models (the first five models) statistically significantly outperform non-autoregressive models (the next two).
K
K 100
Model
Parameters
Accuracy f
L2 loss |
Tol. var. loss
Accuracy f
L2 loss |
Tot, var. loss 1
GrT-2
124M
61.4% : 1.5%
0.106 i 0.002
0.211 = 0.001
42.3% + 2.4%
0.097 + 0
0.192 + 0.001
GPT-2-MEDIM
355M
62.6% + 1.7%
0.104 + 0.002
0.209 + 0.002
42.9% + 2.4%
0.096 + 0
0.19 + 0.001
GPT-2-LARGE
774M
62.4% = 188%
0,102 0.002
0,208 = 0.002
43.1%
= 2.3%
0.095 0.001
0.189 +
LLAMA
7B
62.6% = 1.7%
0.101 0.002
0.206 = 0.002
43.3%
= 2.4%
0.095 0.001
0.189 0.001
LLAMA 2-CHAT
62.9% 1.7%
0.102 0.002
0.207 1 0.002
43.2%
2.5%
0.095 0.001
0.189 1
BERT
110M
56.3% + 1.5%
0.113 = 0.003
0.222 + 0.003
38.6% + 2.5%
0.1 = 0.001
0.191 + 0.001
BRRT-LARGE
336M
55.2% + 1.2%
0.116 + 0.002
0.226 + 0.003
38.9% + 2.9%
0.1 - 0.001
0.191 + 0.001
Null GPT-2
124M
27.3% + 1%
0.209 0.003
0.322 = 0.005
13.8%
1.7%
0.145 0.001
0.248 0.003
One AT and one BERT is trained on each dataset. On each model, five probes are used to prodict topics from each of the five datasets. Results are shown in Figure 5 AT shows a strong distinction between predicting its own dataset versus datasets from other topic models, whereas this distinction is present but weaker for BERT. These results are evidence that AT and BERT indeed encode topic information because, the topic information were instead constructed by the probe, the probe would work equally well on mismatched datasets matched ones.
4.4 From HMM-LDA to Natural Corpora
In this section, we start with a syuthetic dataset with more realistic assumption than topic models that topics exchangeable component in partially exchangeable sequence. We generate dataset based both LDA and a hidden Markov model (HMM). Then, we move to evaluate the extent to which LLMs recover topic mixture two natural corpora, 20Neusgroups (20NG) and WikiTent-1083
4.4.1 Topic distribution can be recovered when LDA is sequence sub-component
Most natural texts not exchangeable. For example, consider this sentence: "The sediment found in the quartz includes silicon." This sentence is neither fully exchangeable nor generated purely by latent state model. However, if we only consider the words that contribute to the sentence's topic (geology), we are left with sediment, quartz, silicon. These words can be plausibly generated by an exchangeable topic model as a sub-component: the sentence "The silicon found in the sediment includes quartz" also has the topic of geology even though word order changes. While we do not exhaustively list possible factors behind language that can be embedded by LLMs, we study to what extent can topics be encoded by model embeddings under settings where topical words occupy varied proportion with respect to the whole sequence.
12

---

Published in Transactions on Machine Learning Research (07/)025)
Model Similar to the LDA experiments, we generate synthetic documents with HMM--DA model [Griffiths et al., 2004) (Figure 7 HMM-LDD combines semantics and synlax, and posits that each word comes from a latent class and that class transitions are governed by an HMM. Using the HMM-LDA model allows one to manipulate the degree to which topic cxists as a sub-component of the sequence.
When p is 1, all the words are generated from the LDA model. As p decreases, more words are generated from the syntactic classes. We can then examine how well topic distributions are recovered from the embeddings of models that are trained on data generated from HMM-LDA models that vary p.
Results Results in Figure 7 suggest that without exchangeability, the autoregressive transformer still learns the same semantic latent variables provided there is a large enough proportion of LDA-generated component (i.e., p 2 0.4). This suggests that the theory of finding explainable quantities in autoregressive models can potentially be extended to real texts where interpretable probabilistic models, such as the topic model, exist as sub-component of language.
4.4.2 Recovering latent topic distributions in natural corpora
This section turns to analyzing LLMs pretrained on natural language. Although an exhaustive list of predictive sufficient statistic is difficult to identify, topical words do form a proportion of factors underlying the text. This allows us to hypothesize that lopic mixtures can be decoded from the transformer residual streams.
Datasets We use 20Newsgroups (20NG) and WikiText-103 (Merity et al. 2016 20NG is a collection of 18,000 posts written in a style similar to informal emails, divided into 20 subjects. Contrasting with the informal language style in 20NG, FikitText-103 consists of over 100 million tokens sourced from the set of verified articles on Wikipedia that are classified as Good and Featured.
Setup We train LDA models across three random seeds on each dataset, and use pretrained large language models as our LLM. Specifically, the LLMs are GPT-2, GPT-2-MEDIUM, and GPT-2-LARGE (Radford et al, 2019); LLAMA 2 and LLAMA 2-CHAT (Touvron el al. 2023); and BERT, and BERT-LARGE (Devlin el al. 2019]. The probes target LDA-learned topic mixtures with K E {20, 100}. For the models, we additionally include a randomly initialized GPT-2, called Null GPT-2, as control that can differentiate the contribution of the LLM from the contribution of the probe. Considering that the first token may lead to useful embedding that stands for the (CLS) token in BERT series, we searched across the first, last, and average embeddings.
Topic prediction results Performance is shown in Table 4 for 20NG and in Table 7 in the Appendix for Wikitext-103 To predict mixtures of twenty topics, random guessing would yield a 5% accuracy for K 20 or 1% accuracy for K = 100. The best-performing LLM on each dataset demonstrates success at encoding topic distributions by achieving 88.5%/74.2%% accuracy on WikiText- and 02.9%/43.27 2% on 20NG. The probes here lake averaged loken embeddings the last layer in each LLM. While having direct information over whole sequences is advantageous compared to over an arbitrary token, we show that using only the last token preserves strong performance, with 73.7% 27%5%%%59 58.99% accuracy on WikiTexxt and 22.2%%34.77 on 20NG (Table 14 and 13 in Appendix). These results suggest that topic mixtures are directly encoded in embeddings, rather than resulting from high-quality word representations.
Topics encoded by inner layers We also explored the possibility of decoding topic mixtures from other LLM layers, as opposed to only the last layer (Table g in Appendix A.5 In the LLAMA-2 series, the decodability of topic mixtures progressively increases as we move from the word embedder to the intermediate layers. Although word embedder can potentially achieve high accuracy demonstrated by the synthetic dataset, we observe that LLAMA-2 relies on its inner layers to encode topic mixtures. This phenomenon can be caused by LLAMA-2's stronger ability to model the next word, which is correlated with the ability to capture latent structures, which we discuss in more detail in the next paragraph.
Probe performance and LLM perplexity To show that the the posterior predictive on topics corresponds to autoregressive prediction, we analyze 100 probes trained on 100 different token positions, along with their corresponding I.L.M perplexity (Figure 5e). Each position is defined based on the corresponding percent
13

---

Published in Transactions on Machine Learning Research (07/2025)
of the total document (each document has different length). We expect that as perplexity token increases, probe performance based on the embedding taken from that token would decrease. This hypothesis is supported by the linear trend in Figure 5c A possible confound in this result is that perplexity tends to decrease as the position in a document increases, so the results could be driven by position rather than perplexily; however, we ran a linear mixed-effects model and found that perplexity continued to have a statistically significant effect controlling for position (sce Appendix A.4).
5 Discussion
In this paper, we have shown that autoregressive modeling is connected to Bayes optimal agents, where the sufficient statistic is encoded by the model in two major cases: 1) exchangeable distributions that emcompass a wide range of generative models including the topic model; and 2) latent state models that include the HMM. We have also seen that when these latent variables form a significant subcomponent of the data, they explanatory approach on why lopic mixtures embedded in transformer
can still be encoded, leading residual streams.
There are several potential implications and future directions of our findings.
Interpretability. Understanding the inner workings of LLMs is important for Al safety and trustworthiness. However, the sheer size of LLMs makes it challenging to analyze. Our results suggest that interpretable models of document structure such as topic models provide useful guidance about what mechanistic interpretability should look for. This in turn means that one goal for interpretability work should be enhancing models like LDA because such advances will in turn sharpen our ability to interpret LLM representations.
Reverse direction- training LLMs to capture inductive biases. We have shown how autoregressive models capture latent structures when predicting the next token. Another direction to explore whether it is possible to explicitly incorporate inductive biases during training more reliably and efficiently capture these latent structures, have been done in modeling language (McCoy & Griffiths 2025 and vision
Carballo-Castro et al. 2024).
Extensions to masked language models. It has been shown that the autoregressive objective can be broken down into the Bayesian predictive distribution in several cases. It is possible to extend this breakdown to masked language models as well (Appendix A. Whereas the autoregressive model learns a different posterior distribution p(f) for each n as it progresses through tokens in a sequence, the masked model learns that same p(0 m where U denotes the set of unmasked indices, resulting in a less expressive objective. On synthetic topic model bags-of-words. the masked model underperforms, but it is worth exploring whether its performance is consistently worse than that of the autoregressive model in a wider range of cases.
Applications. Our results help inform construction and use of embeddings in practice, and can be applied to generative models that satisfy the conditions discussed by the paper. The analyses can be naturally extended lo other latent variables that do not depend heavily on word order, such as the author type of the document (Adreas 2022) or the author's sentiment (Radford et al. 2017
As another example, our results can also inform time series modelling. A practitioner might have a time series dataset with underlying factors that are informed by human experts. Tf the practitioner wants to use deep autoregressive models construct an embedding that contains a certain kind of information from the input, they may need ensure that this information is predictive sufficient statistic for the task.
6 Conclusion
We have developed a general framework for analyzing what the embeddings of an autoregressive language model should represent. Our analyses suggest that such embeddings should represent latent structures such that the next token In+1 is independent from previous tokens Fin when conditioned on that structure, a property possessed by predictive sufficient. statistics. We confirmed this hypothesis with probing experiments
14

---

Published in Transactions on Machine Learning Research (07 (07/2025)
cases where predictive sufficient statistics can be identified. We hope that our findings contribute to bridging the gap between Bayesian probabilistic models and deep neural networks.
Acknowledgments
This work was supported by ONR grant number N00014-23-1-2510. TRS is supported by an NDSEG Fellowship. RTM was supported by the NSF SPRF program under Grant No. 2204152.
References
Yossi Adi, Einat Kermany, Yonatan Belinkov, Ofer Lavi, and Yoav Goldberg, Fine-grained analysis of sentence embeddings using auxiliary prediction tasks. In International Conference Learning Representations, 2017.
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 Represenlations, 2023.
Jacob Andreas. Language models as agent models. In Findings of the Associalion for Compulational Linguistics: EMNLP 2022, pp. 5769 5779, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics.
Leonard E Baum and Ted Petric. Statistical inference for probabilistic functions of finite state markov chains.
The annals of mathematical stalistics, (7(6):1554-1555) 1966.
Yonatan Belinkov. Probing classifiers: Promises, shortcomings, and advances. Computational Linguistics, 48 (1):207-219, 2022.
Yonatan Belinkov, Nadir Durrani, Fahim Dalvi, Hassan Sajjad, and James Glass. What do neural machine translation models learn about morphology? In Proceedings of the 55th Annual Meeting of the Association for Computational Linguástics (Volume 1: Long Papers), pp. 861-872, Vancouver, Canada, July 2017. Association for Computational Linguistics.
Yoshua Bengio, Réjean Ducharme, and Pascal Vincent. A neural probabilistic language model. Advances in Neural Information Processing Systems, 13, 2000.
Jose M Bernardo and Adrian Smith. Bayesian Theory, Wiley series in probability and statistics. John Wiley & Sons Ltd., Chichester, 2000.
David Blei, Andrew Ng, and Michael Jordan. Latent dirichlet allocation. Advances in Newral Information Processing Systems, 14, 2001.
Nick Cammarata, Shan Carter, Gabriel Goh, Chris Olah, Michael Petrov, Ludwig Schubert, Chelsea Voss, Ben Lgga,, and Swee Kiat Lim. Thread: Circuits. Distill, 2020. doi: 10.23915/disttil.00022 000024.
Alba Carballo-Castry Sonia Laguna, Moritz Vandenhirtz, and Julia Vogt. Exploiting interpretable capabilities with concept-onhanced diffusion and prototype networks. 10 2024. doi: 4448550///Xiv.2/// 2410.18055.
Hoagy Cunningham, Aidan Ewart, Logan Riggs, Robert Huben, and Lee Sharkey. Sparse autocncoders find highly interpretable features in language models. ArXiv, abs/2309.08600, 2023.
Jacob Devlin, Ming- Wei Chang, Kenton Loe, and Kristina Toutanova. BERT: Pro-training of deep bidiroctional transformers for language understanding. Minneapolis, Minnesota, June 2019. Association for Computational
Tinguistics.
Allyson Ettinger, Ahmed Elgohary, and Philip Resnik. Probing for semantic evidence of composition by means of simple classification tasks. In Proceedings of the 1st Workshop on Evaluating Vector-Space Representations for NLP, pp. 134-139, Berlin, Germany, August 2016. Association for Computational Linguistics.
15

---

Published in Transactions on Machine Learning Research (07 (07/2025)
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? case study of simple function classes. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022.
Andrew Gelman, John B. Carlin, Hal S. Stern, and Donald B. Rubin. Bayesian Data Analysis. Chapman and Hall/CRC, 2nd ed. edition, 2004.
Thomas Griffiths, Mark Steyvers, David Blei, and Joshua Tenenbaum. Integrating topics and syntax. In L. Saul, Y. Weiss, and L. Botton (eds.), Advances in Neural Information Processing Systems, volume 17. MIT Press, 2004.
Nate Gruver, Marc Finzi, Shikai Qiu, and Andrew Gordon Wilson. Large language models are zero-shot time series forecasters, 2024.
Abhiject Gupta, Gemma Boloda, Marco Baroni, and Scbastian Padó. Distributional vectors encode referential attribules. In Proceedings of the 2015 Conference on Empirical Melhods in Natural Language Processing, pp. 12 21, Lisbon, Portugal, September 2015. Association for Computational Linguistics.
John Hewitt and Christopher D. Manning. A structural probe for finding syntax in word representations. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Compulational Linguisties: Human Language Technologies, Volume 1 (Long and Short Papers), pp. 4129-44 4138,, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.
Dieuwke Hupkes, Sara Veldhoen, and Willem Zuidema. Visualisation and' diiagnostic classifiers' reveal how recurrent and recursive neural networks process hierarchical structure. Journal of Artificial Intelligence Research, 61:907-922 926, 2018.
Daniel Jurafsky and James Martin. Speech and Language Processing: An Introduction to Natural Language Processing, Computational Linguistics, and Speech Recognition. Prentice Hall, 2008
Rudolph Emil Kalman. new approach to linear filtering and prediction problems. ASME Journal of Basic Engineering, 82(1):35 45, March 1960.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2014.
Arne Köhn. What's in an embedding? analyzing word embeddings through multilingual evaluation. Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing, pp. 2067-2073, Tisbon, Portugal, September 2015. Association for Computational Linguistics.
Iryna Korshunova, Jonas Degrave, Ferenc Huszár, Yarin Gal, Arthur Gretton, and Joni Dambre. Bruno: A
deep recurrent model for exchangeable data. Advances in Neural Information Processing Systems, 31, 2018.
Yuchen Li, Yuan-Fang Li, and Andrej Risteski. How do transformers learn topic structure: Towards mechanistic understanding. ArXiv, hts/23333.0245,, 2023.
Yongjie Lin, Yi Chern Tan, and Robert Frank. Open sesame: Getting inside BERT's linguistic knowledge. In Proceedings of the 2019 ACL Workshop BlackborNLP: Analyzing and Interpreting Neural Networks for .LP, pp. 241-253, Florence, Italy, August 2019. Association for Computational Linguistics.
Nelson F. Liu, Matt Gardner, Yonatan Belinkov, Matthew E. Peters, and Noah A. Smith. Linguistic knowledge and transfcrability of contextual representations. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pp. 1073-1094, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.
Thomas McCoy and Thomas Griffiths. Modeling rapid language learning by distilling bayesian priors into artificial neural networks. Nature Communications, 16(1):4676, 2025. doi: 10. 10.1038/5411477-055--000
16

---

Published in Transactions on Machine Learning Research (07/2025)
R. Thomas McCoy, Shunyu Yao, Dan Friedman, Matthew Hardy, and Thomas L Griffiths. Embers of autoregression: Understanding large language models through the problem they are trained to solve. arviv preprint XrXr::2009 13638, 2023.
Yu Meng, Yunyi Zhang, Jiaxin Huang, Yu Zhang, and Jiawei Han. Topic discovery via latent space clustering of pretrained language model representations. Proceedings of the ACM Web Conference 2022, 2022.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016.
Vladimir Mikulik, Grégoire Delétang, Tom McGrath, Tim Genewein, Miljan Martic, Shane Logg, and Pedro A. Ortega. Mela-trained agents implement bayes-optimal agents. ArXiv, abs/2010 201101223,, 2020.
Samuel Müllrr, Noah Hollmann, Scbastian Pineda Arango, Josif Grabocka, and Frank Hutter. Transformers can do bayesian inference. In International Conference on Learning Represenlations, 2022.
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. In The Fleventh International Conference on Learning Represenlalions, 2023.
Madhur Panwar, Kabir Ahuja, and Navin Goyal. In-context. learning through the bayesian prism. In The Twelfth International Conference Learning Representations, 2024.
L.R. Rabiner. A tutorial on hidden markov models and selected applications in speech recognition. Proceedings of the IEEE, (7(2)2257 286, 1989. doi: 10.1109/5.18622
Alec Radford, Rafal Jozefowicz, and Nya Sutskever. Learning to generate reviews and discovering sentiment. arviv preprint arxi nen:170440111 0144,, 2017.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenA/ I blog, 1(8):9, 2019.
Anna Rogers, Olga Kovaleva, and Anna Rumshisky. A primer in BERTology: What we know about how BERT works. Transactions of the Associalion for Compulational Linguistics, 8:842 866, 2020.
Adam S. Shai, Sarah E. Marzen, Lucas Teixeira, Alexander Gietelink Oldenziel, and Paul M. Riechers. Transformers represent belief state geometry in their residual stream, 2024.
Xing Shi, Inkit Padhi, and Kevin Knight. Does string-based neural MT learn source synlax? In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pp. 1526 15661134, Austin, Texas, November 2016. Association for Computational Linguistics.
Joshua Tenenbaum. Bayesian modeling of human concept learning. Advances in Neural Information Processing Systems, 11, 1998.
lan Tenney, Patrick Xia, Berlin Chen, Alex Wang, Adam Poliak, R Thomas McCoy, Najoung Kim, Ben- jamin Van Durme, Sam Bowman, Dipanjan Das, and Ellie Pavlick. What do you learn from context? probing for sentence structure in contextualized word representations. In International Conference on Learning Representations, 2019.
Naftali Tishby and Noga Zaslavsky, Deep learning and the information bottleneck principle. 2015 IEEE Information Theory Workshop (TTW), pp. 1-5, 2015.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein,
17

---

Published in Transactions on Machine Learning Research (07 (07/2025)
Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina. Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023.
Tulia Ture, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Well-read students learn better: On the importance of pre-training compact models. arXiv: Computation and Lenguage, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Usækoreit, Llion Jones, Aidan N Gomez, f. ukasz Kaiser, and Mia Polosukhin. Attention is all you need. Advances in Newral Information Processing Systems, 30, 2017.
Hengyi Wang, Shiwei Tan, Zhiqing Hong, Desheng Zhang, and Hao Wang. Variational language concepts for interpreting foundation language models. In Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen (eds.), Findings of the Association for Computational Linguistics: EMNLP 2024, pp. 8645 8671, Miami, Florida, USA, November 2024a. Association for Computational Linguistics. doi: 11/6653/v12222/ /t2.///../...ddnnneeeeemmm findings-emnlp. 505.
Xinyi Wang, Wanrong Zhu, Michael Saxon, Mark Steyvers, and William Yang Wang. Large language models are implicitly latent variable models: Explaining and finding good demonstrations for in-context learning, 2024b.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning implicit Bayesian inference. orXiv preprint arXfv:2111.00880 2021.
Takateru Yamakoshi, Thomas Griffiths, and Robert Hawkins. Probing BERT's priors with serial reproduction chains. In Findings of the Association for Computational Linguistios: ACL 2022, pp. 3977-3992, Dublin, Ireland, May 2022. Association for Computational Linguistics.
Naimeng Yo, Hanming Yang, Andrew Siah, and Hongseok Namkoong. Pre-training and in-context learning is bayesian inference la de finetti, 2024.
Tom Young and Yang You. On the inconsistencies of conditionals learned by masked language models. ai preprint 200000200..0008,, 2022.
Carolina Zheng, Keyon Vala, and David M. Blei. Revisiting topio-guided language models, 2023.
18

---

Published in Transactions on Machine Learning Research (07 (07/202) (025)
A Appendix
A.1 Masked language models
Masked language models (MLMs) are another class of LLMs that have been successful at modcling text, and we explore whether the objective of MLMs also equivalent to implicit. Bayesian inference under the exchangeability assumption. We find that it is, but differs from autoregressive language models a way that results in less expressivity.
MLMs, notably including the BERT series (Devlin et al. 2019), randomly mask certain tokens in the input sequence, and the model's goal is to predict the original words at the masked indices only by considering their surrounding context. Unlike autoregrossive models, the MLM does not model coherent joint distribution of the data Yamakoshi et al. 2022: Young & You 2022 However, the log objective can be extended as follows,
(MLM(FEN)
Σ (ogp(anlaasec))
nEM
Σ log / p(ann0)p(()assee)d) nEM
(4)
where M denotes the set of masked indices, and U denotes the set of unmasked indices. The proof given below, and is a simple exlension of the derivation for the autoregressive version. The difference between the MLM objective and the autoregressive objective that in the summation, the prediction of each token En the same posterior over the latent variable p(@) In other words, cach token In is predicted independently from the latent variable 0. As a result, MLM forms less expressive Bayesian inference objective than autoregressive models. Therefore, we aim to empirically evaluate both the ability of MLM to recover latent variables, and whether its performance differs from that of autoregressive models.
We now prove the equivalence between the MLM objective and Bayesian inference Equation statement is that, given an exchangeable process T1, ..,,... Fn.
The
(MLM(FEN)= = Σ
(5)
nEM
/
Σ log
(6)
nEM
where M is the set of masked indices and U is the set of unmasked indices, and by construction M n =
0.
Proof. We first prove the autoregressive version of the equivalence, which Korshunova et al. (2018 proposes
for each individual term in the summation but briefly mentions why it is equivalent,
logp(ZinN) = log p(x,) ) +
log 1
=logp(x1) +
Σ log p(mm+=100)p())m...)
(7)
To do so, we shall prove the equivalence in each term in the summation, that is, the statement that
p(Zn+1)FIm) S, 100(
(8)
for each > 1. First, de Finetti's theorem states that, under the same exchangeability condition,
p(#im+1) [ p(0) n+1
Il 1lo,10)d0.
99
19

---

Published in Transactions on Machine Learning Research (07 (025))25)
and also that each 2; is conditionally independent given 0 for all t. To show the equivalence, we first divide each side of Equation by p((1m)) assuming that p(uim)) > 0. The left hand side becomes,
p(1in+1)) p(2n42l21mn)
(10)
p(z1m)
The right hand side becomes,
I p())
I p(ar10)
p(0) 10)p++10)((1111
d0
(11)
p(z1n))
p(11:n)
p p 10) p(0))((z1:1 d0
(12)
p((1m))
p(5n+110)p(( (11:
(13)
Line 11 uses conditional independence to combine the product on through n, and line 13 uses Bayes
rule. Therefore, because of Equation g we prove the stalement of Equation
The proof for MLM in Equation can be shown by simple extension. We shall use zo short hand for Ti Using de Finetti's thcorem,
p(o(mju))= = p(0) II p(ax00)d0, VE(nJUU
(14)
Dividing each side of the above equation by p(xc), , the loft hand side becomes,
p(z())ur) p(znlov)
(15)
p(xv)
The right hand side becomes,
/ p(0)
[nj p(a,10) p(xv) d8
J p(0) p(z, (1)p0)((z)l0 p(ac)
(16)
p( (0) (()))1)))(0)((( p(0)p(xz)00 (0)
de
(17)
plzc)
p( p(mm10p((000))
(18)
Same as in the proof for the autoregressive version, line 16 uses conditional independence to combine the product U, and line 18 uses Bayes rule. Therefore, have that p( p((2_())p((( Thus, we have shown the equivalence for each term in the summation of the original statement Equation
A.2 Definition of additional exchangeable conjugate models
We consider the generative process for sequence Fi, where žij are i.i.d. across j.
A.2.1 Topic model
The generative model for LDA is For each topic & in (I,.., K),
For each document i,
1. Draw topic 3g Dirichlety
1. Draw topic mixture 0 Dirichlet(a).
2. For each word j document i,
(a) Draw topic assignment tij Categorical(()),
(b) Draw word mz ~ (3
where V is the vocabulary size, and and n pre-initialized hyperparameters.
20

---

Published in Transactions on Machine Learning Research (07 2025)
A.2.2
model
0, BETA(a, 3)
Fij BERNOULLI(O)),
where 3 are fixed hyperparameters.
A.2.3
Gamma-Expones model
0: GAMMA(a, B)
Tij EXPONENTIAL(0))
where 3 are fixed hyperparameters.
A.2.4 HMM-LDA
There are C' classes (cI...,,cc) CC) which follow an HMM transition mairix. One class is the semantic class Cı, and a word coming from 61 generated by an LDA model. Words of other classes are 'syntactic' and are directly generated by a distribution over words with a Dirichlet prior, one distribution for each .c., cc). The probability of generating from 61, the LDA class, is determined by a parameter p which specifies the transition probability into that class from all other classes. The remaining transition probabilities generated from Dirichlet(y) prior; lower corresponds to more sparse transition matrices.
A.3 Implementational details
All computations for syuthetic dalasets are run on single Tesla T4 GPUs, and those for natural corpora are run on single A100 GPUs.
Experimental process Each dataset is split into three sets: set 1, set 2, and set 3. Set 1 is used for training the transformer. Set 2 is used for validating the transformer and getting embeddings from transformer that are used to train the probe. Set 3 is used for validating the probe.
Except discrete hypothesis space datasets and natural corpora, the sizes for the three sets are: 10000, 3000, 1000, and each sequence 500-tokens long. In the discrete hypothesis space datasets, we experimented with different sequence lengths (detailed in our results), and the sizes for the three sets are: 20000, 19000, 1000.
In HMM--DA, sequence lengths 400, and the sizes for the three sets are 10000, 1000, 1000.
On 20NG, probe training and validation are run on 11,314 and 7,532 documents, respectively. On WikiText- 103, probe training and validation are run on 28,475 and 60 documents, respectively. Both splits are derived directly from train-validation split provided by the dataset sources. Note that set not used for natural corpora because we use pretrained LLMs.
Transformer Except for topic models, we use three-layer transformer decoder with hidden-size = 128 and number of attention heads = 8. If the input is categorical (similar to tokens in natural corpus), we employ the standard word embedder layer before the decoder layers. the input is contimuuus, we use a Linear layer to map inputs to dimension 128 in place of the word embedder layer.
Dropout 0.1 applied, and learning rate 0.001, batch-size 64.
Transformers on topic models Autoregressive transformer (AT) and BERT hyperparameters for training are given in Table g Configurations of BERT are identical those of BERT-TINY from Ture et al. (2019). AT hidden sizes and final layer embedding sizes are 128, same as BERT-TINY, and it uses four layers, resulting in 655,336 parameters. Bert has 608,747 747 parameters.
21

---

Published in Transactions on Machine Learning Research (07 (07/2025)
Probe The probe is linear layer with softmax activations. Learning rate is tuned in [0.001, 0.01], and batch-size = 64.
Hyperparameters for exchangeable conjugate models The Gaussian-Gamna hyperparameters do, 90, Ml, λo are 15,1, The OOD hyperparameters are {2, 1,
On the Beta-Bernoulli model, we use a = 2, 3 8. In the OOD case, a = 8, = 2. On the Gamma--Expnnntiin model, we use a = 2, / Tn the OOD case, a 2.3
A.4 Linear mixed-effects model
We want to validate our hypothesis that the LLM's latent topic representation helps it predict individual tokens. While Fig. 5c is suggestive of this relationship, here we use statistical testing to confirm it.
Concretely, we use a linear mixed-effects model to predict the per-token perplexity. We analyze 701,243 individual tokens from 20NG test corpus using GpT-2. Perplexity naturally decrcases as the LLM processes the document, so we include a fixed effect of token position and random effect for the document itself; finally, include the topic decoding accuracy (a binary / outcome based on the topic probe) as the variable of interest. We extract 100 tokens per document, stratified so they are evenly spaced, and represent. the token position as the percent into the document,
perplexity ~ fokken_posithe + accuracy + (1 documentt__ id).
(19)
We find significant effects for both token position and topic accuracy,
Effect
Group
Term
Estimate
Std. Error
Statistic
DF
p-valuc
fixed
(Intercept)
4.65
0.01
413.28
21078.63 63
<2e- 16
fixed
aopii__accraayy
-0.15
0.01
-16.51
355354.64
<2e-16
fixed
token position
-0.78
0.01
-58.47
696289.CK .69
<2e-16
ran pars
document id Residual
sd (Intercept)
0.63
ran pars
sd Observation 3.22
Finally, we obtain a Variance Inflation Factor of 1.014742 between accuracy and token position, suggesting an acceptable degree of colinearity between the two variables.
A.5 Additional results
Hyperparameter sweep Table 5c shows hyperparameter sweep in the Bayesian conjugate models setting, across transformer embedding size in {8,32, 128}, number of layers in (2,3,4), and number of attention heads in {4,8). In general, we observe that embedding size affects performance most significantly. In the Gaussian and Bernoulli datasets, performance improves with higher embedding size. In the Gamma- Exponnntia dataset, performance is best with embedding size = 32.
Hyperparametors for probes on AT and BERT on LDA are given in Table 10
Hyperparamelers for probes on the LLMs in natural corpora are given in Table 11 and Table 12
Moments Figure Sa shows that the probe decodes the moments of the posterior distribution of the Gaussian-Gamma model. Because the higher distribution moments more volatile in value and less directly related to estimating the parameters f, we examine whether existing discrepancies are caused by an overly simple probe. We perform a second set of experiments where the probe has a hidden layer with ReLU activations (Figure 8b), showing stronger alignment.
Figure shows results on posterior distribution moments on Beta-Bernoulli and Gaama--xxonenin models. On the Gamma-Exee Exponenta model, divide the target second, third, and fourth moments by factors of 10, 100, 1000, respectively that each moment is given roughly equal importance.
22

---

Published in Transactions on Machine Learning Research (07 (07/2025)
Table 5: Scaled MSE across hyperparameter settings where the probe targets sufficient. statistics, along with standard error across three random seeds. In the Caussian-Gamma case, we report only the second sufficient statistic, i.e. the standard deviation of the seen sequence (which is more challenging than the mean), to avoid clutter.
(a) Gaussian- Gamma.
mum heads 4
mmm heads =
d= 8
d d= 32
d= 128
d=8
d= 32
d d= 128
layers
0.168
0.022
0.076
0.069 0.013
0.303 0.039
0.045 + 0.010
0.050 0.005
layers
0.566
0.017
0.066
0.048 0.005
0.321
0.014
+ 0.005
0.046 0.010
layers
0.451
0.059
0.056
0.047 0.009
0.490
1: 0.006
0.033 0.004
(b) Beta-Bernoulli.
heads
heads
128
d
d 128
layers
0.010 = 0.002
0.002 + 0.000
0.001 = 0.000
0.007 + 0.002
0.002 + 0.001
0.001 0,000
layers
0.008 = 0.002
0.001 + 0.000
0.000 = 0.000
0.010 + 0.001
0.001 + 0.000
0.001 + 0.000
layers
0.007 + 0.001
0.001 + 0.000
0.000 = 0.000
0.010 + 0.002
0.001 + 0.000
0.000 = 0.000
(c) Gamma- Expammaa-Epoonniiall.
mum heads
mum heads
d= 8
d= 32
d= 128
d= 8
d= 32
d = 128
layens
0.004 0.001
0.001 1: 0.000
0.009 0.006
0.002 0.001
0.001 1: 0.001
0.067 0.006
layers
0.003 0.001
0.001 1: 0.000
0.006 0.004
0.002 0.000
0.001 1: 0.000
0.058 0.006
layers
0.002 0.000
0.000 1: 0.000
0.021 0.008
0.002 0.000
0.001 1: 0.000
0.033 0.005
23

---

Published in Transactions on Machine Learning Research (07 (075)
(a) Maussian-Cammnn (linear probe).
(b) Maussian-Cammin (two-layer probe).
Figure 8: Probe recovery of transformer-learned posterior distribution moments (blue) and ground truth moments (red) across 1000 test datapoints. The first row shows parameters probed in the general non-OOD case. The second row shows corresponding information in the OOD case.
Memorization Figure 10 conducts the same memorization experiments for the Gamma- Laamma-xxponentall case as in the Gaussian-Gamma case in the main text (Figure
Parsimonious embeddings Table 0 shows results from training an MLP on sufficient statistics to target transformer embeddings. To compute these results, get mean-squared error, which is then divided by the variance of the transformer embeddings along each dimension separately. The error is averaged across the embedding dimension and subtracted from to yield the mean proportion of variance accounted for in the embedding. These results establish that the relationship runs in both ione:tions: we can decode sufficient. statistics from embeddings, and we can predict embeddings from sufficient statistics.
Table 6: Validation variance transformer embeddings accounted for by training an MLP only on sufficient statistics of training sequences.
Transformer size
Gunseien-Crmma
Beta-Bemanlli
Camman Ereinn--tpponnttalll
71.1%
28.4%
72.2% 45.6%
84.6%
79.6%
Figure II shows Gammaaa probing results on sufficient statistics on the 10th token, complementing our exploration of token memorization transformer (Figure
Probing results on Wikitext-103 Table shows probing results on Wikitext-103.
Table 7: Wikitext-103 topic prediction performance based different LLMs. On accuracy, the TLAMA2-CHAT performance is not statistically significantly different from that of the MLMs; otherwise, the autoregressive models statistically significantly outperform the non-eutoregressive models.
K
Model
Parameters
Accuracy f
L2 loss |
Tol. var. loss
Accuracy f
1.2 loss 4
Tol. var. losss 1
GrT-2
124M
86.7% + 0.5%
0.025 1
0.098 1
73.9% 1: 2.3%
0.026 0.001
0.089 0.002
GPT-2-MRDUM
355M
88.2%
0.6%
0.097 0.001
74.2% + 1.3%
0.025 +
0.097 0.002
GPT-2-1ARGE
774M
88.5% 0.8%
0.094 0.001
74.2%
1.4%
0.025 +
0.088 0.001
LLAMA
87.3% -
0.023 =
0.091 = 0.001
70.4% = 1.1%
+
0.09 + 0.001
LLAMA 2-CHAT
7B
85.3% + 0.7%
0.024 +
0.094 0
69.9% +: 1%
0.026 1: 0
0.09
BERT
110M
84.9% + 1.1%
0.027 1
0.103 1. 0.001
72.4% 1 1.4%
0.029 1.
0.097 1 0.002
BERT-LARE LARGE
336M
85.4% + 1.7%
0.03 +
0.111 +
72.1% + 0.9%
0.031 +
0.104 + 0
Null GPT-2
124M
58.1%
0.121 0.003
0.247 = 0.006
32.9% + 3.2%
0.099 0.003
0.195 0.007
24

---

Published in Transactions on Machine Learning Research (07 (07/202) 2025)
RPP1
0.3
0.4
0.2
14
0.00
03
03
14-
62-
21
04
12-
033
0.1
206
406
6%
305
1009
266
406
04
906
1089
206
489
061
306
1000
200
489
041
906
3300
(a) Beta-Beraaa Bernoulli.
[UP))
200
10-
150
1200
treat
14
51A)
6PP)
170
2403
100
2003
5633
066
1008
20
800
100
1200
(b) Gamma- Exponential
22.
1103
e171
6A7
10-
3300
(c) Glamma-Exxoneee (non-linear probe with factors).
Figure 9: Probe recovery of centformer-eearnee posterior distribution moments (blue) plotted with ground truth moments (red). The first row shows parameters probed 1000 test datapoints in the general, i..,, non-OOD, case. The second row shows corresponding information in the OOD case.
25

---

Published in Transactions on Machine Learning Research (07/2025)
Figure 10: Probing over the first 10 tokens themselves using the 10th token embedding of the transformer. Aside from perfectly encoding the 10th token, this embedding does not show memorization over the other 9 tokens suggested by the noise in probe recovery. Visually, the noise is less than in the Gaussian-Gamme case (Figure 4). This is partly due to spuriously large magnitude datapoints in this distribution, and the noise here is still significantly higher than in capturing sufficient statistics (Figure
Figure 11: Probing over mean (left) and standard deviation (right) of the first 10 tokens using the 10th token embedding of the transformer in the Gaussian- Gammma dataset. The first row corresponds to the same generation process, and the second row corresponds to the OOD case.
26

---

Published in Transactions on Machine Learning Research (07 (07/2025)
Inner layer performance natural corpora Table shows topic distribution recovery from LLM inner layers.
Table 8: Across-layer performance in the format of (accuracy, L2-loss), where move from the word embedder to the final layer (prior to word emission) from left to right.
Model
Parameters
Layer
Layer
Layer
Layer 75%
Layer 90%
Layer
Layer -1
GPT-2
124M
87.4%, 0.027
87.1%, 0.025
86.1%, 0.025
86.7%, 0.026
86.3%, 0.026
85.9%, 0.026
86.7%, 0.025
--E-2MMEIUMM
355M
88.2%.0.026
0.026
86.9%, 0.024
85.5%, 0.027
85, 0.028
83.9%. 0.03
81.8%, 0.04
88.2%, 0.024
GPT-2-LARGE
771M
85.1%, 0.026
87.6%, 0.024
86.5%, 0.025
85.4%, 0.026
84.5%, 0.028
82.9%, 0.03
88.5%, 0.023
LLAMA 2
7B
59.2%, 0.192
74.4%, 0.104
87.3%, 0.025
86.6%, 0.023
86.9%, 0.024
86.7%, 0.024
87,3%, 0.023
LLAMA 2-CHAT
7B
38.8%, 0.193
74.9%, 0.098
87.2%, 0.025
87.5%, 0.023
86.5%, 0.024
85.8%, 0.025
85.3%, 0.024
BERT
110M
88.4%, 0.028
88.6%, 0.024
86.2%, 0.028
85.9%, 0.03
86.4%, 0.028
86.4%, 0.028
84.9%, 0.027
BERT-LARGE
336M
88.2%, 0.029
87.8%, 0.026
86.4%, 0.027
85%, 0.03
85.9%, 0.029
85.3%,, 0.03
86.4%, 0.03
Using averaged tokens as document embeddings performs better across models both natural corpora. Here, we also report on results from using only last. tokens as document emboddings (Table 13 and Table 14
Using only last token embedding on natural corpora
Table 9: Autoregressive transformer (AT) and BERT hyperparameters for training on the synthetic datasets.
Parameter
Tuning range
Chosen value
Batch-size
[8, 128]
16
Learning rate
[ 10 10-3)
10-4
27

---

Published in Transactions on Machine Learning Research (07 (075))
Table 10: Probe hyperparameters for training on top of synthetic datasct language models.
Parameter
Tuning range
Chosen value
Batch-size
[8. 64]
16
Learning rate
[10-1,0.00 0.03]
10-3
Weight-decay
[0, 3.4 10-4-1044
0
Embedding choice
{First, Last, Average)
Last for AT
Average for BERT
Table 11: Probe hyperparameters for training on top of GPT-2, GPT-2-MEDIUM, GPT-2-LARGE, BERT, and BERT-LARGE.
Parameter
Tuming range
Chosen value
Batch-size Learning rate
{128}
128
[10-5, 10-3] [0, 3.4]
3 10
Weight-decay Embedding choice
3.4 10-3 Average
(First, Last, Average}
Table 12: Probe hyperparameters for training on top of LLAMA 2 and LLAMA 2-CHAT.
Parameter
Tuning range
Chosen value
Batch-size Learning rate
{128}
128
[10-3, 10
10-4
Weight-decay Embedding choice
3.4]
0.34
(First, Last, Average]
Average
Table 13: 20NG topic prediction performance based on different LLMs using the last token as document embedding.
= 20
K 100
Model
Parameters
Accuracy
L2 loss
Tot. var. loss
Accuracy 1
L.2 loss
Tot. var. loss
GPT-2
124M
17.8% +: 1.8%
0.143 =:2 0.001
0.26 = 0.002
31% +: 3%
0.115 = 00.001
0.217 : 0.002
GPT-2-MEDIM
355M
47.6% + 1.2%
0.144 + 0.001
0.26 + 0.001
30.8% + 3%
0.116 0.001
0.217 + 0.002
GPI-2-LARGE
774M
47.8% :: 1.8%
0.144 10
0.261 4 0.002
31.1% + 2.9%
0.116 0
0.217 : 0.002
LLAMA 2
7B
45.4% +: 2%
0.152 0
0.267 2 0.001
27.95
+ 2.9%
0.122 + 0.001
0.223 = 0.003
LLAMA 2 2-CHAT
7B
46.1%
1.8%
0.15 0.001
0.265 0.001
28.35
2.5%
0.122 0.001
0.222 + 0.003
BERT
110M
51.7% +: 1.4%
0.126 + 0.002
0.236 + 0.002
34.9% + 2.8%
0.108
0.204 = 0.001
BERT-LANGE
336M
52.2% 4: 2.1%
0.125 + 0.001
0.235 + 0.002
34.7% + 2.7%
0.107 + 0.001
0.201 + 0.001
Table 14: WikiText-1009 topic prediction performance based on different LLMs using the last token document embedding.
K 20
K - 100
Model
Parameters
Accuracy 1
L22 loes J.
Tot. var. loes
Accuracy 1
L.2 loss J
Tot.
var.
loss
GPT-2
124M
69.2% + 1.4%
0.079 + 0.003
0.191 +: 0.005
50.9% +: 1.3%
0.055
0.001
0.146 + 0.003
GPT-2-MEDIUM
355M
70.7% + 0.6%
0.077 = 0.002
0.188 : 0.004
51.6% + 1.9%
0.056
0.001
0.146 = 0.003
GPT- SPAG2-
774M
72% + 0.6%
0.074 + 0.001
0.184 + 0.003
52.8% + 2%
0.053
0.001
0.143 + 0.002
LLAMA 2
7B
69.7% + 2.3%
0.075 : 0.002
0.157 = 0.003
51.9% + 1.4%
0.054
0.001
0.144 = 0.003
LLAMA 2-CHAT
7B
71% 4: 1.8%
0.073 + 0.001
0.181 + 0.002
54.2% + 0.8%
0.052
0.002
0.141 + 0.005
BERT
110M
73.7% + 1.8%
0.067 +
0.177 + 0.002
58.9% + 0.9%
0.05 + 0.002
0.14 0.001
BERI-LARGE
336M
63.9% +: 3%
0.1 = 0.001
0.219 + 0.001
50% + 1%
0.065 0.002
0.164 0.003
28


Reviews:
[{"review_id": "agNhsAFD59", "reviewer": "Reviewer_a9NK", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1. The paper provides a principled Bayesian framing of what information the embeddings of LMs should contain: sufficient statistics. Theoretically, this is also the only information that the embedding should contain.\n\n2. The authors analyze diverse data generation models, including exchangeable models, LDA, and HMMs, and also provide probing-based empirical evidence for each case.\n\n3. The authors also provide some ablation studies to show that observed effects are likely due to the LM’s representations rather than overly powerful probes.\n\n4. The hypothesis is verified across different model architectures used in the experiments: BERT, GPT2, LLAMA 2. \n\n5. The authors also attempt to extend the hypothesis to real-world text data.\n\nWeaknesses:\n\n1. The paper's theoretical and empirical results focus on simple data generative models, which are quite far from the real-world text data distribution. A more sophisticated data distribution model that takes syntax and semantics into consideration would greatly strengthen the paper.\n\n2. The memorization defined in the last paragraph of page 7 is a bit weird. It is only natural that the embedding of the last token does not memorize the previous tokens, as it is supposed to predict the next token. In my opinion, memorization should represent whether the LM is memorizing the training data or learning the distribution/sufficient statistics.\n\n3. The experiment with real-world text data (20NG) is a bit of a far stretch. The assumption is still that 20NG follows an LDA generative process, which clearly is not. It is only natural that a trained LM can catch the semantics of a piece of text, thus being predictive of its topic, compared to a randomly initialized untrained baseline. I don't see how the sufficient statistics hypothesis is supported through this experiment.\n\n4. It would be nice if there were some real-world implications of the conclusion: would training LM embeddings to better encode the sufficient statistics improve LM performance? The real-world LMs likely do not perfectly learn the autoregressive factorization of the underlying data distribution.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper provides a principled theoretical and empirical investigation into what embeddings in autoregressive language models (LMs) represent. The authors argue that these embeddings serve as predictive sufficient statistics—compressing all the necessary information from the past to predict the next token. They ground their claims in two canonical data-generation processes: exchangeable models (e.g., topic models like LDA) and latent state models (e.g., HMMs). Theoretical analyses are supported by probing studies on synthetic and real-world datasets. The paper shows that embeddings capture latent generating distributions, generalize out-of-distribution, and avoid trivial token memorization. The authors further validate their framework on LDA-generated data, mixed HMM-LDA synthetic data, and natural corpora (20 Newsgroups and WikiText-103) using pretrained LLMs (GPT-2, LLaMA-2, BERT).", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Please clarify the memorization experiment shown in Table 4 and the last paragraph of page 7. To me, it is not clear what the authors are trying to show.\n\n2. Please clarify/strengthen the last real-world experiment. Please see weaknesses for detailed comments.\n\n3. Many figures' x and y axes are not clearly labeled, and the ticks are too small.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "htQXon3vEa", "reviewer": "Reviewer_DJUX", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Pros: \n\n[1] The paper is well-written and easy to follow. The figures are well-designed. \n\n[2] I like the way to do experiments with both synthetic data and real data. Sometimes with synthetic data, it is easier for controlled experiments. \n\n[3] This paper focuses on an important and interesting problem, which will benefit the LLM community significantly. \n\nCons: \n\nWhile the paper is well-written, some additional figures and concrete examples might further improve the readability to the general audience in the community. For example, what topics are encoded by different layers of LLMs. \n\nIn Section 4.4.2,  “we observe that Llama-2 relies on its inner layers to encode topic mixtures.” Is there a potential to explain why there is such phenomenon? \n\nThere are multiple interesting findings in this paper. I would like to suggest adding a section or subsection, to summarize the Lessons Learned or Observations with bullet forms. This will benefit the readers and future studies in this area. \n \nCould you also discuss the limitations and potential future directions? \n\nWhile I aggree with the importance of this work, could you also discuss what potential implications or what applications could benefit from understanding the embeddings? For example, could we train a better LLM or LLM could be explainable. \n\nCould you also discuss the relations to studies about embedding understanding in image data (e.g., R1), as well as embedding interpretation in other language models (e.g., R2)? \n\nTypos: In section 4.3, the first sentence might contain some repeated words, e.g., “draw”. \n\nR1: Exploiting Interpretable Capabilities with Concept-Enhanced Diffusion and Prototype Networks, 2024\n\nR2: Variational Language Concepts for Interpreting Foundation Language Models, 2024", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "In this paper, the authors investigate what embeddings represent in autoregressive language models. The authors demonstrate that these embeddings can be viewed as predictive sufficient statistics, capturing key latent structures underlying text generation. By examining interpretable scenarios such as exchangeable models and latent state models, they empirically show that transformer embeddings encode meaningful latent distributions, generalize well to out-of-distribution data. Extending their analysis to more realistic settings, the authors specifically focus on topic modeling, revealing that large language models implicitly encode topic mixtures similar to those inferred by LDA. Experiments are conducted on both synthetic adn real datasets. Overall, this paper studies an interesting topic and is well-written, the experiments are also well-designed.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See above comments", "broader_impact_concerns": "I have no concerns about the impact", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "OPSlHBAUxN", "reviewer": "Reviewer_iADS", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**:\n- The experiments are extensive and convincing. The authors thought carefully about appropriate controls to ensure validity of the results (controlling for the probe's performance, as well as the two baselines: the end-to-end word embedder and Null GPT-2). The findings of this paper are relevant to TMLR audience members interested in the connection between LLMs and Bayesian inference, as well as those involved in the field of mechanistic interpretability.\n\n**Weaknesses**:\n- While the experiments are thorough, a large part of the paper’s 12 pages is devoted to explaining the experiments’ setup/details, and I feel that the paper would benefit from a more organized discussion of the findings. A suggestion is to move Section 4.2 to the appendix (the results on HMMs are interesting, but similar to the exchangeable results) and to expand Section 4.4.\n- The organization and polish of the experiments section could be improved. There are additional changes that could be made to improve readability, such as in the beginning of Section 4.1.2, implementation details on the OOD experiment are provided before the experiment itself is described. I also find some of the tables a little hard to read due to formatting (e.g., Table 3).\n- I have provided some open questions, changes, and suggestions for the authors to address.\n\n**Questions**:\n- For the topic model experiments, I am not sure about the justification for why you chose to probe only the last layer embedding. For example, Thompson and Mimno (2020) [1] find that using GPT-2’s final layer as the representation for a clustering-based topic model performs very poorly, with significant improvement from using either the $(L-1)$th or $(L-2)$th layer. The theory outlined in Section 3 would still hold if the sufficient statistics were present in earlier, but not later, layers.\n- Section 4.1.2: how important is ruling out memorization of the inputs in the last layer embedding to the claims of the paper? Would it be possible for the model to use memorized inputs to reconstruct the predictive distribution in the last linear layer?\n- Table 3: The performance of the true model (LDA) in this experiment seems a little low, do you have an idea of why that is?\n\n**Stylistic suggestions:**\n- While Figure 1 shows three generative processes, the discrete hypothesis model is a special case of the exchangeable data model, as you describe in Section 3.1.1. It would be more clear if this relationship was shown in the figure. Also, in the sentence “The relevant predictive sufficient statistics are the sufficient statistics for $\\theta$”, drawing the parallel between the sufficient statistic and the posterior would be a bit more clear if you listed the latent variables in question, ”$\\theta$, $z_{n+1}$, $h$”.\n- Section 4.1.1: I suggest putting the full generative process for the Gaussian-Gamma model in the main text, as a reader who doesn’t know this model wouldn’t be able to infer what it is from just the section text and equation. You might also want to say something like “We describe the Gaussian-Gamma model as an example of a Bayesian conjugate model, with full details on the remaining two in Appendix A.1.”\n- Section 4.1.2: The OOD abbreviation should be defined where it first appears in the text.\n- Section 4.1.2: What do you mean by “suggesting a similar robustness with a Bayes-optimal agent”?\n- Section 4.1.2: It would be more clear to use the mathematical formula for the captured variance in embeddings instead of writing it in words.\n- Section 4.2.1: typo: “hyperparamters”\n- Section 4.2.2: missing word: “should encode the predictive sufficient [statistics]”\n- Table 2: In the caption, do you mean standard deviation or standard error?\n- Section 4.3: typo: “that draws draws words”\n- Section 4.3.2: I understood the general idea of this section, but found the explanation of the experiment details confusing. How do the five distinct topic models differ (e.g., can you provide hyperparameters)? Can you use some mathematical notation to explain how the probe training vs. prediction works?\n- Section 4.4.2: “20NG can be naturally grouped into twenty topics” — you should specify that this is because each post is from a newsgroup about a specified subject.\n- Section 4.4.2: “LDA-learned topic mixtures with $K = 20, 100$” $\\to$ “with $K \\in ${$20, 100$}”\n\n[1] Topic Modeling with Contextualized Word Representation Clusters (Thompson and Mimno, 2020) (https://arxiv.org/abs/2010.12626)", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The main contribution of this paper is to draw the connection between the activations from autoregressive language models and the concept of predictive sufficient statistics, and to validate this through empirical studies on Transformer-based models. The authors point out that if a language model captures the predictive distribution $p(x_{n+1} | x_{1:n})$ perfectly, then it must have learned to represent the predictive sufficient statistics of the data distribution. To test to what extent learning to represent sufficient statistics occurs in these models in practice, they focus on data generated from two families of sequential probabilistic models containing instances with known and tractable sufficient statistics: exchangeable models and latent state models. Their methodology for recovering the sufficient statistics from the trained autoregressive LM is to train a linear probe on the last layer embedding.\n\nThe empirical studies show the following: the sufficient statistics can be recovered in simple synthetic settings for exchangeable and latent state models, and that as the main experimental context, topic modeling, transitions from easier (still synthetic) settings to natural text corpora, a significant amount of the latent topic mixture can still be recovered.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "**Minor**: In Section 4.1.2, did you find similar results on memorization for the other studied models (Beta-Bernoulli and Gamma-Exponential)? You could consider adding this to the appendix.\n\n**Recommended**: Can you add some discussion on how the BERT topic model in Section 4.3 fits with the Section 3 framework, and additional interpretation of the BERT results in this light? Section 3 focuses on next-word prediction, whereas BERT is a masked language model (i.e., it predicts the current word), and is not autoregressive.\n\n**Recommended**: The conclusion could benefit from more discussion on the intended takeaways — additionally, while the experiments focus heavily on topic models, the paper also hints that the findings are more broadly applicable (such as in the intriguing last paragraph of the introduction). It would be interesting to revisit this in the conclusion, with connections to the empirical results.", "broader_impact_concerns": "n/a", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)