Source-linked AI summary

ReMA: Learning to Meta-think for LLMs with Multi-Agent Reinforcement Learning

Ziyu Wan, Yunxiang Li, Xiaoyu Wen, Yan Song, Hanjing Wang, Linyi Yang, Mark Schmidt, Jun Wang, Weinan Zhang, Shuyue Hu, Ying Wen

arXiv:2503.09501v3cs.AIcs.CLcs.LGcs.MA

TL;DR

Existing construction-based and single-agent approaches provide limited flexibility for acquiring meta-thinking in LLM reasoning. ReMA uses multi-agent reinforcement learning to separate strategic oversight from detailed execution, and it reports stronger performance across mathematical and LLM-as-a-Judge benchmarks, especially on out-of-distribution datasets. The framework also extends to multi-turn reasoning, while the paper identifies sensitivity and instability in that setting.

  • Problem

    Construction-based approaches lack sufficient flexibility for exploring suitable meta-thinking patterns, while RL on base LLMs can suffer from readability, language-mixing, and large-search-space challenges.

  • Method

    ReMA separates meta-thinking and detailed reasoning into high- and low-level agents trained with aligned multi-agent reinforcement-learning objectives.

  • Results

    ReMA consistently achieves the highest average performance across three backbone models and outperforms baselines by up to 6.68% on mathematical benchmarks and 8.49% on LLM-as-a-Judge benchmarks.

  • Takeaways & Limitations

    ReMA particularly improves out-of-distribution performance and provides insights into how distinct agents’ evolving interactions relate to LLM reasoning capabilities.

  • Takeaways & Limitations

    Multi-turn ReMA is highly sensitive to response-length and turn-count hyperparameters, with some configurations causing repetition collapse or empty responses.

Abstract

from arXiv · show

Recent research on Reasoning of Large Language Models (LLMs) has sought to further enhance their performance by integrating meta-thinking -- enabling models to monitor, evaluate, and control their reasoning processes for more adaptive and effective problem-solving. However, current single-agent work lacks a specialized design for acquiring meta-thinking, resulting in low efficacy. To address this challenge, we introduce Reinforced Meta-thinking Agents (ReMA), a novel framework that leverages Multi-Agent Reinforcement Learning (MARL) to elicit meta-thinking behaviors, encouraging LLMs to think about thinking. ReMA decouples the reasoning process into two hierarchical agents: a high-level meta-thinking agent responsible for generating strategic oversight and plans, and a low-level reasoning agent for detailed executions. Through iterative reinforcement learning with aligned objectives, these agents explore and learn collaboration, leading to improved generalization and robustness. Empirical results from single-turn experiments demonstrate that ReMA outperforms single-agent RL baselines on complex reasoning tasks, including competitive-level mathematical benchmarks and LLM-as-a-Judge benchmarks. Additionally, we further extend ReMA to multi-turn interaction settings, leveraging turn-level ratio and parameter sharing to improve efficiency. Comprehensive ablation studies further illustrate the evolving dynamics of each distinct agent, providing valuable insights into how the meta-thinking reasoning process enhances the reasoning capabilities of LLMs. Our code can be found in https://github.com/ziyuwan/ReMA-public

1 Introduction

ReMA addresses limitations in existing approaches to LLM meta-thinking by separating strategic oversight from detailed reasoning within a multi-agent reinforcement-learning system. The framework is evaluated across mathematical and LLM-as-a-Judge tasks, with reported generalization insights from ablations.

  • LLMs can develop meta-thinking abilities that generalize to out-of-distribution tasks.
  • Construction-based approaches teach predefined meta-thinking patterns with supervised fine-tuning or preference optimization but offer limited flexibility for exploration.
  • ReMA separates strategic meta-thinking from detailed reasoning using high- and low-level agents trained with aligned reinforcement-learning objectives.
  • ReMA formally defines and optimizes a multi-agent meta-thinking reasoning process through multi-agent reinforcement learning.
  • Experiments cover mathematical reasoning and LLM-as-a-Judge tasks, while multi-turn extensions use turn-level ratios to optimize trajectory returns and stabilize training.
  • Ablation studies examine evolving agent dynamics and interaction patterns, including role reversals under different reward settings.

2 Preliminaries

