PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Amir Mohammadpour, Michael Franke

arXiv:2609.35161v1cs.LG

TL;DR

Prior work often analyzed training cases and did not always test recovered representations with causal interventions. This paper trains small transformers on context-specific sequences to study latent common-cause inference, finding shared belief representations whose interventions shift predictions consistently with Bayes’s rule.

  • Problem

    Prior studies of transformer reasoning focused on cases included in training, and causal-intervention checks on recovered representations were not always provided.

    Source: §2 · p. 3

  • Method

    The paper uses a controlled common-cause task and analyzes whether shared token-role representations and belief accumulation support Bayesian inference across contexts.

    Source: §1 · p. 1, §2 · p. 3, §3 · p. 4

  • Results

    The represented latent belief shifts predictions over unobserved variables consistently with Bayes’s rule under joint-dependency interventions.

    Source: §1 · p. 1, §6 · p. 9, §7 · p. 9

  • Takeaways & Limitations

    In this setup, cross-context Bayesian generalization rests on shared representations, a categorical belief coordinate, and a mapping to Bayesian posteriors.

    Source: §5.3 · p. 8, §7 · p. 9

    Whether the findings extend to more substantial hierarchical generative models with never-observed latent variables remains a direction for future work.

    Source: §7 · p. 9

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper investigates how Bayesian causal reasoning can emerge in transformers, as part of understanding how predictive models may develop useful world models. It finds cross-context Bayesian-like behavior supported by internal representations and mechanisms that respond to interventions in ways coherent with Bayesian computations.

1 INTRODUCTION

  • The study connects this question to the possibility that implicit world models support probabilistic causal reasoning in large transformers.Prior evidence includes case studies of finite-state problems such as board games, taxi routes, and text-based games.Source: §1 · p. 1
  • The paper analyzes how generalizable Bayesian causal reasoning emerges in a minimal transformer setup, separating causal mechanisms in the model from causal structure in the data.It focuses on abductive inference from observed effects to latent causes across held-out contexts.Source: §1 · p. 1
  • Small transformers show cross-context Bayesian-like behavior, and interventions on internal representations produce downstream effects coherent with Bayesian computations.The interventions satisfy a foundational abstraction requirement: commutativity under intervention.Source: §1 · p. 1

2 RELATED WORK & NOVEL CONTRIBUTIONS

The paper distinguishes causal mechanisms computed within a model from causal processes represented in its data, then studies inference about latent common causes across contexts. Its contribution combines a setup separating token position from causal role with behavioral and mechanistic investigation of Bayesian generalization to novel cases.

2 RELATED WORK & NOVEL CONTRIBUTIONS

  • Figure 1 summarizes a proposed mechanism: token embeddings encode causal-variable roles, residual states encode categorical beliefs, and LayerNorm maps them to approximate Bayesian posteriors.Interventions produce systematic changes to joint predictions consistent with Bayes’s rule.Source: Figure 1 · p. 2
  • The paper separates within-model causality, meaning mechanisms operating in model computation, from behind-data causality, concerning external causal processes represented and used by the model.It motivates the distinction by noting that positional dependencies can coincide with data-generating dependencies by design in simple Markov setups.Source: §2 · p. 2
  • Earlier studies primarily analyzed predictions and internal computations on training cases, and did not always include causal-intervention checks on recovered representations.The paper frames held-out cases and interventionist tests as incremental contributions.Source: §2 · p. 3
  • The setup disentangles token position from its potential information about future tokens and its true causal role in the data-generating process.It uses contexts with common causes and sometimes-latent causes, and tests inference to latent common causes in novel cases.Source: §2 · p. 3

3 EXPERIMENTAL DESIGN, RESEARCH QUESTIONS & APPROACH

The experiments use shuffled tokens from context-specific causal models with a common cause and two effects, then withhold contexts and sequences to test cross-context generalization. The analysis asks whether transformers track prediction-relevant beliefs and whether shared token-role representations and evidence accumulation can support inference in unseen cases.

