The Information Geometry of Softmax: Probing and Steering

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

REVIEW 4 major objections 4 minor 1 cited by

The Information Geometry of Softmax: Probing and Steering

T0 review · 4 major / 4 minor · reviewed 2026-08-02 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗ Review JSON

Pith's one-line read Dual steering moves a softmax model's concept probability with minimal off-target change, and its formula is exact.

desk verdict A clean conditional result: the geometry story is sound, the theorem is correctly derived, but the paper's own Figure 7 shows the key assumption fails in practice, so the abstract overstates what is proven. read the letter →

arxiv 2602.15293 v2 pith:LAVQBX6P submitted 2026-02-17 cs.LG cs.AIcs.CLstat.ML

Kiho Park , Todd Nief , Yo Joong Choe , Victor Veitch This is my paper · ORCID

classification cs.LGcs.AIcs.CLstat.ML

keywords informationgeometrysoftmaxlinearrepresentationsteeringKLdivergenceBregmanprobeinterpretability

verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

This paper argues that the geometry of softmax representation spaces is not Euclidean but information-geometric, and that this distinction matters for steering model behavior. The authors prove a specific claim: to change a binary concept (like dog to cat) while disturbing everything else as little as possible, you should update the representation in the dual coordinate — the expected unembedding vector — rather than by directly adding the probe vector in the raw space. They call this dual steering, show it is the solution of a constrained KL-minimization problem, and verify empirically on language and image models that it preserves off-target probabilities and rankings better than standard Euclidean steering. If correct, this gives a principled geometric explanation for why additive steering leaks probability mass and provides a robust recipe for concept manipulation.

What carries the argument

The central object is the dual coordinate map φ(λ)=∇A(λ)=E[γ|λ], the gradient of the log-normalizer (equivalently, the mean of the unembedding vectors under the softmax). The key identity is that KL divergence between softmax distributions equals the Bregman divergence of A, which turns steering into a projection problem: moving along βW in the dual space is the forward-KL projection onto the probe hyperplane, whereas adding βW in the primal space is an unprincipled Euclidean step. The proof decomposes the KL objective into an on-target term and an off-target term via a concept-factorization of the output distribution.

What would settle it