The preliminaries formulate vanilla autoregressive reasoning and its reinforcement-learning training objective. They also describe token-level policy optimization and readability, language-mixing, and exploration challenges for RL on base LLMs.

  • 2.1 Vanilla Reasoning Process (VRP): Vanilla reasoning generates each response token autoregressively from the prompt and preceding tokens.
  • 2.1 Vanilla Reasoning Process (VRP): Intermediate reasoning steps precede the extracted final answer, forming a chain-of-thought response.
  • 2.2 Training VRP via Reinforcement Learning: Reinforcement learning frames decoding as a deterministic, token-level Markov decision process.
  • 2.2 Training VRP via Reinforcement Learning: The reward function compares a generated answer with the golden answer for a question sampled from the dataset.
  • 2.2 Training VRP via Reinforcement Learning: GRPO and REINFORCE++ are adopted to compute policy gradients, with GRPO optimizing a token-level objective over grouped responses.
  • 2.2 Training VRP via Reinforcement Learning: RL on base LLMs can produce poor readability and language mixing, while the large search space makes efficient meta-thinking learning difficult.

3 Method

ReMA integrates meta-thinking into LLM reasoning through two collaborating agents: a high-level planner and a low-level executor. Its MARL training supports single-turn and multi-turn settings, including parameter sharing and turn-level ratio clipping for more stable optimization.

  • 3.1 Deploying Meta-Thinking Reasoning Process for LLMs: ReMA decomposes problem solving into sequential meta-thinking and reasoning phases handled by separate high-level and low-level agents.The high-level agent plans, monitors, or revises strategy, while the low-level agent produces the detailed solution.
  • 3.1 Deploying Meta-Thinking Reasoning Process for LLMs: The high-level agent generates and summarizes meta-thoughts from the prompt and interaction history, while the low-level agent executes detailed problem-solving under those instructions.
  • 3.2.2 Scaling up to Multi-turn MAMRP: Multi-turn MAMRP enables interleaved cognitive switching across turns for complex, long-horizon planning tasks.The multi-turn extension is introduced because single-turn interaction lacks immediate and fine-grained cognitive switching.
  • 3.2.1 Optimizing Single-turn MAMRP: Single-turn MAMRP is optimized as a Markov Game with a joint hierarchical policy that samples meta-thinking output before reasoning output.The joint policy factors into the high-level policy for m and the low-level policy for y conditioned on x and m.
  • 3.2.1 Optimizing Single-turn MAMRP: Training alternates between freezing one agent and updating the other, while each policy maximizes its respective reward under task- and model-specific objectives.The individual rewards can include the final reward, regularization, and format rewards.
  • 3.2.2 Scaling up to Multi-turn MAMRP: Multi-turn training shares parameters between agents through role-specific system prompts and uses a turn-level ratio to stabilize reinforcement-learning updates.Parameter sharing avoids frequent GPU model swapping and enables larger training batches; turn-level clipping treats each turn as a single action.

4 Experiments

Experiments evaluate ReMA on mathematical reasoning and LLM-as-a-Judge tasks in single- and multi-turn settings, alongside ablations of meta-thinking behavior and training configurations. ReMA generally improves benchmark performance and out-of-distribution generalization, while multi-turn results reveal both gains and sensitivity to rollout choices.

  • Experimental setup: Experiments compare ReMA with VRP, VRPRL, and MRPRL baselines across mathematical reasoning and LLM-as-a-Judge benchmarks using greedy pass@1 performance.Mathematical evaluation includes MATH500 and several out-of-distribution datasets; judge evaluation includes RewardBench970 and JudgeBench.
  • Single-turn results: 6.68% maximum improvement on mathematical benchmarks and 8.49% on LLM-as-a-Judge benchmarks are reported for ReMA across different LLMs.The comparison uses greedy decoding and reports average performance across models.
  • Single-turn results: 20% improvement on AMC23, 13.33% on AIME24, and 14.23% on RewardBench970 are reported for selected ReMA and model combinations.ReMA achieves the highest performance on most benchmarks, particularly out-of-distribution datasets.
  • Meta-thinking effects: RL under Meta-thinking achieves better learning dynamics and harder-problem generalization than RL from Base or RL from SFT, especially on AIME24.RL from SFT performs best initially on in-distribution and easier sets but does not improve on harder ones.
  • Meta-thinking behavior: Llama-3.2-1B-Instruct converges to EMPTY, whereas Llama-3.1-8B-Instruct adapts meta-thinking strategies to problem difficulty.The structured action space contains DECOMPOSE, REWRITE, and EMPTY; smaller models tend toward EMPTY to avoid formatting penalties.
  • Multi-turn results and ablations: Multi-turn ReMA increases training accuracy from approximately 55% to 70% and yields about 5% average improvement across seven test benchmarks, but performance is highly configuration-sensitive.Observed failures include repetition within a turn and empty responses after a few turns; ablations also find better sample efficiency with turn-level ratio and faster convergence with shared parameters.

