Source-linked AI summary

Look Inward to Explore Outward: Learning Temperature Policy from LLM Internal States via Hierarchical RL

Yixiao Zhou, Yang Li, Dongzhou Cheng, Hehe Fan, Yu Cheng

arXiv:2602.13035v1cs.LGcs.AIcs.CL

TL;DR

RLVR lacks reward-coupled temperature control despite temperature shaping trajectory exploration and policy learning. IntroLLM addresses this with a hierarchical, hidden-state-conditioned temperature policy jointly optimized with token generation, and it outperforms fixed and heuristic methods on mathematical reasoning benchmarks while learning interpretable exploration patterns.

  • Problem

    Existing RLVR methods use static or heuristic temperature schedules that are not directly optimized by task-level rewards and cannot adapt across reasoning contexts, prompts, or training stages.

  • Method

    IntroLLM models temperature selection as a hierarchical reinforcement-learning policy conditioned on hidden states and jointly trains it with the token policy using shared task rewards.

  • Results

    Learned temperature policies consistently outperform fixed-temperature and heuristic adaptive baselines on mathematical reasoning benchmarks, with higher temperatures used for uncertain multi-step reasoning and lower values for execution-heavy tokens.

  • Takeaways & Limitations

    Temperature can be learned as a context-aware exploration strategy that adapts across prompts, token positions, and training stages while yielding interpretable behavior.

  • Takeaways & Limitations

    The framework is most appropriate where correctness can be automatically checked and may produce incorrect or unsafe intermediate outputs without reliable verification signals.

Abstract

from arXiv · show

Reinforcement Learning from Verifiable Rewards (RLVR) trains large language models (LLMs) from sampled trajectories, making decoding strategy a core component of learning rather than a purely inference-time choice. Sampling temperature directly controls the exploration--exploitation trade-off by modulating policy entropy, yet existing methods rely on static values or heuristic adaptations that are decoupled from task-level rewards. We propose Introspective LLM, a hierarchical reinforcement learning framework that learns to control sampling temperature during generation. At each decoding step, the model selects a temperature based on its hidden state and samples the next token from the resulting distribution. Temperature and token policies are jointly optimized from downstream rewards using a coordinate ascent scheme. Experiments on mathematical reasoning benchmarks show that learned temperature policies outperform fixed and heuristic baselines, while exhibiting interpretable exploration behaviors aligned with reasoning uncertainty.

1. Introduction

RLVR makes decoding strategy part of training, but fixed and heuristic temperatures cannot adapt exploration to changing prompts, reasoning stages, or model capabilities. IntroLLM learns context-aware temperature control jointly with token generation from task rewards, improving mathematical reasoning while producing interpretable exploration patterns.

  • Motivation: RLVR depends on trajectory diversity for policy improvement, making sampling temperature a training-time control over exploration and exploitation.Temperature modulates policy entropy by scaling token-distribution logits.
  • Limitations of existing methods: Heuristic adaptive rules encode assumptions about exploration and are decoupled from task-level rewards, limiting adaptation through experience.Such rules use signals including token entropy, KL divergence, or predefined token categories.
  • IntroLLM: IntroLLM formulates temperature selection as hierarchical reinforcement learning, conditioning a temperature policy on hidden states and a token policy on selected temperature.The two policies make coupled decisions during generation.
  • IntroLLM: Temperature and token policies are jointly optimized from the same task-level reward, allowing exploration to vary across prompts, token positions, and training stages.The framework uses coordinate ascent and a mixed discrete–continuous temperature policy that learns when and how much to adjust.
  • Results: Learned temperature policies outperform fixed and heuristic methods on mathematical reasoning benchmarks, assigning higher temperatures to uncertain multi-step reasoning and lower values to execution-heavy tokens.The reported behaviors include numerical computation, factual retrieval, and final answer synthesis as lower-temperature settings.

2. Preliminary

