GeoSteer: Faithful Chain-of-Thought Steering via Latent Manifold Gradients

AI summary

Overview

Research area: Natural Language Processing, specifically large language model reasoning (Chain-of-Thought) and activation-level steering of hidden representations.

Technical level: Advanced. The paper assumes familiarity with variational autoencoders, latent-variable models, Kullback-Leibler divergence, Jacobians and the chain rule, and activation steering methods such as Activation Engineering and representation engineering.

One-sentence scope: The paper proposes GeoSteer, a method that learns a low-dimensional VAE latent manifold of Chain-of-Thought hidden states together with a differentiable quality function, then steers a target LLM's hidden states along the gradient of that quality function during inference, and evaluates the result on GSM8k with four Qwen3 models.

What This Paper Is About

Existing LLMs frequently produce Chain-of-Thought reasoning traces that contain logical inconsistencies, unnecessary sentences, or outright errors in intermediate steps, even when the final answer is correct. This gap means final-answer accuracy alone cannot tell you whether the reasoning process itself is trustworthy.

The goal of GeoSteer is to improve the intermediate reasoning of a target model at inference time without retraining it and without requiring large preference datasets. It does so by treating hidden states as points on a curved latent manifold and moving them along a gradient that points toward higher-quality reasoning regions, rather than applying the linear, Euclidean-style shifts used by prior activation steering methods.

Key Contributions

1. A manifold-based steering framework (GeoSteer). The method adjusts hidden states inside a VAE-learned latent space so that the model generates more logically consistent reasoning processes.

2. A geometry-aware activation steering method. Instead of linear manipulation in Euclidean space, GeoSteer computes a quality gradient in latent space and pulls it back through the encoder's Jacobian into hidden-state space, giving geometrically coherent updates.

3. A three-part pipeline combining data construction, manifold learning, and inference-time control. Step-level quality scores are produced by a teacher model, a VAE plus a differentiable regressor estimate the manifold and quality function, and the pullback gradient steers generation.

4. An empirical demonstration across multiple model scales. Experiments on GSM8k with Qwen3-0.6B, Qwen3-1.7B, Qwen3-4B, and Qwen3-8B show effects on both exact-match accuracy and judged reasoning quality.

Main Findings

• Average improvement as stated in the abstract: GeoSteer improved accuracy by 0.9 points and enhanced reasoning quality by 4.5 points on average compared with the original LLMs. The tables in the paper report per-model, per-strength results rather than a single aggregate, so these averages are the abstract's summary figures.

• Accuracy was largely preserved, with occasional gains: In Table 1, for Qwen3-1.7B the exact-match score improved from 82.3 to 84.9 at β = 150, which the paper calls the largest gain observed. Qwen3-8B showed moderate improvements when β was between 125 and 150 (90.7 to 91.3 at β = 125 and β = 300, and 90.7 to 91.4 at β = 200). Qwen3-4B remained near-stable across β values.

• Small models are sensitive to steering strength: Qwen3-0.6B dropped sharply for β > 100, falling from 60.0 at baseline to 55.0 (β = 100), 52.0 (β = 125), 50.9 (β = 150), 46.2 (β = 200), and 28.7 (β = 300).

• Statistical significance on accuracy was limited and concentrated: Significant exact-match changes at p < 0.05 on the McNemar test appear for Qwen3-1.7B at β = 150, 200, and 300 (all improvements), and for the Qwen3-0.6B baseline and the Qwen3-4B baseline at β = 300.

• Judged reasoning quality improved for most settings: In Table 2, for Qwen3-0.6B the pairwise win rate rose from 10.01% to 15.31% at β = 10, and reached 23.20% versus a 21.53% baseline at β = 100. For Qwen3-1.7B the steered model improved at every tested β, for example 1.90% to 4.78% at β = 10 and 5.00% to 8.49% at β = 300. Qwen3-4B and Qwen3-8B also showed higher win rates, with smaller gains than the smaller models.

