PaperlayerView exact revision on arXiv

Source-linked AI summary

Large Language Models Explore by Latent Distilling

Yuanhao Zeng, Ao Lu, Lufei Li, Zheng Zhang, Yexin Li, Kan Ren

arXiv:2604.24927v2cs.CLcs.AIcs.LG

TL;DR

Test-time scaling is limited when sampled candidates differ lexically but repeat the same reasoning strategy. ESamp uses latent-representation novelty to steer decoding toward semantically under-explored generations, matching or outperforming baselines across tasks with less than 5% throughput overhead.

  • Problem

    Standard sampling often produces lexical variation without distinct reasoning strategies, limiting candidate diversity and downstream Pass@k scaling.

  • Method

    ESamp trains an online Latent Distiller to estimate novelty in hidden representations and reweight decoding toward semantically under-explored regions.

  • Results

    Less than 5% throughput overhead accompanies ESamp’s superior or comparable baseline performance across model families and tasks, especially for reasoning-model Pass@k efficiency.

  • Takeaways & Limitations

    ESamp provides a practical decoding strategy for increasing semantic exploration and Pass@k scaling efficiency across diverse LLM generation tasks.

  • Takeaways & Limitations

    Under fully saturated GPU utilization, asynchronous overlap may provide fewer overhead benefits.

Abstract

from arXiv · show

Generating diverse responses is crucial for test-time scaling of large language models (LLMs), yet standard stochastic sampling mostly yields surface-level lexical variation, limiting semantic exploration. In this paper, we propose Exploratory Sampling (ESamp), a decoding approach that explicitly encourages semantic diversity during generation. ESamp is motivated by the well-known observation that neural networks tend to make lower-error predictions on inputs similar to those encountered before, and incur higher prediction error on novel ones. Building on this property, we train a lightweight Distiller at test time to predict deep-layer hidden representations of the LLM from its shallow-layer representations to model the LLM's depth-wise representation transitions. During decoding, the Distiller continuously adapts to the mappings induced by the current generation context. ESamp uses the prediction error as a novelty signal to reweight candidate token extensions conditioned on the current prefix, thereby biasing decoding toward less-explored semantic patterns. ESamp is implemented with an asynchronous training--inference pipeline, with less than 5% worst case overhead (1.2% in the optimized release). Empirical results show that ESamp significantly boosts the Pass@k efficiency of reasoning models, showing superior or comparable performance to strong stochastic and heuristic baselines. Notably, ESamp achieves robust generalization across mathematics, science, and code generation benchmarks and breaks the trade-off between diversity and coherence in creative writing. Our code has released at: https://github.com/LinesHogan/tLLM.

1. Introduction

ESamp addresses the limited semantic diversity of naive test-time sampling by steering LLM generation toward less-explored latent representation regions. It combines an online Latent Distiller with novelty-based token reweighting and an asynchronous implementation that improves Pass@k efficiency with low overhead.

  • Motivation: Test-time scaling is constrained when lexically different candidates share the same reasoning structure, limiting downstream selection mechanisms.Naive token-level stochasticity often produces surface variation without changing the underlying reasoning strategy, causing redundant candidates and diminishing returns.
  • Method: ESamp encourages semantic exploration by penalizing tokens associated with predictable latent representations during generation.The method steers generation toward less-explored semantic regions rather than relying on lexical variation alone.
  • Method: A lightweight Latent Distiller learns shallow-to-deep hidden-representation mappings online, using prediction error as a novelty signal.The approach is grounded in the observation that neural networks predict familiar mappings more accurately than previously unseen ones.
  • Method: ESamp reweights the base model’s next-token distribution through a KL-regularized objective, suppressing redundant continuations and favoring under-explored semantic behaviors.In practice, distiller prediction errors are projected onto candidate tokens conditioned on the current prefix.
  • Results and efficiency: ESamp improves Pass@k scaling efficiency across benchmarks and model families while incurring less than 5% throughput overhead in standard serving scenarios.Its asynchronous pipeline decouples distiller training and inference from main LLM generation, supporting practical deployment.

