DiffAdapt: Difficulty-Adaptive Reasoning for Token-Efficient LLM Inference

AI summary

Overview

Research area: Efficient inference for reasoning large language models, specifically difficulty-adaptive allocation of test-time compute (chain-of-thought length control).

Technical level: Intermediate. The paper is readable without deep mathematics, but assumes familiarity with chain-of-thought reasoning, token-level probability entropy, hidden-state representations, and lightweight classification probes.

Scope: The paper diagnoses a "U-shaped" entropy pattern across reasoning models, proposes a three-stage framework (DiffAdapt) that predicts question difficulty from a frozen LLM's hidden states, and empirically evaluates whether routing each question to an Easy, Normal, or Hard inference strategy can preserve or improve accuracy while cutting token usage.

What This Paper Is About

Reasoning LLMs typically generate long thinking traces for every question, spending the same compute on trivial problems as on genuinely hard ones. The authors first analyze token-probability entropy across difficulty levels and find a consistent U-shaped curve — high entropy on easy problems despite high accuracy — which they interpret as evidence of overthinking. They then build DiffAdapt, a framework that trains a small probe on the model's own hidden states to classify each incoming question as Easy, Normal, or Hard and apply a corresponding inference strategy, without fine-tuning the base LLM.

Key Contributions

1. An empirical characterization of overthinking via entropy. Across three models, the authors report a consistent U-shaped entropy pattern: high entropy on easy problems (despite high accuracy), low entropy on medium-difficulty problems, and high entropy on hard problems. They quantify a 22–25% entropy reduction from the easy to the medium difficulty region.

2. A three-tier inference strategy design. Three strategies (Easy, Normal, Hard) each pair a fixed prompt with a sampling temperature and a maximum token length. The hyperparameters were chosen by a grid search over 125 parameter combinations (5 Normal × 5 Hard × 5 Easy) on MATH500 rather than by heuristic selection, with a selection criterion of accuracy ≥ 95% followed by minimal average token consumption.

3. The DiffAdapt framework. A three-stage pipeline: (1) a proxy model samples responses (10 iterations per problem, temperature 0.6, max length 32K tokens) and heuristically labels data with Easy/Normal/Hard; (2) a small MLP probe is trained on the last-layer hidden state after prefilling the question; (3) at inference, the probe selects a strategy. The base LLM weights stay frozen, so only the probe is trained (100 epochs, AdamW, learning rate 1e-3).

4. Evaluation across five models and eight benchmarks, including comparison to fixed strategies and the training-free dynamic baseline DEER, an orthogonality study with length-control RL models, efficiency and latency measurements, ablations, and a blind LLM-as-a-Judge study of reasoning quality.

Main Findings

• U-shaped entropy across difficulty: Difficulty was measured on DeepMath-103K (difficulty ratings 1–10 assigned by GPT-4o), using 300 questions per difficulty level and 10 sampling iterations per problem at temperature 0.6. Entropy is computed as the average per-token entropy over the generated sequence. The pattern is reported for DeepSeek-R1-Distill-Qwen-7B on DeepMath-103K and said to be consistent across multiple architectures.

• Three difficulty regions map to three strategies: the "Overthinking Region" (easy, high correctness with high entropy), the "Certainty Region" (normal, low entropy with optimal performance), and the "Capability Limit Region" (hard, high entropy with low accuracy). The hard region motivates a "Fail Fast" strategy that cuts losses rather than reasoning longer.

• Oracle headroom is large: in oracle experiments with Qwen3-4B across eight benchmarks (max token limit 32K), selecting the correct strategy with minimal token use yields a 7.2% average accuracy gain over the best fixed baseline. Token allocation in the oracle ranged from 198 tokens on GSM8K to 4,675 on AIME25. Earlier oracle experiments reported by the authors achieved 50% token savings while improving accuracy by over 10%.

• DiffAdapt beats fixed strategies across models and domains: in-domain results (GSM8K, MATH500, AIME24&25, OlympiadBench) show the largest improvements for Qwen3-4B at higher token budgets, stable gains for DeepSeek-R1-Qwen-7B, and benefits across model families for DeepSeek-R1-Llama-8B. Out-of-domain (Minerva, GPQA, MMLU-Pro) gains become more pronounced as the maximum token budget increases. The paper states that models with larger inter-strategy performance differences show greater DiffAdapt improvements.

