PaperlayerView exact revision on arXiv

Source-linked AI summary

Keypoint-based Progressive Chain-of-Thought Distillation for LLMs

Kaituo Feng, Changsheng Li, Xiaolu Zhang, Jun Zhou, Ye Yuan, Guoren Wang

arXiv:2405.16064v1cs.CL

TL;DR

CoT distillation aims to transfer reasoning abilities from large, costly LLMs to smaller models, but existing methods treat rationale tokens equally and do not order learning by reasoning difficulty. KPOD addresses these issues with token-significance weighting and in-rationale progressive distillation, and experiments on four reasoning benchmarks show significant improvements over baselines.

  • Problem

    Existing CoT distillation methods may inadequately mimic keypoint tokens and fail to distinguish the learning order of reasoning steps.

    Source: §1 · pp. 1–2

  • Method

    KPOD combines mask-learned rationale token weighting with progressive distillation that expands from final reasoning steps to the full rationale in an easy-to-hard order.

    Source: §1 · p. 2, §3.3 · p. 4, §3.4 · p. 6

  • Results

    KPOD shows significant performance improvements over baselines on four reasoning benchmarks.

    Source: §1 · p. 2

  • Takeaways & Limitations

    KPOD provides a unified CoT distillation framework that emphasizes keypoint tokens and progressively teaches reasoning steps from easy to hard.

    Source: §1 · p. 2, §5 · p. 8

Abstract

from arXiv · show

Chain-of-thought distillation is a powerful technique for transferring reasoning abilities from large language models (LLMs) to smaller student models. Previous methods typically require the student to mimic the step-by-step rationale produced by LLMs, often facing the following challenges: (i) Tokens within a rationale vary in significance, and treating them equally may fail to accurately mimic keypoint tokens, leading to reasoning errors. (ii) They usually distill knowledge by consistently predicting all the steps in a rationale, which falls short in distinguishing the learning order of step generation. This diverges from the human cognitive progression of starting with easy tasks and advancing to harder ones, resulting in sub-optimal outcomes. To this end, we propose a unified framework, called KPOD, to address these issues. Specifically, we propose a token weighting module utilizing mask learning to encourage accurate mimicry of keypoint tokens by the student during distillation. Besides, we develop an in-rationale progressive distillation strategy, starting with training the student to generate the final reasoning steps and gradually extending to cover the entire rationale. To accomplish this, a weighted token generation loss is proposed to assess step reasoning difficulty, and a value function is devised to schedule the progressive distillation by considering both step difficulty and question diversity. Extensive experiments on four reasoning benchmarks illustrate our KPOD outperforms previous methods by a large margin.

1. Introduction

KPOD addresses two weaknesses in chain-of-thought distillation: equal treatment of tokens and undifferentiated learning of rationale steps. It combines keypoint-aware token weighting with progressive easy-to-hard distillation.

1. Introduction

  • CoT distillation transfers large-model reasoning to smaller students by having them mimic generated step-by-step rationales.This targets the deployment cost and resource requirements of very large reasoning models.Source: §1 · p. 1
  • Rationale tokens differ in importance, but prior methods treat them equally, risking errors when key calculation tokens are not precisely mimicked.Reasoning-irrelevant words can be replaced, whereas deviations in keypoint tokens can alter the reasoning outcome.Source: §1 · p. 1
  • Prior methods also learn all rationale steps simultaneously, rather than following the human-like progression from easier to harder reasoning.Direct curriculum-learning applications may misjudge step difficulty because inessential tokens dominate token-level difficulty estimates.Source: §1 · pp. 1–2
  • KPOD learns token significance through mask learning, using answer prediction and mask-ratio losses to identify inessential tokens.The probability that a token is not masked becomes its significance weight.Source: §1 · p. 2
  • Its progressive strategy starts with later rationale steps, then expands toward the full rationale while a value function schedules difficulty and selects diverse questions.A weighted token-generation loss estimates step difficulty while the value function controls progression and question diversity.Source: §1 · p. 2
  • KPOD outperforms baselines on four reasoning benchmarks, according to the reported experiments.The method is presented as a unified framework for token significance and reasoning difficulty.Source: §1 · p. 2

