Geometric Scaling of Bayesian Inference in LLMs
Naman Aggarwal, Siddhartha R. Dalal, Vishal Misra
TL;DR
This work tests whether the geometric substrates that enable Bayesian inference in wind-tunnel experiments—namely entropy-ordered value manifolds, orthogonal hypothesis frames, and entropy-driven attention refinement—persist in production-scale LLMs. Across Pythia-410M, Phi-2, Llama-3, and Mistral, the authors show that last-layer value representations align along a dominant entropy axis, and domain-restricted prompts collapse the geometry to wind-tunnel-like low-dimensional manifolds. In-context learning experiments (SULA) reveal that model states move along the entropy axis as evidence is accumulated, indicating inference-time engagement with the geometric substrate. Causal interventions perturbing the entropy axis disrupt local geometry but do not proportionally degrade Bayesian-like calibration, suggesting the geometry is a privileged readout of distributed uncertainty rather than a single bottleneck. Overall, the paper demonstrates that Bayesian-like geometry is preserved at scale, with static representations (value manifolds and key frames) universal across architectures, while dynamic focusing hinges on architectural routing and training data quality. These findings offer a principled geometric lens for interpreting uncertainty processing in modern transformers and inform future work on causal probing and scalable interpretability.
Abstract
Recent work has shown that small transformers trained in controlled "wind-tunnel'' settings can implement exact Bayesian inference, and that their training dynamics produce a geometric substrate -- low-dimensional value manifolds and progressively orthogonal keys -- that encodes posterior structure. We investigate whether this geometric signature persists in production-grade language models. Across Pythia, Phi-2, Llama-3, and Mistral families, we find that last-layer value representations organize along a single dominant axis whose position strongly correlates with predictive entropy, and that domain-restricted prompts collapse this structure into the same low-dimensional manifolds observed in synthetic settings. To probe the role of this geometry, we perform targeted interventions on the entropy-aligned axis of Pythia-410M during in-context learning. Removing or perturbing this axis selectively disrupts the local uncertainty geometry, whereas matched random-axis interventions leave it intact. However, these single-layer manipulations do not produce proportionally specific degradation in Bayesian-like behavior, indicating that the geometry is a privileged readout of uncertainty rather than a singular computational bottleneck. Taken together, our results show that modern language models preserve the geometric substrate that enables Bayesian inference in wind tunnels, and organize their approximate Bayesian updates along this substrate.