• DiffAdapt outperforms DEER: DEER performs comparably to the strongest fixed strategy in-domain, but exhibits limited generalization and larger degradation under distribution shift. The introduction states that relative to DEER, DiffAdapt achieves up to 62% token reduction and 18% performance improvement across eight mathematical reasoning benchmarks on five models.

• Token savings versus a fixed Normal strategy: Qwen3-4B 22.4%, DeepSeek-R1-Qwen-7B 9.7%, ThinkPrune-7B 10.1%. DEER, by contrast, uses more tokens than the baseline: −53.3% for DS-R1-Qwen-7B and −27.5% for Qwen3-4B. The authors attribute this mechanistically to DEER's confidence-based continuation decisions frequently driving generation to the maximum token cap rather than adapting by difficulty.

• Latency: under identical settings (Qwen3-4B, first 40 OlympiadBench problems, batch size 10, single A800 GPU, max token limit 32K, temperature 0.6, DEER think threshold 0.9), end-to-end wall-clock time was 64 minutes for the vLLM baseline, 57 for +DEER, and 10 for +DiffAdapt — reported as 6× faster than the vLLM baseline and 5× faster than DEER.

• Orthogonality with length-control RL: on Nemotron-1.5B and ThinkPrune-7B, Easy strategies perform best in most settings because these models are trained to solve problems with low compute. DiffAdapt is slightly below Easy strategies at low token limits but reaches state-of-the-art results under high token budgets.

• Robustness ablations (Qwen3-4B, in-domain): transferring thresholds (α=0.85, β=0.35, γ=0.60) from the DeepSeek-R1 family yields a negligible average difference of approximately 0.3% (average 71.2 vs. 70.9 for the default). Replacing the 2-layer MLP with a Linear Head drops accuracy by roughly 3.2% on average (67.7). Training on only 30% of the data gives an average of 68.5. Across token budgets of 33.3%, 50.0%, 66.7%, 83.3%, and 100%, the default configuration scores 59.3, 67.2, 72.4, 75.8, and 76.8.

• Reasoning integrity is largely preserved: in a blind pairwise LLM-as-a-Judge study on 50 GSM8K queries judged by Qwen3-30B, DiffAdapt won 76% (38/50), the baseline won 12% (6/50), and 12% (6/50) were ties. Failure analysis of baseline wins attributed 10% (5/50) to subjective preference and 2% (1/50) to truncation error, meaning actual logic failure due to truncation occurred in only 2% of cases.

• Zero-shot transfer to MMLU-Pro: using probes and thresholds trained only on DeepMath, DiffAdapt outperforms fixed-strategy baselines by 3–7% across token budgets and for DeepSeek-R1-Qwen-7B and DeepSeek-R1-Llama-8B, indicating the difficulty signal transfers to unseen domains without re-training.

Methodology in Plain English

The authors begin diagnostically. They take a large, decontaminated math dataset with difficulty ratings, sample 300 questions per difficulty level, generate 10 answers per question, and compute two things for each: how often the model is correct, and the average entropy (uncertainty) of the token probability distribution across the generated trace. Plotting accuracy and entropy against difficulty reveals the U-shape.

From that shape they hand-design three inference modes. Easy mode uses a low temperature (0.5) and only 0.4× the maximum token budget with a prompt telling the model the problem looks straightforward. Normal mode uses temperature 0.8 and the full 1.0× budget with a step-by-step prompt. Hard mode uses temperature 0.4 and 0.5× the budget with a "fail fast" prompt that asks the model to outline a method while being mindful of resources. They first test the theoretical ceiling by running all three strategies on every question and picking the best one per question (the oracle).

The actual system has two components. First, a labeling step: a proxy model (typically the same LLM) generates 10 responses per training problem, entropy and correctness are computed, and a heuristic rule assigns a label — Normal if correctness is at least α and entropy at most β, Hard if correctness is below γ, and Easy otherwise. Second, a probe: after the model reads (prefills) the question, its last-layer hidden state is extracted and fed to a small multi-layer perceptron trained with cross-entropy loss to predict the three-way difficulty label. At inference, the probe runs once on the prefill state and picks the strategy; because it does not intervene in prefilling or decoding, it composes with batching, KV cache, and prefix caching.

Validation covers three reasoning models (Qwen3-4B, DeepSeek-R1-Qwen-7B, DeepSeek-R1-Llama-8B) and models trained with length-control RL (Nemotron-1.5B, ThinkPrune-7B), on five math benchmarks and three out-of-domain benchmarks, with each experiment run three times and the mean reported. Comparisons include fixed Easy/Normal/Hard strategies and DEER, a training-free early-exit method run with its default think threshold of 0.9.

