# What Should Embeddings Embed? Autoregressive Models Represent Latent
Generating Distributions

Liyi Zhang
Department of Computer Science
Princeton University
Princeton, NJ 08540
<EMAIL_ADDRESS>
&Michael Y. Li
Department of Computer Science
Stanford University
Stanford, CA 94305
<EMAIL_ADDRESS>
&Thomas L. Griffiths
Department of Psychology and Computer Science
Princeton University
Princeton, NJ 08540
<EMAIL_ADDRESS>

###### 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 connect the autoregressive prediction
objective to the idea of constructing predictive sufficient statistics to
summarize the information contained in a sequence of observations, and use
this connection to identify three settings where the optimal content of
embeddings can be identified: independent identically distributed data, where
the embedding should capture the sufficient statistics of the data; latent
state models, where the embedding should encode the posterior distribution
over states given the data; and discrete hypothesis spaces, where the
embedding should reflect the posterior distribution over hypotheses given the
data. We then conduct empirical probing studies to show that transformers
encode these three kinds of latent generating distributions, and that they
perform well in out-of-distribution cases and without token memorization in
these settings.

##  1 Introduction

Autoregressive language models (LMs) are trained to predict the next token in
a sequence (e.g., Bengio et al., 2000). Many large language models (LLMs) use
the autoregressive objective for pretraining (e.g., Radford et al., 2019), and
their document-level embeddings have been shown to capture elements of latent
structure that appear in text, such as agent properties (Andreas, 2022) and
syntax (Hewitt and Manning, 2019). However, it remains unclear why embeddings
represent these elements, and, more formally, what information an embedding
learned by a capable autoregressive predictor should capture. An understanding
of what LLMs represent and why is important to the design and evaluation of
architectures and optimizers, and to supporting transparency and safety.

The representations formed by LLMs are challenging to analyze partly due to
polysemanticity, where a neuron may activate for several distinct concepts
Cunningham et al. (2023). Previous work has developed methods for probing LLM
internal representations for specific concepts (Li et al., 2023; Meng et al.,
2022; Zheng et 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. By linking the autoregressive objective to finding predictive
sufficient statistics, we show that optimal content of embeddings can be
identified in 1) independent identically distributed data, where the embedding
should capture the sufficient statistics of the data; 2) latent state models,
where the embedding should encode the posterior distribution over states given
the data; and 3) discrete hypothesis spaces, where the embedding should
reflect the posterior distribution over hypotheses given the data. We use
probing methods to confirm that the relevant information can be decoded from
LM embeddings. We also show that content that appears to be of a similar
difficulty but not expected to be captured via predictive sufficient
statistics is more challenging to recover from LM embeddings by probing.

Our analysis suggests that LLMs should represent latent structure that
captures the posterior distribution over the generative process underlying
text. Specifically, latent structure such that the next word $x_{n+1}$ is
independent from previous words $x_{1:n}$ when conditioned on that structure
should be recoverable from LLM embeddings. Additionally, by linking LM
representations to Bayesian inference, our approach suggests that it may be
effective to evaluate and interpret the behavior of LMs through comparison to
more interpretable Bayes-optimal agents.

##  2 Related work

The embeddings produced by language models have been investigated in detail
(Gupta et al., 2015; Köhn, 2015; Ettinger et 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 different aspects of the latent
structure of text, including part of speech (Shi et al., 2016; Belinkov et
al., 2017), syntactic number (Giulianelli et al., 2018; Conneau et al., 2018),
sentence structure (Tenney et al., 2019; Hewitt and Manning, 2019; Liu et al.,
2019; Lin et al., 2019), entity attributes (Gupta et al., 2015; Grand et al.,
2022), sentiment (Radford et al., 2017), semantic roles (Ettinger et al.,
2016; Tenney et al., 2019), world states (Li et al., 2021), and agent
properties (Andreas, 2022). However, our work is motivated from a
fundamentally different perspective. Instead of focusing on what is captured
in the embeddings of these models, in this paper we explore why these
particular kinds of structure might be represented as a consequence of
statistical properties of the training data.

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.
(2024) 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. Zhang et al. (2023) and Zheng et al. (2023) also
connect LLM embeddings to Bayesian inference, but they focus on topic models
embedded in LLMs, while we extend the connection to more general cases.

Metalearned RNNs have also been shown to encode information equivalent to a
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., 2022). 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.

Figure 1: Three data generation processes where prediction of the next token
$x_{n+1}$ is independent from previous tokens $x_{1:n}$ given a 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 $\theta$ (or
$p(\theta|x_{1:n})$, $p(z_{n+1}|x_{1:n})$, and $p(h|x_{1:n})$ respectively).
We show the embeddings learned by autoregressive transformers represent this
information.

##  3 Three cases where optimal embeddings can be identified

