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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Chen, Weize, Yuan, Jiarui, Jin, Tailin, Ding, Ning, Chen, Huimin, Liu, Zhiyuan, Sun, Maosong
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866915304160886784
author Chen, Weize
Yuan, Jiarui
Jin, Tailin
Ding, Ning
Chen, Huimin
Liu, Zhiyuan
Sun, Maosong
author_facet Chen, Weize
Yuan, Jiarui
Jin, Tailin
Ding, Ning
Chen, Huimin
Liu, Zhiyuan
Sun, Maosong
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2505_19217
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle The Overthinker's DIET: Cutting Token Calories with DIfficulty-AwarE Training
Chen, Weize
Yuan, Jiarui
Jin, Tailin
Ding, Ning
Chen, Huimin
Liu, Zhiyuan
Sun, Maosong
Computation and Language
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.
title The Overthinker's DIET: Cutting Token Calories with DIfficulty-AwarE Training
topic Computation and Language
url https://arxiv.org/abs/2505.19217