📊 statistics

The Information Geometry of Softmax: Probing and Steering

Original authors: Kiho Park, Todd Nief, Yo Joong Choe, Victor Veitch

Published 2026-06-01

📖 5 min read🧠 Deep dive

Original authors: Kiho Park, Todd Nief, Yo Joong Choe, Victor Veitch

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to steer a massive, complex ship (an AI model) to a specific destination (a new concept, like changing a sentence from "he" to "she" or an image of a "dog" to a "cat").

For a long time, researchers have tried to steer these ships by simply pushing the steering wheel in a straight line. They assumed the ocean was flat and uniform, like a giant sheet of ice. If they wanted to go North, they just pushed the wheel North. This is called Euclidean steering.

However, this paper argues that the ocean AI models sail on isn't flat ice. It's actually a curved, bumpy landscape where the "closeness" of two points depends on how similar the ship's behavior is, not just how far apart they look on a map. The authors call this Information Geometry.

Here is the breakdown of their discovery and their new method, Dual Steering, using simple analogies:

1. The Problem: The "Flat Map" Trap

The paper starts by saying that most current methods treat the AI's internal thoughts (representations) as if they exist in a flat, straight-line world.

The Analogy: Imagine you are mixing two colors of paint. If you mix red and blue in a straight line, you get purple. But if you are mixing probabilities (like the chance of rain vs. the chance of sun), the math is different.
The Issue: When researchers push the AI in a straight line to change one thing (e.g., "dog" to "cat"), they accidentally spill paint everywhere else. The AI might suddenly start talking about "bicycles" or "clouds" just because the straight-line push disturbed the balance of the whole system. This is called "off-target leakage."

2. The Solution: The "Curved Map" (Information Geometry)

The authors realized that the AI's "brain" works like a probability machine. When the AI decides what word to say next, it uses a mathematical tool called Softmax to turn numbers into percentages (probabilities).

The Insight: Because the AI deals in probabilities, the space it lives in is curved. In this curved space, the "straightest" path isn't a straight line; it's a curve that respects the rules of probability.
The Analogy: Think of the Earth. If you want to fly from New York to London, the shortest path isn't a straight line through the ground; it's a curve along the surface (a great circle). If you tried to fly in a straight line through the Earth, you'd crash. Similarly, if you try to steer an AI in a "straight line" through its curved probability space, you crash into unwanted concepts.

3. The New Method: "Dual Steering"

The paper introduces a new way to steer the AI called Dual Steering. Instead of pushing the AI in a straight line (Euclidean), they push it along the natural curves of the probability space.

How it works:
- Euclidean Steering (The Old Way): You grab the steering wheel and yank it straight toward "Cat." In doing so, you accidentally knock over the "Dog" and "Bird" statues on the dashboard.
- Dual Steering (The New Way): You navigate using a special map that knows the terrain is curved. You steer the ship along a path that changes the "Dog" to "Cat" without touching the "Bird" or "Bicycle" statues.
The "Dual" Concept: The paper explains that there are two ways to look at the AI's data:
1. Primal Space: The raw numbers the AI sees.
2. Dual Space: The actual probabilities (the behavior) the AI produces.
  The authors found that to change the behavior cleanly, you need to do your steering in the "Dual Space" (the world of probabilities) and then translate that back to the raw numbers.

4. What They Proved

The authors proved mathematically that Dual Steering is the "Goldilocks" method:

It successfully changes the target concept (e.g., makes the AI say "cat" instead of "dog").
It minimizes damage to everything else. It keeps the probability of unrelated things (like "friend" or "bicycle") exactly the same as they were before.

5. Real-World Tests

They tested this on real AI models (like Gemma-3 for text and MetaCLIP for images).

Text Example: When they asked the AI to change a sentence from "He is my..." to "She is my...", the old method accidentally made the AI start talking about "friends" or "uncles" in weird ways. The new method changed "He" to "She" perfectly while keeping the rest of the sentence structure and meaning intact.
Image Example: When changing an image of a "cat" to a "dog," the old method sometimes accidentally added a "dog-cat hybrid" or a "bicycle" into the picture. The new method swapped the cat for a dog cleanly, leaving the background and other objects untouched.

Summary

The paper claims that AI models don't live on a flat map; they live on a curved, probability-based landscape. If you try to steer them with a straight line, you cause chaos. By using Dual Steering, which respects the natural curves of that landscape, you can change the AI's mind about one specific thing without accidentally breaking everything else.

Note: The paper focuses strictly on the geometry of how these models work and how to steer them using linear probes. It does not claim to solve general AI safety issues, nor does it discuss clinical or medical applications. It is purely about making the "steering wheel" of AI more precise and less likely to cause accidental side effects.

Technical Summary: The Information Geometry of Softmax: Probing and Steering

Problem Statement

Current approaches to interpreting and manipulating AI models often rely on the linear representation hypothesis, which posits that high-level concepts correspond to specific directions in a model's vector representation space. While this hypothesis has enabled methods for identifying and steering concepts (e.g., sentiment, truthfulness), existing techniques are frequently brittle and less effective than direct fine-tuning. A primary limitation identified by the authors is the implicit assumption that representation spaces possess a flat, Euclidean geometry. The paper argues that for models utilizing softmax distributions (including Transformers, LLMs, and contrastive models like CLIP), the natural geometry of the representation space is not Euclidean but rather information geometry. Consequently, standard Euclidean operations (like adding a probe vector directly to a representation) may fail to preserve semantic structure or cause unintended "off-target" drift, where probability mass leaks to unrelated concepts.

