Test-time Diverse Reasoning by Riemannian Activation Steering

AI summary

Overview

Research area: Machine learning / large language model inference — specifically test-time activation steering for mathematical reasoning, formulated as Riemannian optimization.

Technical level: Advanced. The paper assumes familiarity with transformer hidden states, best-of-N decoding, matrix determinants, and Riemannian manifold optimization (tangent spaces, exponential maps, parallel transport, block-coordinate descent).

Scope: The paper proposes SPREAD (SPherical intervention for REAsoning Diversity), an unsupervised test-time activation-steering method that injects diversity into multiple simultaneous reasoning trajectories by solving a log-determinant volume-maximization problem over a product of spheres, and evaluates it on three mathematical reasoning benchmarks with two 1.5B-parameter Qwen models.

What This Paper Is About

Best-of-N reasoning — generating N candidate solutions and picking the best — is bottlenecked by diversity collapse, where stochastic sampling still yields near-identical reasoning paths that repeat the same mistakes. The authors aim to fix this by intervening directly in the model's hidden activation space at test time, pushing the internal representations of concurrently generated sequences apart so that they explore different reasoning trajectories, without any fine-tuning or extra neural networks.

Key Contributions

1. The SPREAD steering method. An unsupervised activation-steering scheme that, at a "synchronization anchor" during batched generation, extracts hidden activations from all N sequences and computes N additive steering vectors. These vectors are added to the activations of all subsequent tokens until the next anchor, and the process repeats until an end-of-sequence (EOS) token.

2. A geometric reformulation. Instead of maximizing only the volume of the full set of N steered activations, the method maximizes the sum of squared parallelepiped volumes over every subset of the N steered vectors — which prevents degenerate solutions where N−1 vectors cluster together and one is pushed far away. Proposition 2 shows this is equivalent to a log-determinant problem: minimize −log det[I + (H+V)ᵀ(H+V)] subject to ‖v_i‖² ≤ α_i.

3. A Riemannian optimization algorithm with theory. Because the problem is non-convex (Example 3) and globally NP-hard in general, the authors prove the norm constraints are binding (Proposition 4) and turn the problem into one over a product of spheres. They derive the Riemannian gradient (Lemma 5), propose a Riemannian block-coordinate descent algorithm using exponential maps (Algorithm 1), and prove smoothness (Propositions 8, 9) and convergence (Theorem 6): with step sizes η_i = 1/L_i, the minimum gradient norm decays as O(1/√k) and any limit point is a stationary point.

4. A single-hyperparameter, architecture-free design. Setting α_i = C·‖h_i‖₂/p reduces tuning to one relative parameter C > 0, and the method needs no additional neural architecture to measure quality or diversity.

Main Findings

• Pass@N accuracy: On AIME24 with Qwen2.5-1.5B at temperature 1.0, SPREAD reports 3.3 (C=1) and 6.7 (C=10) versus 0.0 for vanilla temperature sampling. The paper states SPREAD "consistently performs at least as well as vanilla temperature sampling, and often achieves improvements of several percentage points," with C = 1 giving the strongest results under most conditions.

• Unique Solution Count and Diversity Score: The authors report that SPREAD consistently outperforms temperature sampling on both metrics across the three benchmarks, using GPT-4.1-mini as an LLM judge that returns a diversity score in [0,1] and an integer count of distinct approaches.

• Pareto dominance on AIME24: Figures 3 and 4 show accuracy–diversity frontier plots for AIME24 with Qwen2.5-1.5B and Qwen2.5-Math-1.5B-Instruct. SPREAD (red circles) dominates vanilla sampling at different temperatures (blue squares). Corresponding Pareto plots for MATH and OlympiadBench are placed in the appendix.

• Computational efficiency: On synthetic activations with dimension p ranging from 1536 up to 2¹⁴ = 16384 (the hidden size of models such as LLaMA-3.1-405B), with K = 20 and averaged over 30 independent runs, execution time for N = 32 stays under 1.8 seconds even at the largest dimension. The authors describe this as demonstrating scalability and practical efficiency.

• Convergence guarantee: Theorem 6 gives both an asymptotic result (gradient norm tends to zero) and a non-asymptotic sublinear rate, with the step sizes chosen as η_i = 1/L_i so that each coordinate-descent update provably decreases the objective — removing the need for extensive learning-rate tuning.

Methodology in Plain English

The authors generate N answers to the same question at once. At fixed token positions during generation — the paper uses τ ∈ {100, 600, 1100, 1600}, with K = 20 iterations of the optimizer each time — they grab the hidden state vector for each sequence. These vectors form a matrix H.

The intuition is geometric: a set of vectors spans a parallelepiped, and the volume of that shape is a natural measure of how spread out the vectors are. The authors want to nudge each hidden state by a small vector v_i so that the resulting shape has maximum volume — and not just as a whole, but for every possible subset of sequences, so no single sequence is doing all the work of being different.

Maximizing the sum of squared subset volumes turns out to be mathematically identical to maximizing a determinant (the D-optimal design quantity). It is equivalent to minimizing −log det[I + (H+V)ᵀ(H+V)] — a smooth log-determinant objective — under the requirement that each nudge is not too large. Bounding the nudge size matters because a huge perturbation would overwrite the useful information in the hidden state and destroy the generation.

