Small transformers track Bayesian evidence for latent common causes via a context-invariant mechanism

Amir Mohammadpour; Michael Franke

cs.LGSep 28, 2026

Small transformers track Bayesian evidence for latent common causes via a context-invariant mechanism

Authors: Amir Mohammadpour, Michael Franke

Organizations: Department of Linguistics, University of Tübingen

Abstract

We present an in-depth investigation of how a form of Bayesian reasoning about common causes can emerge as a cross-contextual generalization in small, tractable transformers. Incrementing on recent work, our set-up (i) disentangles causal mechanisms in the model from the causal structure of the true data-generating process, (ii) orients more towards natural language prediction by considering inference of latent common causes, and (iii) considers whether and how Bayesian evidence accumulation for latent common causes can be implemented in representations and mechanisms that allow for cross-context generalization to novel test cases.

Figures & tables

Figure 1: We train small transformers on scrambled token sequences from different contexts, each of which are causal models with a common cause and two effect variables. Held-out contexts require the models to extrapolate Bayesian inference of the causal variable’s value from contexts seen in training. We find that “Bayesian generalization” emerges, and we detail its mechanics, finding that (1) token embeddings geometrically encode the relevant roles of the underlying causal variables, that (2) the residual stream encodes a categorical belief representation, which is (3) transformed via the LayerNorm into an approximate Bayesian posterior. Surgical interventions yield systematic downstream effects on the joint distribution over relevant variables, as predicted by Bayes rule.

Figure 2: Training success and stochastic dependencies in models’ behavior. (a) Training success measures; (global) averaged over all positions, and (held-out) averaged over the masked positions in held-out sequences . (b) Total variation distances at second token of supervised (blue) and held-out sequences (orange). The three columns represent variable exchange error

\Delta_{\mathrm{ex}}

(Eq. 5 ), child’s predictive sensitivity to sibling

S_{Y}

(Eq. 6 ) and child’s predictive sensitivity to parent

S_{X}

(Eq. 7 ). Open markers: mean with 95% two-level bootstrapped interval. Shaded bands:

\bar{\epsilon}_{1}

(Eq. 8 ), model’s predictive error at the same position.

\Delta_{\mathrm{ex}}

and

S_{Y}

, are in the range of models’ error margin.

S_{X}

reaches

0.56

, ground-truth is

|2p_{c}-1|

(mean

0.60

): an order of magnitude above the floor.

Figure 3: Context-centered embeddings exhibit a common geometry. (a) Variable plane encodes the variation among variable identities; $d_{r_{1}}$ separates parent from children, $d_{r_{2}}$ children from each other. (b) Value plane separates tokens by the regime value; $d_{v}$ and $d_{r}$ mostly captures a common value shift among all variables and parent’s deviation from children, respectively. (c) Child values are exchangeable; their distance in the value plane tracks contexts dependency power (i.e., furthest for $p_{c}=0.9$ , and shortest for $p_{c}=0.7$ ). (d) The design axes align with the natural directions that explain the variance of the raw token embeddings.

Figure 4: Geometry of the residual stream. Residual streams after functional attention, pooled over 45 models and all contexts, projected onto the design planes. Dots are pooled cells; markers are pool medians. (a) Each variable and each pair occupies its own ray, a pair lying between its members. (b) Value plane; in the absence of parent (gray dots), belief about it lies on the anti-diagonal. A concordant second child leaves the median in place and a discordant one returns it to the origin. Belief encodes the sign of the evidence and is insensitive to its count. (c) Value plane (complementary); sequences containing the parent lie off the belief axis.

Figure 5: Accumulation gain. (a) Bootstrap mean with 95% CI of absolute parent log-odds in units of

|\lambda_{c}|

. Held-out two-child cases are not exact. (b) Amplification of the parent’s log-odds from one child to two concordant children, split into the factor from the components,

F_{w}

, and from the final LayerNorm,

F_{\sigma}

. Grey diagonals mark total gains of

1.25

1.5

and

1.75\times

; the dashed diagonal marks

2\times

. (c) Shrinkage of the final LayerNorm input from one child to two concordant children,

