📊 statistics

Semi-Supervised Conditional Diffusion via Label Augmentation

This paper introduces Label-Augmented Conditional Diffusion (LACD), a method that leverages unlabeled data by assigning them trivial labels to achieve faster convergence in total variation distance and improved generative performance compared to purely supervised conditional diffusion models.

Original authors: Jin Su, Yuan Gao, Yong Zhou, Jian Huang

Published 2026-07-21

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jin Su, Yuan Gao, Yong Zhou, Jian Huang

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 paint pictures of cats, dogs, and birds. In the world of artificial intelligence, this is called "conditional generation"—teaching a machine to create specific things based on a label you give it. For years, the best tools for this job have been "diffusion models." Think of these models like a master sculptor who starts with a block of noisy, static-filled marble and slowly chips away the noise to reveal a perfect statue. To learn how to carve a specific animal, the sculptor needs to see thousands of examples of that animal being carved. But here's the catch: getting those labeled examples is expensive and hard. You might have a million photos of animals lying around, but only a few hundred have tags saying "cat" or "dog." The rest are just unlabeled blobs of pixels.

This is where the problem gets tricky. If you only show the sculptor the few hundred labeled photos, they might get confused or start memorizing the few examples they have, resulting in weird, repetitive art. But if you could somehow use those millions of unlabeled photos to help the sculptor understand the general shape of "animals" before teaching them the specific differences between cats and dogs, the result would be much better. This paper tackles exactly that challenge: how to use a mountain of unlabeled data to improve the training of these AI artists, even when the labeled data is scarce.

The authors of this paper, Jin Su and colleagues, propose a clever new method called Label-Augmented Conditional Diffusion (LACD). Their big idea is to stop treating the unlabeled data as "useless" and instead give it a temporary, "trivial" label—let's call it "Unknown." Imagine you have a classroom of students. You have a few students who know exactly what a "cat" looks like, and a huge crowd who just know they are looking at "an animal" but don't know the species. Instead of ignoring the second group, the authors tell the teacher (the AI model) to treat "Unknown" as a valid category.

Here is how the magic happens: The AI learns to generate images for "cat," "dog," and "bird" by looking at the labeled students. Simultaneously, it learns what a generic "animal" looks like by studying the huge crowd of unlabeled students. Because the AI is trained to understand that "cat," "dog," and "bird" are all variations of that generic "animal" concept, the massive amount of unlabeled data helps it learn the shared features of all animals (like having fur, ears, or tails) much faster and more accurately. This shared understanding acts as a strong foundation, allowing the model to distinguish between specific animals with far fewer labeled examples than usual.

The paper doesn't just guess that this works; they prove it with math and test it with computers. They showed that when you have a lot of unlabeled data, this new method creates images that are statistically closer to the real thing than the old methods that only used labeled data. In their experiments, they tested this on everything from simple computer-generated shapes to real-world photos of cities and even brain wave data.

For example, on a famous image dataset called CIFAR-10, they found that if they only had 1,000 labeled images, adding 10,000 unlabeled images made the AI's pictures of birds and cars look significantly more realistic and diverse. The old method, which ignored the unlabeled photos, produced images that looked a bit blurry or repetitive. The new method, using the "Unknown" label trick, produced images that were sharper and more varied, almost as good as if they had been trained on all 50,000 images with labels.

The authors also ran simulations on synthetic data to see exactly how the math works. They found that as they added more unlabeled data, the error in the generated images dropped steadily. In some cases, the improvement was strict and guaranteed by their mathematical proofs, meaning the new method is mathematically proven to be better than the old way when unlabeled data is available. They even tested this on a tabular dataset of EEG brain signals for epilepsy detection, showing that the method works not just for pretty pictures, but for complex data tables too.

However, the paper is careful not to overhype the results. They note that while the method is a huge improvement, it doesn't completely replace the need for labeled data; it just makes the labeled data you do have go much further. They also point out that their current mathematical proofs rely on certain assumptions about how the data is distributed, and real-world data can sometimes be messier than their simulations. But overall, the message is clear: by treating unlabeled data as a "generic" category, we can teach AI models to be better artists with fewer expensive labels, turning a pile of untagged data into a powerful learning tool.

Technical Summary: Semi-Supervised Conditional Diffusion via Label Augmentation

Problem Statement

The paper addresses the challenge of estimating conditional distributions $P_{X|Y}$ in high-dimensional settings where labeled data $(X, Y)$ is scarce, but unlabeled data $X$ is abundant. While conditional diffusion models have achieved minimax optimality in supervised settings, they suffer from overfitting and poor performance when paired data is limited. Existing semi-supervised learning (SSL) methods primarily focus on discriminative tasks or estimating marginal densities, leaving a gap in utilizing unlabeled covariates $X$ to improve the estimation of the conditional distribution $P_{X|Y}$ in generative modeling.

Methodology: Label-Augmented Conditional Diffusion (LACD)