2. Related Work

Prior work explores diversity through stochastic sampling, structured search, and controlled logit reweighting. ESamp differs from these approaches by targeting redundancy in continuous representation space rather than relying solely on token- or vocabulary-level interventions.

  • Stochastic Sampling: Stochastic sampling methods such as Top-p, Min-p, and entropy-based sampling restrict candidate pools to inject randomness efficiently.These methods mitigate deterministic-decoding degeneration through heuristic distribution truncation.
  • Structured Search: Structured search methods including Diverse Beam Search, Stochastic Beam Search, and Tree of Thoughts explicitly traverse generation trees to find reasoning trajectories.Multiple branches or backtracking introduce substantial computational overhead.
  • Controlled Generation: Contrastive Decoding modifies logits heuristically, while Controlled Decoding formulates steering as token-level KL-regularized reinforcement learning with a learned value function.The framework shows that optimal steering can be achieved by reweighting logits with the learned value function.
  • Controlled Generation: DeRa and OverRIDE adopt formulations similar to Controlled Decoding for controlled generation.These works follow the described theoretical framework for controlled generation.
  • Representation-Space Exploration: ESamp is conceptually closest to OverRIDE but estimates redundancy in continuous representation space rather than vocabulary space, where token-repetition penalties may miss semantically equivalent sequences.OverRIDE uses online adaptation to suppress redundancy and penalizes token repetition; ESamp uses LD for representation-space redundancy estimation.

3. Problem Formulation

The paper formulates LLM generation as an MDP and seeks semantic exploration through an intrinsic novelty reward over parallel trajectories. It optimizes a KL-regularized policy while assuming redundancy vanishes after semantic regions are explored.

  • MDP formulation: LLM generation is modeled as an MDP whose state is the token prefix, actions are vocabulary tokens, and policy probabilities come from the LLM.Standard likelihood-based or heuristically truncated decoding can produce lexical variation without equivalent semantic diversity.
  • Novelty objective: A batch of K trajectories receives a per-step intrinsic reward for steering generation toward semantic regions not yet explored by the batch.The reward is incorporated through a KL-regularized policy optimization framework.
  • Novelty objective: The KL-regularized objective uses a frozen reference LLM and a positive regularization strength α, and has a closed-form optimal policy.The paper states that πref is frozen and α > 0 controls regularization strength.
  • Assumption: The per-step formulation relies on vanishing redundancy: after one trajectory explores a semantic region, later visits receive near-zero reward, so Q∗(st, zt) ≈ r(st, zt).The paper formalizes this condition and discusses a practical relaxation in Appendix A.2.

4. Methodology

Exploratory Sampling (ESamp) promotes semantic exploration by using an online Latent Distiller to model depth-wise hidden-representation transitions and identify novel contexts. It reweights candidate tokens according to latent prediction errors and semantic alignment, with parallel generation coordination and asynchronous execution to reduce overhead.

  • Semantic representation-based exploration: ESamp grounds exploration in continuous hidden-layer representations, encouraging contexts to occupy diverse semantic regions instead of producing superficial lexical variation.The method uses hidden representations of the generated prefix as a semantic embedding of the current context.
  • Online latent distillation: The Latent Distiller is a lightweight MLP trained online to map shallow-layer hidden representations to deep-layer representations encountered during generation.Training minimizes mean squared error, allowing the Distiller to capture mappings associated with previously generated contexts.
  • Error-guided token reweighting: ESamp uses the latent prediction error as an intrinsic reward, increasing probability for candidate tokens whose vocabulary-head directions align with unexplained representation components.The reward is r(s, z) = log πref(z|s) − log qdist(z|s), and β controls exploration intensity.
  • Error-guided token reweighting: The resulting policy discourages redundant semantic responses rather than merely frequent tokens, biasing generation toward unexplored semantic regions.The latent error norm measures context novelty, while cosine similarity provides a selective semantic direction.
  • Parallel coordination: In parallel generation, the shared online Distiller coordinates sequences through an implicit first-come, first-served mechanism that suppresses redundant exploration.Mappings learned from one sequence reduce the probability of similar regions for other sequences.
  • Asynchronous implementation: The asynchronous pipeline overlaps lightweight Distiller computation with the host LLM’s middle-layer execution, but overlap benefits may diminish under fully saturated GPU utilization.The Distiller is removed from the critical path when sufficient execution slack exists.

