Reproduce “RLKD: Distilling LLMs' Reasoning via Reinforcement Learning” on OpenResearch

About this paper

Shicheng Xu, Liang Pang, Yunchang Zhu, Jia Gu, Zihao Wei, Jingcheng Deng, Feiyang Pan, Huawei Shen and 1 more

Researchers at the Chinese Academy of Sciences and Huawei Inc. developed RLKD, a reinforcement learning framework that utilizes a Generative Structure Reward Model to distill the implicit multi-branch reasoning structure from large language models into smaller ones. This method enables student models to achieve more authentic and generalizable reasoning capabilities, outperforming SFT-based methods with 0.1% of the data.

Key insights

• Authentic reasoning in LLMs involves a dynamic interplay of meta-reasoning (selecting sub-problems) and solving (executing chosen sub-problems), which can be conceptualized as an implicit multi-branch structure.
• Distilling the underlying structural reasoning at a step-level, rather than merely imitating surface-level token sequences, is essential for fostering robust and generalizable reasoning abilities in smaller models.
• A generative reward model can effectively quantify the alignment of reasoning *structures* between teacher and student, providing precise, step-level feedback for reinforcement learning that moves beyond traditional outcome-based or surface imitation rewards.

Summary by alphaXiv.