To keep the nudges from being too large, each steering vector is restricted to lie on a sphere of radius √α_i. A product of N spheres is a Riemannian manifold, so the authors use manifold optimization: instead of moving in a straight Euclidean line, they compute the gradient, project it onto the tangent space of the sphere, and then move along the sphere's surface using the exponential map — a formula involving cos and sin that guarantees the updated vector still has exactly the right length.

Because the manifold is a product of independent spheres, one expensive joint update can be replaced by N cheap sequential updates — Riemannian block-coordinate descent. Convergence is guaranteed if the step size for block i is set to 1/L_i, where L_i is derived from the data (the Frobenius norm of H, the total radii, and α_i), so no hand-tuning of learning rates is required. Initialization uses each hidden vector's offset from the centroid of all hidden vectors, plus small Gaussian noise, normalized to the correct sphere radius.

Experiments use Qwen2.5-1.5B (base) and Qwen2.5-Math-1.5B-Instruct on AIME24 (30 problems), MATH500 (500 problems) and OlympiadBench (675 problems), with p = 1536, steering applied at layer 28 (the final layer), C ∈ {1, 10}, temperature ∈ {0.2, 0.4, 0.6, 0.8, 1.0}, maximum generation length 2048 tokens, random seed 42, and NVIDIA RTX A5000 24GB hardware.

Why This Matters

Impact on research. The paper connects two usually separate literatures: activation steering (typically applied with fixed, contrastively derived direction vectors for attributes like truthfulness or toxicity) and test-time reasoning diversity. It argues that mathematical reasoning is a poor fit for contrastive steering because mathematical correctness has no simple positive/negative exemplar, and instead treats diversity itself as the optimization objective. It also offers a case study in reformulating an intractable determinant-maximization problem as a tractable Riemannian one with convergence guarantees.

Real-world applications (as framed by the paper's own task list):

• Mathematical problem solving, where a system must try several solution strategies before committing to an answer.
• Code generation, where alternative implementations of the same specification have different robustness and maintainability properties.
• Symbolic reasoning tasks, where different derivations of the same result can validate each other.
• Any best-of-N deployment where a reward model scores candidates but the candidates themselves are too similar to be useful.

Industry relevance. The method is a test-time, inference-only intervention: it modifies no model weights, needs no fine-tuning data, adds no extra neural networks, and requires only one new hyperparameter (C). It operates on hidden states already computed during generation, and the efficiency experiment suggests the per-anchor optimization cost remains modest even at the 16384-dimensional hidden size used by large models such as LLaMA-3.1-405B.

Future Directions

• Larger and more diverse models. The evaluation covers two 1.5B-parameter Qwen variants only; whether the gains transfer to larger or non-Qwen models is not established.
• Layers and anchors. Steering is applied exclusively at layer 28, the final layer, with anchors at token positions 100, 600, 1100, and 1600. The effect of steering earlier layers, or of different anchor densities and values of the iteration budget K, is not reported.
• Activation diversity versus reasoning diversity. The authors explicitly acknowledge there is no one-to-one equivalence between hidden-activation diversity and reasoning-path diversity, so the causal link between the objective being optimized and the outcome being measured remains an open question.
• Beyond mathematics. Whether the approach generalizes to code generation, symbolic reasoning, or other domains with more subjective correctness criteria, and how sensitive it is to the single scaling constant C, are natural extensions.

Target Audience

Researchers and engineers working on LLM inference-time techniques — particularly those interested in best-of-N sampling, decoding diversity, and activation steering. It will also appeal to readers with optimization backgrounds, since the central contribution is a log-determinant maximization problem over a product of spheres solved by Riemannian block-coordinate descent with convergence guarantees. Practitioners deploying reasoning models in production may find the efficiency measurements relevant, while readers without manifold-optimization background will find Sections 4.1–4.2 demanding.

Authors’ abstract

Best-of-$N$ reasoning improves the accuracy of language models in solving complex tasks by sampling multiple candidate solutions and then selecting the best one based on some criteria. A critical bottleneck for this strategy is the output diversity limit, which occurs when the model generates similar outputs despite stochastic sampling, and hence recites the same error. To address this lack of variance in reasoning paths, we propose a novel unsupervised activation steering strategy that simultaneously optimizes the steering vectors for multiple reasoning trajectories at test time. At any synchronization anchor along the batch generation process, we find the steering vectors that maximize the total volume spanned by all possible intervened activation subsets. We demonstrate that these steering vectors can be determined by solving a Riemannian optimization problem over the product of spheres with a log-determinant objective function. We then use a Riemannian block-coordinate descent algorithm with a well-tuned learning rate to obtain a stationary point of the problem, and we apply these steering vectors until the generation process reaches the subsequent synchronization anchor. Empirical evaluations on popular mathematical benchmarks demonstrate that our test-time Riemannian activation steering strategy outperforms vanilla sampling techniques in terms of generative diversity and solution accuracy.