5. Experiments

Experiments evaluate ESamp across mathematics, science, code generation, and creative writing using diverse model families and decoding baselines. Results show robust test-time exploration, improved semantic diversity without sacrificing quality, and low serving overhead.

  • Benchmarks: ESamp is evaluated on competition mathematics, expert-validated science questions, code generation, and BookCorpus story continuation tasks.The benchmarks include AIME 2024/2025, GPQA-Diamond, LiveCodeBench v5, and 512-token creative-writing continuations.
  • Models: Experiments span instruction-tuned, reasoning, and cross-family models, including Qwen2.5-7B/32B-Instruct, Qwen3-8B, and GPT-OSS-20B.Qwen3 uses “no thinking” mode and GPT-OSS-20B uses “medium thinking effort” mode to limit context exhaustion.
  • Test-Time Exploration: ESamp achieves superior or comparable Pass@k performance to baselines and robustly generalizes across benchmarks, with GPT-OSS-20B matching baseline Pass@64 performance using only Pass@8.Its largest gains occur on AIME24/25, while FIRE’s benchmark-specific gains do not generalize to LiveCodeBench v5.
  • Semantic Diversity: ESamp breaks the coherence–diversity trade-off by achieving the highest diversity, lowest semantic similarity, and best generation quality in creative writing and math reasoning.In math reasoning, it also combines the highest diversity scores with superior Pass@k, indicating distinct valid reasoning paths.
  • Generation Dynamics: Baseline generation diversity plateaus as sequences lengthen, whereas ESamp continues driving trajectory divergence during decoding.Figure 4(a) measures average pairwise cosine similarity across parallel Qwen2.5-7B generations on BookCorpus.
  • Serving Efficiency: Overhead is less than 2% for standard serving and approximately 4.25% at K = 16, while the distiller and hidden-state buffer use less than 200MB of VRAM for an 8B model.The asynchronous distiller computation overlaps with LLM execution; optimization could reduce memory use to around 50MB.

6. Conclusion

ESamp addresses the limited semantic diversity of standard decoding by estimating novelty in internal representations and steering generation toward under-explored semantic regions. It matches or outperforms baselines while maintaining negligible latency overhead through an asynchronous pipeline.

  • 6. Conclusion: ESamp estimates novelty in internal representations to steer generation toward under-explored semantic regions.This targets standard decoding’s tendency to produce surface-level lexical variation without genuine semantic diversity.
  • 6. Conclusion: ESamp matches or outperforms baselines while its asynchronous pipeline ensures negligible latency overhead.The combination supports practical, efficient LLM decoding.

A. Derivations … B.3. Mixed-Batch Guardrails

