Concise Reasoning, Big Gains: Pruning Long Reasoning Trace with Difficulty-Aware Prompting

TL;DR

The paper tackles inefficiencies in chain-of-thought distillation caused by verbose reasoning traces and lack of adaptation to problem difficulty. It introduces Difficulty-Aware Prompting (DAP), a two-step approach that first generates long CoTs with a strong teacher and then rewrites them into concise, difficulty-appropriate traces, yielding LiteCoT—a 100K dataset with an average of ~720 tokens per sample. Trained on LiteCoT, the Liter family (1.5B, 7B, 32B) demonstrates superior accuracy and notably lower inference token requirements across 11 benchmarks, including 74.2% Pass@1 on AIME24 with only about 5K inference tokens. The method generalizes well to diverse tasks and offers substantial training and inference efficiency gains over prior long-CoT distillation approaches and competing CoT optimization techniques.

Abstract

Existing chain-of-thought (CoT) distillation methods can effectively transfer reasoning abilities to base models but suffer from two major limitations: excessive verbosity of reasoning traces and inadequate adaptability to problem difficulty. Long reasoning traces significantly increase inference costs, and uniform-length solutions prevent base models from learning adaptive reasoning strategies. To address these issues, we propose a difficulty-aware prompting (DAP) method to dynamically shorten reasoning traces without performance loss. In our approach, a large teacher model first judges each problem's difficulty and then rewrites its reasoning traces to an appropriate shorter length, yielding concise yet complete reasoning traces. Leveraging the DAP pipeline, we curate a distilled dataset called LiteCoT consisting of 100K concise reasoning examples, with solutions averaging only 720 tokens (an order of magnitude shorter than typical CoTs). Using LiteCoT, we distilled a new family of reasoning models called Liter (1.5B, 7B, and 32B) based on the Qwen2.5 architecture. Experiments show that a student model fine-tuned on just 100K of these difficulty-pruned CoT samples outperforms a model distilled on 800K original Long CoT samples, while significantly reducing training and inference costs. Our method also generalizes well: across 11 diverse benchmarks, the shorter difficulty-aware CoTs achieve equal or better accuracy than Long chains, using far fewer tokens. For example, on the challenging AIME24 exam, our approach reaches $74.2\%$ Pass@1 using only about 5K inference tokens, surpassing other methods that consume many more tokens. Our code and data are available at https://github.com/Evanwu1125/LiteCoT.

Figures (7)

• Figure 1: (a) Training tokens comparison of LiteCoT with prior distillation methods (32B model); (b) Inference efficiency and accuracy on the challenging AIME24 benchmark (32B model). Our Liter models trained with LiteCoT achieve higher accuracy (74.2% Pass@1) at substantially reduced inference token counts compared to prior methods.
• Figure 2: Overview of the DAP pipeline. First, the DeepSeek R1 model generates initial Long CoT data from collected user queries. Then, a difficulty-aware prompt guides the model to rewrite concise CoT outputs, adapting reasoning length dynamically based on problem difficulty.
• Figure 3: Token count distribution of LiteCoT compared with previous CoT distillation datasets.
• Figure 4: Inference time vs. accuracy across eight benchmarks. Models trained with short CoT (filled symbols) demonstrate significantly faster inference times compared to Long CoT (hollow symbols) while maintaining or exceeding accuracy levels.
• Figure 5: Radar chart show the performance comparison of the Vanilla, DeepSeek-R1-Distilled, and Liter versions of the Qwen2.5 model with different parameter sizes (1.5B, 7B, 14B, 32B) on multiple math and reasoning benchmarks (Math500, AIME24, OlympiadBench, AIME25, GPQA, AMC23).