PaperlayerView exact revision on arXiv

Source-linked AI summary

Improving Reasoning Capabilities in Small Models through Mixture-of-Layers Distillation with Stepwise Attention on Key Information

Yao Chen, Jiawei Sheng, Wenyuan Zhang, Tingwen Liu

arXiv:2604.15701v1cs.CL

TL;DR

Existing CoT distillation methods transfer rationales but do not adequately model teachers’ evolving attention to critical information. MoLSAKI transfers stepwise critical-token attention and uses adaptive Mixture-of-Layers alignment, achieving reported gains across reasoning benchmarks, while its evaluation remains limited in model diversity and task complexity.

  • Problem

    Existing CoT distillation methods mainly transfer teacher rationales and do not adequately explore teachers’ dynamic attention toward critical information during reasoning.

  • Method

    MoLSAKI transfers teacher stepwise attention on critical tokens to students and uses Mixture-of-Layers for adaptive alignment between teacher and student layers.

  • Results

    MoLSAKI yields average relative improvements of 7.5% for GPT2-Large and 11.3% for TinyLlama over baselines.

  • Takeaways & Limitations

    Stepwise critical-token attention provides a framework for transferring progressive focus on key information during small-model reasoning distillation.

  • Takeaways & Limitations

    The study uses relatively simple reasoning tasks with short chains of thought and has limited exploration of model sizes and architectures.

Abstract

from arXiv · show

The significant computational demands of large language models have increased interest in distilling reasoning abilities into smaller models via Chain-of-Thought (CoT) distillation. Current CoT distillation methods mainly focus on transferring teacher-generated rationales for complex reasoning to student models. However, they do not adequately explore teachers' dynamic attention toward critical information during reasoning. We find that language models exhibit progressive attention shifts towards key information during reasoning, which implies essential clues for drawing conclusions. Building on this observation and analysis, we introduce a novel CoT distillation framework that transfers the teacher's stepwise attention on key information to the student model. This establishes structured guidance for the student's progressive concentration on key information during reasoning. More importantly, we develop a Mixture of Layers module enabling dynamic alignment that adapts to different layers between the teacher and student. Our method achieves consistent performance improvements across multiple mathematical and commonsense reasoning datasets. To our knowledge, it is the first method to leverage stepwise attention within CoT distillation to improve small model reasoning.

1 Introduction

Large language models support complex reasoning but are costly to deploy, motivating CoT distillation into smaller models. MoLSAKI addresses the underuse of critical information by transferring stepwise attention and adaptively aligning teacher and student layers.

  • Large language models’ computational demands make deployment on edge devices prohibitively expensive, motivating reasoning distillation into smaller models.
  • Existing CoT distillation transfers teacher-generated rationales but often treats tokens equally, limiting students’ use of key information across multi-step reasoning.
  • Numerical tokens receive significantly more attention than non-numerical tokens during mathematical reasoning in both teacher and student models.
  • MoLSAKI transfers the teacher’s stepwise attention on critical tokens to guide the student’s progressive focus during reasoning.
  • An SVAMP example shows the distilled student failing to use numerical information adequately, while the teacher uses all numerical information to reach the correct result.
  • MoL adaptively aligns teacher and student layers with different depths through weighted, dynamic correspondence.

2 Related Work

Related work distills reasoning through teacher-generated rationales and has also transferred self-attention patterns using fixed or restricted layer mappings. MoLSAKI instead focuses on critical tokens and dynamic layer alignment.

  • CoT distillation methods transfer reasoning abilities into smaller models using teacher-generated rationales, including zero-shot extraction and separate rationale-answer prediction.
  • Prior attention-transfer methods use uniform, identical-count, or final-layer mappings and require matched attention dimensions or fixed correspondences.
  • MoLSAKI focuses distillation on critical tokens in reasoning steps rather than full attention matrices.
  • MoLSAKI uses dynamic layer alignment to address mismatched teacher and student layer counts.

3 Methodology

MoLSAKI combines CoT supervision, stepwise attention extraction, and adaptive Mixture-of-Layers alignment. It represents reasoning steps and critical-token attention so the student can match the teacher’s evolving focus.

  • MoLSAKI first prepares teacher-annotated CoT data, then extracts teacher and student stepwise attention, and finally performs adaptive MoL layer alignment.
  • Teacher prompting produces question-rationale-answer records, with the rationale and answer used for subsequent student CoT distillation.
  • CoT distillation trains two tasks: final-answer prediction and rationale generation, both using cross-entropy losses.
  • The attention loss guides the student to progressively focus on key information by distilling teacher attention on critical tokens.
  • Reasoning inputs are segmented into steps using periods, with the question included in the stepwise representation.
  • Tokenization and regular-expression matching identify reasoning-step tokens and critical-token index sets for attention extraction.

