Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Learning to Think: Information-Theoretic Reinforcement Fine-Tuning for LLMs

Jingyao Wang, Wenwen Qiang, Zeen Song, Changwen Zheng, Hui Xiong

TL;DR

Learning to Think (L2T) tackles the problem of balancing reasoning quality and token efficiency in LLMs by introducing a universal dense process reward rooted in information theory. The reward combines a fitting information gain term with a compression penalty and is estimated efficiently via PAC-Bayes bounds and the Fisher information matrix, enabling episodic reinforcement fine-tuning without task-specific annotations. By recasting reasoning as an episodic RL problem and optimizing with GRPO, L2T achieves improved reasoning accuracy under fixed token budgets while reducing token usage across diverse benchmarks. The approach demonstrates strong gains in reasoning efficiency and robustness, offering a principled path to dynamic budget allocation in practical LLM deployments.

Abstract

Large language models (LLMs) excel at complex tasks thanks to advances in their reasoning abilities. However, existing methods overlook the trade-off between reasoning effectiveness and efficiency, often encouraging unnecessarily long reasoning chains and wasting tokens. To address this, we propose Learning to Think (L2T), an information-theoretic reinforcement fine-tuning framework for LLMs to make the models achieve optimal reasoning with fewer tokens. Specifically, L2T treats each query-response interaction as a hierarchical session of multiple episodes and proposes a universal dense process reward, i.e., quantifies the episode-wise information gain in parameters, requiring no extra annotations or task-specific evaluators. We propose a method to quickly estimate this reward based on PAC-Bayes bounds and the Fisher information matrix. Theoretical analyses show that it significantly reduces computational complexity with high estimation accuracy. By immediately rewarding each episode's contribution and penalizing excessive updates, L2T optimizes the model via reinforcement learning to maximize the use of each episode and achieve effective updates. Empirical results on various reasoning benchmarks and base models demonstrate the advantage of L2T across different tasks, boosting both reasoning effectiveness and efficiency.

Learning to Think: Information-Theoretic Reinforcement Fine-Tuning for LLMs

TL;DR

Learning to Think (L2T) tackles the problem of balancing reasoning quality and token efficiency in LLMs by introducing a universal dense process reward rooted in information theory. The reward combines a fitting information gain term with a compression penalty and is estimated efficiently via PAC-Bayes bounds and the Fisher information matrix, enabling episodic reinforcement fine-tuning without task-specific annotations. By recasting reasoning as an episodic RL problem and optimizing with GRPO, L2T achieves improved reasoning accuracy under fixed token budgets while reducing token usage across diverse benchmarks. The approach demonstrates strong gains in reasoning efficiency and robustness, offering a principled path to dynamic budget allocation in practical LLM deployments.

Abstract

Large language models (LLMs) excel at complex tasks thanks to advances in their reasoning abilities. However, existing methods overlook the trade-off between reasoning effectiveness and efficiency, often encouraging unnecessarily long reasoning chains and wasting tokens. To address this, we propose Learning to Think (L2T), an information-theoretic reinforcement fine-tuning framework for LLMs to make the models achieve optimal reasoning with fewer tokens. Specifically, L2T treats each query-response interaction as a hierarchical session of multiple episodes and proposes a universal dense process reward, i.e., quantifies the episode-wise information gain in parameters, requiring no extra annotations or task-specific evaluators. We propose a method to quickly estimate this reward based on PAC-Bayes bounds and the Fisher information matrix. Theoretical analyses show that it significantly reduces computational complexity with high estimation accuracy. By immediately rewarding each episode's contribution and penalizing excessive updates, L2T optimizes the model via reinforcement learning to maximize the use of each episode and achieve effective updates. Empirical results on various reasoning benchmarks and base models demonstrate the advantage of L2T across different tasks, boosting both reasoning effectiveness and efficiency.
Paper Structure (49 sections, 5 theorems, 37 equations, 10 figures, 4 tables, 1 algorithm)

This paper contains 49 sections, 5 theorems, 37 equations, 10 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Reasoning of LLMs
  4. Process Reward Models
  5. Problem Settings and Analysis
  6. Problem Settings
  7. Empirical Evidence
  8. Motivation Analysis
  9. Learning To Think
  10. Problem Reformulation
  11. Learning Information-Theoretic Dense Process Reward
  12. Optimizing LLM with Reinforcement Fine-Tuning
  13. Practical Implementation with GRPO
  14. Experiments
  15. Experimental Settings
  16. ...and 34 more sections

Key Result

Theorem 4.2

Given the low-rank parameter proxy $\tilde{\theta}_k$ and $\tilde{\theta}_{k-1}$ for parameters $\theta_k$ and $\theta_{k-1}$, assume that $\tilde{\theta}_k$ and $\tilde{\theta}_{k-1}$ follow Gaussian distribution, e.g., $p(\tilde{\theta}_k) = \mathcal{N}(\tilde{\theta}_k | \mu_k, \Sigma_k)$ where $

Figures (10)

  • Figure 1: Results of DeepScaleR-1.5B-Preview across different tasks on Omni-MATH. We partition the generated reasoning chain into episodes, measuring accuracy $\mathrm{Acc}(k)$ and average token consumption $\overline{T}(k)$ at different episode depths. More details and results are shown in Appendix \ref{['sec_app:experiments_motivating']}.
  • Figure 2: Efficiency comparison across different benchmarks. We compute the token budget required for each benchmark and treat the budget of the base model w/o fine-tuning as reference ($1\times$).
  • Figure 3: Pass@1 vs. token budget of different methods on AIME. We record the model reasoning accuracy under different maximum token budgets to evaluate the ability of using test-time compute.
  • Figure 4: Effect of L2T components.
  • Figure 5: Parameter sensitivity of $\alpha$ and $\beta$
  • ...and 5 more figures

Theorems & Definitions (6)

  • Definition 4.1
  • Theorem 4.2
  • Lemma B.1
  • Proposition D.1
  • Theorem D.2
  • Theorem D.3