Assume we have a sequence $x_{1:n}$ and an autoregressive language model (LM)
that predicts the next item in the sequence, $p(x_{n+1}|x_{1:n})$. We denote
the LM embedding for sequence $x_{1:n}$ as $\phi_{n}=f(x_{1:n})$. The
distribution $p(x_{n+1}|x_{1:n})$ is some function $g(\phi_{n})$ of this
embedding, with that function implemented by the final layers of the neural
network instantiating the LM. That is, the probability of the next element in
the sequence $x_{n+1}$ only depends on $\phi_{n}$. This establishes our basic
question: what should $\phi_{n}$ represent in order to accurately predict
$x_{n+1}$?

In this section we describe three cases where this question can be answered
analytically (Figure 1). First, when $x_{1:n}$ are independently sampled
conditioned on an unknown parameter, $\phi_{n}$ needs only represent the
sufficient statistic of this sequence. Second, when $x_{1:n}$ are generated by
a state space model (in the discrete case, a hidden Markov model), $\phi_{n}$
need only represent the posterior distribution over states given $x_{1:n}$.
Finally, when $x_{1:n}$ are sampled independently from one of a discrete set
of latent distributions, $\phi_{n}$ need only represent the posterior
distribution over hypotheses about that distribution given $x_{1:n}$. In each
case we explain how $p(x_{n+1}|x_{1:n})$ factorizes to make it possible for
$x_{1:n}$ to be summarized by some $\phi_{n}$ and identify the form of the
corresponding $g(\phi_{n})$.

###  3.1 General approach

The key idea behind our approach is that we can identify situations where the
embedding $\phi_{n}$ contains all of the information from $x_{1:n}$ required
to predict $x_{n+1}$. This idea is directly related to the notion of a
sufficient statistic (Gelman et al., 2004). Given a distribution $p(x)$ with
parameters $\theta$, a statistic $s(x)$ is sufficient for $\theta$ if the
conditional distribution of $x$ given $s$ does not depend on $\theta$. In
other words, if we only know $s$, we can estimate $\theta$ just as well as if
we know the actual value of $x$. For example, if $x$ is a set of independently
and identically distributed draws from a Gaussian distribution and $\theta$ is
the mean of that distribution, then the mean of the sample $x$ is a sufficient
statistic for $\theta$. The generative process can be rewritten to generate
the sample mean $s$ given $\theta$ and then generate the variation around that
sample mean independently from $\theta$, and $s$ contains all the information
relevant to estimating $\theta$.

In the autoregressive setting, we care about predictive sufficiency (Bernardo
and Smith, 2000). A statistic $s(x_{1:n})$ is predictive sufficient for the
sequence $x_{1:n}$ if

$\displaystyle p(x_{n+1}|x_{1:n})=p(x_{n+1}|s(x_{1:n})).$ (1)

If a model performs autoregressive modeling perfectly, its embedding should
represent a predictive sufficient statistic. Our three cases thus correspond
to settings where predictive sufficient statistics can be easily identified
and could plausibly be represented by a neural network.

###  3.2 Case 1: Exchangeable models

Predictive sufficiency is particularly straightforward to establish in
exchangeable 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(x_{1:N})=p(x_{\pi(1:N)})$ for some permutation $\pi$. Any
exchangeable model can be re-expressed in terms of the $x_{i}$ being sampled
independently and identically distributed according to a latent distribution
$p(x|\theta)$ parameterized by $\theta$, with
$p(x_{1:N})=\int_{\theta}\prod_{i}p(x_{i}|\theta)p(\theta)\,d\theta$ (Gelman
et al., 2004). This idea leads to the following proposition:

Proposition. Given an exchangeable sequence $x_{1:N}$ where each $x_{i}$ is of
dimension $d_{x}$, and given functions
$f:\mathbb{R}^{nd_{x}}\mapsto\mathbb{R}^{d_{m}}$,
$g:\mathbb{R}^{d_{m}}\mapsto\mathbb{R}^{d_{x}}$ such that, for each $1\leq
n\leq N$, $g\circ f(x_{1:n})=p(x_{n+1}|x_{1:n})\;\forall x_{n+1}$,
$f(x_{1:n})$ is a sufficient statistic for $x_{1:n}$.

In other words, if we have a perfect autoregressive predictor that is
composable into $g\circ f$, the output of $f$ is a 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 and Smith,
2000). Because $p(x_{n+1}|x_{1:n})=g(f(x_{1:n}))\;\forall n$, $f(x_{1:n})$ is
a predictive sufficient statistic for the sequence, and it is also a
sufficient statistic. ∎

The resulting sufficient statistic also fully specifies the posterior on the
parameters of the generating distribution, $p(\theta|x_{1:n})$. Sufficient
statistics are easily identified for a wide range of distributions, including
all exponential family distributions (Bernardo and Smith, 2000), and are easy
to represent. Since the LM’s predictive distribution decomposes into the form
above, this theoretical result gives us strong predictions about the contents
of embeddings for models trained on exchangeable data.

###  3.3 Case 2: Latent state models

In a latent state model, each $x_{i}$ is generated based on a latent variable
$z_{i}$. These $z_{i}$ are interdependent, with $z_{i}$ being generated from a
distribution conditioned on $x_{i-1}$. Common latent state models include
Kalman filters (where $x_{i}$ and $z_{i}$ are continuous and $p(x_{i}|z_{i})$
and $p(z_{i}|z_{i-1})$ are linear-Gaussian) (Kalman, 1960) and hidden Markov
models (where the $z_{i}$ are discrete) (Baum and Petrie, 1966). In a latent
state model, the posterior predictive distribution is