• The printed Qwen3-0.6B table shows a diverging trend at high β: In Table 2, the baseline column for Qwen3-0.6B rises with β (up to 59.89% at β = 300) while the steered column falls (15.62% at β = 300). For the other three models the steered values exceed the baseline values at most tested β values, with exceptions such as Qwen3-4B at β = 125 and β = 300.

• Optimal steering strength depends on model capacity: Qwen3-1.7B and Qwen3-8B achieved their best exact match at β ∈ {125, 150}. Qwen3-0.6B improved on pairwise evaluation at moderate β but degraded sharply at high β. The strongest pairwise gains for Qwen3-0.6B appeared at β ∈ {100, 125, 150, 200}.

• Steering changes internal representations even when the output text barely changes: In the case study, the baseline latent trajectory stayed in a region with an average quality score around 0.65, while the steered trajectory showed a sharp transition around token 63 (the word "take"; the figure caption refers to token 64), where the average quality score increased from approximately 0.65 to around 0.78. The transition occurred at the boundary between problem description and arithmetic computation, and both generations produced nearly identical surface text and the same final answer.

• The high-quality region in the case study sat around 0.75–0.80 quality: The steered CoT became slightly more concise while preserving the essential reasoning steps, whereas the baseline included extra explanatory phrases.

• Reasoning quality and final correctness are related but not equivalent: Figure 3 shows the step-wise quality curves for correct (EM = 1) and incorrect (EM = 0) solutions have similar shapes—quality rises early then stabilizes—with substantially overlapping variance bands. Incorrect examples tended to contain fewer reasoning steps, giving smaller variance toward the end of the trajectory. The paper reports cases where high step-wise quality still ended in a wrong answer due to a late arithmetic mistake, and correct answers that passed through low-quality regions.

• The quality estimator captures local coherence, not global correctness: The paper concludes that EM depends heavily on the correctness of the final computation, whereas the latent-space score evaluates the structural coherence of partial reasoning, so intermediate quality is useful but not sufficient for predicting final-answer correctness.

Methodology in Plain English

The approach has three stages.

Stage 1 — Build a scored reasoning dataset. A teacher model, gpt-oss-20b, generates two kinds of Chain-of-Thought for each training problem in GSM8k: high-quality trajectories that are coherent and typically reach correct answers, and flawed trajectories that contain arithmetic or logical errors while still sounding natural. Each prefix of each trajectory is given a scalar quality score in [0, 1] by prompting the teacher with a scoring prompt that asks for a single numeric value.

Stage 2 — Learn a manifold and a quality function. The target (student) model reads the reasoning tokens, and the hidden state at the last token of each prefix is extracted. A variational autoencoder is trained to compress these hidden states into a lower-dimensional latent space (k < d) and reconstruct them, with an objective combining reconstruction error and a KL divergence term against a prior. This regularization makes the latent space closer to an isotropic Gaussian, reducing the strong anisotropy of the raw hidden-state space and making distances and directions more uniform. A separate differentiable regressor R_ψ is fit to predict the teacher's quality score from latent vectors by minimizing squared error. Its gradient with respect to the latent vector indicates the direction that improves reasoning quality.

Stage 3 — Steer during inference. At each generation step, the model's hidden state is encoded into the latent space. The latent quality gradient is mapped back into hidden-state space using the transpose of the encoder's Jacobian (a pullback through the chain rule). The hidden state is then updated by adding a normalized version of this pullback gradient scaled by a hyperparameter β, and the updated state is used to generate the next token. Repeating this over all steps produces a trajectory guided toward higher-quality regions while trying to preserve the model's linguistic behavior.

Evaluation. Experiments use Qwen3-0.6B, Qwen3-1.7B, Qwen3-4B, and Qwen3-8B on the GSM8k test set in a 4-shot setting, with the four examples randomly chosen from the GSM8k test set. Two metrics are used: Exact Match (EM) for final answers, and a Pairwise Win Rate in which GPT-4o compares baseline and steered CoTs for clarity, coherence, and reasoning quality. Statistical significance is tested with the McNemar test at p < 0.05.