RLVR optimizes an autoregressive token policy against verifiable rewards, with GRPO providing group-relative advantages and clipped updates for efficient, stable optimization. IntroLLM extends this pipeline by inserting a hidden-state-conditioned temperature policy before token sampling.

  • RLVR: RLVR trains a stochastic token policy to maximize expected task reward over sequences generated from prompts.The framework uses deterministic, rule-based signals to provide reliable feedback and mitigate reward hacking.
  • RLVR: GRPO samples groups of trajectories per prompt and computes sequence-level advantages without separate value models.This avoids the computational overhead of value networks required by PPO.
  • RLVR: The policy uses a clipped surrogate objective with an importance ratio to improve training stability and enable efficient credit assignment.The importance ratio compares current and old token-policy probabilities.
  • Sampling temperature: Temperature rescales vocabulary logits: higher values flatten token distributions for exploration, while lower values sharpen them for exploitation.In RLVR, this directly changes trajectory diversity during training.
  • IntroLLM: IntroLLM sequentially selects temperature from the hidden state before the token policy generates the next token from the temperature-modulated distribution.Both policies are trained jointly using verifiable task rewards.

3. Problem Formulation

The paper formulates adaptive temperature selection as a sequential hierarchical decision process because exploration needs differ across reasoning steps, prompts, difficulty levels, and training stages. A hidden-state-conditioned temperature policy and token policy jointly maximize expected reward while separating exploration control from content generation.

  • Problem formulation: Static or heuristic temperature schedules cannot adapt to heterogeneous exploration needs across reasoning steps, prompts, difficulty levels, and training stages.Reasoning pivots may benefit from high temperature, whereas factual tokens require low temperature.
  • Hierarchical process: The temperature policy selects τt from hidden state ht and the previous temperature, while the token policy generates yt conditioned on ht and τt.This defines temperature control as a learnable sequential decision process.
  • Objective: The hierarchical process maximizes expected reward over jointly generated token and temperature sequences.The token sequence is y and the temperature sequence is τ.
  • Design properties: The formulation provides reward-driven exploration, hidden-state-conditioned adaptation, and a principled separation between content generation and exploration control.These properties motivate the subsequent architectural and joint-optimization design.

4. Introspective LLM

Introspective LLM uses coupled temperature and token policies in a hierarchical rollout, with temperature decisions conditioned on hidden states and token generation conditioned on selected temperatures. Coordinate ascent updates the two policies separately while sharing trajectory-level rewards.

  • Hierarchical formulation: The hierarchical formulation couples a hidden-state-conditioned temperature policy with a token policy that generates from a temperature-modulated distribution.At each decoding step, the temperature policy selects τt before the token policy samples yt.
  • Temperature policy design: Temperature selection uses a two-stage mixed discrete–continuous action: first decide whether to update, then sample a bounded intensity when updating.The update decision uses a Bernoulli variable, while the new intensity is sampled from a Beta distribution and transformed to [τmin, τmax].
  • Temperature policy design: When the update decision is inactive, temperature remains unchanged, reducing variance from redundant adjustments while preserving adaptive control when changes are useful.This design learns both whether to adjust temperature and how much adjustment to make.
  • Joint policy optimization: Coordinate ascent alternates token-policy and temperature-policy updates, fixing one policy while optimizing the other under shared trajectory-level advantages.The token policy is updated with realized temperatures fixed, and the temperature policy is updated with token actions fixed; both use GRPO objectives and the same final reward-derived advantage.
  • Joint policy optimization: Rollouts retain token and temperature actions so both policies’ likelihoods can be recomputed and optimized with GRPO.Cached temperatures preserve the original exploration context for token likelihoods, while recorded temperature actions support reinforcement of the exploration strategy.

5. Experiments