$\displaystyle p(x_{n+1}|x_{1:n})=\int
p(x_{n+1}|z_{n+1})p(z_{n+1}|x_{1:n})\,dz_{n+1}.$ (2)

In this case, the posterior distribution $p(z_{n+1}|x_{1:n})$ captures all of
the information in $x_{1:n}$ relevant to predicting $x_{n+1}$, rendering
$x_{n+1}$ independent of $x_{1:n}$ when conditioned on the posterior, and acts
as a predictive sufficient statistic in this model.

More formally, we can consider an embedding a representation $\phi_{n}$ such
that a fixed operator $g(\phi_{n})$ can produce $p(x_{n+1}|x_{1:n})$. The
posterior distribution $p(z_{n+1}|x_{1:n})$ satisfies this characterization,
with Equation 2 showing that the relevant operator is the integral of
$p(x_{n+1}|z_{n+1})$ over $z_{n+1}$. That operator can be easily approximated
linearly and hence by a single layer of a neural network. The embedding
$\phi_{n}$ thus need only represent $p(z_{n+1}|x_{1:n})$.

###  3.4 Case 3: Discrete hypothesis spaces

In a more specific version of Case 1, assume each $x_{i}$ is generated
independently from some unknown generative model. Let $\mathcal{H}$ denote the
set of hypotheses $h$ about the identity of this model. In this case,
$x_{1:n}$ are exchangeable, but any sufficient statistics might be difficult
to identify. The posterior predictive distribution can be written as

$\displaystyle
p(x_{n+1}|x_{1:n})=\sum_{h\in\mathcal{H}}p(x_{n+1}|h)p(h|x_{1:n}).$

By an argument similar to Case 2, $p(h|x_{1:n})$ is a predictive sufficient
statistic in this model and an embedding thus need only capture this posterior
distribution.

###  3.5 Probing embeddings to recover predictive sufficient statistics

The three cases identified above specify information that should be encoded in
the embeddings formed in neural networks. This sets up the second element of
our approach, which is building probes to check whether this information is
actually encoded in the representations of specific networks. We focus on
transformers (Vaswani et al., 2017), since they are widely used as language
models. We denote the target to be decoded for each sequence $x_{1:n}$ by a
vector $t$. For example, in the discrete hypothesis space case the target $t$
is $p(h|x_{1:n})$, which is a vector on the simplex with vertices
corresponding to different values of $h$. Exact choices of $t_{n}$ in each
dataset are given in the next section. Given a trained transformer, we decode
$t$ by by training a second model (a probe) to predict the target $t$ from the
embedding $\phi_{n}$ of the corresponding document $x_{1:n}$. The probe $g$
maps from the document embedding to the target. To ensure that the relevant
statistical information is contained in the LM, not in the probe, we keep the
probe simple by defining it as a linear layer with softmax activations:
$g(\phi_{n})=\text{Softmax}(\text{Linear}(\phi_{n}))$.

##  4 Empirical analysis of embeddings

We have identified three cases where a predictive sufficient statistic is
expected to be contained in an autoregressive model’s embeddings. In this
section, we conduct probing on transformers trained on different datasets to
empirically verify this hypothesis. We also conduct ablation studies to
further test this idea. In all experiments, we use a linear probe unless
otherwise mentioned, three-layer transformer decoder, and the Adam optimizer
(Kingma and Ba, 2014) (more implementational details in Appendix A.3).

###  4.1 Exchangeable models

####  4.1.1 Generative model

We evaluate our approach on three exchangeable models: the Gaussian-Gamma,
Beta-Bernoulli, and Gamma-Exponential models. Due to space limitations we only
describe the Gaussian-Gamma model in detail, but the full generative process
underlying the other two models is given in Appendix A.1.

The Gaussian-Gamma model generates from a Gaussian distribution with two
unknown parameters. For $M$ sequences $\\{x^{(k)}\\}_{1\leq k\leq M}$ with $N$
tokens each, each sequence is i.i.d. generated by a mean and precision
parameter sampled from a prior. Across the $M$ sequences, $M$ mean and
precision parameters are sampled to generate these sequences. For sequence
${x}^{(k)}$, the generative process is,

$\displaystyle\tau_{k}$
$\displaystyle\sim\textsc{Gamma}(\alpha_{0},\beta_{0})$
$\displaystyle\mu_{k}|\tau_{k}$
$\displaystyle\sim\mathcal{N}(\mu_{0},(\lambda_{0}\tau_{k})^{-1})$
$\displaystyle x^{(k)}_{n}|\mu_{k},\tau_{k}$
$\displaystyle\sim\mathcal{N}(\mu_{k},\tau_{k}^{-1})\;\text{for}\;n\in\\{1,...,N\\},$