\Delta\sigma^{2}

, split into the part inside the variable plane and the rest. Grey diagonals mark constant total shrinkage; the dashed line marks the pooled in-plane share,

\sum\Delta\sigma^{2}_{\mathrm{plane}}/\sum\Delta\sigma^{2}

. (d)

\Delta\sigma^{2}

in the held-out context against the mean over the two supervised contexts, one point per model. Dashed: identity; solid: pooled ratio

\sum\Delta\sigma^{2}_{\text{held-out}}/\sum\Delta\sigma^{2}_{\text{supervised}}

; open markers are clipped at the axis limit.

Figure 6: Joint dependency intervention.

n

unit beliefs are added along the belief coordinate after the functional attention; 135 (model, context) cells, means with

95\%

bootstrap CIs. (a) Parent log-odds shift against

n

at the first (solid) and second (dashed) child positions. (b) Slope of the two-child over the one-child line from (a) as intervention ratio against the accumulation gain

G(|\Omega|)

, supervised and held-out cluster; right: intervention ratio over accumulation gain is close to one, i.e., intervention follows the functional form. (c) Sibling against parent in

\tanh

coordinates, colored by context index

p_{c}

; right: residual mean squared error from the Bayes predicted identity line.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Layers / heads per layer	2 / 1
$d_{\text{model}}$ / head dimension	16 / 16
MLP hidden size, activation	16, GELU
Dropout	none
Biases	QKV and output projections
Positional encoding	sinusoidal, max length 1000
Vocabulary	19 (1 shared delimiter $+$ 3 contexts $\times$ 6)

Appendix

Table 1: Architecture configuration common to all 45 models.

Optimizer	AdamW, $\beta=(0.9,0.999)$ , weight decay $0.1$
Learning rate	$0.05$ , constant (no warmup or decay)
Gradient clipping	norm $1.0$ , no accumulation
Batch size	384 sequences (92,160 tokens)
Training budget	20k steps (7.68M sequences)
Validation set	4,800 fixed sequences
Evaluation / checkpoint interval	every 10 steps

Appendix

Table 2: Optimization and checkpoint selection.

SparseYour attention	Total	$p_{0}=0.7$	$p_{0}=0.8$	$p_{0}=0.9$
Layer 1	31	8	13	10
Layer 2	14	7	2	5

Appendix

Table 3: Sparse attention layer by held-out context.

Figure 7: For each model and context

C_{c}

, the centered table

\tilde{E}_{c}\in\mathbb{R}^{6\times 16}

, one row per variable and value, is split into the variable plane, the value axis and the interaction plane (Appendix C ). (a) Share of

\|\tilde{E}_{c}\|_{F}^{2}

in each subspace, by context

p_{c}

. Bars are means over the 45 models, and error bars are 95% bootstrap CIs with the model as resampling unit. This split uses no parent/child structure. Most of the table lies in the variable plane. The value share grows with

p_{c}

, and the interaction share stays small. (b) Share of the parent-vs-children axis within each plane,

\|d_{r_{1}}^{\top}\tilde{E}_{c}\|^{2}/(\|d_{r_{1}}^{\top}\tilde{E}_{c}\|^{2}+\|d_{r_{2}}^{\top}\tilde{E}_{c}\|^{2})

for the variable plane and

\|d_{p}^{\top}\tilde{E}_{c}\|^{2}/(\|d_{p}^{\top}\tilde{E}_{c}\|^{2}+\|d_{a}^{\top}\tilde{E}_{c}\|^{2})

for the interaction plane. Dots are single (model, context) values, and open markers are means with 95% bootstrap CIs. The dashed line at

0.5

marks a plane with no parent/child asymmetry. The variable plane lies near it, and the interaction plane is almost entirely the parent deviation. Supervised and held-out contexts agree throughout. (c) Share of each axis before and after the LayerNorm that precedes the functional attention, for the first, second and third token of a triplet.

d_{r_{1}}

separates the parent from the children,

d_{r_{2}}

the two children,

d_{v}

is the common value shift,

d_{p}

the parent’s departure from it, and

d_{a}