Measure, along an actual dual-steering path on a large language model, the concept logit logit P(W=1|λt) versus the probe projection βWᵀλt/||βW||, and compare to test contexts with the same projection. If dual-steering paths systematically give lower concept logits at matched projections than test contexts (as the paper's Figure 7 shows), then the probe-constancy assumption is violated and the proven off-target optimality no longer holds — this observation would settle whether the theorem's premise is met.

Watch

Extended reading notes

Core claim

The central claim is that the softmax family P(γ|λ)=exp(λᵀγ−A(λ)) carries a Bregman (dually flat) geometry, with primal coordinates λ and dual coordinates φ(λ)=E[γ|λ]. Within this geometry, the correct way to steer a representation onto the hyperplane defined by a linear probe βW (i.e., to set βWᵀλ=c) is to move the dual coordinate along the probe direction: φ(λ̂)=φ(λ0)+tβW. The paper proves (Theorem 3) that any minimizer of the forward KL divergence D_KL(Pλ0||Pλ) over that hyperplane has exactly this form, and that if the distribution is factorizable into a concept part and an off-target part, the same minimizer also minimizes the KL divergence of the off-target distribution. Thus dual stee

Load-bearing premise

The theorem assumes the linear probe exactly satisfies P(W=1|λ)=σ(βWᵀλ+bW) for every λ on the steering hyperplane (so the concept probability is constant there) and that the full output distribution factorizes into a concept part and an off-target part; the paper itself demonstrates in Section C.2 that this constancy fails on real models.

Editorial extensions

If this is right

Euclidean steering (adding the probe vector to λ) is a type error whenever the primal and dual spaces differ; the paper predicts it will leak probability mass to off-target outputs, which the experiments confirm.
Dual steering provides a concrete, implementable algorithm (regularized Newton updates on the covariance matrix) that works even when the Hessian is rank-deficient, because regularization nudges the distribution toward higher entropy.
The theory says dual steering preserves not just probabilities but also the ranking of off-target components, which is why it is more reliable than Euclidean steering for practical use.
The analysis explains when Euclidean steering can work: when the total counterfactual probability mass stays roughly constant along the path, so the on-target term in the KL decomposition is effectively fixed.
The geometry of softmax layers is shown to be dually flat, implying that interpolation between representations has two natural paths with distinct semantics — an AND-like primal interpolation and an OR-like dual interpolation.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

The same dual-coordinate argument likely extends to other exponential-family output layers, where the dual map is the expectation parameter; a testable extension is to apply dual steering to attention logits or mixture-of-experts gates.
If the probe-constancy assumption fails (as the paper's own Figure 7 suggests it does in practice), then dual steering may trade incomplete target steering against off-target distortion; a natural follow-up is to measure and correct for the deviation of the concept logit along the steering path.
The rank-preservation property of dual steering could be exploited beyond steering, e.g., in controlled generation or debiasing where maintaining a stable ordering of alternatives matters.
The paper focuses on the output softmax layer, but a direct empirical test for intermediate layers would assess whether the same dual-coordinate update remains optimal when representations are transformed by later layers.

Share X Bluesky LinkedIn Reddit HN

Signed reviews

No signed human review yet.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Desk Editor's Note

read the letter

You should know this paper has a genuine theoretical core and an honest disclosure of its own soft spot, but the abstract sells the guarantee too hard. The Bregman-duality treatment of softmax is a nice unification, and the primal vs. dual interpolation experiments are intuitive. Theorem 3 is correctly derived: the KL projection onto a probe hyperplane yields dual steering, and with concept-factorizability it minimizes off-target KL. The proof checks out. The regularized Newton implementation is pragmatic, and the empirical results consistently show dual steering beating Euclidean steering on off-target preservation. Credit also goes to Section 5.3, which openly says the probe-constancy assumption is violated and points to C.2 and Figure 7 — that is real intellectual honesty.

The soft spot is load-bearing, not cosmetic. Theorem 3's off-target optimality (3.4) requires the concept probability to be constant on the hyperplane. Figure 7 shows dual steering paths have lower logits than test contexts at the same projection, so the first term in (A.30) is not constant on the actual steered path. That means the proven guarantee does not apply to the experiments as run. The headline comparisons in Figure 4 align paths by measured target probability, not by probe projection, so the plotted curves are not even on the hyperplane the theorem talks about. Concept-factorizability is also assumed without direct verification. The implemented method uses a tuned regularized-Newton approximation without formal convergence guarantees. These are addressable, but the abstract's claim that dual steering is 'proven' to minimize off-target changes is too strong for what is actually established.

Who gets value from this: anyone working on interpretability, activation steering, or the linear representation hypothesis. The framework is clear, the method is practical, and the empirical evidence suggests it works even where the theory is conditional. I would cite it. It deserves a serious referee, but with a request to align the claims with the assumptions — either soften the abstract or add experiments that probe the invariance assumption directly, e.g., test whether off-target preservation holds conditional on projection value rather than on achieved target probability.

Send it to peer review. I would take it as a conditional: the core is sound, and the gap is known and disclosed. The right fix is disciplined claims plus a bit more evidence that the assumption approximately holds in the tested regime.

Referee Report

4 major / 4 minor

Summary. The paper develops an information-geometric view of softmax-based representation spaces, arguing that the natural geometry is Bregman/dually flat. It introduces primal and dual interpolation paths, and on that basis proposes "dual steering": instead of adding the probe direction β_W in the primal space as in Euclidean steering, one moves the dual coordinate along β_W. The central result, Theorem 3, states that the KL minimizer over the hyperplane Λ_W(c)={λ : β_W^⊤ λ=c} has dual coordinate φ(λ̂)=φ(λ_0)+tβ_W, and that, under a concept-factorizability condition, the same minimizer also minimizes the off-target KL divergence D_KL(P^Z_{λ_0}||P^Z_λ). The method is implemented via a regularized Newton update on the covariance (Hessian) and tested on Gemma-3-4B and MetaClip-2, with metrics measuring target-concept probability, counterfactual mass, off-target KL, and rank stability. The empirical section reports that dual steering preserves off-target distributions better than Euclidean steering across several binary concepts.

Significance. If the theorem applied to the experimental regime, this would be an important contribution: it gives a principled geometric reason for the failure of Euclidean steering and provides a concrete, falsifiable alternative. The proof of Theorem 3 is internally coherent, and the KL decomposition in Appendix A.2 checks out. Proposition 1 is also correctly proven. The paper usefully distinguishes primal and dual interpolation and supports the distinction with clear figures. The main weakness is not the internal derivation but the gap between the theorem's assumptions and the tested setting: the paper's own Appendix C.2 shows that the exact linear-probe assumption is violated on the tested models, and concept-factorizability is never verified. As a result, the unqualified claim in the abstract that dual steering is proven to minimize off-target changes overstates what is established. The empirical results are suggestive, but the theoretical guarantee currently covers an idealized regime rather than the actual implementation.

major comments (4)

[§3.1 (Theorem 3), §A.2 (Eq. A.30), §C.2 (Fig. 7)] The proof of (3.4) relies on the assumption that, on the hyperplane Λ_W(c), the concept distribution P^W_λ is constant, so that the first term in Eq. (A.30) is independent of λ. Section C.2 and Figure 7 directly test this premise and find that dual-steering paths yield lower logit P(W=1|λ) than test contexts at the same β_W^⊤λ/||β_W||. Thus P^W_λ is not a function of the probe projection, the first term in (A.30) is not constant, and the constrained KL minimizer is not equivalent to the minimizer of D_KL(P^Z_{λ_0}||P^Z_λ). Section 5.3 acknowledges this, but the abstract and Section 3.1 state the optimality result without this qualification. The theorem is correct under its assumptions; the load-bearing issue is that the experiments are run where the assumptions are known to fail.
[§5.2 (Figure 4) vs. §3.1 (Theorem 3)] The empirical comparison in Figure 4 aligns steering paths by the measured target-concept probability P^W_{λ_t}(1) on the x-axis, not by the probe projection β_W^⊤λ. Theorem 3 concerns the hyperplane Λ_W(c) where β_W^⊤λ is fixed. These are different level sets when the probe is miscalibrated, and Figure 7 shows that they do differ on the tested models. Consequently, even the qualitative comparison in Figure 4 is not directly covered by Theorem 3. The authors should either align paths by β_W^⊤λ in the main experiments or explicitly state that the theorem applies only to an idealized probe and that the empirical comparison is heuristic.
[§4.2 (Algorithm 1), §5.2] The implemented method is not the exact dual-coordinate update φ(λ_t)=φ(λ_0)+tβ_W. Algorithm 1 solves (Σ_t+αI)v=β_W and takes a normalized step; the actual dual displacement is approximately Σ_t(Σ_t+αI)^{-1}β_W, which is not colinear with β_W when β_W is in the approximate null space of Σ_t. Section C.1's Figure 5 confirms the cosine similarity between the dual step and β_W is substantially below 1, especially initially. Moreover, for LLMs the covariance is truncated to Top-K tokens (Appendix B.2), changing the Hessian. No error bound or convergence guarantee is given for this approximation. The theorem does not establish the optimality of the regularized, truncated path actually evaluated.
[Definition 2, §5.1, Eq. (3.4)] The off-target optimality result (3.4) requires concept-factorizability of P_λ for all λ in Λ_W(c) ∪ {λ_0}. This is a strong structural assumption: it requires the output space to split into counterfactual pairs with P_λ(y_i^w)=P^W_λ(w)P^Z_λ(z_i). The paper never tests this assumption. In Appendix B.2 the metrics define P^Z_λ(z_i)=P_λ(y_i^0)+P_λ(y_i^1), but that identity does not by itself establish the factorization used in the proof. Without factorizability, the decomposition (A.30) is not valid, and the interpretation of the KL term as an off-target divergence is not justified. The authors should either provide empirical evidence for factorizability on their counterfactual pairs or present (3.4) as a conditional statement that is not directly applicable to the reported experiments.

minor comments (4)

[§5.1] The rank-difference metric is defined as a weighted sum of absolute differences of inverse ranks, but the notation 'ranking_{P^Z_{λ_t}}(z)' is not formally defined. Please define how ties and out-of-vocabulary tokens are handled.
[§B.2] The Top-K covariance truncation is described as using 'e.g., 20,000' tokens for LLMs, but no K is reported for the CLIP experiments or in the main figures. Since truncation affects the Hessian and hence the steering path, the value of K should be reported per experiment.
[§3.2, Eq. (3.5)] The decomposition of reverse KL in Eq. (3.5) is stated without proof and without explicitly defining the domain of the sum over z. For consistency with Appendix A.2, a short derivation or reference would help.
[Throughout] The model name appears as both 'MetaClip-2' and 'MetaClip' in different places; please standardize. Also, Figure 2's caption is dense and would benefit from a legend explaining the colored lines in the probability plots.

Circularity Check

0 steps flagged · score 0.0 of 10

No material circularity: Theorem 3 follows from Bregman projection and an explicit KL decomposition under stated assumptions; the paper's own caveats concern external validity, not circularity.

full rationale

The derivation chain is self-contained. The key result, Theorem 3, is proved directly in Appendix A.2: Eq. (3.3) follows from the first-order optimality condition for minimizing D_KL(P_λ0 || P_λ) over the affine set {λ : β_W^T λ = c}, and the stationarity equation (A.21) gives φ(λ̂) − φ(λ0) parallel to β_W. Eq. (3.4) follows from the explicit KL decomposition (A.29)–(A.30) under the stated hypotheses: probe constancy (3.1) makes the first term constant on the hyperplane, and concept-factorizability (Definition 2) justifies the decomposition. No fitted parameter is renamed as a prediction: the probe β_W is an input supplied by assumption (3.1), and the regularization α and probe directions used in experiments are implementation/training choices that do not enter the theorem's logic. The self-citations ([PCV24], [Par+25]) are background for the linear representation hypothesis and related steering work, not load-bearing substitutes for the proof; the proof itself uses standard external results ([Ama16], [Ban+05]) and direct calculation. The paper's own Appendix C.2 and Figure 7 candidly state that the probe-constancy premise is not strictly satisfied in practice; that is a limitation on the external validity of the theorem, not a circularity, because violating an assumption is not equivalent to defining the conclusion into the premises. No load-bearing step reduces to its own input by construction.

Assumptions & free parameters 4 free parameters · 6 assumptions · 0 invented entities

The central theoretical results rest on the softmax exponential-family assumption, the exactness of a linear probe, and concept-factorizability. The first is a domain assumption, the latter two are ad-hoc idealizations that the paper itself shows are violated by real models. The practical method adds tuned hyperparameters (α, η, T, Top-K) and fitted probe directions. No new physical or independently testable entities are introduced; the counterfactual-pair partition ZW is a bookkeeping construct rather than an invented entity.

free parameters (4)

Regularization parameter α = 5e-3
Tuned for both LLM and CLIP experiments to make the Hessian invertible; directly shapes the approximate dual-steering path and off-target preservation.
Probe direction βW (Primal MD / Dual MD) = Mean difference of training context embeddings
Estimated from training data; the theorem treats βW as given, but the experimental steering directions are fitted vectors.
Step size η and max iterations T = Not reported
Algorithm 1 hyperparameters; the paper says 'sufficient number of steps' without reporting values, so the practical path is under-specified.
Top-K covariance truncation = 20000 tokens
Chosen to approximate the LLM covariance matrix; affects the regularized Newton step and thus the empirical steering path.

assumptions (6)

domain assumption Outputs follow the softmax form P(γ=γy|λ)=exp(λ^T γy − A(λ)) over a finite set of fixed unembedding vectors.
The entire geometry and steering theory are built on this exponential-family form (Eq. 1.1).
standard math The log-normalizer A is sufficiently convex that ∇A defines a dual map and ∇A* is its inverse on the image.
Invoked in Section 2.1 to define the dual coordinate system and used throughout the proof of Theorem 3.
ad hoc to paper An exact linear probe exists: P(W=1|λ)=σ(βW^T λ + bW) for all λ.
Equation (3.1); needed so P^W_λ is constant on the steering hyperplane. The paper later reports this is not satisfied in practice (Section C.2).
ad hoc to paper Concept-factorizability of Pλ with respect to W for all λ in ΛW(c)∪{λ0}.
Definition 2 and the second part of Theorem 3; required to decompose the KL divergence into on-target and off-target terms.
standard math Existence of the KL minimizer on the hyperplane ΛW(c).
The theorem is conditional on the minimizer existing; the paper does not prove existence in the non-strictly-convex boundary cases relevant to rank-deficient Hessians.
domain assumption For CLIP, the temperature scaling from the final training step is correctly re-applied before defining the softmax distribution.
Section B.1.2; necessary for the CLIP embeddings to define a valid softmax distribution of the required form.

how reviews work

0 comments

Cite this review

Pith. "Pith review of The Information Geometry of Softmax: Probing and Steering." pith.science (2026). https://pith.science/paper/LAVQBX6P

@misc{pith2026260215293,
  author       = {Pith},
  title        = {Pith review of: The Information Geometry of Softmax: Probing and Steering},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/LAVQBX6P}},
  note         = {Machine review of arXiv:2602.15293}
}