where $\alpha_{0},\beta_{0},\mu_{0},\lambda_{0}$ are fixed hyperparameters.
The predictive distribution for the next token in sequence ${x}$ is,

$\displaystyle p(x_{n+1}|x_{1:n})=\int
p(x_{n+1}|\mu,\tau)p(\mu,\tau|x_{1:n})d(\mu,\tau).$

The optimal Bayesian agent uses the same prior distributions as the Gamma and
Normal that generate $\mu$ and $\tau$. It will analytically infer the ground
truth posterior $p(\mu,\tau|x_{1:n})$ for any stream of data it sees, and use
this posterior for predicting next tokens. A Bayesian agent can also use other
suitable priors and converge to the optimal posterior. To be consistent with
other exchangeable conjugate models, we denote $\theta=(\mu,\tau)$ to indicate
latent variables whose posterior distribution is predictive sufficient.

(a) Gaussian-Gamma.

(b) Beta-Bernoulli.

(c) Gamma-Exp.

Figure 2: Probe recovery of transformer-learned sufficient statistic (blue)
and ground truth sufficient statistic (red), across 1000 test datapoints. The
first row shows parameters probed in the non-OOD case (from left to right:
Gaussian mean $\mu$, Gaussian precision $\tau$, Bernoulli mean, and
Exponential mean). The second row shows the corresponding information in the
OOD case.

####  4.1.2 Probing experiments

##### Experimental setup

We hypothesize that a transformer trained on this dataset should come to
represent the sufficient statistics of the corresponding distribution. We use
three probing setups.

Sufficient statistic. As suggested by the Proposition in Section 3.2, we probe
the sufficient statistic for $\theta$. The probe uses the last token embedding
of the transformer as input.

Moments of the posterior distribution. A consequence of knowing the sufficient
statistic is finding the true posterior $p(\theta|x_{1:N})$, so we decode the
moments of this posterior from the transformer embedding. The moments are
functions of the mean and variance of the stream of data seen so far by the
autoregressive predictor. In these conjugate models, it might be unsurprising
for a transformer to encode the mean of the sequence it processes, because the
optimal strategy for its loss function is to always predict the mean of the
sequence that it sees so 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 posterior based on the data. Thus, we probe the
transformer on out-of-distribution (OOD) datasets that are generated from a
distribution of the same form but with distinct hyperparameters.

##### Implementation details

For Gaussian-Gamma hyperparameters $\alpha_{0},\beta_{0},\mu_{0},\lambda_{0}$
we choose $\\{5,1,1,1\\}$. The OOD dataset is instead generated with
hyperparameters $\\{2,1,5,1\\}$, so it is centered on a different mean (5 vs.
1) and has a higher spread compared to the original data. An analogous design
was used for the other exchangeable conjugate models, with details in Appendix
A.2.

(a) Gaussian-Gamma (linear probe).

(b) Gaussian-Gamma (two-layer probe).

Figure 3: 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. 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.

##### Results

Results for all three models are shown in Figure 2, where the probe decodes
the sufficient statistics. The remaining figures focus on the Gaussian-Gamma
case (see Appendix A.3 for the other models). Figure 3(a) shows that the probe
decodes the moments of the posterior distribution. Because the higher
distribution moments are more volatile in value and less directly related to
estimating the parameters $\theta$, 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 3(b)),
showing stronger alignment.

##### Analyzing whether the transformer memorizes tokens

The transformer embedding (of size 128) can easily memorize a set of tokens,
so we tested whether memorization is leading to the successful recovery of the
sufficient statistics. We started by taking the 10th token embedding and
probing the first 10 tokens’ sufficient statistics in each sequence, with high
performance shown in Figure 8 in Appendix A.3. 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, a correlation exists between probe results and true token values, but
some level of correlation is expected because even seeing a single token can
reveal information about this sequence’s generating distribution. However, the
noise suggests that memorization is less of a goal than finding sufficient
statistics.

###  4.2 Hidden Markov model

####  4.2.1 Generative model

For an HMM generated data with $M$ sequences with $N$ tokens each, we
formulate the generative process as, for a sequence ${x}$,

$\displaystyle A_{c}$
$\displaystyle\sim\textsc{Dirichlet}_{C}(\gamma)\;\text{for}\;c\in\\{1,...,C\\}$
$\displaystyle B_{c}$
$\displaystyle\sim\textsc{Dirichlet}_{V}(\delta)\;\text{for}\;c\in\\{1,...,C\\}$
$\displaystyle{z}_{0}$ $\displaystyle\sim\textsc{Categorical}({\pi})$
$\displaystyle{x}_{i}$ $\displaystyle\sim\textsc{Categorical}(B_{{z}_{i}})$
$\displaystyle{z}_{i+1}$ $\displaystyle\sim\textsc{Categorical}(A_{{z}_{i}}),$

where $C,V,{\pi}$ 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. $\gamma,\delta$ are scalar
hyperparamters that are also initialized and fixed, and they represent the
evenness of the samples from the Dirichlet distributions. $A$, $B$, and ${z}$,
as a result, represent the transition matrix, the emission matrix, and the
latent states, respectively.