The appendix derives the KL-regularized optimal policy and formalizes when per-step novelty rewards are sequentially optimal. It also describes an asynchronous GPU implementation with preallocated memory and safeguards for mixed prefill/decode batches.

  • A. Derivations / A.1. Closed-Form Optimal Policy: The appendix provides a closed-form derivation for the KL-regularized reinforcement-learning objective, maximizing reward while keeping π close to πref.The regularization coefficient α controls the KL penalty and acts as a temperature.
  • A.1. Closed-Form Optimal Policy: The Lagrangian enforces probability normalization, and differentiating it with respect to π(z|s) yields the stationarity condition used to solve the policy.The derivation introduces multiplier λ for the constraint that probabilities sum to 1.
  • A.2. Reward Structure and Per-Step Optimality: Under vanishing redundancy, once a semantic region is explored, all later continuations receive zero novelty reward, making Q*(s_t, z_t) = r_t(s_t, z_t).The resulting sequential policy reduces to π* ∝ πref exp(r/α).
  • A.2. Reward Structure and Per-Step Optimality: The Distiller approximates vanishing redundancy through rapid fitting and local generalization, assigning near-zero novelty to explored regions and their neighborhoods.Its adaptive prediction error can still identify genuinely novel mappings encountered later, rather than permanently excluding previously explored regions.
  • B. Engineering the Parallel Pipeline / B.1. Asynchronous CUDA Streams: The parallel pipeline uses a dual-stream runtime in which the main vLLM graph and high-priority Distiller MLP computations synchronize through CUDA Events.GPU-side event synchronization lets the CPU dispatch both instruction chains without halting, hiding kernel execution latency.
  • B.2. Pre-allocated “In-Graph” Memory: The implementation uses a static GPU ring buffer, writing first-layer hidden states with a custom kernel and reading them asynchronously without CPU transfers.Distillation loss computation and weight updates remain in GPU HBM, avoiding dynamic allocation and PCIe bottlenecks.
  • B.3. Mixed-Batch Guardrails: A metadata check activates the Distiller only for decode-phase tokens, dynamically bypassing it for mixed batches or prefill-heavy steps.This guardrail reflects the Distiller’s optimization for generation rather than prompt processing.

B.4. Latency Analysis … C.7. Surface Entropy Versus Latent Prediction Error

Additional analyses show that ESamp is computationally lightweight, robust across model scales and random seeds, and effective because its latent prediction error captures structured novelty rather than arbitrary noise. ESamp generally preserves single-sample accuracy while improving multi-sample coverage and outperforming surface- or vocabulary-space alternatives.

  • B.4. Latency Analysis: 15–20ms of layer execution leaves ample slack for Distiller inference, which takes less than 0.5ms on Llama-3-8B with an A100.The margin ensures Distiller logits are available before logit fusion begins.
  • C.1. Hyperparameter Sensitivity Across Model Scales: ESamp uses a fixed first-to-final-layer Distiller and β = 0.25 as a robust default across Qwen3 model scales without per-model tuning.The implementation exposes β as the primary exploration-strength hyperparameter and reports low sensitivity to small changes.
  • C.2. Pass@1 Accuracy and the Exploration–Grounding Trade-off: ESamp matches or exceeds Vanilla in most Pass@1 settings, with no systematic accuracy decline despite occasional small decreases.Its objective is candidate-set coverage rather than optimizing the single most likely sample, while some settings show substantial Pass@1 gains.
  • C.3. Does Latent Prediction Error Encode Structured Novelty?: The true Distiller error vector yields substantial gains, whereas matched-magnitude Gaussian noise collapses to approximately Vanilla-level performance.This supports structured novelty information in the error direction rather than generic perturbation.
  • C.4. Latent-Space Versus Vocabulary-Space Novelty: The vocabulary-space Distiller is unstable and substantially underperforms latent-space ESamp because online KL learning over high-dimensional discrete distributions is noisy.ESamp instead estimates novelty in compact continuous representation space, improving online stability and usefulness for exploration.
  • C.5. Multi-Seed Stability: Across seeds 41, 42, and 43, ESamp remains stable and on AIME25 trades lower Pass@8 for stronger Pass@32 and Pass@64.The trade-off reflects its intended use for improving coverage when multiple candidates are sampled.
  • C.6. Baseline Hyperparameter Tuning: Even the weakest ESamp setting matches or exceeds the best-tuned baseline settings, while the default β = 0.25 achieves the best result.The comparison uses three-point grids for the main decoding baselines on Qwen3-8B AIME24 Pass@64.
  • C.7. Surface Entropy Versus Latent Prediction Error: ESamp achieves higher diversity than entropy-adaptive decoding, indicating that latent representation novelty captures semantic variation beyond token-level entropy.The entropy baseline adjusts candidates according to their contribution to normalized token-distribution entropy.

C.8. Composability with FIRE and Self-Consistency … D.3. Open-Source Throughput