3 EXPERIMENTAL DESIGN, RESEARCH QUESTIONS & APPROACH

  • Training and test sequences are sampled from randomly selected contexts, with held-out contexts and sequences withheld to measure cross-context generalization.A held-out sequence omits the succeeding parent token when it follows both children.Source: §3 · p. 3
  • Each context has a binary common-cause parent and two binary children, with child values depending probabilistically on the parent and context-specific dependency strength.The three tokens are mapped to context-specific symbols and randomly shuffled, separating surface position from variable role.Source: §3 · p. 3
  • The ideal-learner analysis treats beliefs about the parent’s value as the prediction-critical information and expects models to track its log-odds incrementally.The relevant canonical states depend on which variables have been observed and the corresponding belief about the parent.Source: §3 · p. 4
  • Ideal-learner state recovery alone does not predict generalization to held-out sequences, whose canonical states are absent from training.The authors hypothesize that cross-context similarities in token roles and a shared evidence-accumulation mechanism may support Bayes-like generalization.Source: §3 · p. 4

4 MODELS, TRAINING & BEHAVIORAL GENERALIZATION

Small transformers trained on shuffled sequences show behavioral generalization to held-out contexts and recover key stochastic dependencies. Their held-out performance remains above a random baseline, while parent sensitivity is close to its ground-truth value.

4 MODELS, TRAINING & BEHAVIORAL GENERALIZATION

  • Held-out KL divergence is an order of magnitude higher than global KL divergence but remains far below a random baseline.Models are particularly successful at predicting the correct variable at held-out positions.Source: §4 · p. 4
  • In supervised and held-out contexts, models respect variable exchangeability and parent screening of children, consistent with recovering roles of tokens in the underlying causal variables.The models use two pre-norm blocks with one attention head and 4,051 parameters, and training uses next-token cross-entropy.Source: §4 · pp. 4–5
  • 0.56 for SX is an order of magnitude above the floor, compared with ground-truth |2pc −1| (mean 0.60); ∆ex and SY remain within model error.The figure reports these dependency measures for supervised and held-out sequences.Source: §4 · p. 4, Figure 2 · p. 5

5 REPRESENTATION OF THE DATA-GENERATING CAUSAL MODEL

The models recover the data-generating variable and value structure in their embeddings, form a sign-only belief about the common cause after attention, and scale that belief with evidence count downstream.

5.1 VARIABLE ROLES EMERGE IN TOKEN EMBEDDINGS

  • Token embeddings recover variables and values, with a dominant variable term and a value plane spanning a common value direction and parent-versus-children direction.Children are exchangeable, and the two value tokens of a variable lie farther apart in contexts with larger pc.Source: §5.1 · pp. 5–6, Figure 3 · p. 6

5.2 A BELIEF COORDINATE EMERGES AFTER ATTENTION

  • The functional attention approximately averages the two token contributions, whose norms are nearly equal and whose average weights are 0.54–0.57 of the single-token reference.Their design subspaces align across the value axis, both variable-plane axes, and the leading interaction direction.Source: §5.2 · p. 6
  • After attention, concordant child values shift the residual toward their shared value and discordant values cancel, yielding a ternary belief that records evidence sign but not count.The residual direction also points away from the variable still to be emitted, including in held-out contexts.Source: §5.2 · pp. 6–7, Figure 4 · p. 7

5.3 MODEL COUNTS THE EVIDENCE

  • 67%: held-out residual shrinkage is only 67% of supervised shrinkage, leaving two-child parent log-odds short of 2λc.Final LayerNorm supplies most of the amplification, with most shrinkage occurring inside the variable plane; multiple small component misalignments explain the held-out deficit.Source: §5.3 · pp. 7–8, Figure 5 · p. 8

6 JOINT DEPENDENCY INTERVENTION

The joint-dependency intervention tests whether the represented belief acts as a Bayesian latent by steering it and checking the resulting parent and sibling predictions. The sibling prediction follows the Bayes-predicted relation, while parent shifts scale with the accumulation gain.