We implement a forward-backward algorithm (Rabiner, 1989) to compute the
posterior $p({z}_{n+1}|{x}_{1:n})$, and explore whether this distribution can
be decoded from the transformer embedding on ${x}_{1:n}$.

####  4.2.2 Probing experiments

Like our experiments for exchangeable sequences, we use probes to explore
whether an autoregressive transformer represents a predictive sufficient
statistic for non-exchangeable sequences.

##### Decoding latent measures other than the predictive sufficient statistic

To further examine the role of the autoregressive objective in encoding a
predictive sufficient statistic, we find other possible targets and test on
whether they are more difficult to decode. We implement the Viterbi algorithm,
which finds the most likely sequence of latent states given the whole observed
sequence, and we use Viterbi’s output as our target. This is related to but
distinct from the predictive sufficient statistic, $p(z_{n+1}|x_{1:n})$.

##### Implementational details

We choose $C=4,V=64,\gamma=0.5$, and set ${\pi}$ to be uniform in our
experiments. We also vary $\delta$ to control the level of difficulty: how
distinct is one class from another.

##### Results

Based on the autoregressive objective, we hypothesized that a transformer
trained on HMM sequences should encode the posterior distribution on latent
states $p({z}_{n+1}|{x}_{1:n})$. Table 1 suggests that this is the case,
reflected in terms of both accuracy of decoding and error in reproducing this
posterior distribution (as measured by squared loss). Furthermore, decoding
performance is better for the posterior than for several related quantities.
It is more difficult to decode other quantities using the same embedding, or
even to decode any of these quantities using the embedding on one token ahead,
${x}_{n+1}$. Several other quantities exhibit high volatility across random
seeds, since the transition matrix has a greater influence on results when the
tokens conditioned in the target posterior and the token used in the embedding
have a mismatch.

Table 1: Probing target quantities in HMM dataset using different transformer
token embeddings..

|  | $\delta=0.5$ |  | $\delta=1$ |
---|---|---|---|---|---
Target Quantity | Embedding | Accuracy $\uparrow$ | Squared Loss $\downarrow$ | Accuracy $\uparrow$ | Squared Loss $\downarrow$
$p({z}_{n+1}|{x}_{1:n})$ | ${x}_{1:n}$ | $90.8\pm 5.5\%$ | $0.011\pm 0.012$ | $90.4\pm 5.9\%$ | $0.011\pm 0.01$
$p({z}_{n+1}|{x}_{1:n})$ | ${x}_{1:n+1}$ | $66.6\pm 20.9\%$ | $0.072\pm 0.042$ | $65.9\%\pm 21.6\%$ | $0.058\pm 0.035$
$p({z}_{n+1}|{x}_{1:n+1})$ | ${x}_{1:n}$ | $53.2\pm 11.2\%$ | $0.356\pm 0.043$ | $50.8\pm 11.3\%$ | $0.278\pm 0.038$
$p({z}_{n+1}|{x}_{1:n+1})$ | ${x}_{1:n+1}$ | $86.5\pm 4\%$ | $0.066\pm 0.034$ | $82.2\pm 5.8\%$ | $0.067\pm 0.036$
$\hat{{z}}_{n+1}$ | ${x}_{1:n}$ | $59.8\%\pm 12.1\%$ | / | $61.4\pm 14.7\%$ | /
$\hat{{z}}_{n+1}$ | ${x}_{1:n+1}$ | $80.8\%\pm 7\%$ | / | $77.5\pm 8.8\%$ | /

Note: The target $p({z}_{n+1}|{x}_{1:n})$ is a simplex vector found by running
the forward-backward algorithm, and the target $\hat{{z}}_{n+1}$ is a scalar
standing for the most likely latent class found by Viterbi algorithm. Average
and standard deviation across 10 random seeds are reported.

###  4.3 Discrete hypothesis spaces

####  4.3.1 Generative model

We simulate a dataset where the data consist of a sequence of two-dimensional
points uniformly sampled from a rectangular region in 2D space (cf. Tenenbaum,
1998). In this case, the hypothesis space $\mathcal{H}$ is defined as the set
of all rectangles whose corner points are pairs of integers in
$\\{0,1,2,...,7\\}$. For each sequence ${x}$, the generative process is,

$\displaystyle h_{\text{rect}}$ $\displaystyle\sim\mathcal{H},$ $\displaystyle
x_{i}$ $\displaystyle\sim\textsc{Uniform}(h_{\text{rect}})\;\forall i.$

The generative process uniformly samples a rectangle $h_{\text{rect}}$ from
the set of all rectangles $\mathcal{H}$. Then, each token in a sequence is
sampled uniformly from the region defined by $h_{\text{rect}}$.

(a) A hypothesis space with equal-width rectangles.

(b) A hypothesis space with unequal-width rectangles.

Figure 5: Two discrete hypothesis spaces $\mathcal{H}$ used in our
experiments. Any continuous rectangle contained within the axes (e.g., the red
or the orange rectangle) is a valid hypothesis $h\in\mathcal{H}$. This results
in 784 distinct hypotheses in each hypothesis space. The data consist of a
sequence of points sampled uniformly from the target rectangle.

