The Information Geometry of Softmax: Probing and Steering

Kiho Park; Todd Nief; Yo Joong Choe; Victor Veitch

cs.LGFeb 17, 2026

The Information Geometry of Softmax: Probing and Steering

Authors: Kiho Park, Todd Nief, Yo Joong Choe, Victor Veitch

Organizations: University of Chicago · INSEAD

Abstract

This paper concerns the question of how AI systems encode semantic structure into the geometric structure of their representation spaces. The motivating observation is that the natural geometry of these representation spaces should reflect the way models use representations to produce behavior. We focus on the important special case of representations that define softmax distributions. In this case, we argue that the natural geometry is information geometry. Our focus is on the role of information geometry on semantic encoding and the linear representation hypothesis. As an illustrative application, we develop "dual steering", a method for robustly steering representations to exhibit a particular concept using linear probes. We prove that dual steering optimally modifies the target concept while minimizing changes to off-target concepts. Empirically, we find that dual steering enhances the controllability and stability of concept manipulation.

Figures & tables

Figure 1: Dual steering (bottom) effectively modifies the target concept (e.g., $\texttt{verb}\Rightarrow\texttt{third-person}$ or $\texttt{cat}\Rightarrow\texttt{dog}$ ) while preserving off-target distributions (e.g., $P(\text{``maintain''})+P(\text{``maintains''})$ or $P(\text{``cat + bicycle''})+P(\text{``dog + bicycle''})$ ), whereas Euclidean steering (top) fails to maintain off-target distributions despite reaching the target concept probability. Left: Token probability changes in Gemma-3-4B when steering the context “Author gives an insight into what it costs US taxpayers to build and” using a linear probe for $\texttt{verb}\Rightarrow\texttt{third-person}$ . Euclidean steering leaks significant mass to off-target tokens (e.g., “to”) during intermediate steps, whereas dual steering directly shifts probability from base tokens (e.g., “maintain”, “operate”) to target tokens (e.g., “maintains”, “operates”). Center & Right: Steering MetaCLIP-2 on the context “a photo of one cat” for the concept $\texttt{cat}\Rightarrow\texttt{dog}$ . Dual steering transfers probability from base images (e.g., “cat”, “cat + bicycle”) directly to targets (e.g., “dog”, “dog + bicycle”). In contrast, Euclidean steering unintentionally promotes the off-target “cat + dog” image (green frame in the right column), which becomes the Top-1 result during intermediate steps. In the probability plots, Top- $K$ tokens (LLM) or images (CLIP) are shown explicitly, with the remainder grouped as “others.”

Figure 2: Primal interpolation emphasizes the shared structure (intersection) of distributions, whereas dual interpolation results in a linear mixture. We visualize output probability changes along interpolation paths between two context embeddings $\lambda(x_{0})$ and $\lambda(x_{1})$ . The dual interpolation (right, $m$ -geodesic: $\phi_{t}=(1-t)\phi(\lambda_{0})+t\phi(\lambda_{1})$ ) corresponds to a weighted average of the endpoint distributions. In contrast, near the midpoint, the primal interpolation (left, $e$ -geodesic: $\lambda_{t}=(1-t)\lambda_{0}+t\lambda_{1}$ ) upweights shared components (e.g., “the”, “called” in LLM or “black-and-white dog” in CLIP), while suppressing endpoint-specific outputs (e.g., “Paris” vs. “Berlin,” or “black dog” vs. “white dog”). Top- $K$ tokens (LLM) or images (CLIP) are shown explicitly, with the remainder grouped as “others.”

Figure 3: Illustration of an ideal steering process. The blue, red, and gray boxes represent partitions of the total probability mass. The goal is to shift only the boundary (purple bar) governing the concept distribution

P^{W}_{\lambda}

, while preserving the vertical partitions defined by the off-target distribution

P^{Z}_{\lambda}

Figure 4: Illustration of Euclidean versus dual steering in the primal and dual spaces. Euclidean steering (red) drifts into a region where the unrelated token “friend” gains more probability mass, whereas dual steering (blue) remains within the region where the off-target distribution is preserved. The background color indicates the top-1 token at each point, e.g., green for “father” and purple for “friend”.

Figure 5: Dual steering (red) consistently preserves off-target distributions better than Euclidean steering (blue), while both boost the target concept probability. We plot three robustness metrics (y-axes) against the target concept probability (x-axis) achieved via steering along the dual mean difference. As the target concept probability approaches 1 (moving right), Euclidean steering degrades the off-target distribution, whereas dual steering maintains it. Columns: Tasks include LLM steering for $\texttt{English}\Rightarrow\texttt{French}$ (left), CLIP on synthetic objects for $\texttt{yellow}\Rightarrow\texttt{green}$ (middle), and CLIP on real images (COCO) for $\texttt{carrot}\Rightarrow\texttt{broccoli}$ (right). Rows: The top row shows the total probability mass on counterfactual pairs (constant is better). The middle and bottom rows show the KL divergence and rank difference of off-target distributions (lower is better). Lines represent the mean, and shading indicates the standard error of the mean (SEM) across test contexts.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Top- $K$	Step 0	Step 5	Step 10	Step 50	Step 100
1	0.2031	0.2031	0.2031	0.2926	0.3321
20	0.0120	0.0133	0.0151	0.0006	0.0000
2000	0.0001	0.0002	0.0002	0.0000	0.0000
20000	0.0000	0.0000	0.0000	0.0000	0.0000
200000	0.0000	0.0000	0.0000	0.0000	0.0000

