🤖 machine learning

SMOPD: Multi-Reward Reinforcement Learning via Specialize-and-Merge Online Policy Distillation

The paper proposes SMOPD, a two-stage "specialize-and-merge" online policy distillation framework that trains reward-specialized teacher models and consolidates them into a single student policy to effectively balance multi-reward signals of varying granularities, outperforming existing methods like GDPO across various model backbones.

Original authors: Wen Wang, Jiahua Bao, Tu Yongsiqi, Yihao Liu, Haotian Zhou, Haoxuan Ma, Mengyu Zhou, Wenkui Fan, Junwei He, Xiaoxi Jiang, Guanjun Jiang

Published 2026-08-05

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Wen Wang, Jiahua Bao, Tu Yongsiqi, Yihao Liu, Haotian Zhou, Haoxuan Ma, Mengyu Zhou, Wenkui Fan, Junwei He, Xiaoxi Jiang, Guanjun Jiang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to be the perfect assistant. You want it to be helpful, harmless, and also able to follow strict formatting rules, like using specific tags or calling the right tools. This is the world of Reinforcement Learning from Human Feedback (RLHF). Think of it like training a dog: you give it a treat (a "reward") when it does something good, and it learns to repeat that behavior. But here's the tricky part: what if you want the dog to sit and stay and fetch, all at once? If you just give it one big treat for doing everything, the dog might get confused. It might realize that "fetching" is the easiest way to get a treat, so it ignores "sitting" and "staying" entirely. In the world of AI, this is called reward masking, where one strong signal drowns out the others.

Recently, scientists tried to fix this by giving the AI separate scores for each task before combining them, a method called GDPO. It was like giving the dog a scorecard for sitting, another for staying, and a third for fetching, then adding them up. But there was a catch: some tasks are easy to measure (like "did you fetch the ball?"), while others are rare and hard to spot (like "did you stay perfectly still for 10 seconds?"). If the AI gets a million "fetch" scores but only one "stay" score, the "stay" signal gets lost in the noise. The AI learns to fetch perfectly but forgets to stay still. This paper, SMOPD, asks a simple question: How do we make sure the AI learns the rare, hard-to-spot skills without forgetting the easy ones?

The Problem: The "Noisy Classroom"

Imagine a classroom where a teacher is trying to teach students two things at once: how to solve complex math problems (a "dense" reward because there are many steps to get right) and how to raise their hands before speaking (a "sparse" reward because it only happens once per class).

If the teacher just shouts out a total score at the end of the day, the students will focus entirely on the math. They'll ignore the hand-raising because, statistically, the math points are everywhere, while the hand-raising points are rare. Even if the teacher tries to weigh the hand-raising equally, the students still don't hear it clearly enough over the noise of the math scores. The AI models were facing this exact problem: they were great at the "dense" tasks but terrible at the "sparse" ones, like formatting their answers correctly or being harmless when it mattered most.

The Solution: Specialize, Then Merge

The authors of this paper propose a clever two-step strategy called SMOPD (Specialize-and-Merge Online Policy Distillation). Instead of trying to teach one student to be perfect at everything at the same time, they create a team of specialists and then combine their skills.

Step 1: The Specialization (The "Specialist Teachers")
First, the researchers train several different "teacher" models. But here's the twist: each teacher is told to focus only on one specific skill.

Teacher A is told: "Ignore the math! Your only job is to make sure the students raise their hands." To do this, the teacher gets a massive bonus for every hand raised, making that signal impossible to ignore.
Teacher B is told: "Ignore the hand-raising! Your only job is to solve the math problems."

Because each teacher is hyper-focused, they become experts at their specific job. Teacher A learns the rare "hand-raising" skill perfectly because it's the only thing that matters to them. Teacher B becomes a math wizard. They don't try to balance the two; they just go all-in on one.

Step 2: The Merge (The "Student" Model)
Now, the researchers bring in a new "student" model. This student doesn't just listen to one teacher; it listens to all of them at the same time.

When the student is writing an answer, it looks at what Teacher A (the hand-raiser) would do and what Teacher B (the math whiz) would do.
The student learns to copy the best parts of both. If Teacher A is confident about the formatting, the student copies that. If Teacher B is confident about the math, the student copies that too.

Crucially, the student also has a "safety net" (called an anchor) that checks the final answer to make sure it's still a good overall answer, not just a copy-paste job. This ensures the student learns to balance the skills naturally, rather than just being a confused mix of two different people.

What They Found

