Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Unsupervised Decomposition and Recombination with Discriminator-Driven Diffusion Models

Archer Wang, Emile Anand, Yilun Du, Marin Soljačić

TL;DR

This work introduces a discriminator-driven diffusion framework for unsupervised decomposition and recombination of data into factorized latent components. By training a discriminator to distinguish single-source from recombined generations and adversarially updating the encoder and generator, the method yields more disentangled and recombinable latent factors, improving both perceptual quality and factorization on images. The approach also extends to robotics, where factorized video representations enable recombined rollouts that broaden state-space exploration in LIBERO while maintaining physical plausibility. Overall, the discriminator feedback provides a scalable regularizer that enhances compositional generalization in both image synthesis and sequential robotic planning, with measurable gains in FID, MIG, MCC, and exploration coverage.

Abstract

Decomposing complex data into factorized representations can reveal reusable components and enable synthesizing new samples via component recombination. We investigate this in the context of diffusion-based models that learn factorized latent spaces without factor-level supervision. In images, factors can capture background, illumination, and object attributes; in robotic videos, they can capture reusable motion components. To improve both latent factor discovery and quality of compositional generation, we introduce an adversarial training signal via a discriminator trained to distinguish between single-source samples and those generated by recombining factors across sources. By optimizing the generator to fool this discriminator, we encourage physical and semantic consistency in the resulting recombinations. Our method outperforms implementations of prior baselines on CelebA-HQ, Virtual KITTI, CLEVR, and Falcor3D, achieving lower FID scores and better disentanglement as measured by MIG and MCC. Furthermore, we demonstrate a novel application to robotic video trajectories: by recombining learned action components, we generate diverse sequences that significantly increase state-space coverage for exploration on the LIBERO benchmark.

Unsupervised Decomposition and Recombination with Discriminator-Driven Diffusion Models

TL;DR

This work introduces a discriminator-driven diffusion framework for unsupervised decomposition and recombination of data into factorized latent components. By training a discriminator to distinguish single-source from recombined generations and adversarially updating the encoder and generator, the method yields more disentangled and recombinable latent factors, improving both perceptual quality and factorization on images. The approach also extends to robotics, where factorized video representations enable recombined rollouts that broaden state-space exploration in LIBERO while maintaining physical plausibility. Overall, the discriminator feedback provides a scalable regularizer that enhances compositional generalization in both image synthesis and sequential robotic planning, with measurable gains in FID, MIG, MCC, and exploration coverage.

Abstract

Decomposing complex data into factorized representations can reveal reusable components and enable synthesizing new samples via component recombination. We investigate this in the context of diffusion-based models that learn factorized latent spaces without factor-level supervision. In images, factors can capture background, illumination, and object attributes; in robotic videos, they can capture reusable motion components. To improve both latent factor discovery and quality of compositional generation, we introduce an adversarial training signal via a discriminator trained to distinguish between single-source samples and those generated by recombining factors across sources. By optimizing the generator to fool this discriminator, we encourage physical and semantic consistency in the resulting recombinations. Our method outperforms implementations of prior baselines on CelebA-HQ, Virtual KITTI, CLEVR, and Falcor3D, achieving lower FID scores and better disentanglement as measured by MIG and MCC. Furthermore, we demonstrate a novel application to robotic video trajectories: by recombining learned action components, we generate diverse sequences that significantly increase state-space coverage for exploration on the LIBERO benchmark.
Paper Structure (50 sections, 3 theorems, 30 equations, 16 figures, 4 tables, 2 algorithms)

This paper contains 50 sections, 3 theorems, 30 equations, 16 figures, 4 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Background
  4. Compositional Generation with Discriminator
  5. Factorized Generative Model
  6. Compositional Recombination
  7. Discriminator Feedback
  8. Experiments
  9. Synthetic Recombination Experiment: Isolating Discriminator Feedback
  10. Image Settings
  11. Perceptual Quality of Recombined Images.
  12. Disentanglement Evaluation.
  13. Videos for Robotics Setting
  14. Factorized video representation.
  15. Detailed Background on Decomp Diffusion and COMET
  16. ...and 35 more sections

Key Result

Lemma 5.1

Let $\{E_k(x)\}_{k=1}^K$ be a collection of energy functions and $\{\alpha_k\}_{k=1}^K$ be a collection of non-negative coefficients. Define the weighted aggregate energy $E(x) \coloneqq \sum_{k=1}^K \alpha_k E_k(x)$. Next, for $i\in\{1,\dots,K\}$, we define the sequence of unnormalized densities $\

Figures (16)

  • Figure 1: (a) We encode two input images into $K=4$ latent components. The middle panels show component-wise reconstructions: for each $k$, we decode using only $z_k$, with all other components set to zero. Red dashed boxes mark which components are chosen from each source to form a hybrid latent code $\tilde{z}$. Decoding $\tilde{z}$ yields the final image (right), which merges appearance/scene attributes from both sources. (b) During training, the discriminator learns to distinguish between predictions from latents from a single input and latents from multiple sources, while the model attempts to fool the discriminator. At inference, recombined latents are sampled via the standard diffusion denoising process.
  • Figure 2: Synthetic recombination experiment. Left: real samples. Middle: naive recombination in entangled coordinates produces off-manifold samples. Right: discriminator feedback restores the geometry of real data. Top row shows $(x_0, x_1)$ projections; bottom row shows $(x_2, x_3)$ projections.
  • Figure 3: For each subfigure, the first two images are source samples drawn from the dataset, and the third image is a recombined generation obtained by mixing latent factors from the two sources. Across all datasets (Virtual KITTI, CLEVR, CelebA-HQ, and Falcor3D; left-to-right, top-to-bottom), our method produces recombined samples that remain visually coherent and structurally consistent, while Decomp Diffusion often exhibits recombination artifacts.
  • Figure 4: Recombined video plan from Scene 5 and Scene 6 from Libero. The two source scenes are "Red mug on right plate" from Scene 5 and "Chocolate pudding left of plate" from Scene 6. The generated video plan shows the robot attempting to place the pudding on the edge of the plate. This specific action sequence does not appear in the original demonstrations, which only include placements beside the plate. The influence of Scene 5 can be observed in the trajectory as the robot arm is pulled sideways, suggesting interactions between the recombined behaviors.
  • Figure 5: Mean ± std joint-state coverage over 10 paired runs for Scene5 (top) and Scene6 (bottom). Recombined videos yield action trajectories with consistently higher exploration than non-compositional baselines under identical warm-up conditions.
  • ...and 11 more figures

Theorems & Definitions (10)

  • Lemma 5.1: Additive energies form a PoE distribution
  • proof
  • Remark 5.2: Naïve recombination artifacts
  • Remark 5.3: Additive diffusion
  • Proposition 5.4: Closure under subset recombination implies Cartesian-product support
  • proof
  • Remark 5.5
  • Definition 5.6: Mutual information
  • Lemma 5.7: Pairwise mutual information contracts under recombination
  • proof