1. CoT Distillation …

The framework extracts stepwise attention on critical tokens from teacher and student layers, then uses Mixture-of-Layers alignment to handle mismatched architectures. Teacher layer weights emphasize attention variation, while student routing generates adaptive layer weights for attention-loss alignment.

  • Stepwise Attention: Stepwise attention on critical tokens is computed across reasoning steps and critical-token positions from each model’s self-attention matrices.The framework extracts relevant columns and aggregates their rows at each reasoning step.
  • Stepwise Attention: The teacher and student attention representations align reasoning-step and critical-token dimensions, without requiring shared tokenizers or vocabularies.Both representations use matching counts for reasoning steps and critical tokens despite architecturally distinct models.
  • Adaptive Layer Alignment: Mixture-of-Layers dynamically aggregates attention across layers through trainable router weights, addressing the lack of complete teacher–student layer correspondence.This replaces rigid single-layer alignment, which is described as inflexible and suboptimal for mismatched architectures.
  • Adaptive Layer Alignment: The largest variation in the teacher’s stepwise attention occurs in intermediate layers, based on column-gradient analysis across Llama3-8B layers.The framework uses these gradients to derive temperature-controlled teacher layer weights.
  • Adaptive Layer Alignment: The student MoL routes normalized value representations into compact layer embeddings, concatenates them, and produces adaptive weights through an affine transformation and temperature-controlled softmax.These student layer weights are used for the subsequent attention alignment.
  • Attention Alignment Loss: The stepwise attention loss compares teacher and student attention distributions after layer weighting and temporal-step normalization using averaged KL divergence.The overall objective combines prediction, explanation, and stepwise-attention losses.

4 Experiments

Experiments evaluate MoLSAKI across mathematical and commonsense reasoning datasets, model scales, teacher configurations, hyperparameters, and layer-alignment strategies. The method consistently improves reasoning performance and supports adaptive alignment between teacher and student layers.

  • 4.1 Setup: MoLSAKI is evaluated on five reasoning datasets, including in-domain and out-of-domain mathematical settings and an in-domain commonsense setting.The mathematical evaluation uses SVAMP in-domain and SingleEq, Asdiv, and GSM8K out-of-domain; CSQA is used for commonsense reasoning.
  • 4.2 Main Results: MoLSAKI achieves average relative improvements of 7.5% for GPT2-Large and 11.3% for TinyLlama over baselines.The comparison is reported against baseline approaches in the main accuracy results.
  • 4.2 Main Results: MoLSAKI improves in-domain accuracy over DSS and MMIloss across two student-model scales.The authors attribute this to guiding student attention at each reasoning step through stepwise attention alignment.
  • 4.2 Main Results: MoLSAKI consistently outperforms baselines on out-of-domain reasoning benchmarks for two different student models.The result is presented as evidence of strong out-of-domain generalization from distilling attention focused on critical tokens.
  • 4.3 Hyperparameter Analysis: SVAMP performance peaks at β = 1.0, while increasing teacher temperature τ1 reduces performance and student temperature τ2 = 0.5 maximizes results on SVAMP and AsDiv.Extreme τ2 values degrade performance, whereas out-of-domain advantages persist across most β settings.
  • 4.5 SL Alignment vs. MoL Alignment: Adaptive weighted MoL alignment outperforms fixed single-layer alignment, while mixed teacher configurations yield 5.1% and 7.9% relative accuracy improvements.The configuration results are reported for Unified and Hybrid settings, respectively, using GPT2-Large as the student.

5 Conclusion

The conclusion identifies stepwise attention on critical tokens as a source of reasoning cues and presents MoLSAKI as a framework for transferring those cues through adaptive layer alignment. This addresses critical-information underutilization in CoT distillation.

  • 5 Conclusion: Stepwise attention on critical tokens is proposed as implicitly encoding essential reasoning cues in large models.The perspective motivates using attention patterns in addition to teacher-generated rationales.
  • 5 Conclusion: MoLSAKI transfers the teacher’s stepwise attention on critical tokens to the student through Mixture-of-Layers adaptive alignment.The framework targets critical-information underutilization during reasoning distillation.

Limitations

The study’s scope is constrained by limited computational resources and experiments focused on relatively simple reasoning tasks with short chains of thought.

  • Limitations: Limited computational resources restricted exploration of diverse teacher and student model sizes and architectures.The authors identify broader scale and architecture coverage as future work.
  • Limitations: Experiments focused on relatively simple reasoning tasks with short chains of thought, leaving longer and more intricate reasoning paths for future study.The authors explicitly question whether the findings generalize to more complex problems.

