PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning Data Representations with Joint Diffusion Models

Kamil Deja, Tomasz Trzcinski, Jakub M. Tomczak

arXiv:2301.13622v2cs.LGcs.CVstat.ML

TL;DR

Joint generative–predictive models can perform unevenly across synthesis and classification or be unstable to train. This paper shares a diffusion model’s UNet encoder with a classifier for stable end-to-end joint training, achieving stronger generation and classification quality while enabling visual counterfactual explanations.

  • Problem

    Joint models often provide uneven synthesis and classification performance or are unstable to train.

  • Method

    The paper extends diffusion models with a classifier sharing the UNet encoder and trains both objectives end to end.

  • Results

    The joint diffusion model outperforms recent joint and hybrid methods in both generation quality and classification performance.

  • Takeaways & Limitations

    Shared generative and discriminative representations also enable conditional generation and visual counterfactual explanations.

Abstract

from arXiv · show

Joint machine learning models that allow synthesizing and classifying data often offer uneven performance between those tasks or are unstable to train. In this work, we depart from a set of empirical observations that indicate the usefulness of internal representations built by contemporary deep diffusion-based generative models not only for generating but also predicting. We then propose to extend the vanilla diffusion model with a classifier that allows for stable joint end-to-end training with shared parameterization between those objectives. The resulting joint diffusion model outperforms recent state-of-the-art hybrid methods in terms of both classification and generation quality on all evaluated benchmarks. On top of our joint training approach, we present how we can directly benefit from shared generative and discriminative representations by introducing a method for visual counterfactual explanations.

1 Introduction

Joint models aim to synthesize data and predict labels using shared representations, but existing hybrid approaches often favor one task. The paper analyzes diffusion-model representations and introduces a jointly trained diffusion model with a classifier, reporting benefits across generation, classification, and counterfactual explanation.

  • Shared representations could benefit uncertainty calibration, semi-supervised learning, unsupervised domain adaptation, and continual learning.
  • Existing hybrid models often concentrate on synthesizing data or predicting rather than performing both tasks simultaneously.
  • The proposed model shares the UNet encoder between diffusion generation and classification, producing representations trained by both objectives.
  • The method uses classifier-optimized diffusion representations for conditional sampling and image alteration.
  • Shared generative and discriminative representations support visual counterfactual explanations that identify image regions requiring change for a specified output.
  • The joint diffusion model outperforms other joint and hybrid methods on both generation quality and classification performance.

2 Background

The background formulates joint modeling through shared generative and predictive distributions and describes diffusion models as noisy forward processes with learned reverse denoising. Its training objective predicts injected noise using a shared UNet, whose encoder–decoder structure supports this work.

  • A joint distribution can factorize as p(x, y) = p(x|y)p(y) or p(y|x)p(x), combining prediction with modeling the marginal data distribution.
  • Diffusion models create latent variables by progressively adding Gaussian noise according to a variance schedule β1, ..., βT.
  • The learned reverse process models pθ(xt−1|xt) with a parameterized denoising distribution optimized through a variational lower bound.
  • Rather than estimate the previous latent variable directly, diffusion training can predict the added noise ϵ using a neural network ϵθ.
  • Most diffusion models use a shared UNet for noise prediction, and its encoder–decoder structure is central to the proposed joint model.

3 Related Work

Related work combines generative and discriminative modeling through classifier guidance, diffusion representations, energy-based models, and shared parameterizations. These approaches motivate a joint diffusion model that shares representations while avoiding instability or separately trained components.

  • Classifier guidance adds gradients from an independently trained classifier during backward diffusion to steer samples toward a target class.
  • Diffusion models have also been used to generate visual counterfactual explanations, while this work simplifies that setup through shared classifier–diffusion representations.
  • Prior work uses pretrained diffusion UNet activations for tasks such as image segmentation, motivating end-to-end training with generative and discriminative losses.
  • Joint Energy-based Models use one parameterization for classification or generation, but their gradient estimators can be unstable and prone to divergence when optimization is poorly tuned.

4 Diffusion models learn data representations