####  4.3.2 Probing experiments

##### Implementational details

We use rectangles with seven unit blocks on each side, resulting in a size-784
hypothesis space, i.e., 784 different possible rectangles from which the
sequences are drawn. The unit lengths are either 1 (Figure 5(a)), or 0.4 and
0.16 alternating (Figure 5(b)).

##### Probing results

We probe for the distribution $p(h|x_{1:n})$, which is a length-784 simplex
vector, using the last token embedding of each sequence. Results are shown in
Table 2. In general, the true hypothesis can be found with high accuracy, even
though the number of classes is high.

##### Embedding visualization

We visualize the transformer embeddings to show what they represent in terms
of this discrete set of hypotheses (Figure 6). We use Principal Components
Analysis (PCA) to reduce the 128-dimensional embeddings to 2D, and color the
points by properties of each rectangle. Embeddings are clustered into
different regions based on the mean of the center of the true generating
rectangle (Figure 6(a)), where, for instance, the top right yellow cluster
corresponds to the upper right-most rectangle, and the middle teal cluster
corresponds roughly to the many rectangles whose center have a mean of 3.5.
Additionally, the embeddings encode the position along individual axes (Figure
6(b)), as well as distance between corners of the generating rectangle (Figure
6(c)).

(a) Colored by mean of the generating rectangle.

(b) Colored by mean along the y-axis of the generating rectangle.

(c) Colored by spread along the y-axis of the generating rectangle.

Figure 6: Two-dimensional representation of embeddings of all validation
datapoints in the discrete hypothesis space dataset (the setup is unequal
width and sample size $=50$). Subfigures show the same embeddings, colored by
different properties of the true rectangle used to generate the data. Table 2:
Probing results on discrete hypothesis spaces. In general, the probe achieves
strong performance in recovering a 784-length vector, and performance
increases as task difficulty decreases. Average and standard deviation across
10 random seeds are reported.

| Equal Width |  | Unequal Width |
---|---|---|---|---
Sample Size | Accuracy $\uparrow$ | Squared Loss $\downarrow$ | Accuracy $\uparrow$ | Squared Loss $\downarrow$
20 | $87.3\pm 1.2\%$ | $0.173\pm 0.016$ | $66\pm 1.7\%$ | $0.29\pm 0.012$
50 | $99.5\pm 0.2\%$ | $0.008\pm 0.004$ | $88.5\pm 0.9\%$ | $0.159\pm 0.011$

##  5 Discussion

##### Limitations and Future Directions

Several limitations of our analysis motivate future work.

Extending to more complex Bayesian probabilistic models. We focused on three
cases that consist of a one-step generation process and a simple likelihood
$p(x|z)$ that can be approximated by a linear layer. Our analysis can be
extended to hierarchical models or likelihoods parameterized by, for instance,
neural networks, while simultaneously exploring more choices of transformer
embeddings.

Interpreting neural networks with Bayesian probabilistic models. A similar
Bayesian approach could be used to interpret the representations learned by
other autoregressive neural networks. A direction for future work is extending
our analyses to deep learning models for other modalities.

Task difficulty. As seen in the hypothesis space experiments, the ability to
decode predictive sufficient statistics decreases as task difficulty
increases. Developing a better understanding of the relationship between task
difficulty and transformer embeddings is an important direction for future
work.

##### 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
$x_{n+1}$ is independent from previous tokens $x_{1:n}$ when conditioned on
that structure, a property possessed by predictive sufficient statistics. We
confirmed this hypothesis with probing experiments on three 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.

## Acknowledgements

This work was supported by ONR grant number N00014-23-1-2510.