Why This Matters

Impact on research. The paper reframes test-time scaling from "more tokens is better" to "the right number of tokens per problem." The entropy analysis offers a measurable diagnostic for overthinking, and the finding that a linear probe on prefill hidden states suffices for difficulty routing suggests difficulty may already be linearly decodable inside a frozen model — a claim relevant to interpretability and to efficiency research more broadly. The orthogonality result with length-control RL indicates that inference-time routing is complementary to, rather than competing with, training-time efficiency methods.

Real-world applications:

• Serving reasoning models at scale, where per-query token spend is the dominant cost driver and uniform budgets cause latency and throughput degradation.
• Cost-sensitive deployments such as educational tutoring or coding assistants, where many user queries are easy and do not require long deliberation.
• Batch or latency-sensitive pipelines (for example, evaluation harnesses or agent loops) where a lightweight routing decision avoids regeneration of long traces.
• High-stakes query triage where the "Fail Fast" mode can cap spend on questions that are likely beyond the model's capability, rather than burning the full budget on a likely-wrong answer.

Industry relevance. The method requires no retraining of the base LLM, adds only a small classifier on hidden states, and is explicitly stated to be compatible with batching, KV cache, and prefix cache. Reported reductions of up to 22.4% in token usage and a 6× wall-clock speedup against a vLLM baseline on one configuration translate directly into lower serving cost at comparable accuracy, which is the metric most infrastructure teams optimize.

Future Directions

• Calibration-free deployment. Thresholds α, β, γ are tuned per model; although transferring DeepSeek-R1 thresholds to Qwen3 changed average accuracy by only about 0.3%, the authors note that extremely distinct domains may still need a lightweight calibration phase and call for completely calibration-free mechanisms.

• Using generation-time signals. The probe deliberately relies only on prefill hidden states to avoid disturbing decoding, which means it ignores dynamics that emerge during generation. Incorporating lightweight generation signals such as early step-wise entropy is proposed as future work, with the caveat that it introduces an efficiency–precision trade-off.

• Label stability under distribution and sampling shifts. Because difficulty labels come from a proxy model with a fixed sampling protocol (10 samples at temperature 0.6), the authors recommend a light validation phase to re-establish thresholds when transferring across domains or changing sampling configurations.

• Fail-safes under tight budgets. Aggressive budget reduction can hurt accuracy on particularly hard or error-prone cases; a practical fallback is to revert to the Normal strategy when probe confidence is low, the prefill signal is out of distribution, or the selected strategy has underperformed recently.

• Composition with orthogonal efficiency work. The paper notes that weight pruning and sparsity allocation, zero-cost proxy search for compact architectures, and efficient attention designs reduce per-token cost and can be composed with DiffAdapt's token-count reduction.

Target Audience

Researchers and engineers working on LLM inference efficiency, test-time compute scaling, and reasoning-model serving systems. It is also relevant to practitioners who need to cut serving cost without retraining a base model, and to interpretability researchers interested in what frozen hidden states encode about task difficulty. Readers should be comfortable with chain-of-thought prompting, entropy as an uncertainty measure, and the idea of training a small classifier head on top of a frozen LLM.

Authors’ abstract

Recent reasoning Large Language Models (LLMs) demonstrate remarkable problem-solving abilities but often generate long thinking traces whose utility is unclear. Our work aims to improve their efficiency, enabling them to reach high performance without overthinking. First, we analyze the entropy of token probabilities in reasoning traces. Across three models, we observe a consistent U-shaped entropy pattern: high entropy on easy problems despite high accuracy, low entropy on problems with medium difficulty, and high entropy on hard problems reflecting uncertainty. Specifically, we notice 22--25\% entropy reduction from easy to medium difficulty regions, suggesting an {overthinking} phenomenon on easy instances. Building on these insights, we introduce \textbf{DiffAdapt}, a lightweight framework that selects Easy/Normal/Hard inference strategies per question based on their difficulty and reasoning trace entropy. Each inference strategy consists of a fixed prompt, temperature and maximum token length. In contrast to existing efficiency optimization methods, our approach does not fine-tune base LLM but a small probe that classifies LLM's final hidden state, allowing inexpensive adaptation. We comprehensively evaluate our method on five models and eight benchmarks. Our method achieves comparable or improved accuracy while reducing token usage by up to 22.4\%, establishing a practical path toward compute-efficient reasoning.