Enhancing Reasoning Abilities of Small LLMs with Cognitive Alignment

TL;DR

This work tackles the gap between large and small LLM reasoning by introducing a cognitive-alignment pipeline (CRV) and a gap-aware optimization method (CogPO). The CRV system uses Critic, Rethinker, and Verifier agents to generate and refine CoT data tailored to the cognitive capabilities of small models, while CogPO extends Direct Preference Optimization to enforce reasoning alignment via per-gap beta scheduling. Empirical results across challenging reasoning benchmarks show that CRV+CogPO significantly improves small-model performance over direct SFT and other preference methods, with strong generalization across backbones and data scales. The framework offers a scalable path to deploy efficient, capable reasoning agents in resource-constrained settings and opens avenues for domain-specific applications.

Abstract

The reasoning capabilities of large reasoning models (LRMs), such as OpenAI's o1 and DeepSeek-R1, have seen substantial advancements through deep thinking. However, these enhancements come with significant resource demands, underscoring the need for training effective small reasoning models. A critical challenge is that small models possess different reasoning capacities and cognitive trajectories compared with their larger counterparts. Hence, directly distilling chain-of-thought (CoT) rationales from large LRMs to smaller ones can sometimes be ineffective and often requires a substantial amount of annotated data. In this paper, we first introduce a novel Critique-Rethink-Verify (CRV) system, designed for training smaller yet powerful LRMs. Our CRV system consists of multiple LLM agents, each specializing in unique tasks: (i) critiquing the CoT rationales according to the cognitive capabilities of smaller models, (ii) rethinking and refining these CoTs based on the critiques, and (iii) verifying the correctness of the refined results. Building on the CRV system, we further propose the Cognitive Preference Optimization (CogPO) algorithm to continuously enhance the reasoning abilities of smaller models by aligning their reasoning processes with their cognitive capacities. Comprehensive evaluations on challenging reasoning benchmarks demonstrate the efficacy of our CRV+CogPO framework, which outperforms other methods by a large margin.

Figures (6)

• Figure 1: A motivating example. Large models (right) apply vector-based algebraic abstraction to solve the problem, while small models (left) employ simple formulaic geometric decomposition. This trajectory mismatch underscores the inefficacy of direct CoT distillation across models with substantial capacity gaps.
• Figure 2: Overview of our CRV+CogPO framework, consisting of two synergistic phases: (1) SFT training with cognitively aligned data generated by the CRV system, and (2) CogPO: dynamic $\beta$ adjustment preference optimization training using cognitive reasoning pairs with different quality gaps. Disclaimer: We use the Qwen logo as our backbone; however, any LLMs with sufficient capabilities can serve as agents as well.
• Figure 3: An illustration of the proposed CogPO algorithm, showing the different preference gaps between CoT pairs and the corresponding mini-tasks.
• Figure 4: Experimental results for different sizes of Qwen2.5 models on AIME2024, MATH500, GPQA Diamond, and LiveCodeBench V2.
• Figure 5: Experimental results for other model series (Llama3.1-8B-Instruct, Mistral-7B-V0.3) beyond Qwen2.5, on AIME2024, MATH500, GPQA Diamond, and LiveCodeBench V2.