## References

  * Bengio et al. (2000) Yoshua Bengio, Réjean Ducharme, and Pascal Vincent.  A neural probabilistic language model.  _Advances in Neural Information Processing Systems_ , 13, 2000.
  * Radford et al. (2019) Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al.  Language models are unsupervised multitask learners.  _OpenAI blog_ , 1(8):9, 2019.
  * Andreas (2022) Jacob Andreas.  Language models as agent models.  In _Findings of the Association for Computational Linguistics: EMNLP 2022_ , pages 5769–5779, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics.
  * Hewitt and Manning (2019) 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 Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_ , pages 4129–4138, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.
  * Cunningham et al. (2023) Hoagy Cunningham, Aidan Ewart, Logan Riggs, Robert Huben, and Lee Sharkey.  Sparse autoencoders find highly interpretable features in language models.  _ArXiv_ , abs/2309.08600, 2023.
  * Li et al. (2023) Yuchen Li, Yuan-Fang Li, and Andrej Risteski.  How do transformers learn topic structure: Towards a mechanistic understanding.  _ArXiv_ , abs/2303.04245, 2023.
  * Meng et al. (2022) 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.
  * Zheng et al. (2023) Carolina Zheng, Keyon Vafa, and David M. Blei.  Revisiting topic-guided language models, 2023.
  * Tenney et al. (2019) Ian Tenney, Patrick Xia, Berlin Chen, Alex Wang, Adam Poliak, R Thomas McCoy, Najoung Kim, Benjamin 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.
  * Gupta et al. (2015) Abhijeet Gupta, Gemma Boleda, Marco Baroni, and Sebastian Padó.  Distributional vectors encode referential attributes.  In _Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing_ , pages 12–21, Lisbon, Portugal, September 2015. Association for Computational Linguistics.
  * Köhn (2015) Arne Köhn.  What’s in an embedding? analyzing word embeddings through multilingual evaluation.  In _Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing_ , pages 2067–2073, Lisbon, Portugal, September 2015. Association for Computational Linguistics.
  * Ettinger et al. (2016) 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_ , pages 134–139, Berlin, Germany, August 2016. Association for Computational Linguistics.
  * Adi et al. (2017) Yossi Adi, Einat Kermany, Yonatan Belinkov, Ofer Lavi, and Yoav Goldberg.  Fine-grained analysis of sentence embeddings using auxiliary prediction tasks.  In _International Conference on Learning Representations_ , 2017.
  * Hupkes et al. (2018) Dieuwke Hupkes, Sara Veldhoen, and Willem Zuidema.  Visualisation and’diagnostic classifiers’ reveal how recurrent and recursive neural networks process hierarchical structure.  _Journal of Artificial Intelligence Research_ , 61:907–926, 2018.
  * Rogers et al. (2020) Anna Rogers, Olga Kovaleva, and Anna Rumshisky.  A primer in BERTology: What we know about how BERT works.  _Transactions of the Association for Computational Linguistics_ , 8:842–866, 2020.
  * Belinkov (2022) Yonatan Belinkov.  Probing classifiers: Promises, shortcomings, and advances.  _Computational Linguistics_ , 48(1):207–219, 2022.
  * Shi et al. (2016) Xing Shi, Inkit Padhi, and Kevin Knight.  Does string-based neural MT learn source syntax?  In _Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing_ , pages 1526–1534, Austin, Texas, November 2016. Association for Computational Linguistics.
  * Belinkov et al. (2017) 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 Linguistics (Volume 1: Long Papers)_ , pages 861–872, Vancouver, Canada, July 2017. Association for Computational Linguistics.
  * Giulianelli et al. (2018) Mario Giulianelli, Jack Harding, Florian Mohnert, Dieuwke Hupkes, and Willem Zuidema.  Under the hood: Using diagnostic classifiers to investigate and improve how language models track agreement information.  In _Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP_ , pages 240–248, Brussels, Belgium, November 2018. Association for Computational Linguistics.
  * Conneau et al. (2018) Alexis Conneau, German Kruszewski, Guillaume Lample, Loïc Barrault, and Marco Baroni.  What you can cram into a single $&!#* vector: Probing sentence embeddings for linguistic properties.  In _Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 2126–2136, Melbourne, Australia, July 2018. Association for Computational Linguistics.
  * Liu et al. (2019) Nelson F. Liu, Matt Gardner, Yonatan Belinkov, Matthew E. Peters, and Noah A. Smith.  Linguistic knowledge and transferability of contextual representations.  In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_ , pages 1073–1094, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.
  * Lin et al. (2019) Yongjie Lin, Yi Chern Tan, and Robert Frank.  Open sesame: Getting inside BERT’s linguistic knowledge.  In _Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP_ , pages 241–253, Florence, Italy, August 2019. Association for Computational Linguistics.
  * Grand et al. (2022) Gabriel Grand, Idan Asher Blank, Francisco Pereira, and Evelina Fedorenko.  Semantic projection recovers rich human knowledge of multiple object features from word embeddings.  _Nature human behaviour_ , 6(7):975–987, 2022.
  * Radford et al. (2017) Alec Radford, Rafal Jozefowicz, and Ilya Sutskever.  Learning to generate reviews and discovering sentiment.  _arXiv preprint arXiv:1704.01444_ , 2017.
  * Li et al. (2021) Belinda Z. Li, Maxwell Nye, and Jacob Andreas.  Implicit representations of meaning in neural language models.  In _Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers)_ , pages 1813–1827, Online, August 2021. Association for Computational Linguistics.
  * Xie et al. (2021) Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma.  An explanation of in-context learning as implicit Bayesian inference.  _arXiv preprint arXiv:2111.02080_ , 2021.
  * McCoy et al. (2023) 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.  _arXiv preprint arXiv:2309.13638_ , 2023.
  * Wang et al. (2024) 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, 2024.
  * Zhang et al. (2023) Liyi Zhang, R. Thomas McCoy, Theodore R. Sumers, Jian-Qiao Zhu, and Thomas L. Griffiths.  Deep de finetti: Recovering topic distributions from large language models, 2023.
  * Mikulik et al. (2020) Vladimir Mikulik, Grégoire Delétang, Tom McGrath, Tim Genewein, Miljan Martic, Shane Legg, and Pedro A. Ortega.  Meta-trained agents implement bayes-optimal agents.  _ArXiv_ , abs/2010.11223, 2020.
  * Panwar et al. (2024) Madhur Panwar, Kabir Ahuja, and Navin Goyal.  In-context learning through the bayesian prism.  In _The Twelfth International Conference on Learning Representations_ , 2024.
  * Garg et al. (2022) Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant.  What can transformers learn in-context? a case study of simple function classes.  In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, _Advances in Neural Information Processing Systems_ , 2022.
  * Akyürek et al. (2023) Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou.  What learning algorithm is in-context learning? investigations with linear models.  In _The Eleventh International Conference on Learning Representations_ , 2023.
  * Müller et al. (2022) Samuel Müller, Noah Hollmann, Sebastian Pineda Arango, Josif Grabocka, and Frank Hutter.  Transformers can do bayesian inference.  In _International Conference on Learning Representations_ , 2022.
  * Gelman et al. (2004) Andrew Gelman, John B. Carlin, Hal S. Stern, and Donald B. Rubin.  _Bayesian Data Analysis_.  Chapman and Hall/CRC, 2nd ed. edition, 2004.
  * Bernardo and Smith (2000) Jose M Bernardo and Adrian Smith.  _Bayesian Theory_.  Wiley series in probability and statistics. John Wiley & Sons Ltd., Chichester, 2000.
  * Kalman (1960) Rudolph Emil Kalman.  A new approach to linear filtering and prediction problems.  _ASME Journal of Basic Engineering_ , 82(1):35–45, March 1960.
  * Baum and Petrie (1966) Leonard E Baum and Ted Petrie.  Statistical inference for probabilistic functions of finite state markov chains.  _The annals of mathematical statistics_ , 37(6):1554–1563, 1966.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _Advances in Neural Information Processing Systems_ , 30, 2017.
  * Kingma and Ba (2014) Diederik P. Kingma and Jimmy Ba.  Adam: A method for stochastic optimization, 2014.
  * Rabiner (1989) L.R. Rabiner.  A tutorial on hidden markov models and selected applications in speech recognition.  _Proceedings of the IEEE_ , 77(2):257–286, 1989.  doi: 10.1109/5.18626.
  * Tenenbaum (1998) Joshua Tenenbaum.  Bayesian modeling of human concept learning.  _Advances in Neural Information Processing Systems_ , 11, 1998.