read the original 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

Figures reproduced from arXiv: 2602.15293 by the authors.

**Figure 2.** Primal interpolation emphasizes the shared structure (intersection) of distributions, [PITH_FULL_IMAGE:figures/full_fig_p004_2.png] view at source ↗

**Figure 3.** Ideal steering modifies the concept distribution [PITH_FULL_IMAGE:figures/full_fig_p006_3.png] view at source ↗

Figures from the paper (7 more)

**Figure 4.** Figure 4: Dual steering (red) consistently preserves off-target distributions better than Euclidean [PITH_FULL_IMAGE:figures/full_fig_p011_4.png]

**Figure 5.** Figure 5: 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 effectivel…

**Figure 6.** Figure 6: Histogram of projections of test set context embeddings onto primal and dual mean differences β ⊤ W λi/∥βW ∥2 . Both primal and dual mean differences effectively separate the base and target context embeddings, indicating their utility as linear probes for the respecti…

**Figure 7.** Figure 7: Projection of Euclidean (green) and dual (purple) steering paths onto the linear probe (β ⊤ W λt/∥βW ∥2 ) against the logit of the target concept probability (logit P(W = 1 | λt )), using primal or dual mean differences as the linear probe. Blue and red dots represent …

**Figure 8.** Figure 8: 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 “l…