the difference between the children’s shifts. Points lie on the identity line, so the LayerNorm leaves the decomposition intact. (d) Partner term relative to query term,

\|b\|_{F}/\|a\|_{F}

, at the second token of a triplet, where

a

and

b

are the query’s and the partner’s contributions to the cell means of the residual stream. Markers as in (b). Before the functional attention ( ln1 input) the partner term is essentially zero. After it ( ln2 input), the ratio is close to 1 in both model groups, so attention adds the partner’s embedding to the query’s own at equal weight.

Figure 9: Functional form of the attention block, read at the residual node after the functional attention for the second token of a triplet. Each dot is one model and context; open circles are medians. (a) Share of the residual stream that a sum of one query term and one antecedent term cannot explain (log scale). The observed share is compared to that of random tables fitted in the same way. (b) Size of the query and antecedent terms relative to the residual of the same token observed alone. Dashed lines mark the values expected if the node summed the two tokens (1) or averaged them (0.5). (c) Alignment of the query and antecedent terms within each design subspace: the value axis

d_{v}

, the variable plane

(d_{r_{1}},d_{r_{2}})

and the interaction plane

(d_{p},d_{a})

. For each plane, 1st and 2nd are its two principal angles, from most to least aligned; a cosine of 1 means the two terms point in the same direction.

Figure 10: Writers carrying the change in the readout. For each writer

m\in\{e,a_{1},f_{1},a_{2},f_{2}\}

, the share

s_{m}=\big(\mathbb{E}_{2}[\langle\gamma\odot\tilde{m},w_{c}\rangle/\sigma]-\mathbb{E}_{1}[\langle\gamma\odot\tilde{m},w_{c}\rangle/\sigma]\big)/\mathbb{E}_{2}[A]

is the writer’s contribution to the change in the readout from one observed child (

\mathbb{E}_{1}

) to two concordant children (

\mathbb{E}_{2}

), normalized by the two-child readout

A=v-\langle b,w_{c}\rangle

. The layer-norm bias term

\langle b,w_{c}\rangle

is constant and excluded. For each model and context, the five shares sum to

1-1/F

, where

F=\mathbb{E}_{2}[A]/\mathbb{E}_{1}[A]

is the total readout gain. Columns: contexts

p_{c}

. Rows: models with sparse attention in layer 1 (31 models) or layer 2 (14 models). Dots: individual models, blue for supervised contexts, orange for the held-out context. Black points: panel mean with 95% percentile-bootstrap CIs (10,000 resamples). The attention writer of the sparse layer carries the largest share, followed by the embedding

e

. In layer-2 models, the layer-1 writers

a_{1}

and

f_{1}

contribute almost nothing.

Figure 11: Held-out shortfall in the drop of

\sigma^{2}

, by term. For each model and term,

\Delta

is the drop from one observed child to two concordant children (

\mathbb{E}_{1}-\mathbb{E}_{2}

) in the held-out context, minus the mean drop over the model’s two supervised contexts. Negative

\Delta

means the held-out context shrinks the term less. Rows are the terms of the exact expansion of

\sigma^{2}

over the five writers: five variances

\mathrm{var}(m)=\tfrac{1}{d}\lVert\tilde{m}\rVert^{2}

and ten covariances

2\,\mathrm{cov}(m,m^{\prime})=\tfrac{2}{d}\langle\tilde{m},\tilde{m}^{\prime}\rangle

. The bottom row is their sum,

\Delta\sigma^{2}

. Panels: models with sparse attention in layer 1 (

n=31

) and layer 2 (

n=14

). Grey dots are single models. Black points are means with 95% percentile-bootstrap CIs over models (2,000 resamples). A point is filled if its CI excludes zero and open otherwise. The x-axis spans the 2nd–98th percentile of all dots in both panels, padded by 5%. The 24 of 720 dots beyond it are drawn open at the edge. In both families the held-out context shrinks

\sigma^{2}

less. In layer-1 models the shortfall is spread over the covariances of

e

and

a_{1}

with the MLP writers

f_{1}

and

f_{2}

. In layer-2 models it sits mainly in

2\Delta\mathrm{cov}(e,a_{2})

