The Overthinker's DIET: Cutting Token Calories with DIfficulty-AwarE Training

TL;DR

DIET introduces a difficulty-aware training framework that reduces token usage in large language models while improving reasoning performance, enabling better scaling and more efficient inference.

Contribution

The paper proposes DIET, a novel method integrating task difficulty into RL training to optimize token compression and model efficiency.

Findings

Significantly reduces token counts without sacrificing reasoning quality.

Enhances inference scaling by maintaining high per-sample quality with fewer tokens.

Preserves the positive correlation between response length and difficulty, improving verbosity management.

Abstract

Recent large language models (LLMs) exhibit impressive reasoning but often over-think, generating excessively long responses that hinder efficiency. We introduce DIET ( DIfficulty-AwarE Training), a framework that systematically cuts these "token calories" by integrating on-the-fly problem difficulty into the reinforcement learning (RL) process. DIET dynamically adapts token compression strategies by modulating token penalty strength and conditioning target lengths on estimated task difficulty, to optimize the performance-efficiency trade-off. We also theoretically analyze the pitfalls of naive reward weighting in group-normalized RL algorithms like GRPO, and propose Advantage Weighting technique, which enables stable and effective implementation of these difficulty-aware objectives. Experimental results demonstrate that DIET significantly reduces token counts while simultaneously improving reasoning performance. Beyond raw token reduction, we show two crucial benefits largely overlooked by prior work: (1) DIET leads to superior inference scaling. By maintaining high per-sample quality with fewer tokens, it enables better scaling performance via majority voting with more samples under fixed computational budgets, an area where other methods falter. (2) DIET enhances the natural positive correlation between response length and problem difficulty, ensuring verbosity is appropriately allocated, unlike many existing compression methods that disrupt this relationship. Our analyses provide a principled and effective framework for developing more efficient, practical, and high-performing LLMs.