**Figure 9.** Figure 9: 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…

**Figure 10.** Figure 10: 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…

Discussion (0). Continue with ORCID to comment.

Forward citations

Cited by 1 Pith paper

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score. Full citation record

ToxiREX: A Dataset on Toxic REasoning in ConteXt
cs.CL 2026-06 unverdicted novelty 6.0 of 10

ToxiREX is a new dataset of 128k Reddit comments in six languages with hierarchical annotations for implicit toxicity in conversational context based on an existing reasoning schema.

Reference graph

Works this paper leans on

2 extracted references · 1 linked inside Pith · cited by 1 Pith paper

[1]

Understanding intermediate layers using linear classifier probes

[AB16] G. Alain and Y. Bengio. “Understanding intermediate layers using linear classifier probes”.arXiv preprint arXiv:1610.01644(2016) (cit. on p. 5). [Ama16] S.-i. Amari.Information geometry and its applications. 2016 (cit. on pp. 2, 4, 12, 19). [AN00] S.-i. Amari and H. Nagaoka.Methods of information geometry. 2000 (cit. on p. 2). [Ard+24] A. Arditi, O...

arXiv 2016
[2025]

The geometry of refusal in large language models: concept cones and representational independence

2025 (cit. on p. 1). [Wol+25] T . Wollschläger, J. Elstner, S. Geisler, V . Cohen-Addad, S. Günnemann, and J. Gasteiger. “The geometry of refusal in large language models: concept cones and representational independence”.arXiv preprint arXiv:2502.17420(2025) (cit. on p. 12). [Wu+25] Z. Wu, A. Arora, A. Geiger, Z. Wang, J. Huang, D. Jurafsky, C. D. Manning...

arXiv 2025

Pith tools

Reviewed August 2, 2026 · model on record in the stance chip above.

[1] [1]

Understanding intermediate layers using linear classifier probes

[AB16] G. Alain and Y. Bengio. “Understanding intermediate layers using linear classifier probes”.arXiv preprint arXiv:1610.01644(2016) (cit. on p. 5). [Ama16] S.-i. Amari.Information geometry and its applications. 2016 (cit. on pp. 2, 4, 12, 19). [AN00] S.-i. Amari and H. Nagaoka.Methods of information geometry. 2000 (cit. on p. 2). [Ard+24] A. Arditi, O...

arXiv 2016

[2] [2025]

The geometry of refusal in large language models: concept cones and representational independence

2025 (cit. on p. 1). [Wol+25] T . Wollschläger, J. Elstner, S. Geisler, V . Cohen-Addad, S. Günnemann, and J. Gasteiger. “The geometry of refusal in large language models: concept cones and representational independence”.arXiv preprint arXiv:2502.17420(2025) (cit. on p. 12). [Wu+25] Z. Wu, A. Arora, A. Geiger, Z. Wang, J. Huang, D. Jurafsky, C. D. Manning...

arXiv 2025