Experiments evaluate IntroLLM across mathematical reasoning, out-of-domain tasks, efficiency, learned temperature behavior, and ablations. Learned token-level temperature control improves accuracy and diversity while adapting exploration to difficulty and reasoning dynamics with limited overhead.

  • Evaluation setup: IntroLLM is evaluated on AIME 2024, AMC 2023, MATH-500, Minerva Math, OlympiadBench, and Omni-Math, against fixed and adaptive baselines.
  • Main results: 1.6–2.0% separates the best and worst fixed-temperature GRPO configurations, with no single temperature optimal across benchmarks and model scales.
  • Main results: 30.47%/40.91% average Avg@8 for IntroLLM exceeds EAD’s 29.41%/39.66% and TAMPO’s 28.28%/38.83% across 1.7B/4B models.
  • Main results: 61.79%/73.50% Pass@8 surpasses the second-best method by 2.08%/4.29%, while AMC23 reaches 87.50% on 4B.
  • Out-of-domain generalization: IntroLLM improves over the second-best baseline across GPQA-Diamond, MMLU-Pro, and HumanEval on both model scales.On 1.7B, gains are 1.26%, 0.23%, and 1.82%; on 4B, they are 0.07%, 0.49%, and 0.61%, respectively.
  • Efficiency: The temperature policy adds 21M parameters, a 0.122% increase, with nearly unchanged FLOPs and approximately 2–3 tokens/s throughput degradation at 4k context.
  • Learned behavior: Median temperature increases monotonically from MATH-500 L1 to L5, indicating larger exploration budgets for harder problems.
  • Learned behavior: IntroLLM combines global natural annealing with problem-level reasoning rhythms, using higher temperatures at logical pivots and lower values during computation and answer synthesis.

6. Related Work

Related work positions IntroLLM within RLVR exploration control and adaptive temperature sampling. Prior approaches use localized objectives, diversity signals, supervised control, or heuristics, whereas IntroLLM treats temperature selection as a learned control problem.

  • RLVR and exploration control: RLVR has progressed from outcome-level preference modeling toward process-level reasoning rewards, with exploration control addressing complex-task reward hacking.
  • RLVR and exploration control: Recent RLVR methods localize exploration to decision-critical tokens, optimize Pass@k, or incorporate semantic diversity signals.
  • Adaptive temperature sampling: Static temperature schedules fail to account for token-level heterogeneity and prompt-specific dynamics because exploration needs vary across training stages and token positions.

7. Conclusion

The paper treats exploration control as a learnable component of RLVR, jointly optimizing temperature and token generation under task rewards. This produces context-aware exploration that improves reasoning performance and training stability relative to fixed or heuristic schedules.

  • Sampling temperature is modeled as a learnable policy and optimized jointly with token generation using the same task reward.
  • The framework adapts exploration across prompts, token positions, and training stages through context-aware temperature control.
  • Learning temperature leads to more effective exploration, improved reasoning performance, and more stable training than fixed or heuristically scheduled temperatures.
  • The formulation extends beyond this instantiation toward principled optimization of other generation-time decisions in RL-based language-model training.

Impact Statement

The work proposes adaptive exploration control for reinforcement-learning-based language-model post-training and identifies benefits for reasoning reliability and training efficiency. It also limits its applicability to settings with reliable verification signals and notes broader risks outside verifiable domains.

  • Impact Statement: The framework learns when and how much to explore during generation for reinforcement learning with verifiable rewards.The stated goal is adaptive exploration control during language-model training.
  • Impact Statement: Treating decoding parameters as learnable control variables may reduce reliance on brittle heuristics and manual tuning.The passage connects this design to more transparent and systematic RL training.
  • Impact Statement: Adaptive exploration may improve reasoning diversity but can also increase incorrect or unsafe intermediate outputs without reliable verification signals.The authors therefore position the framework primarily for domains where correctness can be automatically checked.
  • Impact Statement: The work highlights verification, evaluation, and responsible deployment as continuing requirements for RL-based language-model training.This conclusion follows the paper’s stated potential benefits and risks.

B.1. Model Architecture