5 Conclusion

ReMA separates meta-thinking from detailed reasoning through multi-agent reinforcement learning, achieving strong results across reasoning benchmarks and extending to multi-turn settings. The paper’s analyses examine agent coordination and the role of multi-agent training in improving LLM reasoning.

  • Contributions: ReMA uses multi-agent reinforcement learning to elicit meta-thinking by explicitly separating meta-thinking and reasoning into distinct agents.The framework is designed to improve exploration, interpretability, and coordination during reasoning.
  • Empirical findings: ReMA achieves superior results on mathematical reasoning and LLM-as-a-Judge benchmarks, particularly on out-of-distribution datasets.The conclusion reports consistent gains across the evaluated benchmark settings.
  • Multi-turn extension: The framework extends to multi-turn settings, supporting more complex reasoning scenarios that require increased communication between agents.This extension broadens the interaction structure beyond single-turn reasoning.
  • Analysis: Ablation studies show how effective coordination between agents evolves during learning.These analyses support examining structured collaboration as part of the reasoning process.
  • Related work: Multi-agent LLM research includes discussion, role assignment, routing, merging, and training approaches intended to combine capabilities across entities.The related work situates ReMA within broader research on interacting language-model agents.

B Limitation and Future Work

This section describes practical scaling and reward-design considerations for ReMA, while noting that its evaluation is currently limited to mathematics and LLM-as-a-Judge tasks. It also outlines the alternating training procedure for the two agents.

  • Limitations and future work: ReMA has only been tested on mathematics and LLM-as-a-Judge benchmarks, leaving naturally multi-turn tasks as an important future setting.The authors also identify multi-turn MARL dynamics and rollout efficiency as open areas for improvement.
  • Inference scaling: Inference cost grows with the product of the high-level and low-level sampling budgets.If the high-level agent samples k1 responses and each produces k2 low-level samples, total calls scale as k1 × k2.
  • Inference scaling: Allocating more samples to the high-level agent increases strategy diversity, whereas allocating more to the low-level agent can refine solutions for a given plan.The section frames budget allocation as a performance trade-off under fixed computation.
  • Reranking: ReMA considers hierarchical reranking, which aggregates low-level responses within each high-level response before ranking across high-level responses.Flat reranking instead ranks all sampled responses together without preserving hierarchy.
  • Reward design: The primary reward setting combines overall correctness with format rewards for both agents, while alternative designs reward consistency among low-level responses.High-level rewards can use average correctness or the proportion of the most frequent answer, with penalties for explicit answer formatting.
  • Training procedure: Algorithm 1 alternates low-level and high-level updates by generating responses, computing rewards, and optimizing each policy until convergence.The procedure builds separate training datasets for the two agents and updates their policies in sequence.

C.4 Brief convergence analysis

The convergence analysis treats the high- and low-level policies as alternating blocks optimized with TRPO-style updates, while a leader-follower formulation aligns high-level sub-tasks with low-level capabilities. The analysis also identifies conditions under which task decomposition can fail.

  • Block-coordinate optimization: TRPO-style updates optimize a surrogate objective under a KL constraint and guarantee monotonic improvement for each policy update.The paper applies practical PPO or GRPO versions in experiments.
  • Block-coordinate optimization: The high- and low-level policies are updated as two blocks of a single agent through cyclic block coordinate ascent.Each policy is optimized while the other is treated as fixed or stationary.
  • Block-coordinate optimization: Repeated coordinate maximization converges to a fixed point where neither single-policy update can further improve the joint objective.This is the stated convergence condition for the alternating updates.
  • Leader-follower game: The leader-follower formulation uses low-level capabilities to guide high-level task decomposition toward executable sub-task sequences.The reasoning boundary measures the maximum difficulty the low-level policy can reliably handle and informs high-level task selection.
  • Leader-follower game: The loss functions incentivize the high-level policy to select sub-tasks that lead to correct answers while improving the low-level policy’s instruction following.The simplified analysis assumes the high-level policy plans the complete sequence before low-level execution.
  • Leader-follower game: If the high-level policy leads with an overly narrow strategy, the low-level policy may become neither complete nor sound.The paper contrasts this with a low-level-led arrangement intended to preserve completeness and soundness.

D.1.1 Supervised fine-tuning data collection