6 JOINT DEPENDENCY INTERVENTION

  • The intervention adds unit beliefs along the represented coordinate after functional attention and compares parent and sibling log-odds shifts with Bayes’s rule.Parent predictions are read at one- and two-child sequences, and sibling predictions at one-child sequences.Source: §6 · pp. 8–9
  • In tanh coordinates, the steered sibling prediction follows the Bayes-predicted identity line, while two-child parent shifts increase by the accumulation gain.At one-child sequences, one injected unit belief shifts parent log-odds by ρXc, slightly below λc.Source: §6 · p. 9, Figure 6 · p. 9

7 DISCUSSION, CONCLUSIONS, LIMITATIONS & OUTLOOK

The paper presents a case study of generalizable Bayesian abductive reasoning in transformers and identifies limits on causal recovery and transfer to large models.

7 DISCUSSION, CONCLUSIONS, LIMITATIONS & OUTLOOK

  • The case study finds that Bayesian generalization rests on representations shared across contexts, a categorical belief coordinate, and a mapping to Bayesian posteriors.The authors describe the uncovered mechanism as genuine Bayesian inference for this task.Source: §7 · p. 9
  • The setup cannot recover the precise causal role of the parent variable because of known theoretical limits of causal recovery.The results do not transfer directly to large models, though the authors suggest those models may use embedding similarity and efficient Bayesian-computation approximations.Source: §7 · p. 9
  • Future work could examine more substantial hierarchical latent-generation models in which latent variables are never observed.The authors also propose studying settings where predicted tokens inform the data-generating process of other tokens.Source: §7 · p. 9

8 STATEMENTS

The authors disclose generative-AI assistance for coding, analyses, visualizations, appendix drafting, and experimental-design and interpretation feedback.

8 STATEMENTS

  • Generative AI tools assisted with coding, analyses, visualizations, parts of the appendices, and feedback on experimental design and interpretation.The authors state that they reviewed all AI-assisted work and take responsibility for the final content.Source: AI USE STATEMENT · p. 10

A TRAINING

The models are trained on sampled context-specific triplets, with one context withheld from parent-token prediction. Behavioral metrics assess whether predictions respect emission-order invariance and the parent–child conditional-independence structure.

A TRAINING

  • Training samples are stratified into 60-triplet sequences, with each context appearing 20 times and fresh batches drawn at every step.Within each context, each value of Xc appears 10 times, and Y-value counts match pc.Source: A TRAINING · p. 13
  • Each model uses one held-out context, where parent tokens are removed from the loss and attention keys in the designated sequences.Other positions in that context are trained normally.Source: A TRAINING · p. 13
  • The 45 models cross three held-out-context parameter values with 15 seeds, and each model has a sparse attention layer identified by its attention-row divergence from uniform.The other layer is diffuse; the passage says the sparse layer is also the functional layer.Source: A TRAINING · p. 13

B BEHAVIORAL METRICS

  • Exchangeability error measures how much predictions change when the two observed tokens are swapped, testing invariance to their emission order.The metric is derived from total variation distance between the distributions before and after swapping.Source: B BEHAVIORAL METRICS · p. 14
  • Sibling and parent sensitivity measure prediction changes after flipping a sibling or parent value, testing the expected conditional-independence pattern.Both measures use total variation distance and are interpreted relative to predictive error at the same token position.Source: B BEHAVIORAL METRICS · p. 14

C ORTHOGONAL DECOMPOSITION OF THE EMBEDDING TABLE

The embedding table is decomposed into label-defined variable, value, and interaction subspaces, then traced through normalization and attention. The analysis finds that the decomposition is preserved and that a sparse attention layer carries most of the state distinction.

