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

TL;DR

This work tackles the problem of eliciting Chain-of-Thought (CoT) reasoning in base LLMs without retraining by introducing a gradient-based hidden-state optimization framework grounded in probabilistic conditional generation. By modeling CoT elicitation as MAP inference with a likelihood term from a CoT classifier and a prior term enforcing proximity to the original hidden state, the method performs gradient ascent updates $h_{t+1}=h_t+\alpha_t\nabla_h[\log f_\theta(h_t)-\lambda d(h_t,h_0)]+\sqrt{\alpha_t}\cdot z$ to steer internal representations toward reasoning trajectories while preserving linguistic quality. The approach is supported by theoretical bounds on the trade-off parameter $\lambda$ to ensure gradient alignment and stable updates, and it is evaluated across math, commonsense, and logic benchmarks using multiple base LLMs. Results show substantial improvements over vector-arithmetic baselines in many tasks, with moderate gains in others, demonstrating the practical viability of principled hidden-state optimization for unlocking latent CoT in foundation models. The method maintains text coherence and requires no additional training data, highlighting its potential for scalable reasoning enhancement in diverse applications.

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.

Figures (13)

• Figure 1: Our method optimizes hidden states within the base LLM's latent space to elicit Chain-of-Thought reasoning. Before optimization, the model produces direct answers without reasoning steps; after optimization, it generates chain of thought content. Unlike linear activation steering that may disrupt the language prior without regulations, our approach implements guided optimization with principled constraints.
• Figure 2: Impact of different $\lambda$ values on the iterative optimization process, demonstrating how this hyperparameter influences the optimization convergence.
• Figure 3: Upper: Impact of different $\lambda$ values on our method's performance, roughly showing the optimal range. Lower: Comparative analysis of activation steering with different strength. The dashed line in the figure represents the performance of the greedy decoding baseline.
• Figure 4: The gradient norm of $\nabla^*$ and $\nabla^{**}$ in the optimization process
• Figure 5: The prompt for GPT-4o to classify the response as CoT or non-CoT