PaperlayerView exact revision on arXiv

Source-linked AI summary

Marco-o1 v2: Towards Widening The Distillation Bottleneck for Reasoning Models

Huifeng Yin, Yu Zhao, Minghao Wu, Xuanfan Ni, Bo Zeng, Hao Wang, Tianqi Shi, Liangying Shao, Chenyang Lyu, Longyue Wang, Weihua Luo, Kaifu Zhang

arXiv:2503.01461v2cs.LGcs.AIcs.CL

TL;DR

Long-CoT distillation can be difficult for smaller models and may transmit over-thinking biases. Marco-o1 constructs tree-based CoT data with MCTS and combines it with CoT-aware post-training; the authors report improved reasoning performance across math, instruction-following, and planning benchmarks compared with standard distilled models.

  • Problem

    Long-CoT distillation can pose learning difficulties for smaller models and transmit biases such as over-thinking.

    Source: §1 · p. 1

  • Method

    The paper constructs tree-based CoT data with MCTS and applies CoT-aware approaches including length balancing, fine-grained DPO, and a joint post-training objective.

    Source: §1 · pp. 1, 3

  • Results

    The proposed approaches consistently and orthogonally improve reasoning performance over standard distilled models across math, instruction-following, and planning benchmarks.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The findings highlight the importance of well-designed data and post-training strategies for improving the efficiency and reliability of smaller-scale reasoning models.

    Source: §6 · p. 9

    Frequently seen QA prompts can be too easy to produce negative paths, resulting in fewer DPO pairs.

    Source: §2.4 · p. 5

Abstract

from arXiv · show

Large Reasoning Models(LRMs) such as OpenAI o1 and DeepSeek-R1 have shown remarkable reasoning capabilities by scaling test-time compute and generating long Chain-of-Thought(CoT). Distillation--post-training on LRMs-generated data--is a straightforward yet effective method to enhance the reasoning abilities of smaller models, but faces a critical bottleneck: we found that distilled long CoT data poses learning difficulty for small models and leads to the inheritance of biases (i.e. over-thinking) when using Supervised Fine-tuning (SFT) and Reinforcement Learning (RL) methods. To alleviate this bottleneck, we propose constructing tree-based CoT data from scratch via Monte Carlo Tree Search(MCTS). We then exploit a set of CoT-aware approaches, including Thoughts Length Balance, Fine-grained DPO, and Joint Post-training Objective, to enhance SFT and RL on the constructed data. We conduct evaluation on various benchmarks such as math (GSM8K, MATH, AIME). instruction-following (Multi-IF) and planning (Blocksworld), results demonstrate our approaches substantially improve the reasoning performance of distilled models compared to standard distilled models via reducing the hallucinations in long-time thinking. The project homepage is https://github.com/AIDC-AI/Marco-o1.

1 Introduction

Long-CoT distillation can leave smaller models mechanically repeating reasoning patterns without internalizing the logic, sometimes producing no final answer. The paper proposes tree-based CoT data and CoT-aware post-training methods, reporting improvements across math, instruction-following, and planning benchmarks.

1 Introduction

  • Distilled models can repeat reasoning patterns without internalizing their logic, with over-reflection and repetition sometimes preventing a final answer.The paper calls this “formalistic long-time thinking” and asks how to transfer long-CoT reasoning effectively through data construction, SFT, and RL.Source: §1 · p. 1
  • The paper constructs tree-based CoT data from scratch by expanding predefined thought nodes with MCTS, rather than relying on LRMs to generate the reasoning data.The authors say the approach expands the solution space and more closely mimics human-like reasoning patterns.Source: §1 · pp. 1, 3
  • The proposed CoT-aware methods include varying CoT lengths, fine-grained DPO, and a joint DPO-and-SFT objective.The joint objective is intended to mitigate over-optimization observed in DPO.Source: §1 · p. 2
  • The proposed methods consistently and orthogonally improve reasoning performance over standard distilled models across math, eight-language instruction-following, and planning benchmarks.The authors attribute the gains to reduced long-time-thinking hallucinations, including repetition, and fewer cases with no final answer.Source: §1 · p. 2
  • The evaluation covers GSM8K, MATH, and AIME, eight-language Multi-IF instruction-following, and Blocksworld planning.These benchmarks span math at three difficulty levels, instruction-following, and real-world planning tasks.Source: §1 · p. 2

2 Tree-Based CoT Data Construction

The section presents a customizable tree-based CoT construction method that uses MCTS to generate reasoning paths from scratch. Its thought nodes, transition patterns, and path-selection rules support extracting data for SFT and DPO.

2.1 Overall Framework

  • A constrained node-transition tree limits unbounded search; MCTS selects by UCB, expands nodes with LLM prompts, and backpropagates correctness rewards from answer rollouts.Source: §2.1 · p. 3

2.2 Thought Node

  • Thought nodes encode distinct actions: most use role-specific prompts, while Thinking allows unconditioned continuation; alternate-model Reflection can correct errors by backtracking.Source: §2.2 · pp. 3–4

2.3 Reasoning Pattern

  • Reasoning templates vary node order; Hypothesis supports tasks needing early assumptions, while correctness checks can trigger Reflection until a solution is found or the search limit is reached.The framework randomly samples among templates to produce more flexible CoT data.Source: §2.3 · pp. 4–5

