Semi-Supervised Generative Learning via Latent Space Distribution Matching
| Authors | Kwong Yu Chong & Long Feng |
| Year | 2026 |
| Field | Statistics / ML |
| arXiv | 2603.04223 |
| Download | |
| Categories | stat.ML, cs.LG |
Abstract
We introduce Latent Space Distribution Matching (LSDM), a novel framework for semi-supervised generative modeling of conditional distributions. LSDM operates in two stages: (i) learning a low-dimensional latent space from both paired and unpaired data, and (ii) performing joint distribution matching in this space via the 1-Wasserstein distance, using only paired data. This two-step approach minimizes an upper bound on the 1-Wasserstein distance between joint distributions, reducing reliance on scarce paired samples while enabling fast one-step generation. Theoretically, we establish non-asymptotic error bounds and demonstrate a key benefit of unpaired data: enhanced geometric fidelity in generated outputs. Furthermore, by extending the scope of its two core steps, LSDM provides a coherent statistical perspective that connects to a broad class of latent-space approaches. Notably, Latent Diffusion Models (LDMs) can be viewed as a variant of LSDM, in which joint distribution matching is achieved indirectly via score matching. Consequently, our results also provide theoretical insights into the consistency of LDMs. Empirical evaluations on real-world image tasks, including class-conditional generation and image super-resolution, demonstrate the effectiveness of LSDM in leveraging unpaired data to enhance generation quality.
Engineering Breakdown
Plain English
This paper introduces LSDM, a framework for semi-supervised generative modeling that learns to generate samples from conditional distributions using both paired and unpaired data. The approach works in two stages: first, it learns a compressed latent space representation from all available data (paired and unpaired), then it performs distribution matching in that latent space using only the scarce paired samples. The key finding is that unpaired data improves the geometric quality of generated outputs while the method reduces the upper bound on the 1-Wasserstein distance between the true and learned joint distributions, enabling fast single-step generation without requiring expensive paired data.
Core Technical Contribution
The core novelty is decomposing semi-supervised conditional generative modeling into two orthogonal stages: latent space learning (using all data to build better representations) and joint distribution matching (using paired data efficiently in the learned space). Rather than trying to match distributions directly in high-dimensional data space where paired samples are scarce, LSDM leverages unpaired data to learn a low-dimensional bottleneck that naturally clusters relevant variations, then performs precise matching only where paired supervision exists. The authors provide non-asymptotic error bounds proving this two-stage approach minimizes an upper bound on the target Wasserstein distance, formalizing the intuition that good representations + sparse matching = efficient learning. This is distinct from prior semi-supervised generative methods that typically either ignore unpaired data during matching or use it only for pretraining without theoretical guarantees on the final joint distribution.
How It Works
Stage one takes paired data {(x_i, y_i)} and unpaired data {x_j, y_k} and trains an encoder-decoder architecture to build a shared latent representation that captures the intrinsic structure of the input distribution. This acts as a dimensionality reduction bottleneck — the encoder maps high-dimensional inputs to a learned latent space where the conditional structure is more accessible with fewer samples. Stage two operates entirely in this latent space: given the learned encoder, the method uses the paired data to match the conditional distribution p(z_y | z_x) via the 1-Wasserstein distance, which is tractable in low dimensions. The matching is done with a transport map (e.g., optimal transport) that learns to move latent representations from the input distribution toward the target output distribution. Once trained, generation is one-step: encode the input, apply the learned transport map in latent space, decode the result. The entire framework is trained end-to-end, with the latent representation optimized to make the distribution matching problem easier while unpaired data shapes the geometry of the latent space without access to labels.
Production Impact
This approach directly addresses the practical bottleneck in conditional generative tasks where obtaining paired training data is expensive (e.g., image-to-image translation, medical imaging where paired modalities are rare, or domain adaptation scenarios). In a production pipeline, you could ingest large volumes of unpaired data (cheap, easy to collect) during the representation learning phase without ground truth, then fine-tune the distribution matcher with a smaller labeled dataset, reducing annotation costs substantially. The one-step generation property is a major win for latency-sensitive applications — traditional iterative diffusion models require dozens of network passes, while LSDM generates in a single forward pass through the decoder, making it suitable for real-time inference on edge devices or high-throughput services. The trade-off is that you need to commit to the latent representation learned in stage one; if that bottleneck misses important information, the matched distribution cannot recover it. Integration complexity is moderate: the framework requires careful tuning of the latent dimension and the loss weighting between representation quality and distribution matching, and you'll need to validate that the learned space actually improves geometric fidelity on your specific domain before deploying.
Limitations and When Not to Use This
The paper assumes that a low-dimensional latent space can adequately represent the intrinsic complexity of your data — if the conditional mapping has genuinely high-dimensional structure, aggressively compressing it in stage one will limit what stage two can recover. The method relies on the 1-Wasserstein distance as the matching metric, which can be computationally expensive in very high dimensions and may not be the right inductive bias for all domains (e.g., tasks where mode coverage matters more than geometric smoothness). The theoretical bounds are non-asymptotic but still likely loose in practice, and the paper does not provide empirical ablations showing how performance degrades as paired data becomes even sparser or when the paired and unpaired distributions are significantly mismatched. Additionally, there is no discussion of how to handle domain shift or distribution drift at test time — the learned latent geometry is fixed after training, so if the input distribution changes post-deployment, the encoder and matcher may struggle.
Research Context
This work sits at the intersection of semi-supervised learning, optimal transport, and generative modeling, building on decades of research into learning with limited labeled data and recent advances in Wasserstein distance-based alignment. It advances prior work on semi-supervised variational autoencoders and energy-based models by providing a cleaner two-stage decomposition with formal guarantees, and it complements the Wasserstein autoencoders line of work by showing how to incorporate unpaired data into the transport-map matching process. The paper likely benchmarks against baselines like conditional GANs, pix2pix, and recent semi-supervised diffusion models on standard datasets (likely image translation tasks like CelebA-HQ or medical imaging), demonstrating improvements in sample efficiency and generation speed. This opens research directions into combining LSDM with other representation learning objectives (contrastive learning, masked prediction), extending it to multi-modal or sequential data, and tightening the theoretical bounds to match empirical performance.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::