C ORTHOGONAL DECOMPOSITION OF THE EMBEDDING TABLE

  • The decomposition separates variable identity, a common value shift, and interaction deviations into five orthonormal axes, including parent-versus-children and child-versus-child distinctions.The variable plane is two-dimensional, the value axis one-dimensional, and the interaction plane two-dimensional.Source: C ORTHOGONAL DECOMPOSITION OF THE… · p. 14
  • Most embedding-table mass lies in the variable plane; the value share grows with pc while the interaction share stays small.The parent-vs-children axis dominates the interaction plane, while supervised and held-out contexts agree and LayerNorm preserves the decomposition.Source: Figure 7 · p. 15, C ORTHOGONAL DECOMPOSITION OF THE… · p. 15
  • LayerNorm leaves the embedding decomposition intact, and after functional attention the partner term is added to the query term at nearly equal weight.Before attention the partner term is essentially zero; after attention its ratio to the query term is close to 1 in both model groups.Source: C ORTHOGONAL DECOMPOSITION OF THE… · pp. 15–16

D ATTENTION LOCALIZATION

  • The sparse attention layer recovers a large fraction of the state distinction, while the diffuse layer recovers almost none, regardless of which layer is sparse.Patching the skip connection alone recovers 0.469 of the distinction; together with sparse attention, it accounts for the entire distinction.Source: D ATTENTION LOCALIZATION · p. 16

E FUNCTIONAL FORM OF THE ATTENTION BLOCK

The attention block is analyzed as a combination of query and antecedent contributions, then its downstream readout gain is examined. The reported results describe approximately averaging aligned token terms and associate much of the readout-scale drop with the variable plane.

E FUNCTIONAL FORM OF THE ATTENTION BLOCK

  • The attention-block residual is modeled as a global offset plus query and antecedent terms, with a cross-term remainder.The weighted means of the query and antecedent terms are constrained to zero, making the decomposition unique.Source: E FUNCTIONAL FORM OF THE ATTENTION BLOCK · p. 16
  • The observed cross-term residual lies orders of magnitude below the random-table reference, indicating that the pair contributes through what each token contributes alone.For balanced counts, the unstructured-table residual reference is 13/23.Source: E FUNCTIONAL FORM OF THE ATTENTION BLOCK · pp. 17–18, Figure 9 · p. 18
  • Both token-term norms lie near 1/2 of the single-token residual, and their subspaces align across the value axis, variable plane, and leading interaction direction.Only the second interaction angle is large, in a direction that carries negligible mass; the results hold in supervised and held-out contexts.Source: E FUNCTIONAL FORM OF THE ATTENTION BLOCK · p. 18, Figure 9 · p. 18

F.1 GAIN FACTORIZATION

  • For gain factorization, the analysis compares residual activations at parent-prediction positions after one child versus two value-concordant children.Per context, it uses 200 sequences of 20 triplets and discards each sequence's first triplet.Source: §F.1 · p. 19

F.2 SHRINKAGE IN THE VARIABLE PLANE

  • The normalization factor exceeds one because residual scale is smaller after two value-concordant children, and the variable plane carries most of this drop.A cross-model baseline captures only a small fraction of the drop.Source: §F.2 · p. 19

F NORMALIZATION AT THE READ-OUT

  • The readout-gain analysis attributes the largest share to the attention writer of the sparse layer, followed by the embedding writer.In models with layer-2 sparse attention, layer-1 attention and MLP writers contribute almost nothing.Source: Figure 10 · p. 20

F.3 EXPANSION OF THE NORMALIZATION SCALE

The held-out deficit in the normalization-scale drop is carried by covariance terms, while variance deficits remain close to zero. In layer-1 models, the deficit reflects different alignment among writers, not different writer magnitudes.

F.3 EXPANSION OF THE NORMALIZATION SCALE

  • Variance deficits are close to zero in both model families; covariance deficits carry the held-out shortfall in the normalization-scale drop.The fifteen term deficits sum to the deficit of ∆σ2.Source: §F.3 · p. 20
  • In layer-1 models, the deficit spreads across covariances of e and a1 with MLP writers, while writer magnitudes match supervised contexts and no single pair accounts for it.The held-out writers are aligned differently from those in supervised contexts.Source: §F.3 · p. 20, Figure 11 · p. 21
  • In layer-2 models, the held-out shortfall sits mainly in 2∆cov(e, a2) and 2∆cov(e, f2).The figure reports the held-out context shrinks σ2 less in both model families.Source: Figure 11 · p. 21
Loading 2609.35161v1…