Methodology

The paper operationalizes the idea that the geometry of representation vectors $\lambda$ should reflect the closeness of the induced probability distributions $P_\lambda$ .

1. Information Geometric Framework

The authors establish that the geometry induced by the Kullback-Leibler (KL) divergence between softmax distributions is a Bregman (dually flat) geometry. This structure introduces a duality between:

Primal Space ( $\Lambda$ ): The space of representation vectors $\lambda$ .
Dual Space ( $\Phi$ ): The space of expected values (mean embeddings) $\phi(\lambda) = \nabla A(\lambda) = \mathbb{E}[\gamma | \lambda]$ , where $A(\lambda)$ is the log-normalizer.

The paper distinguishes between two natural interpolation paths:

Primal Interpolation (e-geodesic): A straight line in the primal space. This minimizes a weighted sum of reverse KL divergences, behaving like an AND operator that emphasizes shared high-probability regions.
Dual Interpolation (m-geodesic): A straight line in the dual space. This minimizes a weighted sum of forward KL divergences, behaving like an OR operator that creates a linear mixture of distributions, preserving the semantic features of both endpoints simultaneously.

2. Dual Steering

The core methodological contribution is Dual Steering, a technique for modifying representations to exhibit a target concept while minimizing changes to off-target concepts.

The Problem with Euclidean Steering: Standard steering adds a probe vector $\beta_W$ (an element of the dual space) directly to the primal representation $\lambda$ . This is a "type mismatch" in non-Euclidean geometry and often leads to off-target drift.
The Dual Steering Solution: Instead of updating $\lambda$ directly, Dual Steering updates the dual coordinate $\phi(\lambda)$ by adding the probe vector: $\phi(\lambda_t) = \phi(\lambda_0) + t\beta_W$ .
Theoretical Guarantee: Under the assumption of concept-factorizable distributions (where the output distribution can be factored into a target concept distribution and an off-target distribution), the authors prove that Dual Steering identifies the minimizer of the KL divergence that modifies the target concept while minimizing the KL divergence of the off-target distribution. This effectively preserves the relative probabilities of unrelated tokens.

3. Practical Implementation

Since the dual space is a bounded convex set, a linear update in the dual space may exit the feasible region. To address this, the authors propose a Regularized Newton Update:

Approximate the change in dual coordinates using the Hessian of the log-normalizer (which corresponds to the covariance matrix of unembedding vectors).
Solve the linear system $\text{Cov}[\gamma | \lambda] \Delta \lambda = \epsilon \beta_W$ to find the corresponding step in the primal space.
Apply Tikhonov regularization ( $\alpha I$ ) to the covariance matrix to handle rank-deficiency (which occurs when the softmax is highly concentrated on few tokens), ensuring the update remains well-defined and pushes the distribution toward higher entropy regions if necessary.
Use the Conjugate Gradient method to efficiently solve the linear system without explicitly inverting the large covariance matrix.

Key Results

The authors evaluate Dual Steering on Gemma-3-4B (LLM) and MetaClip-2 (Vision-Language) across various binary concepts (e.g., verb $\Rightarrow$ third-person, dog $\Rightarrow$ cat, English $\Rightarrow$ French).

Preservation of Off-Target Concepts: Empirical results show that Dual Steering significantly outperforms Euclidean steering in preserving off-target distributions. While Euclidean steering often causes "leakage" (assigning significant probability to unrelated tokens like "friend" when steering "father" $\Rightarrow$ "mother"), Dual Steering maintains the probability mass of counterfactual pairs and neutral tokens.
Robustness Metrics: Dual Steering achieves lower KL divergence and smaller rank differences for off-target distributions compared to Euclidean methods, even as the target concept probability approaches 1.
Interpolation Behavior: Visualizations confirm that Dual Interpolation (m-geodesic) acts as a linear mixture of distributions, whereas Primal Interpolation (e-geodesic) suppresses endpoint-specific outputs in favor of shared features.

Significance and Claims

The paper claims that the non-Euclidean structure of the intrinsic information geometry is critical for connecting geometric manipulations to semantic tasks like steering.

Theoretical Insight: It demonstrates that the linear representation hypothesis must be understood within the context of information geometry (specifically Bregman duality) rather than flat Euclidean geometry to be robust.
Practical Utility: Dual Steering provides a theoretically grounded method for robustly manipulating representations that minimizes unintended side effects, offering a template for improving the reliability of interpretability and control methods.
Scope: While the results are derived for softmax-based models, the authors suggest the high-level idea of exploiting intrinsic geometry is widely applicable. They explicitly note that the work serves as a template for future research into how geometry influences the robustness of control methods, particularly in intermediate layers of models.

The authors remain modest regarding the "probing assumption" (that a linear probe perfectly captures a concept across the entire hyperplane), acknowledging that in practice, this assumption may not strictly hold, yet Dual Steering still empirically outperforms Euclidean methods.