The model adds a lightweight hidden-state temperature head that predicts temperature-control parameters and is designed as a plug-and-play extension. Training uses hierarchical optimization with separate token and temperature policies, binary verified rewards, and MATH-based data.

  • B.1. Model Architecture: A lightweight MLP branches from the final hidden state to predict temperature parameters.The head is intended to adapt readily to fine-tuned or compressed models of the same family.
  • B.1. Model Architecture: The temperature head uses Linear(dmodel, dmodel/2) → ReLU → Linear(dmodel/2, 3).Its three outputs are the logits [uc, uα, uβ].
  • B.1. Model Architecture: The final-layer initialization sets P(change) = 0.5 and uses Beta(1, 1), producing a uniform initial temperature distribution.This maximum-entropy initialization is intended to facilitate early exploration across the temperature range.
  • B.1. Model Architecture: Hierarchical optimization uses decoupled learning rates of 1 × 10−6 for the token policy and 5 × 10−5 for the temperature policy.Training used 8 NVIDIA H100 GPUs, global batch size 128, and group size G = 8.
  • B.1. Model Architecture: The RLVR reward is binary: symbolic verification assigns R(x, y) = 1 for a correct final answer and R(x, y) = 0 otherwise.No partial rewards or process-supervision rewards were used.
  • B.1. Model Architecture: Training uses 7,500 MATH training problems, with prompts used for reinforcement learning and ground-truth answers used for verification.The dataset spans seven mathematics categories and includes step-by-step solutions.

C.2. Evaluation Benchmarks

Evaluation spans mathematics benchmarks covering balanced difficulty, hard out-of-distribution problems, secondary-school breadth, STEM reasoning, Olympiad deduction, and broad modern mathematical proficiency. Comparisons include fixed-temperature GRPO and adaptive heuristic or meta-policy baselines.

  • C.2. Evaluation Benchmarks: MATH-500 is the primary benchmark for analyzing per-token temperature fluctuations and reasoning uncertainty.It balances difficulty levels L1–L5 and subject categories.
  • C.2. Evaluation Benchmarks: AIME 2024 tests exploration on 15 highly complex, long-context problems with integer answers from 000 to 999.The benchmark is presented as less susceptible to memorization and representative of out-of-distribution hard problems.
  • C.2. Evaluation Benchmarks: AMC 2023 evaluates precision and broad secondary-school problem solving under a 25-question, 75-minute time constraint.The version was selected as recent and potentially unobserved data.
  • C.2. Evaluation Benchmarks: Minerva Math measures formal scientific reasoning and multi-step numerical accuracy using STEM-focused problems from MATH and university-level exams.The suite is used to assess robustness in scientific problem solving.
  • C.2. Evaluation Benchmarks: Fixed-temperature GRPO baselines use τ = 0.6, 1.0, and 1.2, spanning exploitative, default, and exploratory configurations.Temperature remains constant during training, rollout, and every token position.
  • C.2. Evaluation Benchmarks: EAD applies a dynamic simulated-annealing schedule, while TAMPO learns a discrete temperature meta-policy through hierarchical inner- and outer-loop optimization.EAD emphasizes high-entropy early tokens and stable later exploitation; TAMPO updates temperatures using Temperature-Specific Advantages.

D.4. Summary of Differences and Advantages

IntroLLM differs from baseline temperature-control methods by using hidden states and downstream rewards to adapt temperature at reasoning pivots. Its mixed discrete–continuous action space balances flexible exploration with training stability.

  • IntroLLM uses hidden states to detect high-uncertainty reasoning pivots and choose exploration levels, unlike response-wide or hand-scheduled temperature controls.
  • Temperature selection is optimized directly by downstream task rewards rather than relying on manually tuned annealing schedules or discrete bins.
  • A Bernoulli trigger combined with Beta sampling lets the model learn when to keep temperature stable and when to adapt.
  • The mixed action space is presented as improving the balance between exploration flexibility and training stability.

E.1. Analysis of Generation Length

During reinforcement learning, IntroLLM adapts response length rather than maintaining the fixed-length behavior of GRPO. Its responses eventually become approximately 30%–40% longer, which the authors associate with deeper reasoning at critical pivots.

  • GRPO remains near approximately 650 tokens, while IntroLLM moves from an initial convergence phase toward longer reasoning chains.
  • Approximately 30%–40%: IntroLLM’s response length eventually surpasses the fixed-temperature GRPO baseline by this amount.
  • IntroLLM shows a U-shaped adaptation followed by a consistent upward response-length trend during training.
  • The authors interpret the added length as effective thinking depth, with more intermediate steps and alternative logical paths at critical reasoning pivots.
Loading 2602.13035v1…