Distilling the Implicit Multi-Branch Structure in LLMs' Reasoning via Reinforcement Learning
This work addresses the challenge of effectively transferring complex reasoning structures in LLMs for AI and machine learning applications, representing an incremental improvement over existing distillation techniques.
The paper tackles the problem of distilling reasoning abilities from teacher to student LLMs by addressing the limitation of supervised fine-tuning, which collapses the implicit multi-branch structure of authentic reasoning into flat sequences, and proposes RLKD, a reinforcement learning-based framework that improves student reasoning by internalizing this structure, achieving better performance with only 0.1% of data compared to standard methods.
Distilling reasoning paths from teacher to student models via supervised fine-tuning (SFT) provides a shortcut for improving the reasoning ability of smaller Large Language Models (LLMs). However, the reasoning paths generated by teacher models often reflect only surface-level traces of their underlying authentic reasoning. Insights from cognitive neuroscience suggest that authentic reasoning involves a complex interweaving between meta-reasoning (which selects appropriate sub-problems from multiple candidates) and solving (which addresses the sub-problem). This implies authentic reasoning has an implicit multi-branch structure. Supervised fine-tuning collapses this rich structure into a flat sequence of token prediction in the teacher's reasoning path, preventing effective distillation of this structure to students. To address this limitation, we propose RLKD, a reinforcement learning (RL)-based distillation framework guided by a novel Generative Structure Reward Model (GSRM). Our GSRM converts reasoning paths into multiple meta-reasoning-solving steps and computes rewards to measure structural alignment between student and teacher reasoning. RLKD combines this reward with RL, enabling student LLMs to internalize the teacher's implicit multi-branch reasoning structure rather than merely mimicking fixed output paths. Experiments show RLKD surpasses standard SFT-RL pipelines even when trained on 0.1% of data under an RL-only regime, unlocking greater student reasoning potential than SFT-based distillation.