The Information Geometry of Softmax: Probing and Steering

Abstract: This paper concerns the question of how AI systems encode semantic structure into the geometric structure of their representation spaces. The motivating observation of this paper 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.

• The paper presents dual steering, which leverages the Bregman geometry of softmax for robust semantic control in AI models.
• It contrasts primal versus dual interpolation, showing that dual operations preserve target and off-target distributions more effectively compared to Euclidean methods.
• Empirical evaluations on language and vision models confirm that dual steering achieves superior off-target preservation and high concept fidelity.

The Information Geometry of Softmax: Probing and Steering

Motivation and Problem Formulation

This paper analyzes the interplay between high-level concept encoding and the information geometry inherent in the representation spaces of softmax-based AI models. The linear representation hypothesis posits that directions in a model’s embedding space correspond to individual concepts, and that such directions can be manipulated with linear probes to edit or interrogate the model’s behavior. However, traditional methods often presuppose a Euclidean structure on these spaces, which empirically leads to brittle and non-robust edits—particularly for tasks such as steering model outputs toward desired behaviors or concepts. The paper argues that the intrinsic geometry governing these tasks is not Euclidean, but rather the Bregman (information) geometry induced by the softmax operation itself.

Bregman Geometry and Duality in the Softmax Parameterization

The authors identify the softmax-induced geometry of the representation space as a dually flat Bregman geometry. The Kullback-Leibler divergence between two softmax distributions corresponds to a Bregman divergence with respect to the log-normalizer. This geometric insight yields a duality between the “primal” parameter space (representation vectors $\lambda$) and the “dual” space (expected value of unembedding vectors, $\phi(\lambda)$). The mapping between these two spaces, governed by the convex function $A$ (softmax log-normalizer), preserves the structure of probability distributions under model outputs.

This duality implies that concepts—captured by linear probes—should naturally live as directions in the dual space, not the primal one. This motivates interpreting and intervening on model representations using information geometry tools.

Primal vs. Dual Interpolation

The distinction between interpolation in the primal and dual spaces is elucidated. Linear interpolation in the primal space yields intersections of high-probability regions from the endpoint distributions, leading to outputs that accentuate shared structure. In contrast, dual interpolation corresponds to mixture distributions that union the high-probability regions, retaining both endpoint semantics.

Figure 1: Primal interpolation amplifies shared modes (“intersection”), dual interpolation preserves endpoint modes (“union”).

Formally, the primal interpolation minimizes a weighted sum of reverse KL divergences, while dual interpolation minimizes a weighted sum of forward KL divergences. This leads to sharp differences: primal interpolation favors selectivity (AND-like composition), whereas dual interpolation favors inclusivity (OR-like composition) of the underlying semantics.

Dual Steering: Theory and Implementation

The authors propose dual steering: intervening on the dual coordinates rather than directly in the representation (primal) space. Given a linear probe $\beta_W$ representing a concept (e.g., “cat” vs “dog”), the standard Euclidean steering (adding $\beta_W$ to $\lambda$) conflates primal and dual structures, often producing undesirable off-target effects—i.e., modifying concepts or model behaviors not aligned with the steering direction.

Dual steering instead adds the concept direction in the dual space: $\phi(\lambda_t) = \phi(\lambda_0) + t\beta_W$. This operation is proved (under mild concept-factorizability conditions) to be optimal under the information geometry: dual steering moves the target concept as desired while minimizing change to the off-target (neutral/counterfactual) distributions.

Figure 2: Dual steering shifts probability directly from base to target concept pairs, avoiding leakage to neutral or unrelated tokens/images.

Theoretically, this is formalized as a constrained minimization of KL divergence, with dual steering exactly tracing the minimal-KL path given a fixed probe hyperplane. No such guarantee exists for Euclidean steering, which can “leak” probability mass to neutral or semantically unrelated outputs, a phenomenon empirically observed.

Figure 3: Ideal steering modifies only the target concept, keeping off-target distributions unchanged.

Robustness, Practical Challenges, and Regularized Newton Updates

Dual steering faces practical challenges: the dual space is constrained (dual points must stay within the convex hull of unembedding vectors), and the covariance (Hessian) in the softmax can be rank-deficient when probability is highly concentrated. The authors implement a regularized Newton method to trace feasible paths in the primal corresponding to linear movement in the dual, using Hessian regularization to ensure invertibility, and iteratively increase the variance of the target concept.

Figure 4: Dual steering maintains higher cosine alignment in dual space with the concept direction than Euclidean steering, indicating true semantic steering.

This method is robust: it adapts in cases where standard steering would fail due to low entropy or concept entanglement.

Empirical Evaluation

Extensive experiments on language (Gemma-3-4B) and vision-LLMs (MetaCLIP-2) demonstrate the effectiveness of dual steering. The methodology generalizes to multiple binary concept types (e.g., verb forms, language, object attributes), with evaluation conducted both for target concept activation and preservation of off-target (counterfactual) probabilities.

Key empirical findings:

• Dual steering yields strictly better preservation of off-target distributions across all tested metrics than Euclidean steering.
• Dual steering avoids probability “leakage” during intermediate steps, maintains a higher total probability mass on valid counterfactual pairs, and achieves lower KL and rank differences for off-target distributions.
• Effects are robust to the choice of linear probe (primal or dual mean difference).

  Figure 5: Across models and concept types, dual steering preserves off-target distributions (robustness metrics) while effectively controlling the target concept probability.

Implications and Future Directions

Theoretical and empirical analysis supports the main claim: Euclidean approaches to conceptual intervention are fundamentally limited due to ignorance of the true (information-theoretic) geometry of model representations. By leveraging the induced Bregman geometry, dual steering achieves both interpretable and robust control of semantic attributes in generative models.

This framework has implications for:

• Model interpretability: Enables targeted interventions with rigorous guarantees regarding semantic invariance outside the steered direction.
• Control and editing of LM and VLM behavior: Brings improved controllability and stability, critical for high-assurance applications.
• Probe construction and evaluation: Reveals the necessity of geometric compatibility between probes and model parameterizations.

Limitations include the computational complexity of the steering (especially with large vocabularies), and dependence on the richness of the identified probe (breakdowns occur if probe hyperplanes entangle concepts off the training manifold). The paper suggests future work in extending these geometric insights to intermediate representations beyond the output softmax layer and non-dually flat architectures.

Conclusion

This work rigorously establishes the inadequacy of Euclidean assumptions for semantic editing in softmax-based models, deriving and validating a principled framework—dual steering—rooted in information geometry. The results establish dual steering as a method that optimally edits target concepts while provably preserving off-target behaviors, advancing both the theory and practice of model interpretability and control in high-dimensional generative systems.