TLDR: Recall-Extend Dynamics (RED) is a novel method that significantly enhances the reasoning capabilities of small language models (SLMs). It achieves this by dynamically balancing supervised fine-tuning (SFT) for knowledge expansion (‘Extend’) and reinforcement learning (RL) for knowledge refinement (‘Recall’), using entropy regulation. Additionally, RED introduces an accuracy-aware policy shift mechanism for robustly integrating offline distilled data, preventing issues like entropy collapse. Experimental results show RED outperforms existing methods on mathematical reasoning benchmarks, improving both accuracy and reasoning efficiency.
Small language models (SLMs) have shown great promise, but enhancing their reasoning capabilities to match their larger counterparts remains a significant challenge. While large language models (LLMs) have seen substantial improvements through techniques like reinforcement learning with verifiable rewards (RLVR), SLMs often struggle with issues such as ‘overthinking’ and generating redundant content, leading to inefficient training.
Researchers from Tianjin University and Hefei University of Technology have introduced a novel approach called Recall-Extend Dynamics (RED) to address these limitations. RED aims to boost the reasoning abilities of SLMs by intelligently balancing offline data distillation with online reinforcement learning, while also tackling specific challenges in integrating offline data.
Controlled Exploration: Balancing Recall and Extend
The core idea behind RED is to view the integration of reinforcement learning (RL) and supervised fine-tuning (SFT) as a synergy between ‘Recall’ and ‘Extend’ phases. RLVR primarily acts as a ‘Recall’ mechanism, refining existing reasoning paths within the model’s current knowledge and, in doing so, contracting its exploration space. Conversely, SFT serves as an ‘Extend’ mechanism, introducing new reasoning patterns learned from more powerful teacher models, thereby expanding the model’s explorable space.
RED dynamically adjusts the contribution of offline-SFT and RLVR by monitoring the ratio of entropy changes in the model. When the change in RL entropy is small, indicating insufficient exploration, RED increases the weight of offline-SFT to expand the exploration space. If RL entropy shows sufficient change, implying active exploration, the influence of offline-SFT is reduced. This dynamic regulation ensures that the model always has an appropriate exploration space, preventing both under-exploration and unnecessary complexity from distillation.
Adaptive Integration of Offline Data with Accuracy-aware Policy Shifts
A major hurdle in combining offline data (distilled from larger models) with online policy optimization is the potential for distribution discrepancies. Traditional methods of integrating this data can lead to problems like rapid entropy collapse (where the model becomes too deterministic) or a decline in performance.
To overcome this, RED introduces an ‘accuracy-aware policy shift mechanism’. This mechanism dynamically estimates the ‘offline probability’ (how much the model should imitate the distilled data) based on the correctness rate of the samples. For samples with high accuracy, the model is encouraged to learn more from its own policy. However, for samples with low accuracy, the policy offset is adjusted to make the model more inclined to imitate the distillation samples from the larger model. This adaptive approach improves the model’s robustness to varying data quality and optimizes training efficiency, effectively avoiding issues like entropy collapse and performance degradation.
Also Read:
- Enhancing Large Language Model Reasoning Through Contrastive Learning and Reinforced Fine-Tuning
- Proximal Supervised Fine-Tuning: Stabilizing LLM Updates for Broader Capabilities
Experimental Validation and Impact
The RED framework was rigorously tested on a suite of challenging mathematical reasoning benchmarks, including MATH500, AIME, AMC, Minerva, and Olympiad. Using Qwen2.5-1.5B-MATH as the base small model, RED was compared against various state-of-the-art methods, including those that use unified training paradigms (like LUFFY and SRFT) and stage-wise training approaches (like ReLIFT and BREAD).
The results consistently demonstrated RED’s superior performance across all benchmarks, not only in accuracy but also in reasoning efficiency, indicated by shorter average response lengths. An ablation study confirmed that both the dynamic entropy regulation and the accuracy-aware policy shifts are crucial and work synergistically within the RED framework.
Furthermore, a case study revealed that RED fosters a more efficient reasoning process. It increases ‘thinking-related’ probabilities during the initial and intermediate stages of reasoning, allowing for thorough exploration, while decreasing them in the final stage, leading to more decisive and concise conclusions. The training dynamics also showed that RED maintains higher and more stable SFT and RL entropies, signifying better exploration and adaptability throughout the training process.
In conclusion, Recall-Extend Dynamics represents a significant advancement in enhancing the reasoning capabilities of small language models. By intelligently balancing exploration and knowledge expansion, and adaptively integrating distilled data, RED offers a robust and efficient training paradigm that pushes the boundaries of what SLMs can achieve. You can find the full research paper here: Recall-Extend Dynamics: Enhancing Small Language Models through Controlled Exploration and Refined Offline Integration.