The diffusion model’s UNet encoder produces pooled representations that support prediction and encode features differently across noise levels. Coarse features emerge near random noise, while fine-grained features appear earlier in the forward diffusion.

  • Representation architecture: UNet denoising representations are decomposed into multilevel feature tensors, average-pooled by filter, and concatenated into a single vector z_t.The denoising model encodes x_t into Z_t and decodes these features to reconstruct x_{t−1}.
  • Predictive representations: Representations z_0 extracted from original images are predictive, with CIFAR-10 performance sometimes comparable to a stand-alone classifier.An MLP classifier receives the pooled UNet representation z_0.
  • Evaluation: Figure 2 compares stand-alone and representation-based classification accuracy and tracks attribute AUC across ten diffusion timesteps.The figure’s timestep comparison supports the distinction between late-emerging high-grained features and early-emerging low-grained features.
  • Timestep-dependent features: High-grained features such as hair color emerge at late diffusion steps near noise, whereas low-grained features such as necklaces or glasses appear only at earlier steps.The timestep analysis fits logistic regressions to representations from ten diffusion stages across CelebA attributes.

5 Method

The proposed joint diffusion model shares a UNet encoder between denoising and classification, training both objectives end-to-end and optionally applying the classifier to noisy intermediate images. Its shared representations also support classifier-guided conditional sampling by optimizing diffusion representations toward a target label.

  • 5.1 Joint Diffusion Models: DDGMs with classifiers: A shared UNet encoder feeds both the denoising decoder and a classifier, whose average-pooled representations are trained with a single joint objective.The encoder maps x_t to multilevel features Z_t; the decoder reconstructs x_{t−1}, while the classifier predicts the target label from pooled features.
  • 5.1 Joint Diffusion Models: DDGMs with classifiers: The joint model factorizes into a classifier term and a diffusion-model term, combining their training objectives rather than relying on separate parameterizations.The diffusion component uses the simplified diffusion objective, while the classifier uses categorical cross-entropy.
  • 5.1 Joint Diffusion Models: DDGMs with classifiers: Training first noises each labeled image to a random timestep for denoising, then applies the classifier to the original image and computes cross-entropy against its label.The denoising loss is estimated over diffusion timesteps, and both objectives are optimized jointly.
  • 5.2 An alternative training of joint diffusion models: The alternative training reuses the classifier on intermediate noisy images, adding cross-entropy losses across selected timesteps to improve robustness.These noisy classifiers are used only during training, not prediction, and function similarly to noise-based data augmentation.
  • 5.3 Conditional sampling in joint diffusion models: For conditional sampling, the encoder representation is optimized toward a target class before the decoder produces the preceding diffusion state.The same representation supports denoising through dψ(Z_t) and classification through gω(Z_t), replacing reliance on an externally trained classifier for guidance.

6 Experiments

Experiments evaluate whether joint diffusion improves classification, generation, conditional sampling, and visual counterfactual explanations. Across these tasks, the model achieves strong predictive and generative performance while supporting class-directed generation and image explanations.

  • 6.1 Predictive performance of joint diffusion models: The proposed joint diffusion model achieves the best classification performance on FashionMNIST, SVHN, CIFAR-10, and CIFAR-100.It outperforms a stand-alone classifier, a classifier using pretrained diffusion representations, and three joint or hybrid models.
  • 6.3 A comparison to state-of-the-art joint approaches: The joint diffusion model clearly outperforms all compared state-of-the-art joint models in generative performance.Table 3 compares joint models using FID, while Table 4 also places the method alongside state-of-the-art discriminative and generative models on CIFAR-10.
  • 6.2 Generative performance of joint diffusion models: The joint model outperforms standard diffusion models on general FID, while conditional sampling improves generation quality across all evaluated benchmarks.Conditional sampling especially improves Precision, although vanilla diffusion or classifier guidance can be better on particular Precision or Recall components.
  • 6.2 Generative performance of joint diffusion models: Increasing the conditional-sampling step size α produces more precise but less diverse samples, with both measures high around α ∈[100, 250].Larger α values move representations closer to modes, accelerate convergence toward target classes, and reduce background diversity.
  • 6.4 Visual Counterfactual Explanations: On malaria images, conditional sampling changes class labels by removing or adding parasites, and image differences indicate the model’s relevant regions.The examples are described as not cherry-picked, and the method adapts visual counterfactual explanations without changing the original setup.

