Curriculum Learning for Efficient Chain-of-Thought Distillation via Structure-Aware Masking and GRPO

2026-02-05

arXiv.org

AI Summary

This work addresses the challenge that chain-of-thought (CoT) reasoning in large language models is often verbose and difficult to distill effectively into smaller models without compromising interpretability. To tackle this, the authors propose a three-stage curriculum learning framework. First, structure-aware masking combined with shuffled reconstruction enhances the student model’s understanding of CoT structural patterns. Second, grouped relative policy optimization (GRPO) is introduced during masked completion to jointly optimize reasoning accuracy and conciseness. Finally, failure-case-driven targeted rewrites coupled with GRPO reinforcement internalize critical reasoning knowledge. This approach, the first to integrate structure-aware masking with GRPO, achieves an 11.29% absolute accuracy gain and a 27.4% reduction in output length for Qwen2.5-3B-Base on GSM8K, significantly outperforming standard instruction tuning and existing distillation methods.

Abstract

Distilling Chain-of-Thought (CoT) reasoning from large language models into compact student models presents a fundamental challenge: teacher rationales are often too verbose for smaller models to faithfully reproduce. Existing approaches either compress reasoning into single-step, losing the interpretability that makes CoT valuable. We present a three-stage curriculum learning framework that addresses this capacity mismatch through progressive skill acquisition. First, we establish structural understanding via masked shuffled reconstruction. Second, we apply Group Relative Policy Optimization (GRPO) on masked completion tasks, enabling the model to discover its own balance between accuracy and brevity. Third, we identify persistent failure cases and guide the student to internalize teacher knowledge through targeted rewriting, again optimized with GRPO. Experiments on GSM8K demonstrate that our approach enables Qwen2.5-3B-Base to achieve an 11.29 percent accuracy improvement while reducing output length by 27.4 percent, surpassing both instruction-tuned variants and prior distillation methods.

Problem

• Chain-of-Thought Distillation
• Curriculum Learning
• Model Compression
• Reasoning Efficiency
• Interpretability

Innovation

• Curriculum Learning
• Chain-of-Thought Distillation
• Structure-Aware Masking
• Group Relative Policy Optimization
• Reasoning Compression