2. Related Works

The related work spans chain-of-thought reasoning, CoT distillation, and curriculum learning. KPOD is positioned as addressing the missing easy-to-hard organization of reasoning steps in prior distillation methods.

Chain-of-Thought Reasoning

  • CoT prompting enables LLMs to generate intermediate reasoning steps, but strong CoT capabilities typically require models exceeding 100 billion parameters.This makes reasoning resource-consuming for deployment.Source: §2 · p. 2

CoT Distillation

  • CoT distillation transfers LLM-generated step-by-step rationales to smaller students, with prior methods targeting self-rationalization, specialization, rationale consistency, or faithfulness.Examples include SCoTD, Specialized KD, MCC-KD, and SCOTT.Source: §2 · pp. 2–3
  • Existing CoT distillation methods do not account for a reasonable learning order among reasoning steps, which the paper associates with sub-optimal performance.This motivates KPOD’s in-rationale progressive distillation strategy.Source: §2 · p. 3

Curriculum Learning

  • Curriculum-learning methods gradually introduce easy-to-hard samples or organize token sequences, but applying them directly to CoT distillation can be inferior.The paper therefore develops a strategy tailored to stepwise rationale generation.Source: §2 · p. 3

3. Proposed Method

KPOD distills LLM reasoning into smaller students by weighting rationale tokens and progressively ordering reasoning-step learning from easier to harder. Its framework combines mask-based token significance estimation with difficulty- and diversity-aware in-rationale scheduling.

3.1. Preliminaries and Problem Setting

  • KPOD transfers LLM-generated rationales to a student model, training it to generate rationale and answer tokens from each reasoning question.The standard objective combines rationale imitation with final-answer prediction.Source: §3.1 · p. 3

3.2. Framework Overview

  • KPOD addresses uniform token treatment and unordered step distillation with token weighting and progressive in-rationale learning.The token weights both emphasize keypoint-token mimicry and reduce the influence of irrelevant tokens when assessing step difficulty.Source: §3.2 · p. 3

3.3. Rationale Token Weighting

  • The rationale weighting module encodes token embeddings with self-attention and a weight generator, then uses answer prediction and mask-ratio losses to retain keypoint tokens while masking redundant ones.Each token’s non-masking probability serves as its significance weight; Gumbel-Softmax enables discrete mask learning.Source: §3.3 · pp. 3–4

3.4. In-rationale Progressive Distillation

  • The strategy estimates reasoning-step difficulty with weighted token-generation loss and uses a value function to schedule learning order while balancing target difficulty and question diversity.Question diversity is promoted through embedding-based K-means clustering, and the resulting knapsack-constrained optimization is approximately solved with the FTGP submodular maximization algorithm.Source: §3.4 · pp. 5–6

3.5. Training Procedure

  • Training begins by generating the final rationale steps and progressively expands toward complete rationales as scheduled input steps decrease across stages.The framework first learns token significance, derives the schedule from step difficulty, and applies the resulting weighted rationale loss while retaining the answer-prediction term.Source: §3.4 · pp. 5–6, §3.5 · p. 6

4. Experiments

KPOD is evaluated across mathematical and commonsense reasoning benchmarks with multiple student architectures, achieving large gains over prior CoT distillation methods. Ablations, OOD tests, visualizations, and sensitivity analyses support the contributions of token weighting and progressive distillation.

4.2. Overall Performance

  • KPOD exceeds MCC-KD by 5.16%, 5.26%, 4.00%, and 1.48% on GSM8K, ASDiv, SVAMP, and CommonsenseQA, respectively, with LLaMA-7B.CoT distillation methods generally improve smaller student models over models without distillation.Source: §4.2 · p. 6, Table 1 · p. 7