##  Appendix A Appendix

###  A.1 Definition of additional exchangeable conjugate models

We consider the generative process for sequence $x_{i}$, where $x_{ij}$ are
i.i.d. across $j$.

####  A.1.1 Beta-Bernoulli model

$\displaystyle\theta_{i}$ $\displaystyle\sim\textsc{Beta}(\alpha,\beta)$
$\displaystyle x_{ij}$ $\displaystyle\sim\textsc{Bernoulli}(\theta_{i}),$

where $\alpha,\beta$ are fixed hyperparameters.

####  A.1.2 Gamma-Exponential model

$\displaystyle\theta_{i}$ $\displaystyle\sim\textsc{Gamma}(\alpha,\beta)$
$\displaystyle x_{ij}$ $\displaystyle\sim\textsc{Exponential}(\theta_{i}),$

where $\alpha,\beta$ are fixed hyperparameters.

###  A.2 Implementational details

All training processes are run on a T4 or A100 Nvidia GPU.

##### 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.

With the exception of the discrete hypothesis space datasets, the sizes for
the three sets are: 10000, 3000, 1000, and each sequence is 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.

##### Transformer

We use a 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. If the input is continuous, we use a Linear layer to map inputs to
dimension 128 in place of the word embedder layer.

Dropout $=0.1$ is applied, and learning rate $=0.001$, batch-size $=64$.

##### Probe

The probe is a linear layer with softmax activations. Learning rate is tuned
in $[0.001,0.01]$, and batch-size $=64$.

##### Hyperparameters for exchangeable conjugate models

On the Beta-Bernoulli model, we use $\alpha=2,\beta=8$. In the OOD case,
$\alpha=8,\beta=2$.

On the Gamma-Exponential model, we use $\alpha=2,\beta=4$. In the OOD case,
$\alpha=2,\beta=1$.

###  A.3 Experimental results

Figure 7 shows results on posterior distribution moments on Beta-Bernoulli and
Gamma-Exponential models. On the Gamma-Exponential model, we divide the target
second, third, and fourth moments by factors of 10, 100, 1000, respectively so
that each moment is given roughly equal importance.

(a) Beta-Bernoulli.

(b) Gamma-Exponential (linear probe).

(c) Gamma-Exponential (two-layer probe).

Figure 7: Probe recovery of transformer-learned posterior distribution moments
(blue) plotted with ground truth moments (red). The first row shows parameters
probed on 1000 test datapoints in the general, i.e., non-OOD, case. The second
row shows corresponding information in the OOD case.

Figure 8 shows Gaussian-Gamma probing results on sufficient statistics on the
10th token, complementing our exploration of token memorization in transformer
(Figure 4).

Figure 8: Probing over mean (left) and standard deviation (right) of the first
10 tokens using the 10th token embedding of the transformer in the Gaussian-
Gamma dataset. The first row corresponds to the same generation process, and
the second row corresponds to the OOD case.