For RL from SFT experiments, the authors collect expert demonstrations from GPT-4o Mini on 7.5k MATH training problems and use them to fine-tune the LLMs.

  • Data collection: The authors collect demonstrations from GPT-4o Mini on 7.5k MATH training problems for supervised fine-tuning before reinforcement learning.The data-generation process prompts GPT-4o Mini to produce metacognitive reasoning and decompose questions.

D.1.2 Dataset Curation of RewardBench970

RewardBench970 is split into training and test sets under validated loose and strict separation strategies, with balanced representation of original subsets. Under the loose setting, ReMA outperforms RL tuning baselines across models, especially on OOD benchmarks.

  • Dataset split: The dataset contains 5,000 training tuples and 970 remaining test tuples from RewardBench.
  • Dataset split: The strict split prevents instructions from appearing in both training and test sets, while the loose split only prevents tuple overlap.
  • Dataset split: All original RewardBench subsets are evenly represented in both training and test sets.
  • Results: Under the loose split, ReMA significantly outperforms other RL tuning baselines across all models, particularly on OOD benchmarks.

D.2.2 Training on MATH

Multi-turn ReMA is trained on MATH with GRPO and evaluated across rollout configurations, reward designs, and parameter-sharing strategies. The experiments show both measurable training gains and evolving cross-agent behaviors under consistency-based rewards.

  • Training setup: Multi-turn ReMA on MATH uses GRPO, 128 prompts with 16 trajectories each, a 1e-6 learning rate, bfloat16 precision, and four rollout-data mini-batches.
  • Training setup: The main experiment trains for 500 steps with a maximum prompt length of 4096 on 32 NVIDIA A800 GPUs.
  • Training setup: The study compares rollout limits of 30 turns with response lengths of 256 or 1024 tokens and prompt lengths of 4096 or 3072 tokens.
  • Reward effects: Under consistency-based rewards, the lower-level agent increasingly verifies rather than straightforwardly executing the high-level agent’s guidance.
  • Reward effects: Consistency-based rewards produce a self-corrective lower-level behavior that can disrupt the intended separation between planning and execution.

E.2 Detailed Training Curves on Different Datasets of Multi-turn ReMA

Figure 8 provides detailed multi-turn ReMA training curves across datasets. The accompanying example illustrates how high-level policies can change the solving strategy used by the reasoning agent.

  • Training curves: Figure 8 presents detailed training curves for multi-turn ReMA on different datasets.
  • Training curves: The curves should be compared across datasets to examine differences in training behavior and progression.
  • Policy example: A high-level policy improves the solving approach by directing the model to count only coordinates strictly inside a boundary.

F.2 Case study for Experiments in Section E.1

The case study examines how consistency rewards alter the interaction between ReMA’s agents. The high-level agent produces detailed solution attempts, while the low-level agent evolves toward verification.

  • Reward effects: The case study reports that consistency rewards significantly alter the high-level agent’s learning dynamics.
  • Reward effects: The high-level agent generates detailed solution attempts rather than a strategic plan under the consistency reward.
  • Reward effects: The low-level agent consequently evolves to verify the high-level agent’s solutions.

F.3 Case study for Experiments in Section 4.2.2

The case study examines ReMA’s prompt-based separation of high-level metacognitive planning from low-level reasoning, including JSON-formatted decomposition and rewriting. In the JSON-action experiments, outputs converged to shorter and easier sentences, likely because the smaller language model struggled to produce valid JSON while exploring diverse strategies.

  • Interpretability experiment: JSON-action outputs converged to relatively easier and shorter sentences than the main experiments.The passage attributes this pattern to insufficient small-model capacity for simultaneously producing valid JSON and exploring diverse reasoning strategies.
  • Case-study comparisons: The experiments compare reasoning setups with and without high-level metacognition and include case studies of consistency and base-correctness rewards.The cited figures are labeled as case studies for the high-level agent’s consistency reward, base correctness reward, and metacognition comparison.
  • Prompt design: The metacognition prompts ask the high-level agent to rewrite or decompose the original mathematical question before solution generation.The rewrite and decomposition variants require valid JSON outputs labeled REWRITE or DECOMPOSE.
  • Prompt design: ReMA assigns a high-level agent to generate strategic instructions while a low-level agent executes the detailed solution.The high-level prompt forbids directly answering the question, whereas the low-level agent follows the provided instruction and produces the final answer.
  • Prompt design: The MAMRP math prompt instructs the high-level agent to teach a weaker solver through a step-by-step plan without directly providing the answer.The corresponding low-level process uses the question and instruction to construct the final boxed response.
Loading 2503.09501v3…