Probing to Refine: Reinforcement Distillation of LLMs via Explanatory Inversion

TL;DR

This paper introduces a novel distillation framework for large language models that enhances reasoning and generalization by using explanatory probes and reinforcement learning, resulting in significant performance improvements.

Contribution

The paper presents Explanatory Inversion and Explanatory GRPO, innovative methods that improve LLM distillation by fostering deeper understanding and coherent reasoning in smaller models.

Findings

20.39% average performance increase over zero-shot baseline

6.02% improvement over state-of-the-art distillation methods

Models trained with this method require less data and generalize better

Abstract

Distilling robust reasoning capabilities from large language models (LLMs) into smaller, computationally efficient student models remains an unresolved challenge. Despite recent advances, distilled models frequently suffer from superficial pattern memorization and subpar generalization. To overcome these limitations, we introduce a novel distillation framework that moves beyond simple mimicry to instill a deeper conceptual understanding. Our framework features two key innovations. \underline{\textit{First}}, to address pattern memorization, Explanatory Inversion (EI) generates targeted ``explanatory probes'' that compel the student to articulate the underlying logic behind an answer, rather than just memorizing it. \underline{\textit{Second}}, to improve generalization, Explanatory GRPO (\texttt{EXGRPO}) uses a reinforcement learning algorithm with a novel Dialogue Structure Utility Bonus, which explicitly rewards the student for maintaining a coherent reasoning process across these probes. Extensive evaluations on 12 datasets demonstrate significant improvements. Using Gemma-7b as the student model, our method yields an average \textbf{20.39\%} increase over zero-shot performance and a \textbf{6.02\%} improvement over the state-of-the-art distillation baselines. Moreover, models distilled with our method show remarkable training efficiency (e.g., surpassing vanilla fine-tuning with \textbf{10-25\%} training data) and strong generalization to out-of-distribution tasks. Implementation is released at https://github.com/Zhen-Tan-dmml/ExGRPO.git.