the Shape of Beliefs
How language models represent statistics like mean and variance on curved activation manifolds, with steering experiments extended to geographic data.
About this notebook
This notebook replicates and extends the paper Shape of Belief (Sarfati et al, 2026). It studies how a language model represents the mean and the spread of a number sequence. The code generates Gaussian integer sequences with numpy and converts them to text. It also uses a geographic dataset of 90 cities and their monuments. The model is a Llama 3.2 model (1B or 3B) from the transformers library. The notebook trains probes with scikit-learn and draws charts with matplotlib and altair.
A dropdown selects the model. A form sets the mean values, the sigma values, the sequence count, and the random seed. You submit the form and the notebook collects the model activations. A layer slider picks the model layer for the probe. A top-k dropdown changes how the notebook averages the predicted numbers. Charts then show the predicted number against the count of numbers shown so far. In the geographic section, a dropdown or a text box sets a query monument. The notebook then shows the model's belief as points on a 3D globe with leafmap.
The demos need a GPU. The notebook stops with a warning when no CUDA device is present. It loads the model in float16 and reads the softmax over the 1000 integer tokens. It caches the embeddings and activations on disk with marimo persistent_cache. The batch size adapts to the GPU memory. The steering code adds a vector to the residual stream at one layer with a forward hook. The paper stays unclear about some settings, and the authors report that early layers may behave differently.
