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 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.
Lay Summary
AI systems, such as language and vision models, represent concepts internally as directions in a high-dimensional space. This structure has enabled recent work on steering model behavior by manipulating these directions. For example, models can be shifted from producing male-coded outputs to female-coded ones. However, existing methods are often unstable. Steering a target concept frequently leads to unintended changes in many unrelated concepts. We argue that this issue arises from a mistaken geometric assumption. Most methods treat the model's internal space as flat Euclidean space, while the natural geometry of these models is shaped by probability and information theory. By accounting for this geometry, we derive a new method called dual steering, which modifies the intended concept while minimizing interference with unrelated concepts. We validate dual steering on open-source language and vision models, consistently outperforming standard approaches across a wide range of concepts.