4.3. Ablation Study

  • KPOD’s ablations show that token significance weighting, significance-aware step difficulty, progressive distillation, and question diversity each improve performance.KPOD outperforms variants without significance weighting, significance-based difficulty, progressive scheduling, or diversity-based question selection.Source: §4.3 · p. 7, Table 2 · p. 7
  • KPOD surpasses Adaptive CL, SPL, and ICL variants, supporting its in-rationale progressive distillation strategy over these curriculum-learning alternatives.The comparisons replace KPOD’s progressive strategy with each representative curriculum-learning method.Source: §4.3 · pp. 7–8

4.4. OOD Performance

  • KPOD consistently outperforms baselines when trained on GSM8K and tested on ASDiv and SVAMP, indicating stronger OOD generalization.The OOD evaluation uses GSM8K for training and ASDiv and SVAMP for testing.Source: §4.4 · p. 8, Table 3 · p. 8

4.5. Visualizations

  • The token-weighting visualizations assign highest weights to digit and operation tokens, while instructional cues and meaningful subjects receive relatively high weights and less important tokens receive low weights.Figure 2 encodes larger significance weights with deeper red intensity.Source: §4.5 · p. 8, Figure 2 · p. 8

4.6. Parameter Sensitivity Analysis

  • Performance is not sensitive to α or β across relatively large ranges, making these hyperparameters easy to set in practice.α controls the mask-ratio loss, while β belongs to the diversity term for question-set selection.Source: §4.6 · p. 8, Figure 3 · p. 8

5. Conclusion

KPOD combines keypoint-aware token weighting with progressive chain-of-thought distillation so smaller students acquire reasoning capabilities from teacher LLMs in an easy-to-hard manner.

  • KPOD uses rationale token weighting to encourage students to accurately mimic keypoint tokens during distillation.Source: §5 · p. 8
  • Its in-rationale progressive distillation strategy schedules reasoning-step learning from easy to hard.Source: §5 · p. 8
  • Extensive experiments validate KPOD’s effectiveness.Source: §5 · p. 8

A. Implementation Details

The implementation trains LLaMA-7B and FlanT5 students with LoRA, specified optimization settings, and fixed progressive-distillation hyperparameters across four datasets.

B. Training Pseudo-code

KPOD first derives teacher rationales and token significance, then estimates step difficulty, schedules progressively harder question sets, and trains the student in that order before full-rationale generation.

C. Additional Experiments

Additional GSM8K experiments show that KPOD remains relatively stable across broad settings of the learning-difficulty growth rate, cluster count, and initial difficulty.

  • KPOD performance on GSM8K is relatively stable across varying p, K, and C0 settings.Here, p controls learning-difficulty growth, K is the number of question clusters, and C0 is the initial difficulty.Source: C. Additional Experiments · p. 13

D. Proof

This section proves Proposition 3.1 by showing that the value function F is monotone and submodular, enabling application of FTGP’s approximation guarantee. The proof establishes these properties through set inclusions and diminishing marginal gains.

D. Proof

  • Theorem D.1 states that FTGP approximately solves monotone submodular maximization under a knapsack constraint with an approximation-ratio guarantee.Its running time is O(nϵ−1 log ϵ−1), where n is the data scale and ϵ is a hyper-parameter.Source: D. Proof · p. 13
  • The proof reduces Proposition 3.1 to verifying that the value function F is monotone and submodular.Once both properties are established, Theorem D.1 completes the proposition.Source: D. Proof · p. 13
  • Lemma 1 asserts that F is monotone by comparing nested question sets A(t) ⊆ B(t) and analyzing F(B(t)) − F(A(t)).Monotonicity is defined as F(A) ≤ F(B) whenever A ⊆ B.Source: D. Proof · p. 13
  • Lemma 2 asserts that F is submodular by comparing marginal gains from adding x to nested sets A and B.Submodularity requires the marginal gain for A to be at least that for B when A ⊆ B and x ∉ B.Source: D. Proof · p. 14
  • The submodularity algebra uses x_k, y_k, and z_k to express set intersections, then invokes A ⊆ B to obtain y_k ≥ x_k and conclude the required difference is zero.The proof begins with x ∈ N\B and A ⊆ B, ensuring x lies outside both sets before comparing the marginal terms.Source: D. Proof · pp. 14–15
Loading 2405.16064v1…