📊 statistics

Multi-Objective Learning for Diffusion Models: A Statistical Theory under Semi-Supervised Learning

This paper proposes a semi-supervised, two-stage training framework that distills lightweight specialist diffusion models into a generalist model using abundant unlabeled data, thereby achieving efficient multi-objective learning with statistical guarantees that depend only on the complexity of the specialists rather than the generalist.

Original authors: Ziheng Cheng, Yixiao Huang, Hanlin Zhu, Haoran Geng, Somayeh Sojoudi, Jitendra Malik, Pieter Abbeel, Xin Guo

Published 2026-05-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Ziheng Cheng, Yixiao Huang, Hanlin Zhu, Haoran Geng, Somayeh Sojoudi, Jitendra Malik, Pieter Abbeel, Xin Guo

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 or a computer program how to do many different jobs at once. Maybe it needs to draw pictures in the style of anime, in the style of oil paintings, and in the style of realistic photos. Or maybe it needs to help a robot arm pick up a red cube, a blue cube, and a green cube.

In the world of AI, this is called Multi-Objective Learning. The goal is to build one "Generalist" brain that is good at all these tasks.

The Problem: The "Jack of All Trades" is Expensive

Usually, to train this Generalist brain, you need a massive amount of perfect examples. For every task, you need a pair of data: the instruction (like "draw a cat") and the perfect result (the actual picture of the cat).

But here's the catch: getting those perfect pairs is hard and expensive.

In Robotics: You need a human expert to physically demonstrate the task thousands of times.
In Art: You need high-quality, curated image pairs.

However, you do have a lot of the "instructions" (the conditions). You have millions of text prompts or millions of robot starting positions, but you don't have the perfect answers for all of them. This is the Semi-Supervised problem: lots of questions, but very few answers.

If you try to train the big Generalist brain directly on the few answers you have, it's like trying to teach a PhD student to master every subject in the world using only a single textbook. It's inefficient, and the student might not learn well.

The Solution: The "Specialist" Apprentices

The authors propose a clever two-step trick, like a master chef training a new head chef.

Step 1: Train the Specialists (The Apprentices)
Instead of training the big Generalist immediately, you first train small, lightweight "Specialist" models.

You give Specialist #1 the few examples of "Anime style" and teach it only that.
You give Specialist #2 the few examples of "Oil painting" and teach it only that.
You give Specialist #3 the few examples of "Realism" and teach it only that.

Because these specialists are small and focused, they learn very quickly and efficiently from the small amount of data available. They become experts in their tiny, specific corner of the world.

Step 2: The Generalist Learns from the Specialists
Now, you have a problem: you still have millions of "instructions" (unlabeled data) but no "answers" for them.

You take those millions of instructions and feed them to your Specialists.
Specialist #1 generates a fake "Anime" picture for an instruction.
Specialist #2 generates a fake "Oil painting" picture for another instruction.

These are called Pseudo-samples. They aren't perfect, but they are good enough. Now, you have a massive library of "Instruction + Generated Answer" pairs.

You take your big, powerful Generalist model and train it on this massive library of pseudo-samples (plus the original few real examples). The Generalist learns from the Specialists, effectively "distilling" their knowledge into one big brain.

Why This is a Big Deal (The Theory)

The paper provides a mathematical proof (a "statistical theory") showing why this works so well.

Think of it like this:

The Old Way: To learn a complex skill, you need a huge number of practice attempts. The number of attempts needed grows with how complex the skill is.
The New Way: The authors prove that the number of real, expensive practice attempts you need depends only on how complex the Specialists are, not how complex the Generalist is.

Since the Specialists are small and simple, you need very few real examples to train them. The Generalist is huge and complex, but because it learns from the "fake" data generated by the Specialists, it doesn't need a huge number of real examples to learn.

The Result: You get a super-smart Generalist model that handles many tasks well, but you only had to pay the "price" of collecting data for small, simple models.

Real-World Tests

The authors tested this idea in two very different areas:

Robotics: They taught a robot arm to stack and pick up cubes. They trained small specialists for different lighting and camera angles, then used them to train a big generalist. The result? The robot was much better at handling new, unseen environments (like different lights or camera angles) than if they had just trained the big robot on the few real examples they had.
Image Restoration: They tried to fix damaged photos (inpainting). They trained specialists to fix different types of damage (like lines, face masks, or wide scratches). Then, they used those specialists to generate training data for a big generalist. The result was that the generalist fixed photos much better than models trained only on the limited real data.

The Bottom Line

This paper shows that when you are short on expensive, perfect data, don't try to force your big AI to learn everything at once. Instead, hire small, cheap "specialists" to learn the basics quickly, let them generate practice material, and then teach your big "Generalist" to learn from that. It's a smarter, cheaper, and more efficient way to build powerful AI.

Technical Summary: Multi-Objective Learning for Diffusion Models under Semi-Supervised Learning

1. Problem Statement

Diffusion models are increasingly deployed as conditional generators for diverse tasks, such as text-to-image synthesis across heterogeneous prompt domains or robotic control across multiple environments. These scenarios naturally constitute a Multi-Objective Learning (MOL) problem, where a single model must achieve Pareto-optimal performance across multiple target distributions.

A central challenge in this setting is the statistical cost of training a "generalist" model. To handle multiple objectives, the generalist model class ( $\mathcal{F}$ ) typically requires substantially larger capacity than the "specialist" model classes ( $\mathcal{H}_k$ ) needed for individual tasks. In fully supervised settings, the sample complexity for training the generalist scales with its own complexity ( $C_\mathcal{F}$ ), which is often prohibitively high.

Furthermore, real-world deployments often face an asymmetric data regime:

