TLDR: PLAN-TUNING is a new post-training framework that enhances smaller language models’ ability to solve complex problems by teaching them step-by-step planning. It distills high-quality ‘planning trajectories’ from large LLMs and uses these to fine-tune smaller models via supervised and reinforcement learning. This method significantly improves performance on mathematical reasoning benchmarks, including out-of-domain tasks, demonstrating better generalization and structured problem-solving capabilities.
Large Language Models (LLMs) have shown remarkable abilities in solving complex problems by breaking them down into smaller, manageable steps, much like humans plan. This approach, often called ‘natural planning,’ has significantly boosted their performance. However, this capability has mostly been seen in very large, proprietary LLMs, and smaller, open-source models often struggle to effectively use such planning strategies.
Introducing PLAN-TUNING
A new research paper introduces PLAN-TUNING, a novel framework designed to teach smaller LLMs how to plan step-by-step. Instead of just relying on clever prompts during use, PLAN-TUNING embeds this planning ability directly into the model’s knowledge during its post-training phase. The core idea is twofold: first, it ‘distills’ high-quality planning examples (called ‘planning trajectories’) from powerful, large-scale LLMs. Second, it fine-tunes smaller models using these examples through both supervised learning and reinforcement learning techniques, specifically designed to mimic these planning processes.
How PLAN-TUNING Works
The process begins with data generation. For a given problem, a large LLM (like Gemini-2.0-Flash) generates several potential planning trajectories. These plans are then rigorously filtered in a two-stage process. First, a ‘Plan Verifier’ assesses the quality of each plan, checking for coherence and logical soundness. Only plans exceeding a certain quality threshold are kept. Second, because the correct answers are known for the training data, the system executes these selected plans and verifies if they lead to the correct final answer. Only trajectories that pass both quality and correctness checks are used for training, ensuring a high-quality dataset of problem-plan-solution triples.
Once the high-quality planning trajectories are ready, PLAN-TUNING employs two main post-training methods:
-
Supervised Fine-Tuning (SFT): This involves training the smaller LLM to either generate the complete solution (plan, step-by-step execution, and final answer) or to focus solely on generating the high-quality plan. The model learns to map problems to these structured solutions directly.
-
Reinforcement Learning (RL) with GRPO: This method uses a technique called Group Relative Policy Optimization (GRPO). Here, the model generates plans and solutions, and it receives a combined reward based on both the quality of its generated plan (how similar it is to the distilled high-quality plan) and the correctness of its final answer. This encourages the model to produce both good plans and accurate solutions.
Also Read:
- ReMix: Supercharging LLM Reasoning with Smart Data Reuse
- Smarter LLM Training: A Sample-Centric Approach to Enhanced Reasoning
Impressive Results and Generalization
The researchers evaluated PLAN-TUNING on various mathematical reasoning benchmarks. On in-domain tasks like GSM8k and MATH, plan-tuned models consistently outperformed strong baseline models, showing an average improvement of around 7% on GSM8k and a significant boost on MATH. What’s even more compelling is their performance on out-of-domain datasets like OlympiadBench and AIME 2024, where they demonstrated better generalization capabilities, with average performance improvements of approximately 10% and 12% respectively. This indicates that teaching models to plan helps them tackle new, complex problems more effectively.
The study also revealed interesting insights. For instance, training models on a single, specific dataset (like GSM8k or MATH) yielded better results for that domain than mixing the two datasets during training. This suggests that planning styles can be dataset-specific, and a focused approach might be more effective. Furthermore, the two-stage filtering process for data synthesis proved crucial, significantly improving performance by ensuring both plan quality and answer correctness.
In essence, PLAN-TUNING offers an effective strategy for enhancing the complex reasoning abilities of smaller LLMs by systematically incorporating step-by-step planning into their core parameters. This work opens new avenues for making advanced problem-solving capabilities more accessible across a wider range of language models. You can read the full research paper here: PLAN-TUNING: Post-Training Language Models to Learn Step-by-Step Planning for Complex Problem Solving.