ESamp composes with complementary decoding and selection methods, while its architecture, sharing strategies, creative-writing behavior, and open-source runtime implementation support robust and efficient deployment. The optimized open-source path reaches approximately 98.8% of optimized vLLM throughput under the reported benchmark.

  • C.8. Composability with FIRE and Self-Consistency: ESamp combined with FIRE improves Pass@64 beyond either method alone, while remaining compatible with Self-Consistency and yielding slight gains at larger budgets.ESamp modifies candidate-token preferences using representation-space novelty, whereas FIRE changes temperature and Self-Consistency aggregates completed answers by majority vote.
  • C.9. Distiller Architecture Robustness: Pass@k results remain robust across Distiller architectures, motivating the two-layer Gated SwiGLU design as the default for comparable high-budget coverage at lower computational cost.The ablation compares the default two-layer Gated SwiGLU MLP with deeper and simpler alternatives on Qwen3-8B / AIME25.
  • C.10. Shared Versus Per-Prompt Distillers: Per-prompt Distillers perform better on AIME, whereas shared Distillers slightly improve Pass@16 on LiveCodeBench through a larger effective batch and stronger online learning signal.The preferred sharing strategy depends on task structure, with cross-prompt interference potentially harming heterogeneous reasoning problems.
  • C.11. LLM-as-Judge Evaluation for Creative Writing: ESamp achieves the best diversity rank while maintaining quality close to Vanilla in a Gemini 3 Flash evaluation of 16 generations across 2,000 BookCorpus prompts.Lower ranks are better, and the result supports meaningful creative variation rather than merely increased surface-level randomness.
  • D. The tLLM Framework and Practical ESamp Optimizations: The open-source ESamp implementation improves measured throughput over the internal prototype while preserving a general interface for future runtime-adaptation algorithms.The implementation was engineered in the tLLM framework after the main-paper experiments.
  • D.1. tLLM as a Runtime Layer for Test-Time Intervention: tLLM is a vLLM v1 runtime layer supporting internal-state access, asynchronous side computation, and optional sampling guidance without requiring a private vLLM fork.Its producer–consumer abstraction lets external consumers receive hidden states and metadata during generation.
  • D.2. Engineering Optimizations in ESamp: Engineering optimizations include runtime hooks, localized producer–consumer delivery, post-filter candidate intervention, CUDA graphs, Triton grouped prediction, and compatibility with modern vLLM optimizations.These changes reduce hot-path overhead while preserving ESamp behavior and enabling benchmark validation through functional counters.
  • D.3. Open-Source Throughput: 98.8% of optimized vLLM baseline throughput is achieved in the aligned open-source benchmark, corresponding to a measured throughput reduction of about 1.2%.The benchmark uses Qwen2.5-7B, batch size 8, n = 16, active min-p sampling, and an RTX 4090; generality constraints leave room for further improvement.

D.4. Scope of the Open-Source Implementation · E. Experiment Details

D.4 distinguishes ESamp’s decoding contribution from tLLM’s open-source runtime implementation, which supports reproduction, extension, and broader test-time interventions. E. Experiment Details identifies evaluations of ESamp combined with self-consistency and of alternative Distiller architectures.

  • D.4. Scope of the Open-Source Implementation: ESamp uses online latent distillation to encourage representation-space exploration during decoding.
  • D.4. Scope of the Open-Source Implementation: tLLM and ESamp provide a systems path for reproducing and extending the method in a high-throughput inference engine.
  • D.4. Scope of the Open-Source Implementation: The framework supports activation analysis, hidden-state export or editing, online auxiliary-model training, and candidate-level sampling interventions.
  • D.4. Scope of the Open-Source Implementation: ESamp defines the decoding objective and novelty signal, whereas tLLM supplies the runtime substrate for ESamp-like algorithms.
  • E. Experiment Details: Table 16 examines composing ESamp with self-consistency on Qwen3-8B using AIME24.
  • E. Experiment Details: Table 17 reports a Distiller architecture ablation on Qwen3-8B using AIME25.

E.1. Benchmarks and Metrics