Why This Matters

Impact on research. The paper argues that prior activation steering methods treat hidden states as if they live in a flat Euclidean space and can be manipulated with linear arithmetic, ignoring the evidence that internal representations form low-dimensional, structured, curved geometry. GeoSteer is a concrete attempt to replace linear steering with manifold-aware steering. It also reinforces a broader methodological point: evaluating reasoning only by final-answer accuracy hides inconsistent or error-prone intermediate steps, so metrics that capture reasoning quality itself are needed.

Potential real-world applications.

• Tutoring and educational systems that need to show students not just the right answer but a coherent, step-by-step solution path.
• Mathematical and quantitative assistants where an audit trail of intermediate steps must be trustworthy, for example in finance or engineering calculations.
• Verification pipelines where a model's reasoning is checked before downstream action is taken, since step-wise scores can flag weak reasoning regions.
• Inference-time quality control for deployed models, where a small model is steered toward better reasoning without retraining or preference-data collection.

Industry relevance. The method is alignment-free and operates at inference time, requiring no preference datasets and no parameter updates to the target model, which addresses the stated cost and data burden of DPO- and GRPO-style training. It is demonstrated on small open models from the Qwen3 family (0.6B to 8B), which are the scale of models commonly deployed for cost-sensitive tasks. The main practical caveat is that steering strength β must be tuned per model: too small and improvements are marginal, too large and small models degrade sharply—Qwen3-0.6B exact match fell from 60.0 to 28.7 at β = 300.

Future Directions

1. Better quality metrics. The step-wise score used here captures local coherence but not global logical consistency. The authors call for metrics that assess multi-step causal correctness or trajectory-level validity.

2. Testing across model families. Experiments focused on the Qwen series, which explicitly supports CoT generation. The paper proposes evaluating LLaMA, Mistral, and DeepSeek to determine whether the observed trends generalize.

3. Establishing the geometry theoretically. The relationship between gradients in latent space and the true geometry of hidden-state trajectories is not yet theoretically established, and whether the learned manifold faithfully reflects high-quality reasoning structure remains an open question.

4. Reducing dependence on the teacher. Because the latent manifold and quality function are bounded by gpt-oss-20b's coverage, the paper notes this dependency limits the achievable improvement on downstream tasks.

Target Audience

Researchers and engineers working on LLM reasoning, interpretability, and representation-level control benefit most. The paper suits readers already comfortable with VAEs, latent-variable modeling, and activation steering; it is not an introductory treatment of Chain-of-Thought. Practitioners who want an inference-time alternative to preference-based alignment such as DPO or GRPO will find the empirical tables directly useful, but should note that the reported accuracy changes are small, model-dependent, and show significant degradation for the smallest model at high steering strength.

Authors’ abstract

Recent advances in Large Language Models (LLMs) have demonstrated remarkable progress in their reasoning capabilities, such as Chain-of-Thought (CoT). Most approaches rely on CoT rationales. Previous studies have shown that LLMs often generate logically inconsistent reasoning steps even when their final answers are correct. These inconsistencies reduce the reliability of the reasoning process. We propose GeoSteer, a manifold-based framework that improves the quality of intermediate reasoning. The method consists of: (1) constructing a CoT dataset with step-level scores, (2) training a Variational Autoencoder (VAE) model and a quality estimation model to learn a low-dimensional manifold of high-quality CoT trajectories, and (3) steering hidden states of target LLMs toward higher-quality regions in the latent space. This last step enables steering of the hidden states by following gradients along the learned manifold. It facilitates geometrically coherent steering. Evaluation experiments were conducted on the GSM8k dataset using the Qwen3 series. We evaluated performance using two metrics: answer accuracy and overall reasoning quality. GeoSteer improved the accuracy by 0.9 points and enhanced the reasoning quality by 4.5 points on average, compared with those of original LLMs. These results indicate that GeoSteer improves an effective and controllable mechanism for improving the quality of intermediate reasoning in LLMs.