Test-time Diverse Reasoning by Riemannian Activation Steering

TL;DR

This paper introduces a novel test-time activation steering method using Riemannian optimization to enhance output diversity and accuracy in language model reasoning tasks.

Contribution

It proposes an unsupervised Riemannian activation steering strategy that improves diversity in model outputs during inference by optimizing steering vectors at synchronization points.

Findings

Outperforms vanilla sampling in diversity and accuracy on mathematical benchmarks.

Uses Riemannian optimization to find steering vectors that maximize output variance.

Demonstrates effectiveness of test-time activation steering in complex reasoning tasks.

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.