cs.LGSep 28, 2026
SaveSmall transformers track Bayesian evidence for latent common causes via a context-invariant mechanism
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 (Eq. 5 ), child’s predictive sensitivity to sibling (Eq. 6 ) and child’s predictive sensitivity to parent (Eq. 7 ). Open markers: mean with 95% two-level bootstrapped interval. Shaded bands: (Eq. 8 ), model’s predictive error at the same position. and , are in the range of models’ error margin. reaches , ground-truth is (mean ): 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; separates parent from children, children from each other. (b) Value plane separates tokens by the regime value; and 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 , and shortest for ). (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 . 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, , and from the final LayerNorm, . Grey diagonals mark total gains of , and ; the dashed diagonal marks . (c) Shrinkage of the final LayerNorm input from one child to two concordant children, , 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, . (d) in the held-out context against the mean over the two supervised contexts, one point per model. Dashed: identity; solid: pooled ratio ; open markers are clipped at the axis limit.
Figure 6: Joint dependency intervention. unit beliefs are added along the belief coordinate after the functional attention; 135 (model, context) cells, means with bootstrap CIs. (a) Parent log-odds shift against 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 , 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 coordinates, colored by context index ; 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 |
|---|---|
| / 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 6) |
Table 1: Architecture configuration common to all 45 models.
| Optimizer | AdamW, , weight decay |
|---|---|
| Learning rate | , constant (no warmup or decay) |
| Gradient clipping | norm , 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 |
Table 2: Optimization and checkpoint selection.
| SparseYour attention | Total | |||
|---|---|---|---|---|
| Layer 1 | 31 | 8 | 13 | 10 |
| Layer 2 | 14 | 7 | 2 | 5 |
Table 3: Sparse attention layer by held-out context.
Figure 7: For each model and context , the centered table , one row per variable and value, is split into the variable plane, the value axis and the interaction plane (Appendix C ). (a) Share of in each subspace, by context . 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 , and the interaction share stays small. (b) Share of the parent-vs-children axis within each plane, for the variable plane and for the interaction plane. Dots are single (model, context) values, and open markers are means with 95% bootstrap CIs. The dashed line at 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. separates the parent from the children, the two children, is the common value shift, the parent’s departure from it, and 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, , at the second token of a triplet, where and 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 , the variable plane and the interaction plane . 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 , the share is the writer’s contribution to the change in the readout from one observed child ( ) to two concordant children ( ), normalized by the two-child readout . The layer-norm bias term is constant and excluded. For each model and context, the five shares sum to , where is the total readout gain. Columns: contexts . 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 . In layer-2 models, the layer-1 writers and contribute almost nothing.
Figure 11: Held-out shortfall in the drop of , by term. For each model and term, is the drop from one observed child to two concordant children ( ) in the held-out context, minus the mean drop over the model’s two supervised contexts. Negative means the held-out context shrinks the term less. Rows are the terms of the exact expansion of over the five writers: five variances and ten covariances . The bottom row is their sum, . Panels: models with sparse attention in layer 1 ( ) and layer 2 ( ). 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 less. In layer-1 models the shortfall is spread over the covariances of and with the MLP writers and . In layer-2 models it sits mainly in and .