The authors propose Label-Augmented Conditional Diffusion (LACD), a framework that integrates unlabeled data into the conditional diffusion process through a specific augmentation strategy.

1. Augmented Probability Space

The core innovation is the construction of an augmented conditioning variable $\tilde{Y} \in \mathcal{Y} \cup \{\emptyset\}$ , where $\emptyset$ represents the unlabeled state.

Labeled Data: Assigned their true label $Y$ .
Unlabeled Data: Assigned a "trivial" label $\emptyset$ .
Identification: The conditional distribution in this augmented space is defined such that $P_{X|\tilde{Y}=\emptyset} = P_X$ (the marginal distribution), while $P_{X|\tilde{Y}=y} = P_{X|Y=y}$ for $y \in \mathcal{Y}$ . This allows the model to learn the shared baseline component of the distribution from the abundant unlabeled data while using labeled pairs to identify condition-specific variations.

2. Model Specification

The method assumes that conditional densities share a common baseline structure parameterized by a function $h(\tilde{y})$ :
$p_{X|\tilde{Y}=\tilde{y}}(x|\tilde{y}) = p(x; h(\tilde{y}))$
where $h(\tilde{y}) = b_0 + \Phi_0 e_y \cdot \mathbb{1}\{\tilde{y} = y \in \mathcal{Y}\}$ .

$b_0$ captures the shared component (learned from both labeled and unlabeled data).
$\Phi_0 e_y$ captures the condition-specific deviation (learned primarily from labeled data).
This formulation encompasses various families, including location families, exponential families, and latent variable models.

3. Training and Inference

Training: The model employs a score-based diffusion process (Ornstein-Uhlenbeck based). The score network $s_\theta$ is trained via Denoising Score Matching (DSM) on the augmented dataset. The empirical loss is a weighted sum of losses from labeled samples (dependent on $b$ and $\Phi$ ) and unlabeled samples (dependent only on $b$ ).
Sampling: To generate samples for a specific condition $y$ , the reverse diffusion process is simulated conditioned on $\tilde{y} = y$ .

Key Contributions

1. Theoretical Guarantees

The paper establishes finite-sample convergence rates for the LACD estimator in both Total Variation (TV) and Wasserstein-1 (W1) distances.

Strict Improvement in TV: The convergence rate in TV distance is strictly faster than the purely supervised estimator when unlabeled samples are available ( $N_u \gg N_\ell$ ).
Improvement in W1: The convergence rate in W1 distance is at least as fast as the supervised counterpart. Specifically, when the covariate dimension $d_x > 2$ and the number of unlabeled samples is sufficiently large, the LACD rate is strictly faster.
Identifiability: The authors provide sufficient conditions (via measure-theoretic construction) guaranteeing the population-level identifiability of the target conditional distribution under the augmented scheme.

2. Empirical Validation

Extensive experiments were conducted on synthetic, image, and tabular benchmarks:

Synthetic Data: On 2D datasets (e.g., multi-modal Gaussians, rings, spirals), increasing the number of unlabeled samples ( $N_u$ ) consistently reduced TV, MMD, and W1 distances, with the generated samples better recovering target geometry.
Image Benchmarks (CIFAR-10 & Intel Scenes):
- On CIFAR-10, adding unlabeled data significantly improved Fréchet Inception Distance (FID). For instance, with only 1,000 labeled samples, adding 10,000 unlabeled samples reduced the class-averaged FID by ~48%, matching the performance of a model trained with 5,000 labeled samples alone.
- Visual analysis and t-SNE projections indicated that unlabeled data reduced memorization of labeled samples and improved feature-space coverage.
- Similar improvements were observed on the Intel Scenes dataset across various scene categories.
Tabular Data (Bangalore EEG Epilepsy Dataset): On a 16-dimensional EEG dataset, increasing unlabeled samples improved feature-wise and sliced W1 distances, as well as boundary adherence and marginal distribution alignment (KS Complement), demonstrating the method's applicability to non-image data.

Significance and Claims

The paper claims that LACD provides a principled and effective solution for semi-supervised conditional generation. Its primary significance lies in:

Theoretical Justification: It is the first work to rigorously establish that unlabeled observations of the generated object $X$ can strictly improve the estimation of $P_{X|Y}$ in diffusion models, providing convergence rates that outperform supervised baselines.
Practical Efficiency: The method demonstrates substantial gains in sample efficiency, allowing high-quality conditional generation even when labeled data is extremely scarce (e.g., 1-10% of the total data).
Generality: The approach is applicable to diverse data types (synthetic, images, tabular) and distribution families, relying on a shared-parameter architecture that naturally leverages the marginal information contained in unlabeled data.

The authors conclude that while the method shows strong empirical and theoretical promise, future work could explore minimax lower bounds, intrinsic dimensionality adaptations, and relaxing the assumption of missing completely at random (MCAR) for labels.