Eliciting Chain-of-Thought in Base LLMs via Gradient-Based Representation Optimization

TL;DR

This paper introduces a gradient-based optimization method to manipulate hidden states in base LLMs, significantly improving their reasoning abilities across various benchmarks while maintaining text quality.

Contribution

It presents a novel probabilistic optimization framework for eliciting Chain-of-Thought reasoning in base LLMs through hidden state manipulation.

Findings

Outperforms existing hidden state steering methods on multiple reasoning benchmarks.

Effectively balances reasoning enhancement with linguistic coherence.

Provides a theoretically grounded approach for reasoning in base LLMs.

Abstract

Chain-of-Thought (CoT) reasoning is a critical capability for large language models (LLMs), enabling them to tackle com- plex multi-step tasks. While base LLMs, pre-trained on general text corpora, often struggle with reasoning due to a lack of specialized training, recent studies reveal their latent reason- ing potential tied to hidden states. However, existing hidden state manipulation methods, such as linear activation steering, suffer from limitations due to their rigid and unconstrained nature, often leading to distribution shifts and degraded text quality. In this work, we propose a novel approach for elic- iting CoT reasoning from base LLMs through hidden state manipulation grounded in probabilistic conditional generation. By reformulating the challenge as an optimization problem with a balanced likelihood and prior regularization framework, our method guides hidden states toward reasoning-oriented trajectories while preserving linguistic coherence. Extensive evaluations across mathematical, commonsense, and logical reasoning benchmarks demonstrate that our approach con- sistently outperforms existing steering methods, offering a theoretically principled and effective solution for enhancing reasoning capabilities in base LLMs.