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

arXiv: https://arxiv.org/abs/2604.15701 · HTML: https://mlsanity.com/paper/2604.15701/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Placed in Prompting and chain-of-thought reasoning, one of 17 directions in LLM reasoning.
- Builds on 15 of its 40 references that are in the map.
- First posted on arXiv 17 Apr 2026.

## Placed in

- LLM reasoning / [Prompting and chain-of-thought reasoning](https://mlsanity.com/lr/prompting-and-chain-of-thought-reasoning/)
- Knowledge distillation / [Knowledge distillation for LLMs](https://mlsanity.com/distill/knowledge-distillation-for-llms/)

## Builds on (most cited first, 15 of 15)

- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Qwen2.5 Technical Report](https://mlsanity.com/paper/2412.15115/) (2024)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies](https://mlsanity.com/paper/2404.06395/) (2024)
- [TinyLlama: An Open-Source Small Language Model](https://mlsanity.com/paper/2401.02385/) (2024)
- [From System 1 to System 2: A Survey of Reasoning Large Language Models](https://mlsanity.com/paper/2502.17419/) (2025)
- [A Diverse Corpus for Evaluating and Developing English Math Word Problem Solvers](https://mlsanity.com/paper/2106.15772/) (2021)
- [Multi-Step Reasoning with Large Language Models, a Survey](https://mlsanity.com/paper/2407.11511/) (2024)
- [Why do LLMs attend to the first token?](https://mlsanity.com/paper/2504.02732/) (2025)
- [Attention Heads of Large Language Models: A Survey](https://mlsanity.com/paper/2409.03752/) (2024)
- [Knowledge Circuits in Pretrained Transformers](https://mlsanity.com/paper/2405.17969/) (2024)
- [MoH: Multi-Head Attention as Mixture-of-Head Attention](https://mlsanity.com/paper/2410.11842/) (2024)
- [Exploring the System 1 Thinking Capability of Large Reasoning Models](https://mlsanity.com/paper/2504.10368/) (2025)
- [Revealing and Mitigating the Challenge of Detecting Character Knowledge Errors in LLM Role-Playing](https://mlsanity.com/paper/2409.11726/) (2024)
- [Aligning Teacher with Student Preferences for Tailored Training Data Generation](https://mlsanity.com/paper/2406.19227/) (2024)

## Built on by (most cited first, 3 of 3)

- [OISD: On-Policy Internal Self-Distillation of Language Models](https://mlsanity.com/paper/2605.29089/) (2026)
- [ConSA: Controllable Sparsity in Hybrid Attention via Learnable Allocation](https://mlsanity.com/paper/2606.18056/) (2026)
- [Teach Yourself Where to Look: On-Policy Attention Self-Distillation for Reasoning](https://mlsanity.com/paper/2609.33200/) (2026)

## Abstract (arXiv)

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.