Ethics Statement

The supplied passages discuss CoT distillation, its focus on transferring rationales, and the proposed transfer of progressive attention to critical tokens.

  • CoT distillation transfers teacher-generated step-by-step rationales to student models for complex reasoning.
  • Existing methods primarily transfer reasoning results rather than the teacher’s reasoning process.
  • The proposed attention loss encourages students to learn the teacher’s progressive focus on critical tokens during reasoning.
  • Prior self-attention distillation methods may neglect reasoning-specific patterns and require identical teacher-student tokenizers.
  • The paper distinguishes its focus on critical tokens within CoT reasoning from work studying critical tokens during pre-training.

B.1 Mathematical Reasoning

The mathematical analysis examines attention to numerical tokens during reasoning and defines how critical-token attention is extracted and aligned across teacher and student models.

  • The analysis uses 100 randomly sampled instances from each of GSM8K and SVAMP, totaling 200 mathematical reasoning samples.
  • The procedure records attention during generated reasoning, segments steps by periods, and compares numerical with non-numerical tokens.
  • Both teacher and student models assign higher average attention weights to numerical than non-numerical question tokens.
  • Critical words are identified automatically using numerical words for mathematical reasoning or teacher-provided keywords for commonsense reasoning.
  • When a critical word spans multiple tokens, the method sums attention across those tokens to represent the word’s attention.
  • Teacher-derived CoT text and critical words give teacher and student stepwise attention matrices matching shapes despite different tokenizers.
  • Teacher attention toward critical tokens changes dynamically across reasoning steps, indicating shifts in utilized key information.

E.1 Teacher: τ1

The teacher-layer analysis combines visual and quantitative evidence to identify where stepwise attention to critical tokens changes most strongly.

  • The study visually analyzes teacher stepwise attention on critical tokens from qualitative and quantitative perspectives.
  • A selected sample is visualized across all 32 layers of the Llama3-8B teacher, revealing attention variation across layers.
  • The quantitative analysis uses 100 GSM8K and 100 SVAMP samples to compute column gradients for each teacher layer.
  • The largest gradual changes in critical-token attention occur in intermediate layers.
  • The student uses a moderate τ2 so its layers participate substantially but nonuniformly in attention distillation.

F.1 Datasets

The evaluation spans mathematical and commonsense reasoning benchmarks, with dataset transformations, training settings, layer-weight visualizations, and computational-cost analysis described in the supplied passages.

  • The experiments evaluate System 2 reasoning across five commonsense and mathematical reasoning benchmarks.
  • The mathematical datasets are human-authored grade-school word-problem benchmarks, and Asdiv is converted from multiple-choice to open-ended questions.
  • CommonsenseQA measures everyday knowledge and commonsense reasoning about physical and social-world scenarios.
  • Training uses specified learning rates, batch sizes, and maximum steps for GPT2-Medium, GPT2-Large, and TinyLlama, with averages over three random runs.
  • A dual-phase CoT generation pipeline regenerates rationale-answer pairs for incorrect higher-complexity samples.
  • With τ1 = 0.1 and τ2 = 1.0, student layer weights are more evenly distributed.
  • MoLSAKI adds only slight FLOPs because the attention matrix comes from the standard forward pass and MoL adds one linear layer plus RMSNorm.

F.6 Case Study

The case study compares rationale generation and stepwise attention across SVAMP and GSM8K examples. MoLSAKI better preserves and uses relevant numerical conditions, while its distilled student’s attention more closely resembles the teacher’s.

  • Case analysis: The case analysis selects two SVAMP and two GSM8K samples to compare MoLSAKI with baseline methods in rationale generation.The study also visualizes numerical-token attention for GPT2-Medium students and the Llama3-8B teacher.
  • Rationale generation: DSS mentions the number “10” but fails to use it later, producing an incorrect result.This illustrates a mismatch between including a numerical condition in the rationale and using it during subsequent reasoning.
  • Rationale generation: MMIloss overlooks the condition “but he lost 2 of them,” which also leads to an incorrect answer.MoLSAKI, in contrast, uses all relevant numerical conditions in the analyzed examples.
  • Stepwise attention: Compared with baseline methods, the student model distilled by D-SANK shows high similarity to the teacher in stepwise attention on numerical tokens.The passage attributes improved comprehensive attention to numerical conditions and connects this pattern with improved mathematical reasoning.
Loading 2604.15701v1…