Appendix

Table 1: Frobenius norm of the difference between the top-

K

approximation and the full Fisher information matrix along the steering trajectory. The approximation error approaches zero at

K=20{,}000

Figure 6: Cosine similarity between the concept direction and the change in dual coordinates during Euclidean (green and blue) and dual steering (orange and red). Dual steering maintains a higher cosine similarity than Euclidean steering, indicating that dual steering effectively moves along the concept direction in the dual space. In particular, the increase in cosine similarity during the first few steps suggests that the regularized Newton method initially adjusts the trajectory to increase the local variance of the concept, facilitating effective dual steering along the concept direction. Lines represent the mean, and shading indicates the standard error of the mean (SEM) across test contexts.

Figure 7: Histogram of projections of test set context embeddings onto primal and dual mean differences $\beta_{W}^{\top}\lambda_{i}/\|\beta_{W}\|_{2}$ . Both primal and dual mean differences effectively separate the base and target context embeddings, indicating their utility as linear probes for the respective concepts.

Figure 8: Projection of Euclidean (green) and dual (purple) steering paths onto the linear probe ( $\beta_{W}^{\top}\lambda_{t}/\|\beta_{W}\|_{2}$ ) against the logit of the target concept probability ( $\mathrm{logit}\,P(W=1~|~\lambda_{t})$ ), using primal or dual mean differences as the linear probe. Blue and red dots represent the projections and logits of context embeddings from the base and target test groups, respectively. Dual steering typically yields a lower concept probability for any given projection value compared to the test set contexts. This suggests that the probing assumption in Theorem 3.1 is not strictly satisfied in practice.

Figure 9: Total probability mass on counterfactual pairs during Euclidean and dual steering across all experiments. Dual steering consistently preserves a higher probability mass on counterfactual pairs during intermediate steps; this suggests greater robustness, as it avoids “leakage” to neutral tokens more effectively than Euclidean steering. Lines represent the mean, and shading indicates the standard error of the mean (SEM) across test contexts.

Figure 10: KL divergence of off-target distributions during Euclidean and dual steering across all experiments. Dual steering results in lower KL divergence values, indicating better preservation of off-target concepts compared to Euclidean steering. Lines represent the mean, and shading indicates the standard error of the mean (SEM) across test contexts.

Figure 11: Rank differences of off-target distributions during Euclidean and dual steering across all experiments. Dual steering results in lower rank differences, indicating better preservation of off-target concepts compared to Euclidean steering. Lines represent the mean, and shading indicates the standard error of the mean (SEM) across test contexts.

Representation Geometry in Language Models

Representation geometry in language models studies how embeddings and hidden states organize concepts, computation, and behavior.

See all 417 papers

Manifold Steering Reveals the Shared Geometry of Neural Network Representation and Behavior
May 6, 2026Daniel Wurgaft, Can Rager, Matthew Kowal +13Neural Representation Geometry Manifold Learning
Stanford University · University College London · Northeastern University +2
The Information Geometry of Representation Formation
Oct 25, 2026Sanjukta KrishnagopalNeural Representation Geometry Information Geometry
Department of Computer Science, University of California, Santa Barbara
The Cylindrical Representation Hypothesis for Language Model Steering
May 3, 2026Lang Gao, Jinghui Zhang, Wei Liu +7Language Model Steering LLM Interpretability
Department of Natural Language Processing, Mohamed bin Zayed University of Artificial Intelligence, Abu Dhabi, UAE · Department of Computer Science, National University of Singapore, Singapore

May 6, 2026J/K move · Enter open · S save

Manifold Steering Reveals the Shared Geometry of Neural Network Representation and Behavior

Daniel Wurgaft, Can Rager, Matthew Kowal, Vasudev Shyam+12

Stanford University · University College London · Northeastern University+2

Neural Representation Geometry Manifold Learning Causal Interventions in Language Models Activation Steering

Open paper AskPDF arXiv

Neural representations carry rich geometric structure; but does that structure causally shape behavior? To address this question, we intervene along paths through activation space defined by different geometries, and measure the behavioral trajectories they induce. In particular, we test whether interventions that respect the geometry of activation space will yield behaviors close to those the model exhibits naturally. Concretely, we first fit an activation manifold

M_h

to representations and a behavior manifold

M_y

to output probability distributions. We then test the link

M_h \leftrightarrow M_y

via interventions: we find that steering along

M_h

, which we term manifold steering, yields behavioral trajectories that follow

M_y

, while linear steering -- which assumes a Euclidean geometry -- cuts through off-manifold regions and hence produces unnatural outputs. Moreover, optimizing interventions in activation space to produce paths along

M_y

recovers activation trajectories that trace the curvature of

M_h

. We demonstrate this bidirectional relationship between the geometry of representation and behavior across tasks and modalities. In language models, we use reasoning tasks with cyclic and sequential geometries as well as in-context learning tasks with more complex graph geometries. In a video world model, we use a task with geometry corresponding to physical dynamics. Overall, our work shows that geometry in neural representation is not merely incidental, but is in fact the proper object for enabling principled control via intervention on internals. This recasts the core problem of steering from finding the right direction to finding the right geometry.

Abstract

Representation Geometry in Language Models

Explore similar work

Manifold Steering Reveals the Shared Geometry of Neural Network Representation and Behavior

The Information Geometry of Representation Formation

The Cylindrical Representation Hypothesis for Language Model Steering