and

2\Delta\mathrm{cov}(e,f_{2})

Explore similar work

Jun 29, 2026cs.AI

BayesBench: Evaluating LLM Belief Trajectories Under Multi-Turn Evidence Accumulation

Large language models (LLMs) are typically deployed in multi-turn conversations, where each turn provides new evidence that should reduce epistemic uncertainty about their environment. Acting rationally then requires inferring the unobserved quantities that govern it and updating beliefs about them as evidence accumulates. Yet most evaluations only score the model's final-turn answer in a single-turn format, leaving this process unexamined. We ask how closely LLMs' belief updates match those of a rational Bayesian reasoner in multi-turn settings, and introduce BayesBench, a suite of simulation environments that probe this across three progressively complex tasks: (i) Bayesian estimation, where the model infers an unknown parameter from sequential evidence; (ii) Bayesian prediction, where the model turns inferred beliefs about a latent variable into outcome forecasts; and (iii) latent-framed Bayesian prediction, where observations are filtered through a user-persona framing, requiring joint inference over the latent state and the persona. Across seven LLMs (3B--70B), scaling improves latent inference and evidence accumulation, with updates occasionally matching the Bayesian posterior. However, these gains do not reliably carry over to downstream prediction, exposing a gap between inferring latent structure and using it to rationally update beliefs about the target outcome.

Ankur Samanta, Akshayaa Magesh, Tal Lancewicki +7

Meta AI · Columbia University · Work done at Meta +2

LLM Reasoning Strategies

Ask PDF

Jul 20, 2026cs.CL

For What Reason? Interpreting Models' Encoding of Causation and Antithesis

Discourse relations provide document structure, critical to language understanding and enabling language model performance and ethicality. In this work, we investigate how instruction-tuned Transformer models (LLaMA and Mistral) encode discourse relations in English, with a particular focus on the contrasting relations of causation and antithesis. Framing the task as a next-token prediction task and applying a suite of interpretability techniques to test model internals, our findings show that certain early layers make predictive decisions at mid-sequence tokens, while some mid-level layers finalize their decisions closer to the last token. Most of the remaining layers primarily propagate earlier decisions rather than actively influencing them. Additionally, we observe that some layers exhibit a preference for one answer over alternatives, suggesting asymmetric representation of discourse-based reasoning.\footnote{Our code is available at https://github.com/abhidipbhattacharyya/causation_vs_antithesis}

Abhidip Bhattacharyya, Shira Wein

University of Massachusetts Amherst · University of South Florida

Transformer Architectures Causal Reasoning Skills Contrast

Ask PDF

Jul 19, 2026cs.LG

What does a Bayes-filtered transformer believe? A predictive Monte Carlo approach

A Bayes-filtered transformer (BFT) is a transformer trained on sequences that are generated in two steps: first a latent task is drawn from a prior, then observations are drawn conditional on that task. Trained under autoregressive log loss, the BFT's next-token prediction, in the idealized limit, is the Bayesian posterior predictive distribution (PPD) induced by that prior and that conditional law. In practice the trained BFT is only an approximation of this ideal PPD, raising an interpretive question: what prior and posterior over the latent task has the trained BFT actually internalized? Existing work answers this question by comparing the trained BFT's predictions against the predictions of various "reference" posteriors, each standing in for a different candidate algorithm or computation the BFT might be implementing. This prediction-space comparison is fragile: different posteriors can share the same posterior-mean predictions. We use predictive Monte Carlo (PMC) as a general interpretability tool for any BFT: using only next-token generation, PMC returns an approximation to the implicit prior and posterior over the latent task, answering the interpretive question directly in latent space. We apply PMC to three stylized task families spanning 0-Markov and 1-Markov exchangeability. The phenomena previously reported in these settings remain visible in latent space. Code is available at https://github.com/afiq-aswadi/bft-pmc

Afiq Abdillah Effiezal Aswadi, Haotong Ma, Susan Wei

Department of Econometrics and Business Statistics Monash University

Posterior Predictive Distribution Transformer Architectures Next-Token Prediction Monte Carlo Interpretability

Ask PDF