2.4 CoT Data Extraction for Post-Training

  • For SFT, select correct-answer paths; for DPO, pair a correct path with an incorrect path sharing its minimal prefix.SFT paths can be chosen by highest MCTS reward or by desired length.Source: §2.4 · p. 5
  • Familiar, easy QA prompts may yield few negative paths and therefore fewer DPO pairs; harder or less-seen questions can address this constraint.Source: §2.4 · p. 5
  • Shorter CoT paths reduce ineffective outputs and mitigate formalistic long-time thinking, improving reasoning-output quality.Table 3 reports effects of long, medium, and short paths across datasets.Source: Table 3 · p. 5, §2.4 · p. 5

3 CoT-Aware Post-Training

The section addresses formalistic long-time thinking with three CoT-aware post-training methods: balancing reasoning lengths, refining DPO, and combining preference and SFT losses. The methods use tree-derived paths to tailor training signals and reduce emphasis on redundant or noisy preferences.

3 CoT-Aware Post-Training

  • DPO training is prone to formalistic long-time thinking, which the section targets with Thoughts Length Balance, Fine-grained DPO, and a Joint Post-training Objective.The joint objective addresses forgetting during the shift from SFT to preference-based training.Source: §3 · p. 5, §3.3 · p. 6

3.1 Thoughts Length Balance

  • Use the longest CoT data for SFT and shortest for DPO; select correct paths by relative length and form negatives from incorrect paths sharing the shortest common prefix.MCTS samples multiple paths per query to choose representative lengths rather than imposing fixed token thresholds.Source: §3.1 · pp. 5–6

3.2 Fine-grained DPO

  • Conservative DPO adjusts preference targets for potentially noisy labels, while masking shared-prefix tokens focuses loss gradients on the differing output segments.The cDPO gradient weights correct and incorrect paths to support stable training under label noise.Source: §3.2 · p. 6

3.3 Joint Post-training Objective

  • The joint loss adds SFT loss to DPO loss, with α controlling the trade-off between retaining reasoning patterns and preference learning.The paper attributes formalistic long-time thinking partly to forgetting during the switch from SFT to DPO.Source: §3.3 · pp. 6–7

4 Experiments

Across math, planning, and instruction-following experiments, the proposed adjustments improve reasoning performance over the standard distilled model, particularly by reducing formalistic long-time thinking. The experiments also examine the joint-loss trade-off and MCTS at inference.

4.1 Experimental Setup

  • Evaluation covers AIME, GSM8K, MATH500, Blocksworld, and Multi-IF, including multilingual three-turn instruction-following conversations.Multi-IF covers eight languages; AIME uses 60 questions from 2023 and 2024.Source: §4.1 · p. 7

4.2 Experimental Results

  • The combined adjustments improve reasoning tasks, particularly planning and instruction-following, while maintaining competitive math performance and reducing formalistic long-time thinking.The reported methods include data balancing, cDPO, joint SFT loss, and masking shared prefixes; the performance comparison is summarized in Table 5.Source: §4.2 · pp. 7–8, Table 5 · p. 8

4.3 Effects of Joint DPO Loss

  • α=1 provides the best reported trade-off between DPO and SFT losses; smaller values retain some catastrophic forgetting, while larger values weaken preference alignment.The selected configuration is L = LDPO + LSFT, as evaluated across datasets in Table 6.Source: §4.3 · p. 8, Table 6 · p. 8

4.4 MCTS Inference Exploration

  • MCTS inference outperforms inference without MCTS at Test@8 and Test@32 on MATH, demonstrating gains when multiple guesses are allowed.Test@N measures the percentage of problems solved correctly at least once across N guesses; Table 7 reports Test@1, Test@8, and Test@32.Source: §4.4 · p. 9

5 Related Work

Related work describes reasoning models that scale test-time computation, distillation approaches for transferring their capabilities, and MCTS methods for exploring or constructing reasoning paths. It also notes that smaller models may struggle to learn intricate patterns directly from large-model CoT data.

Reasoning Models

  • Reasoning models use CoT generation and increased test-time compute for complex tasks, but smaller models struggle to replicate this behavior.The cited work describes advances in mathematics and coding alongside dependence on large model sizes.Source: §5 · p. 9

Knowledge Distillation

Monte Carlo Tree Search

  • Prior MCTS work uses search to explore reasoning paths and generate structured or diverse training data, including data for training process reward models.Examples include candidate-path synthesis, structured data construction, and MCTS combined with process reward models.Source: §5 · p. 9

6 Conclusion

The paper explores transferring long chain-of-thought reasoning to smaller models through independently generated, flexible tree-based data and CoT-aware post-training. Experiments across math, planning, and instruction-following show robust performance improvements and mitigation of overly formalistic long-time thinking.

6 Conclusion

  • MCTS-generated, flexible tree-based reasoning data reduces reliance on large teacher models, while CoT-aware post-training mitigates overly formalistic long-time thinking.The conclusion frames these as strategies for transferring long chain-of-thought reasoning to smaller models.Source: §6 · p. 9
  • Experiments across math, planning, and instruction-following demonstrate robust performance improvements for the approach.Source: §6 · p. 9
  • Future work will explore reasoning techniques for machine translation, multilingual reasoning, and multimodal reasoning.Source: §6 · p. 9

Related papers

Loading 2503.01461v2…