The evaluation spans mathematical reasoning, code generation, scientific reasoning, and creative writing, using official benchmark data and lighteval evaluation. Performance and diversity are measured with correctness, semantic redundancy, diversity, and fluency/coherence metrics.

  • Benchmarks: Main experiments evaluate AIME 2024, AIME 2025, LiveCodeBench v5, and GPQA-Diamond using official data and lighteval evaluation code.These benchmarks cover mathematics, code generation, and scientific reasoning.
  • Benchmarks: AIME 2024 and 2025 assess advanced mathematical reasoning through multi-step problems requiring integer answers from 0 to 999.The problems target high-performing high school students and evaluate recent complex mathematical tasks.
  • Benchmarks: LiveCodeBench v5 evaluates correct and efficient Python code on post-cutoff competitive-programming problems from platforms including LeetCode, AtCoder, and Codeforces.Its time-sensitive design aims to minimize training-data contamination.
  • Benchmarks: GPQA-Diamond measures deep scientific reasoning on expert-verified biology, physics, and chemistry multiple-choice questions resistant to simple information retrieval.The Diamond subset is the benchmark’s highest-quality tier and requires domain knowledge.
  • Benchmarks: BookCorpus evaluates diversity in creative writing using a cleaned collection containing approximately 51,442 classic literary works.The dataset is sourced from the incredible45/Gutenberg-BookCorpus-Cleaned-Data-English repository on HuggingFace.
  • Evaluation Metrics: Pass@k measures whether at least one of k samples is correct, Embedding Similarity measures semantic redundancy, Vendi Score measures diversity, and PPL proxies fluency and coherence.Vendi Score uses its official implementation.

E.2. Implementation Details … 5. Conclusion

The appendix specifies ESamp’s lightweight implementation and evaluation setup, then contrasts its diverse reasoning traces with the consistency of vanilla generation. The case study concludes that ESamp produces broader mathematical problem-solving strategies while preserving correctness and efficiency in the comparison.

  • E.2. Implementation Details: The Latent Distiller uses a two-layer residual Gated SwiGLU MLP with intermediate dimension 384, while inference exploration strength is β = 0.25.Training uses Adam with learning rate 4 × 10−4 and ϵ = 1 × 10−4.
  • E.2. Implementation Details: Sharing one distiller across an AIME batch produced no significant performance difference from maintaining independent distillers.The authors attribute this observation to the mathematical distinctness of AIME problems.
  • E.3. Beam Search Results: Diverse Beam Search and Stochastic Beam Search performed poorly on Pass@k relative to sampling-based methods, so detailed beam results were omitted.Beam search’s maximum-joint-probability objective conflicts with exploration for reasoning tasks.
  • E.4. OverRIDE and Contrastive Decoding: Contrastive Decoding uses Qwen2.5-0.5B-Instruct or Qwen3-0.6B as the amateur model, while OverRIDE uses λ = 0.8, 10 iterations, adapter rank 16, and learning rate 10−3.GPT-OSS-20B is excluded from Contrastive Decoding because no corresponding small model shares its vocabulary.
  • E.5. Prompts; AIME 2024 / AIME 2025; LiveCodeBench (Code Generation): Evaluations use official lighteval prompts, with dedicated formats for AIME answers, LiveCodeBench programs and multiple-choice outputs.The AIME template requires a boxed final answer, while code prompts specify stdin/stdout behavior and code delimiters.
  • Tree of Thoughts (self-review system prompt): The Tree-of-Thoughts self-review prompt asks models to explore distinct branches, prune wrong ones, and synthesize a step-by-step final answer.It requires the response to end with “Final Answer: <number>” for an integer from 0 to 999.
  • E.6. Case Study; 2. Analysis of File 1: In a single-blind Gemini 3 Flash Preview evaluation of 16-answer sets, ESamp’s file showed significant variety in mathematical modeling and execution strategies.The case study compares vanilla and ESamp generations because evaluating all parallel-generation tokens directly would be impractical.
Loading 2604.24927v2…