The team tested this idea on AI models of different sizes (1.5 billion, 3 billion, and 7 billion parameters) using two very different challenges:

Tool Calling: Making sure the AI uses the right tools (like a calculator or a search engine) and formats the output exactly right (like using specific XML tags).
Safety: Making sure the AI is helpful but also harmless.

The results were impressive. In the "Tool Calling" test, the old methods (GDPO) managed to get the formatting right only about 8.8% of the time because the formatting signal was too sparse. But with SMOPD, the models got the formatting right 97.5% of the time! Even more surprisingly, they didn't lose their math skills in the process. The 1.5 billion model saw a 48% improvement in its overall score compared to the previous best method.

In the "Safety" tests, where the goals of being helpful and being harmless sometimes clash, SMOPD consistently outperformed the other methods across all model sizes. It managed to learn the rare "harmless" behaviors without sacrificing its ability to be helpful.

Why This Matters

The paper suggests that the old way of trying to balance all rewards in a single training session is hitting a wall. You can't just tweak the weights to make everything perfect; sometimes, a skill is just too rare to be learned in a crowd. By letting the AI "specialize" first—letting it get really good at one thing in isolation—and then "merging" those skills, the AI can learn the hard, rare stuff without forgetting the easy stuff.

It's like realizing that to build a perfect orchestra, you shouldn't just tell everyone to play everything at once. Instead, you let the violinists practice their part alone until they are perfect, let the drummers do the same, and then bring them together to play the symphony. The result is a performance where every instrument is heard clearly, and the music sounds better than anyone expected.

Technical Summary: SMOPD

1. Problem Statement

The paper addresses a fundamental limitation in multi-reward Reinforcement Learning from Human Feedback (RLHF) for Large Language Models (LLMs). While existing methods like Group Reward-Decoupled Normalization Policy Optimization (GDPO) have improved upon standard Group Relative Policy Optimization (GRPO) by normalizing reward dimensions separately before aggregation, they still struggle with reward-density imbalance.

In many practical scenarios, models are trained with a mix of dense rewards (providing fine-grained, continuous feedback, e.g., accuracy scores from 0.1 to 1.0) and sparse rewards (providing binary, all-or-nothing feedback, e.g., format compliance of 0 or 1).

The Core Issue: In group-based advantage estimation, dense rewards provide consistent within-group variation, allowing the model to learn effectively. Sparse rewards, however, often assign identical values (e.g., 0) to the majority of responses within a rollout group. This results in zero or near-zero within-group variance for the sparse dimension.
Consequence: Even with per-reward normalization, the sparse reward signal is overwhelmed by the dense reward signal during balanced optimization. The model fails to learn the capability associated with the sparse reward (e.g., strict output formatting) because the optimization signal is insufficient to drive policy updates, regardless of the priority weights assigned in a single-policy setting. Raising the priority of the sparse reward to make it learnable often skews the final policy, degrading performance on other objectives.

2. Methodology: SMOPD

To resolve this tension, the authors propose Specialize-and-Merge Online Policy Distillation (SMOPD), a two-stage training framework that decouples capability acquisition from reward balancing.

Stage 1: Specialize (Reward-Specialized Teachers)

Instead of training a single policy to balance all rewards, SMOPD trains multiple "teacher" policies, each specialized for a specific reward dimension.

Mechanism: The base policy is trained with distinct reward-priority configurations. For a target reward dimension $k$ , the priority weight $w_k$ is set significantly higher than other weights ( $w_k \gg w_j$ ).
Effect: This amplifies the gradient signal for the target reward, allowing the teacher to effectively learn the sparse or hard-to-learn distribution (e.g., a "Format Teacher" that prioritizes binary format compliance) without being drowned out by dense signals.
Outcome: A set of complementary teachers is produced, where each teacher excels at a specific reward dimension, potentially at the expense of others, but collectively covering the full reward space.

Stage 2: Merge (Online Policy Distillation)

The specialized capabilities of the teachers are merged into a single "student" policy through On-Policy Distillation (OPD).

Teacher Mixture: At each token position $t$ in a student-generated sequence, the target distribution is a mixture of the teachers' next-token distributions. SMOPD uses a uniform mixture ( $\alpha_m = 1/M$ ) and approximates the full vocabulary mixture using a top- $\kappa$ construction (retaining the top- $\kappa$ tokens from each teacher) to ensure computational efficiency.
Distillation Loss: The student minimizes the Forward KL divergence between its own distribution and the teacher mixture. Forward KL is chosen because it is "mode-covering," encouraging the student to place probability mass on all high-probability tokens from any teacher (crucial when different teachers are confident about different tokens, e.g., format wrappers vs. function arguments).
GDPO Anchor: To prevent the student from merely imitating the teachers and losing task-level optimization, a GDPO anchor loss is added. This loss uses balanced weights for all rewards, providing a sequence-level signal that guides the student to optimize the original multi-reward objective.
Total Objective: The student optimizes a combination of the distillation loss (token-level guidance) and the anchor loss (sequence-level task direction).

