Cog-Rethinker: Hierarchical Metacognitive Reinforcement Learning for LLM Reasoning

TL;DR

Cog-Rethinker addresses sampling inefficiency in zero-RL reasoning by introducing a hierarchical metacognitive reinforcement-learning framework with two rollout stages: decomposition and reflection. It leverages a metacognitive buffer and BM25-based retrieval to generate diverse subproblem decompositions, while a reflection rollout refines incorrect solutions, all reinforced by a joint training objective that combines RL with supervised fine-tuning on correct samples. Empirical results across multiple math-reasoning benchmarks demonstrate state-of-the-art performance and improved sample efficiency, especially for weaker base models, with ablations confirming the importance of SFT, the metacognitive buffer, and the reflection stage. The approach provides a scalable path for LLMs to acquire advanced reasoning beyond pretraining, achieving faster convergence and reduced data needs in challenging domains.

Abstract

Contemporary progress in large language models (LLMs) has revealed notable inferential capacities via reinforcement learning (RL) employing verifiable reward, facilitating the development of O1 and R1-like reasoning models. Directly training from base models with RL is called zero-RL. However, previous works rely upon activating LLMs' inherent capacities through fixed prompt templates. This strategy introduces substantial sampling inefficiencies for weak LLMs, as the majority of problems generate invalid outputs during accuracy-driven filtration in reasoning tasks, which causes a waste of samples. To solve this issue, we propose Cog-Rethinker, a novel hierarchical metacognitive RL framework for LLM reasoning. Our Cog-Rethinker mainly focuses on the rollout procedure in RL training. After the direct rollout, our Cog-Rethinker improves sample utilization in a hierarchical metacognitive two-stage framework. By leveraging human cognition during solving problems, firstly, it prompts policy to decompose zero-accuracy problems into subproblems to produce final reasoning results. Secondly, with zero-accuracy problems in previous rollout stage, it further prompts policy to refine these answers by referencing previous wrong solutions. Moreover, to enable cold-start of the two new reasoning patterns and maintain train-test consistency across prompt templates, our Cog-Rethinker applies supervised fine-tuning on the policy using correct samples of the two stages with direct rollout template. Experimental results demonstrate Cog-Rethinker's superior performance on various mathematical reasoning benchmarks, we also analyzed its improved sample efficiency that accelerates convergence compared to baseline methods.

Figures (8)

• Figure 1: The case study of our Cog-Rethinker. The Dec Rollout and Ref Rollout denote the policy generation by problem decomposition and answer reflection, respectively. Our Cog-Rethinker can generate more correct samples especially at the beginning of training procedure.
• Figure 2: Overall procedure of our Cog-Rethinker. The upper is the whole rollout stage for different difficulty problems, the lower is the training procedure with token-level policy gradient loss of DAPO and NLL loss of SFT.
• Figure 3: Different prompt templates of our Cog-Rethinker during whole rollout stage.
• Figure 4: Training comparison between our Cog-Rethinker and DAPO.
• Figure 5: Sample utilization efficiency analysis between our Cog-Rethinker and DAPO.

Theorems & Definitions (5)

• Theorem 1: Convergence Rate across Stages
• Lemma 1: Stage-wise Iteration Complexity
• proof
• Theorem 1: Convergence Rate across Stages (Restatement)
• proof