7 Conclusion

The paper introduces a jointly parameterized diffusion model and classifier built on semantically meaningful diffusion representations. It improves classification and generation, achieves state-of-the-art joint-model performance, and supports visual counterfactual explanations without altering the setup.

  • 7 Conclusion: The joint model combines a diffusion model and classifier through shared parameterization.The authors first show that diffusion models learn representations useful for classification, then use those representations in joint training.
  • 7 Conclusion: The approach improves both classification and generative tasks while enabling conditional generations with built-in classifier guidance.The paper reports high-quality generations and state-of-the-art performance among joint models.
  • 7 Conclusion: The joint diffusion model provides visual counterfactual explanations without changes to the original setup.These explanations use the shared generative and discriminative representations.

A Additional experiments

Additional experiments test the joint model under limited labels and domain shift. Shared generative training improves semi-supervised classification and enables stronger unsupervised adaptation using unlabeled target-domain data.

  • A.1 Semi-supervised learning of joint diffusion models: Semi-supervised experiments limit labels to 10%, 5%, or 1% on SVHN, CIFAR-10, and CIFAR-100 and compare standard, pretrained-representation, and joint classifiers.No data augmentation is used in this setup, making full-data performance slightly lower than in Table 1.
  • A.1 Semi-supervised learning of joint diffusion models: With only 5% of labeled CIFAR-10 data, the joint diffusion model performs almost as well as a stand-alone classifier trained on the fully labeled dataset.The diffusion component continues training unsupervised on all available unlabeled data; with extremely few labels, pretraining representations first can be slightly more beneficial.
  • A.2 Domain adaptation: Domain adaptation retrains the source-trained joint model on target data using only the diffusion loss, without altering the basic setup.The experiments compare this procedure with a stand-alone deep neural network classifier.
  • A.2 Domain adaptation: The joint model outperforms a stand-alone classifier in all three domain-adaptation scenarios by a significant margin.It adapts using unlabeled target-domain data after source-task training, while the stand-alone classifier’s accuracy degrades on the target domain.

B Training details and hyperparameters

The implementation shares UNet features between the classifier and generative model through average pooling. Training procedures use diffusion objectives on available data and adjust classifier losses for limited-label settings and domain adaptation.

  • B Training details and hyperparameters: Average pooling converts UNet activations into a lower-dimensional representation but loses spatial information about feature locations.The approach averages each convolutional filter activation across the whole image.
  • B Training details and hyperparameters: Average pooling shares all features across the image between the classifier and generative model.The authors note that max or min pooling could also be used.
  • B Training details and hyperparameters: In semi-supervised training, the diffusion loss uses the full batch while the classifier loss uses only labeled examples.This design addresses settings where up to 99% of labels are omitted.
  • B Training details and hyperparameters: For domain adaptation, the model is first trained jointly on the source task and then retrained on target data using only the DDGM loss.The reported setup produces a performance boost over the baseline classifier without modifying the basic model.

C Additional results: Conditional generations with optimized representations

The joint diffusion model produces conditional samples for Fashion MNIST and CIFAR100, and unconditional generated examples for CIFAR-10, CIFAR-100, and CelebA.

  • Conditional samples are shown for Fashion MNIST and the first 10 CIFAR100 classes, with each row representing one class.
  • Unconditional generated examples are presented for CIFAR-10, CIFAR-100, and CelebA.

D Additional results: Counterfactual image generation

The paper demonstrates counterfactual image generation by altering selected CelebA attributes in both directions on original examples. The examples cover age, smiling, and mustache attributes, with five positive and five negative examples selected for each attribute.

  • Counterfactual generation changes three CelebA attributes—young, smiling, and mustache—on original examples using conditional sampling.
  • The shown transformations include young to old and old to young, plus smiling to no-smiling and no-smiling to smiling.
  • For each attribute, five positive and five negative examples are selected and altered in opposite directions.
Loading 2301.13622v2…