The Information Geometry of Softmax: Probing and Steering

This paper concerns the question of how language models and other 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. We argue that the natural geometry is information geometry, and then show how this interacts with semantic encoding and the linear representation hypothesis. It turns out that the duality structure of information geometry plays a critical role. As an illustrative application, we develop dual steering, a method for robustly steering representations to exhibit a particular concept using linear probes. We formally prove that dual steering optimally modifies the target concept while minimizing changes to off-target concepts. We empirically find that dual steering enhances the controllability and stability of concept manipulation.