Large Language Models Explore by Latent Distilling
| Authors | Yuanhao Zeng et al. |
| Year | 2026 |
| HF Upvotes | 63 |
| arXiv | 2604.24927 |
| Download | |
| Code | https://github.com/LinesHogan/tLLM |
Abstract
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.
Engineering Breakdown
Plain English
This paper addresses a fundamental limitation in test-time scaling of large language models: standard sampling methods like temperature-based decoding produce diverse outputs at a surface level (different words) but fail to explore meaningfully different semantic directions. The authors propose Exploratory Sampling (ESamp), which trains a lightweight 'Distiller' network at test time to model how an LLM's internal representations evolve across layers, then uses this model to guide token selection toward latent space regions the model has high uncertainty about. The core insight is that LLMs make lower-error predictions on familiar inputs and higher-error predictions on novel ones, so exploring uncertain regions in representation space should yield semantically diverse generations that improve test-time scaling performance.
Core Technical Contribution
The key novelty is the depth-wise representation distillation approach: rather than sampling based on output logits alone, ESamp trains a shallow model to predict deep-layer hidden states from shallow-layer representations, creating a measure of representation uncertainty across the network's depth. This uncertainty signal directly guides token selection during decoding—tokens that create representations the Distiller struggles to predict are preferred because they explore the model's latent space more thoroughly. Unlike prior work that relies on output-level diversity metrics (entropy, logit variance), this approach operates in the internal representation space where semantic structure actually lives, enabling fundamentally deeper exploration than surface-level lexical variation.
How It Works
During inference, ESamp operates in three steps: First, at each decoding step, the LLM computes hidden representations at all layers for each candidate token. Second, a lightweight Distiller model—trained on-the-fly using only recent generation context—attempts to predict the deep-layer representations from shallow-layer representations, computing a reconstruction loss across all layers. Third, tokens are selected not just by likelihood but by a combined score that rewards both high probability and high distillation error (high uncertainty). The Distiller uses a simple architecture (likely MLPs or small transformers) trained on just the current sequence context, keeping computational overhead minimal while maintaining an up-to-date model of the LLM's depth-wise latent structure. This creates a feedback loop where exploring uncertain regions gradually shifts the representation manifold the model encounters, preventing collapse into repetitive generation patterns.
Production Impact
For production systems running LLM inference at scale, ESamp offers a concrete way to improve test-time scaling—generating multiple diverse candidate responses and selecting the best one—without expensive fine-tuning or retraining. The approach is orthogonal to model size and architecture, so it integrates cleanly into existing inference pipelines: you add a lightweight distiller module that runs in parallel with the main LLM forward pass. The trade-offs are meaningful: each token requires training the Distiller (typically seconds to minutes of overhead per sequence depending on Distiller size), and you need to store representations at all intermediate layers (moderate memory overhead, roughly proportional to model hidden dimension × sequence length). For applications where semantic diversity directly impacts quality—multi-turn dialogue, code generation with multiple valid solutions, creative writing—the improved exploration often justifies the 5-15% inference latency increase. For latency-critical applications (sub-100ms response time), this approach may be impractical without further optimization.
Limitations and When Not to Use This
The paper assumes that depth-wise representation uncertainty is a reliable proxy for semantic diversity, which may not hold uniformly across all domains or model architectures—some models may have high uncertainty in shallow layers that maps to low semantic variation, or vice versa. Computational cost scales with model depth and Distiller capacity, making it increasingly expensive for very large models; the paper doesn't clearly quantify the overhead or provide ablations on Distiller architecture choices. The approach also requires maintaining representations in memory during decoding, which becomes problematic for very long sequences or batch processing many queries simultaneously. It's unclear how well ESamp generalizes across different model families, scales, and training objectives—the paper would benefit from experiments on a broader range of models (different architectures, sizes, training data) to establish robustness. Additionally, there's limited analysis of failure modes: when does the Distiller's uncertainty signal mislead, or when does high-uncertainty exploration actually degrade generation quality?
Research Context
This work sits at the intersection of test-time scaling research (inference-time improvements without retraining, akin to speculative decoding or ensemble methods) and interpretability work on LLM latent spaces. It builds on the well-established observation that neural networks exhibit higher prediction error on out-of-distribution inputs, extending this principle from input space to internal representation space—a shift that mirrors recent work on uncertainty quantification and latent space exploration in generative models. The paper contributes to the growing body of research on decoding strategies beyond greedy and standard sampling, joining approaches like nucleus sampling, diversity-promoting penalties, and contrastive decoding. This opens a research direction around using test-time learned models (distillers, probes) to guide inference, which could extend to other generation tasks (image, code, structured data) and to multi-stage decoding where early-layer uncertainty informs later-layer predictions.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::