Paired data $(x, y)$ (e.g., curated text-image pairs or expert robot demonstrations) are scarce and expensive to acquire.
Unpaired condition data $y$ (e.g., raw text prompts or state observations) are abundant.

The paper addresses the following question: Can we learn a Pareto-optimal diffusion model in a sample-efficient manner by leveraging weakly supervised (unlabeled condition) data?

2. Methodology

The authors propose a two-stage semi-supervised training framework that reconciles the need for a high-capacity generalist with the scarcity of paired data.

Stage 1: Specialist Training

For each of the $K$ objectives (target distributions $P_k$ ), a lightweight specialist score function $h_k \in \mathcal{H}_k$ is trained using only the limited available paired data $\{(x_i^k, y_i^k)\}_{i=1}^n$ .

The specialist classes $\mathcal{H}_k$ are assumed to have significantly lower complexity ( $C_{\mathcal{H}_k} \ll C_\mathcal{F}$ ) than the generalist class.
The training minimizes the denoising score matching objective on the paired data.

Stage 2: Generalist Training via Pseudo-Sampling

A large generalist model $f \in \mathcal{F}$ is trained to approximate the Pareto-optimal set across all objectives.

Pseudo-Sample Generation: The trained specialists $\{h_k\}$ are used to generate pseudo-paired data. For each abundant unlabeled condition $y_i^k \sim P_Y^k$ , the specialist generates a pseudo-target $\tilde{x}_i^k \sim P_{h_k}(\cdot|y_i^k)$ .
Objective: The generalist is trained on a combination of the original labeled data and the generated pseudo-paired data. The training objective minimizes a scalarization $S$ of the score matching losses across all $K$ objectives:
$\hat{f}_S = \arg\min_{f \in \mathcal{F}} S(\tilde{L}_1(f), \dots, \tilde{L}_K(f))$
where $\tilde{L}_k$ is the empirical loss computed over the pseudo-samples (with a correction term to account for the specialist's bias).

Extension to Sequential Decision Making

The framework is extended to diffusion policies in imitation learning. Here, the "condition" is the state $s$ , and the "target" is the action $a$ . The authors address the distribution shift inherent in on-policy rollouts: since expert state distributions are hard to sample directly, the second stage uses trajectories generated by the learned specialist policies under the environment dynamics to collect pseudo-data for the generalist policy.

3. Key Theoretical Contributions

Sample Complexity Bounds

The paper establishes non-asymptotic generalization bounds for the proposed semi-supervised approach.

Main Result (Theorem 3.2): The number of required paired (labeled) samples to achieve a target error $\epsilon$ $ϵ$ depends only on the complexity of the specialist classes ( $C_{\mathcal{H}_k}$ $C_{H_{k}}$ ), not the generalist class ( $C_\mathcal{F}$ $C_{F}$ ).
- Required labeled samples: $\tilde{O}\left(\frac{K C_\mathcal{H}}{\epsilon^2}\right)$ .
- Required unlabeled conditions: $\tilde{O}\left(\frac{K C_\mathcal{F}}{\epsilon^2}\right)$ .
Comparison: In contrast, vanilla fully supervised training of the generalist would require $\tilde{O}\left(\frac{K C_\mathcal{F}}{\epsilon^2}\right)$ labeled samples. This demonstrates a significant improvement in sample efficiency when $C_\mathcal{F} \gg C_\mathcal{H}$ .
Linear Scalarizations (Theorem 3.5): For linear (weighted-sum) scalarizations and convex function classes, the authors derive sharper rates, further refining the bound on the approximation error.

Distribution Estimation and Policy Guarantees

Distribution Estimation: The paper bridges the gap between score matching error and distribution estimation (Total Variation distance). It proves that the learned generalist achieves Pareto-optimal performance in approximating the target conditional distributions.
Diffusion Policies (Theorem 3.6): The authors provide the first theoretical guarantee on the sub-optimality gap of diffusion policies in imitation learning under semi-supervised settings. They show that the framework successfully handles the distribution shift caused by on-policy rollouts while retaining the sample efficiency benefits of the specialist-to-generalist pipeline.

4. Experimental Results

The authors validate their theory on two distinct domains:

Robotic Manipulation (StackCube & PickCube):
- Setup: Training diffusion policies for tasks under varying domain randomization (scene, lighting, camera pose).
- Result: The semi-supervised MOL method significantly outperforms a labeled-only baseline, particularly in Out-of-Distribution (OOD) settings (e.g., +16% success rate on StackCube Level 2, +14% on Level 3). This suggests the pseudo-rollouts provide crucial supervision for handling distribution shifts.
Image Restoration (CelebA-HQ Inpainting):
- Setup: Inpainting faces with three different mask types (alternating lines, face region, wide stroke).
- Result: The semi-supervised approach improved PSNR, SSIM, and LPIPS across all mask types compared to the labeled-only baseline, demonstrating the method's effectiveness in generative image tasks.

5. Significance and Claims

The paper claims to provide a principled statistical framework for multi-objective learning in diffusion models under semi-supervised regimes. Its primary significance lies in:

Bridging Theory and Practice: It formalizes the "specialist-to-generalist" distillation pipeline, offering a theoretical justification for why this approach is more sample-efficient than direct multi-task learning when labeled data is scarce.
Sample Efficiency: It proves that the statistical cost of learning a complex generalist can be decoupled from its own capacity, relying instead on the complexity of lightweight specialists trained on limited data.
Handling Distribution Shift: It extends these guarantees to sequential decision-making (imitation learning), addressing a critical practical challenge in robotics.

The authors maintain a modest stance, acknowledging that their guarantees rely on a realizability assumption (that the true score functions lie within the specialist classes) and do not currently address Out-of-Distribution (OOD) generalization beyond the specific domain randomization levels tested. They view this work as a foundational step toward understanding the statistical properties of semi-supervised MOL for generative models.