3. Key Contributions

Resolution of Reward-Balancing Tension: SMOPD addresses the inherent conflict in single-policy multi-reward RL where sparse rewards are unlearnable under balanced priorities. By specializing teachers, it makes sparse distributions learnable without permanently skewing the final policy.
Two-Stage Specialize-and-Merge Framework: The method introduces a novel workflow where capability acquisition is separated from reward balancing. Teachers learn in favorable regimes (high priority), and the student merges these capabilities via token-level distillation.
Superiority in Complementary and Conflicting Settings: The paper demonstrates that SMOPD outperforms scalarized baselines (GRPO, GDPO, GD2PO) across different model scales (1.5B, 3B, 7B) and reward structures (complementary: tool accuracy/format; conflicting: helpful/harmless).
Merging Surpasses Teachers: The merged student policy not only balances the objectives but often surpasses the performance of the individual specialized teachers and the scalarized baselines, indicating that the distillation process recovers complementary knowledge unused by single-objective training.

4. Experimental Results

The authors evaluated SMOPD on two primary settings:

A. Complementary Rewards (Tool Calling)

Setup: Training on RLLA-4K with accuracy (dense) and format (sparse binary) rewards.
Findings:
- GDPO Failure: Under balanced priorities, GDPO achieved only 8.8% format compliance on the 1.5B backbone, as the sparse signal was overwhelmed.
- Teacher Success: The Format-specialized teacher (high priority on format) achieved 97.5% compliance but lower accuracy.
- SMOPD Success: The merged student achieved 97.5% format compliance while maintaining high accuracy (RLLA Mean score of 2.740 vs. GDPO's 1.849). This represents a +48% improvement over the GDPO baseline on the 1.5B model.
- Generalization: SMOPD also showed superior performance on held-out benchmarks like API-Bank (LLM-Judge) and BFCL.

B. Conflicting Rewards (Helpful vs. Harmless)

Setup: Training on Safe-RLHF benchmarks (HH-RLHF, PKU-SafeRLHF, Alpaca) with helpfulness and harmlessness rewards.
Findings:
- SMOPD consistently outperformed GDPO and GD2PO across 3B and 7B backbones (Qwen2.5 and Llama-3.2).
- On Qwen2.5-7B, SMOPD achieved an Overall score of 5.646, exceeding the best scalarized baseline (GD2PO: 5.571) and both individual teachers.
- Even when individual teachers degraded (e.g., on Llama-3.2-3B), the merging process recovered performance to approach or exceed scalarized baselines.

C. Ablation Studies

Necessity of the Anchor: Removing the GDPO anchor (using OPD only) resulted in a student that plateaued at the performance level of the teacher mixture. The anchor is essential for pushing the student beyond the "imitation ceiling" to optimize the actual task reward.
Forward vs. Reverse KL: Forward KL was found to be superior to Reverse KL in this multi-teacher setting because it preserves the diverse modes of the teacher mixture, whereas Reverse KL tends to be mode-seeking and may lose specialized capabilities.
Top- $\kappa$ Sensitivity: A small top- $\kappa$ (e.g., 16) was sufficient, as it retained ~99.4% of the teacher probability mass, validating the efficiency of the approximation.

5. Significance and Claims

The paper claims that SMOPD offers a robust solution to the reward-density imbalance problem in multi-reward RL. Its significance lies in:

Decoupling Learning from Balancing: It challenges the paradigm of trying to balance all rewards simultaneously within a single optimization step. Instead, it leverages specialization to ensure all reward signals are learned effectively before merging.
Token-Level Granularity: By utilizing on-policy distillation at the token level, SMOPD provides dense supervision for sparse rewards (like formatting) that group-based advantage estimation cannot provide.
Scalability and Generality: The method is effective across varying model sizes (1.5B to 7B) and reward types (complementary and conflicting), suggesting it is a generalizable approach for multi-objective alignment in LLMs.

The authors conclude that SMOPD produces a unified policy that retains the strengths of specialized teachers while achieving balanced task-level optimization, effectively overcoming the limitations of current scalarized multi-reward methods.