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

Diffusion models are increasingly used as powerful conditional generators, yet real deployments often involve multiple target distributions arising from different tasks, e.g., diverse prompt domains in text-to-image generation, or multiple environments in robotics with diffusion policies. This naturally leads to a multi-objective learning (MOL) problem. A key challenge is that achieving good Pareto trade-offs can require a generalist model class with substantially larger capacity than what suffices for solving any individual task, thereby increasing statistical cost since sample complexity typically scales with the model complexity. To reconcile this, we develop a principled MOL framework for diffusion models with limited data: a semi-supervised regime where paired (labeled) samples are scarce, but (unlabeled) condition data are abundant. We propose a two-stage training procedure that first fits lightweight specialist models from limited paired data, and then distills them into a generalist model by generating pseudo-samples. We establish generalization bounds showing that the required number of paired samples only depends on the complexity of the specialist model classes. We further extend the theory to diffusion policies for sequential decision making to account for distribution shift in on-policy rollouts. Extensive experiments on robotic control and image restoration tasks are conducted to verify our theoretical results.

Modern AI systems that generate images or control robots are typically built on a class of models called diffusion models. In practice, one such model often needs to handle many different tasks at once such as repairing different kinds of damaged photos, or controlling a robot across varied environments. Training a single large model that does all of this well usually requires huge amounts of carefully labeled training data, such as captioned images or expert demonstrations paired with robot states. This paired data is slow and expensive to collect, while unlabeled inputs (raw photos or robot observations on their own) are abundant and cheap. We propose a two-stage training recipe that exploits this gap. First, for each individual task, we train a small "specialist" model on just the limited labeled data as small models need far less data to learn a single task well. Then each specialist generates its own training examples from the abundant unlabeled inputs, and we combine these with the real labeled data to train one larger "generalist" model. We mathematically prove that this dramatically reduces the labeled data required: the cost depends on the size of the small specialists, not the much larger generalist. Experiments on robot manipulation and image restoration confirm clear improvements. Our work explains why a common industry practice, distilling specialists into a generalist, really pays off.