InfoDiffusion: Representation Learning Using Information Maximizing Diffusion Models

Yingheng WangYair SchiffAaron GokaslanWeishen PanFei WangChristopher De SaVolodymyr Kuleshov

article2023ICML64 citations

Proposes InfoDiffusion, a framework that integrates mutual information regularization into diffusion models to extract semantically meaningful, disentangled low-dimensional latent representations without sacrificing generative sample quality.

Listen

Modern generative artificial intelligence models often face a fundamental trade-off between output quality and interpretability. Denoising diffusion models produce state-of-the-art images, audio, and molecular designs, but their internal variables typically lack semantic structure, making them ill-suited for representation learning—the unsupervised discovery of high-level concepts such as facial features, object attributes, or distinct categories. Conversely, traditional frameworks like variational autoencoders provide structured, interpretable representations but generate noticeably lower-quality samples.

The article demonstrates that augmenting diffusion models with low-dimensional auxiliary latent variables regularized by mutual information maximization enables the simultaneous learning of disentangled, human-interpretable representations and high-fidelity data generation. The authors introduce InfoDiffusion, an algorithm designed to prevent powerful diffusion decoders from ignoring compact latent codes while ensuring that learned representations align cleanly with user-specified prior distributions.

To evaluate the proposed method, the researchers conducted extensive empirical benchmarks across five standard image datasets: FashionMNIST, CIFAR10, FFHQ, CelebA, and 3DShapes. The approach uses an encoder to infer compact latent vectors and conditions a multi-step diffusion decoder across all layers using adaptive normalization. The evaluation compared InfoDiffusion against conventional autoencoders, variational baselines, diffusion autoencoders, and leading self-supervised contrastive learning methods on sample quality, representation utility in downstream classification, and factor disentanglement.

The findings show that InfoDiffusion achieves superior or highly competitive representation quality while fully preserving the generation fidelity of diffusion models. On downstream classification tasks, linear models trained on InfoDiffusion's compact representations matched or outperformed existing baselines across all datasets. In disentanglement benchmarks, InfoDiffusion achieved the highest scores on 3DShapes (0.342 DCI) and CelebA (0.299 Total Attribute Disentanglement), outperforming both traditional generative models and 32-dimensional contrastive baselines by wide margins. Furthermore, the model demonstrated strong qualitative control, allowing smooth interpolation between data points and precise manipulation of individual attributes (such as adding smiles or altering hairstyles) without corrupting image fidelity. Finally, the framework successfully accommodated discrete and categorical variables, expanding beyond standard continuous assumptions.

These results indicate that generative systems no longer require separate architectures for high-level data understanding and high-resolution generation. For practical applications such as generative product design, digital content creation, and medical diagnostics, InfoDiffusion reduces system complexity and deployment risks by providing direct, predictable human control over generated content while retaining state-of-the-art visual quality. Unlike earlier diffusion autoencoder variants, it also supports unconditional generation directly from prior distributions without requiring auxiliary latent diffusion models.

Organizations developing controllable generative tools or automated feature extraction pipelines should pilot auxiliary-variable diffusion frameworks in workflows requiring structured manipulation. Future work should evaluate this methodology beyond image generation, applying it to complex non-visual domains such as molecular conformation, material science, and audio synthesis, while investigating scalability to higher-resolution data regimes.

Confidence in these findings is high across standard computer vision benchmarks, backed by rigorous multi-fold evaluations and mathematical proofs of global optimality. However, decision-makers should note that the empirical evaluations were conducted on standardized datasets at moderate image resolutions. Additional validation on complex, out-of-distribution industrial datasets is recommended prior to broad operational deployment.

Cornell University decoders diffusion probabilistic models mutual informations InfoDiffusion information maximizing diffusion models representation learning

arXiv: 2306.08757

Paper: InfoGAN: Interpretable Representation Learning by Information Maximizing Generative Adversarial Nets, Xi Chen et al. (2016). Introduces the foundational paradigm of maximizing mutual information between auxiliary latent codes and generated outputs to learn disentangled representations in deep generative models.
Paper: Denoising Diffusion Probabilistic Models, Jonathan Ho et al. (2020). Establishes the core denoising diffusion probabilistic framework that InfoDiffusion augments with structured auxiliary latent variables.
Paper: beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework, Irina Higgins et al. (2016). Pioneers unsupervised factor disentanglement principles and evaluation protocols that provide standard comparative baselines for InfoDiffusion.
Paper: Fixing a Broken ELBO, Alexander A. Alemi et al. (2018). Provides the theoretical information-theoretic bounds explaining why standard likelihood-based objectives fail to encode meaningful representations into latent codes.
Paper: Disentangling by Factorising, Hyunjik Kim et al. (2018). Develops total correlation regularizers and disentanglement benchmarks directly targeted and compared against in InfoDiffusion's evaluation suite.
Paper: Denoising Diffusion Implicit Models, Jiaming Song et al. (2021). Presents non-Markovian deterministic sampling and latent trajectory mechanics essential for understanding continuous interpolation in diffusion models.

Paper: Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think, Sihyun Yu et al. (2025). Extends the integration of representation learning and diffusion architectures by aligning internal model representations directly with external visual encoders during training.
Paper: Masked Autoencoders Are Effective Tokenizers for Diffusion Models, Hao Chen et al. (2025). Investigates how structured and discriminative latent spaces govern diffusion performance, providing an alternate avenue for combining representation richness with high-fidelity generation.
Paper: Generalization in diffusion models arises from geometry-adaptive harmonic representations, Zahra Kadkhodaie et al. (2024). Provides a theoretical analysis of the internal geometric and harmonic representations that emerge inside diffusion network denoisers.
Paper: An exact information theory of generalization phase transitions in Bayesian diffusion models, Henry Hunt et al. (2026). Formalizes an analytical, information-theoretic perspective on how diffusion models avoid memorization and generalize through restricted Bayesian representations.

Cover for InfoDiffusion: Representation Learning Using Information Maximizing Diffusion Models

Abstract

While diffusion models excel at generating high-quality samples, their latent variables typically lack semantic meaning and are not suitable for representation learning. Here, we propose InfoDiffusion, an algorithm that augments diffusion models with low-dimensional latent variables that capture high-level factors of variation in the data. InfoDiffusion relies on a learning objective regularized with the mutual information between observed and hidden variables, which improves latent space quality and prevents the latents from being ignored by expressive diffusion-based decoders. Empirically, we find that InfoDiffusion learns disentangled and human-interpretable latent representations that are competitive with state-of-the-art generative and contrastive methods, while retaining the high sample quality of diffusion models. Our method enables manipulating the attributes of generated images and has the potential to assist tasks that require exploring a learned latent space to generate quality samples, e.g., generative design.

Topics

decoders diffusion probabilistic models mutual informations representation learning semantics learning objectives latent factor generative model latent variable models

Keywords

InfoDiffusion information maximizing diffusion models representation learning diffusion models variational auxiliary-variable diffusion information regularization mutual information Gumbel-Softmax distribution unsupervised representation learning variational inference

1. Introduction
2. Background
3. Diffusion Models With Auxiliary Latents
3.1. Auxiliary Latent Variables and Semantic Prior
3.2. Auxiliary-Variable Diffusion Decoder
4. Learning and Inference Algorithms For Auxiliary-Variable Diffusion Models
4.1. Variational Inference for Auxiliary-Variable Models
4.2. Inferring Latent Representations
4.3. Discrete Auxiliary-Variable Diffusion
4.4. Sampling Methods
5. InfoDiffusion: Regularizing Semantic Latents By Maximizing Mutual Information
5.1. Regularizing Auxiliary-Variable Diffusion
5.2. A Tractable Objective for InfoDiffusion
5.3. Comparing InfoDiffusion to Existing Models
6. Experiments
6.1. Exploring Latent Representations
6.2. Disentanglement
6.2.1. FINDING DISENTANGLED DIMENSIONS
6.2.2. DISENTANGLEMENT METRICS
6.3. Discrete Latent Priors
6.4. Comparison to Contrastive Methods
6.5. Exploring InfoDiffusion Modeling Choices
7. Related Work
7.1. Representation Learning in Generative Modeling
7.2. Diffusion Models for Representation Learning
8. Conclusion
Acknowledgements
References
A. Proof of Proposition 5.1
B. Proof of Proposition 5.2
C. Additional Experimental Details
D. Additional Sampling Details
D.1. Sampling from Prior
D.2. Sampling from Learned Prior
E. Illustrations of Network Architecture
F. Ablation: Different Approach for Conditioning on z
G. Discrete Latents
H. Regularization Coefficients
I. Qualitative Figures on 3DShapes
J. Assets

Knowls

Knowl 1 — Tractable Objective for Information-Maximizing Diffusion Models
equation
The InfoDiffusion learning objective maximizes an evidence lower bound (ELBO) augmented with a mutual information regularizer $\mathrm{MI}_{x_0, z}$ and an arbitrary strict divergence prior penalty $D(q_\phi(z) \parallel p(z))$ :

$\mathcal{L}_I = \mathbb{E}_{q(x_0, x_1)} \left[ \mathbb{E}_{q_\phi(z \mid x_0)} [\log p_\theta(x_0 \mid x_1, z)] \right] - \mathbb{E}_{q(x_0)} [\mathrm{KL}(q(x_T \mid x_0) \parallel p(x_T))] - \sum_{t=2}^T \mathbb{E}_{q(x_0, x_t)} \left[ \mathbb{E}_{q_\phi(z \mid x_0)} [\mathrm{KL}(q(x_{t-1} \mid x_t, x_0) \parallel p_\theta(x_{t-1} \mid x_t, z))] \right] - (1 - \zeta)\mathbb{E}_{q(x_0)} [\mathrm{KL}(q_\phi(z \mid x_0) \parallel p(z))] - (\lambda + \zeta - 1) D(q_\phi(z) \parallel p(z))$

where $x_0 \in \mathcal{X}_0$ is the clean input data, $x_{1:T}$ are diffusion latent variables over $T$ timesteps, $z \in \mathcal{Z}$ is an auxiliary low-dimensional semantic latent variable, $q_\phi(z \mid x_0)$ is an approximate variational posterior parameterized by an encoder network with parameters $\phi$ , and $q_\phi(z) = \mathbb{E}_{q(x_0)}[q_\phi(z \mid x_0)]$ is the aggregated marginal approximate posterior. The distribution $p_\theta(x_{t-1} \mid x_t, z)$ is a conditional reverse diffusion process parameterized by $\theta$ , $q(x_{1:T} \mid x_0) = \prod_{t=1}^T q(x_t \mid x_{t-1})$ is the fixed forward noising process, and $p(z)$ and $p(x_T) = \mathcal{N}(0, \mathbf{I})$ are prior distributions. Hyperparameters $\zeta \le 1$ and $\lambda \ge 0$ (where $\lambda = \beta - 1$ for a prior regularizer scaling scalar $\beta > 0$ ) control the strength of mutual information maximization and prior matching, respectively.

In practice, the strict divergence $D(q_\phi(z) \parallel p(z))$ is instantiated as the Maximum Mean Discrepancy (MMD) with positive definite kernel $k(\cdot, \cdot)$ :

$\mathrm{MMD}(q_\phi(z) \parallel p(z)) = \mathbb{E}_{z, z' \sim q_\phi(z)}[k(z, z')] + \mathbb{E}_{z, z' \sim p(z)}[k(z, z')] - 2\mathbb{E}_{z \sim q_\phi(z), z' \sim p(z)}[k(z, z')]$

which is optimized using empirical mini-batch averages.
Knowl 2 — Global Optimality of the InfoDiffusion Objective
theoretical result
Let the data space $\mathcal{X}_0$ and latent feature space $\mathcal{Z}$ be continuous spaces, with regularizer weights restricted to $\zeta \le 1$ and $\lambda \ge 0$ . For any fixed target value $I_0$ of the mutual information $\mathrm{MI}_{x_0, z} = \mathbb{E}_{q_\phi(x_0, z)}\left[\log \frac{q_\phi(x_0, z)}{q(x_0)q_\phi(z)}\right]$ , the InfoDiffusion objective $\mathcal{L}_I$ with any strict divergence $D(q_\phi(z) \parallel p(z))$ is globally optimized if and only if:
1. $p_\theta(x_{t-1} \mid x_t, z) = q(x_{t-1} \mid x_t, x_0)$ for all $t \ge 2$ ;
2. $p_\theta(x_0 \mid x_1, z) = q(x_0 \mid x_1)$ ;
3. $q_\phi(z) = p(z)$ .
Under this global optimum, the joint generative model distribution matches the joint data-inference distribution, $p_\theta(x_{0:T}, z) = q_\phi(x_{0:T}, z)$ , which guarantees that the marginal generative distribution recovers the exact true data distribution: $p_\theta(x_0) = q(x_0)$ .
Knowl 3 — Auxiliary-Variable Diffusion Model Generative Formulation
model/method
An auxiliary-variable diffusion model defines a joint probability distribution $p(x_{0:T}, z)$ over input data $x_0$ , diffusion trajectory states $x_{1:T} = (x_1, x_2, \dots, x_T)$ , and a low-dimensional auxiliary latent vector $z \in \mathcal{Z}$ capturing high-level semantic factors of variation. The model factorizes as a conditioned Markov chain:

$p(x_{0:T}, z) = p(x_T) p(z) \prod_{t=1}^T p_\theta(x_{t-1} \mid x_t, z)$

where $p(z)$ is a semantic prior over $z$ , $p(x_T) = \mathcal{N}(x_T; 0, \mathbf{I})$ is the prior over the final diffusion noise, and $p_\theta(x_{t-1} \mid x_t, z)$ is a reverse diffusion step parameterized as:

$p_\theta(x_{t-1} \mid x_t, z) = \mathcal{N}\left(x_{t-1}; \frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t}} \epsilon_\theta(x_t, t, z) \right), \sigma_t^2 \mathbf{I}\right)$

with noise prediction network $\epsilon_\theta(x_t, t, z)$ , variance schedule $\alpha_t$ , and cumulative product $\bar{\alpha}_t = \prod_{s=1}^t \alpha_s$ . The auxiliary latent $z$ is conditionally independent of the forward noising chain $q(x_{1:T} \mid x_0) = \prod_{t=1}^T q(x_t \mid x_{t-1})$ centered at $q(x_t \mid x_{t-1}) = \mathcal{N}(x_t; \sqrt{\alpha_t} x_{t-1}, (1-\alpha_t)\mathbf{I})$ , ensuring that $z$ represents semantic properties of the clean input $x_0$ rather than a control variable for the forward corruption dynamics.
Knowl 4 — Auxiliary Latent Conditioning via Adaptive Group Normalization
model/method
The diffusion noise prediction network $\epsilon_\theta(x_t, t, z)$ parameterizes a U-Net decoder conditioned on both the diffusion timestep $t$ and the auxiliary latent representation $z$ . Conditioning is implemented using Adaptive Group Normalization (AGN):

$\mathrm{AGN}(h, z) = (1 + s(z)) \cdot \mathrm{GroupNorm}(h) + b(z)$

where $h$ is the feature activation map, $\mathrm{GroupNorm}$ is group normalization, and $s(z), b(z)$ are scale and bias vectors predicted from $z$ via a multilayer perceptron $\mathrm{MLP}(z)$ .

Each residual block in the U-Net incorporates two successive AGN layers: the first conditions on the timestep embedding $\mathrm{Embed}(t)$ , and the second conditions on the auxiliary representation embedding $\mathrm{MLP}(z)$ . Conditioning on $z$ across all layers of the U-Net is required; restricting conditioning solely to the U-Net bottleneck layer degrades reconstruction quality, disentanglement, and downstream linear probe accuracy.
Knowl 5 — Two-Phase Prior Sampling and Latent Space Inversion
algorithm
InfoDiffusion performs unconditional generation directly from prior distributions using a two-phase sampling pipeline, and encodes deterministic residual high-frequency details into latent noise variable $x_T$ .
Input: Timesteps $T$ , schedule $\alpha_{1:T}$ , noise prediction model $\epsilon_\theta(x_t, t, z)$ , unconditional diffusion model $\epsilon_\text{uncond}(x_t, t)$ , prior $p(z)$
Output: Generated sample $x_0$
Sample $x_T \sim \mathcal{N}(0, \mathbf{I})$
for $t = T$ down to $T/2 + 1$ do
    Estimate noise $\hat{\epsilon} = \epsilon_\text{uncond}(x_t, t)$
    Sample $x_{t-1} \sim p(x_{t-1} \mid x_t)$ using $\hat{\epsilon}$
end for
Sample auxiliary latent $z \sim p(z)$
for $t = T/2$ down to $1$ do
    Estimate noise $\hat{\epsilon} = \epsilon_\theta(x_t, t, z)$
    Sample $x_{t-1} \sim p_\theta(x_{t-1} \mid x_t, z)$ using $\hat{\epsilon}$
end for
return $x_0$
To deterministically encode an image $x_0$ to preserve high-frequency details not captured by $z = \mathbb{E}[q_\phi(z \mid x_0)]$ , intermediate states $x_t$ are updated recursively via:

$x_{t+1} = \sqrt{\bar{\alpha}_{t+1}} \hat{x}_0(x_t, t, z) + \sqrt{1 - \bar{\alpha}_{t+1}} \epsilon_\theta(x_t, t, z)$

where $\hat{x}_0(x_t, t, z) = \frac{1}{\sqrt{\bar{\alpha}_t}} \left( x_t - \sqrt{1 - \bar{\alpha}_t} \epsilon_\theta(x_t, t, z) \right)$ .
Knowl 6 — Discrete Representation Learning with Relaxed-Bernoulli Distributions
model/method
When latent attributes are inherently categorical or binary, InfoDiffusion parameterizes the approximate posterior $q_\phi(z \mid x_0)$ using the continuous Gumbel-Softmax (Concrete) relaxation. For a categorical variable with $k$ classes (and $k=2$ for Relaxed-Bernoulli) with class probabilities $\pi_1, \dots, \pi_k$ , continuous relaxed latent variables $z_i$ are sampled via:

$z_i = \frac{\exp((\log(\pi_i) + g_i)/\tau)}{\sum_{j=1}^k \exp((\log(\pi_j) + g_j)/\tau)}$

where $g_1, \dots, g_k \sim \mathrm{Gumbel}(0, 1)$ are independent and identically distributed noise draws, and $\tau > 0$ is a temperature parameter.

During training, gradients flow via the reparameterization trick while $\tau$ is annealed linearly from an initial value of $1.0$ down to $0.5$ by reducing $\tau$ by $0.00003$ every 1,000 optimization steps. At test time, the temperature is set to $\tau = 0$ (taking the $\mathrm{argmax}$ ) to yield hard discrete binary/categorical representations.
Knowl 7 — Unification of Autoencoders, VAEs, and Diffusion Models under InfoDiffusion
model/method
The InfoDiffusion framework unifies multiple representation learning and generative modeling architectures as special parametric cases:
1. Single-step diffusion ( $T=1$ ): When diffusion reduces to a single transition step, the architecture matches the InfoVAE framework up to decoder network choices.
2. Single-step with no prior regularization ( $T=1, \lambda=0$ ): Recovers the $\beta$ -VAE framework.
3. Single-step with Jensen-Shannon divergence ( $T=1, D = D_\mathrm{JS}$ ): Recovers Adversarial Autoencoders (AAEs).
4. Zero mutual information and prior regularization ( $\zeta=0, \lambda=0, T > 1$ ): Recovers the non-probabilistic Diffusion Autoencoder (DiffAE).
5. No auxiliary latent ( $z = \emptyset$ ): Recovers standard Denoising Diffusion Probabilistic Models (DDPM).

Knowl 8 — Disentanglement and Representation Quality on 3DShapes and CelebA

data/table

InfoDiffusion was evaluated against deterministic and probabilistic generative baselines on 3DShapes using the DCI disentanglement metric and on CelebA using the Total Attribute Disentanglement (TAD) score, number of captured attributes (Attrs $\ge 0.75$ AUROC), latent linear classification AUROC (Latent Qual.), and sample generation Fréchet Inception Distance (FID).

	3DShapes		CelebA
Model	DCI $\uparrow$	HQ Samples	TAD $\uparrow$	Attrs $\uparrow$	Latent Qual. $\uparrow$	FID $\downarrow$
AE	0.219 $\pm$ 0.001	No	0.042 $\pm$ 0.004	1.0 $\pm$ 0.0	0.759 $\pm$ 0.003	90.4 $\pm$ 1.8
VAE	0.276 $\pm$ 0.001	No	0.000 $\pm$ 0.000	0.0 $\pm$ 0.0	0.770 $\pm$ 0.002	94.3 $\pm$ 2.8
$\beta$ -VAE	0.281 $\pm$ 0.001	No	0.088 $\pm$ 0.051	1.6 $\pm$ 0.8	0.699 $\pm$ 0.001	99.8 $\pm$ 2.4
InfoVAE	0.134 $\pm$ 0.001	No	0.000 $\pm$ 0.000	0.0 $\pm$ 0.0	0.757 $\pm$ 0.003	77.8 $\pm$ 1.6
DiffAE	0.196 $\pm$ 0.001	Yes	0.155 $\pm$ 0.010	2.0 $\pm$ 0.0	0.799 $\pm$ 0.002	22.7 $\pm$ 2.1
InfoDiffusion ( $\lambda=0.1, \zeta=1$ )	0.109 $\pm$ 0.001	Yes	0.192 $\pm$ 0.004	2.8 $\pm$ 0.4	0.848 $\pm$ 0.001	23.8 $\pm$ 1.6
w/ Learned Latent	–	Yes	–	–	–	21.2 $\pm$ 2.4
InfoDiffusion ( $\lambda=0.01, \zeta=1$ )	0.342 $\pm$ 0.002	Yes	0.299 $\pm$ 0.006	3.0 $\pm$ 0.0	0.836 $\pm$ 0.002	23.6 $\pm$ 1.3
w/ Learned Latent	–	Yes	–	–	–	22.3 $\pm$ 1.2

InfoDiffusion achieves the highest disentanglement scores (DCI of $0.342$ and TAD of $0.299$ ) and linear probe latent attribute classification performance ( $0.848$ AUROC) while maintaining low FID scores comparable to DiffAE, outperforming VAE-based autoencoders in visual generation quality.

Knowl 9 — Latent Representation Quality and Image Generation FID on Benchmarks

data/table

Representation quality (measured by 5-fold cross-validated logistic regression accuracy on latent $z$ ) and generation fidelity (FID on 10,000 samples) were evaluated on FashionMNIST, CIFAR10, and FFHQ.

	FashionMNIST		CIFAR10		FFHQ
Model	Latent Qual. $\uparrow$	FID $\downarrow$	Latent Qual. $\uparrow$	FID $\downarrow$	Latent Qual. $\uparrow$	FID $\downarrow$
AE	0.819 $\pm$ 0.003	62.9 $\pm$ 2.1	0.336 $\pm$ 0.005	169.4 $\pm$ 2.4	0.615 $\pm$ 0.002	92.3 $\pm$ 2.7
VAE	0.796 $\pm$ 0.002	63.4 $\pm$ 1.6	0.342 $\pm$ 0.004	177.2 $\pm$ 3.2	0.622 $\pm$ 0.002	95.4 $\pm$ 2.4
$\beta$ -VAE	0.779 $\pm$ 0.004	66.9 $\pm$ 1.8	0.253 $\pm$ 0.003	183.3 $\pm$ 3.1	0.588 $\pm$ 0.002	99.7 $\pm$ 3.4
InfoVAE	0.807 $\pm$ 0.003	55.0 $\pm$ 1.7	0.357 $\pm$ 0.005	160.7 $\pm$ 2.5	0.613 $\pm$ 0.002	86.9 $\pm$ 2.2
DiffAE	0.835 $\pm$ 0.002	8.2 $\pm$ 0.3	0.395 $\pm$ 0.006	32.1 $\pm$ 1.1	0.608 $\pm$ 0.001	31.6 $\pm$ 1.2
InfoDiffusion ( $\lambda=0.1, \zeta=1$ )	0.839 $\pm$ 0.003	8.5 $\pm$ 0.3	0.412 $\pm$ 0.003	31.7 $\pm$ 1.2	0.609 $\pm$ 0.002	31.2 $\pm$ 1.6
w/ Learned Latent	–	7.4 $\pm$ 0.2	–	31.5 $\pm$ 1.8	–	30.9 $\pm$ 2.5

InfoDiffusion achieves higher latent classification accuracy across datasets ( $0.839$ on FashionMNIST, $0.412$ on CIFAR10) compared to all VAE and diffusion baselines, while matching the generation performance of diffusion models (FID $\approx 7.4$ on FashionMNIST, $31.5$ on CIFAR10, and $30.9$ on FFHQ).

Knowl 10 — Representation Learning Comparison to Self-Supervised and Contrastive Methods

data/table

On the CelebA dataset, representations learned by InfoDiffusion (32 dimensions) were compared against self-supervised contrastive baselines (SimCLR, MoCo-v2, DINO) and diffusion autoencoding baselines (PDAE) in terms of Total Attribute Disentanglement (TAD), captured attributes (Attrs), and average AUROC for attribute classification.

Model	Generative	Latent Dim	TAD $\uparrow$	Attrs $\uparrow$	Latent Qual. (AUROC) $\uparrow$
SimCLR	No	2048	0.192 $\pm$ 0.015	–	0.812 $\pm$ 0.003
MoCo-v2	No	2048	0.279 $\pm$ 0.025	–	0.846 $\pm$ 0.001
DINO	No	384	0.000 $\pm$ 0.000	–	0.592 $\pm$ 0.003
SimCLR (fine-tuned)	No	32	0.062 $\pm$ 0.005	2.6 $\pm$ 0.5	0.757 $\pm$ 0.002
PDAE	Yes	32	0.009 $\pm$ 0.001	1.0 $\pm$ 0.0	0.767 $\pm$ 0.003
InfoDiffusion	Yes	32	0.299 $\pm$ 0.006	3.0 $\pm$ 0.0	0.836 $\pm$ 0.002

When restricted to low-dimensional representations ( $32$ dimensions), InfoDiffusion outperforms both contrastive learning (SimCLR: TAD $0.062$ , AUROC $0.757$ ) and pre-trained diffusion representation methods (PDAE: TAD $0.009$ , AUROC $0.767$ ), while also possessing generative sampling capabilities.

Coverage note — None was omitted; all key theoretical formulations, objectives, optimization bounds, architectural modules, sampling algorithms, and benchmark experimental results were captured.

References

1.Burgess, C. and Kim, H. 3d shapes dataset. https://github.com/deepmind/3dshapes-dataset/, 2018.
2.Caron, M., Touvron, H., Misra, I., Jégou, H., Mairal, J., Bojanowski, P., and Joulin, A. Emerging properties in self-supervised vision transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 9650–9660, 2021.
3.Chen, T., Kornblith, S., Norouzi, M., and Hinton, G. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pp. 1597–1607. PMLR, 2020a.
4.Chen, X., Duan, Y., Houthooft, R., Schulman, J., Sutskever, I., and Abbeel, P. Infogan: Interpretable representation learning by information maximizing generative adversarial nets. Advances in neural information processing systems, 29, 2016.
5.Chen, X., Fan, H., Girshick, R., and He, K. Improved baselines with momentum contrastive learning. arXiv preprint arXiv:2003.04297, 2020b.
6.Dhariwal, P. and Nichol, A. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems, 34:8780–8794, 2021.
7.Eastwood, C. and Williams, C. K. A framework for the quantitative evaluation of disentangled representations. In International Conference on Learning Representations, 2018.
8.Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. Generative adversarial networks. Communications of the ACM, 63(11):139–144, 2020.
9.Gretton, A., Borgwardt, K. M., Rasch, M. J., Schölkopf, B., and Smola, A. A kernel two-sample test. The Journal of Machine Learning Research, 13(1):723–773, 2012.
10.Higgins, I., Matthey, L., Pal, A., Burgess, C., Glorot, X., Botvinick, M., Mohamed, S., and Lerchner, A. beta-vae: Learning basic visual concepts with a constrained variational framework. In International conference on learning representations, 2017.
11.Ho, J., Jain, A., and Abbeel, P. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840–6851, 2020.
12.Jang, E., Gu, S., and Poole, B. Categorical reparameterization with gumbel-softmax. arXiv preprint arXiv:1611.01144, 2016.
13.Jing, B., Corso, G., Chang, J., Barzilay, R., and Jaakkola, T. Torsional diffusion for molecular conformer generation. arXiv preprint arXiv:2206.01729, 2022.
14.Karras, T., Laine, S., and Aila, T. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 4401–4410, 2019.
15.Kingma, D. P. and Welling, M. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114, 2013.
16.Kong, Z., Ping, W., Huang, J., Zhao, K., and Catanzaro, B. Diffwave: A versatile diffusion model for audio synthesis. arXiv preprint arXiv:2009.09761, 2020.
17.Krizhevsky, A., Hinton, G., et al. Learning multiple layers of features from tiny images. 2009.
18.LeCun, Y. Phd thesis: Modeles connexionnistes de l’apprentissage (connectionist learning models). 1987.
19.Liu, Z., Luo, P., Wang, X., and Tang, X. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
20.Locatello, F., Bauer, S., Lucic, M., Raetsch, G., Gelly, S., Schölkopf, B., and Bachem, O. Challenging common assumptions in the unsupervised learning of disentangled representations. In international conference on machine learning, pp. 4114–4124. PMLR, 2019.
21.Makhzani, A., Shlens, J., Jaitly, N., Goodfellow, I., and Frey, B. Adversarial autoencoders. arXiv preprint arXiv:1511.05644, 2015.
22.Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Kopf, A., Yang, E., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S. Pytorch: An imperative style, high-performance deep learning library. In Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Buc, F., Fox, E., and Garnett, R. (eds.), Advances in Neural Information Processing Systems 32, pp. 8024–8035. Curran Associates, Inc., 2019.
23.Pedregosa, F., Varoquaux, G., Gramfort, A., Michel, V., Thirion, B., Grisel, O., Blondel, M., Prettenhofer, P., Weiss, R., Dubourg, V., Vanderplas, J., Passos, A., Cournapeau, D., Brucher, M., Perrot, M., and Duchesnay, E. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research, 12:2825–2830, 2011.
24.Preechakul, K., Chatthee, N., Wizadwongsa, S., and Suwajanakorn, S. Diffusion autoencoders: Toward a meaningful and decodable representation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10619–10629, 2022.
25.Ramesh, A., Dhariwal, P., Nichol, A., Chu, C., and Chen, M. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125, 2022.
26.Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models, 2021.
27.Ronneberger, O., Fischer, P., and Brox, T. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention–MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18, pp. 234–241. Springer, 2015.
28.Saharia, C., Chan, W., Saxena, S., Li, L., Whang, J., Denton, E., Ghasemipour, S. K. S., Ayan, B. K., Mahdavi, S. S., Lopes, R. G., et al. Photorealistic text-to-image diffusion models with deep language understanding. arXiv preprint arXiv:2205.11487, 2022.
29.Xiao, H., Rasul, K., and Vollgraf, R. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv:1708.07747, 2017.
30.Xu, M., Yu, L., Song, Y., Shi, C., Ermon, S., and Tang, J. Geodiff: A geometric diffusion model for molecular conformation generation. arXiv preprint arXiv:2203.02923, 2022.
31.Yang, L., Zhang, Z., Song, Y., Hong, S., Xu, R., Zhao, Y., Shao, Y., Zhang, W., Cui, B., and Yang, M.-H. Diffusion models: A comprehensive survey of methods and applications. arXiv preprint arXiv:2209.00796, 2022.
32.Yeats, E., Liu, F., Womble, D., and Li, H. Nashae: Disentangling representations through adversarial covariance minimization. In Computer Vision–ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23–27, 2022, Proceedings, Part XXVII, pp. 36–51. Springer, 2022.
33.Zhang, Z., Zhao, Z., and Lin, Z. Unsupervised representation learning from pre-trained diffusion probabilistic models. Advances in Neural Information Processing Systems, 35: 22117–22130, 2022.
34.Zhao, S., Song, J., and Ermon, S. Infovae: Information maximizing variational autoencoders. arXiv preprint arXiv:1706.02262, 2017.

Citation

MLA

Wang, Y., et al. “InfoDiffusion: Representation Learning Using Information Maximizing Diffusion Models”. International Conference on Machine Learning, vol. 202, 2023, pp. 36336–54, https://proceedings.mlr.press/v202/wang23ah.html.

APA

Wang, Y., Schiff, Y., Gokaslan, A., Pan, W., Wang, F., Sa, C. D., & Kuleshov, V. (2023). InfoDiffusion: Representation Learning Using Information Maximizing Diffusion Models. International Conference on Machine Learning, 202, 36336–36354. https://proceedings.mlr.press/v202/wang23ah.html

Chicago

Wang, Y., Y. Schiff, A. Gokaslan, et al. 2023. “InfoDiffusion: Representation Learning Using Information Maximizing Diffusion Models”. International Conference on Machine Learning 202: 36336–54. https://proceedings.mlr.press/v202/wang23ah.html.

Harvard

Wang, Y. et al. (2023) “InfoDiffusion: Representation Learning Using Information Maximizing Diffusion Models”, International Conference on Machine Learning. PMLR, pp. 36336–36354. Available at: https://proceedings.mlr.press/v202/wang23ah.html.

Vancouver

1. Wang Y, Schiff Y, Gokaslan A, Pan W, Wang F, Sa CD, Kuleshov V (2023) InfoDiffusion: Representation Learning Using Information Maximizing Diffusion Models. In: International Conference on Machine Learning. PMLR, pp 36336–36354

BibTeX

@InProceedings{pmlr-v202-wang23ah,
  title = 	 {{I}nfo{D}iffusion: Representation Learning Using Information Maximizing Diffusion Models},
  author =       {Wang, Yingheng and Schiff, Yair and Gokaslan, Aaron and Pan, Weishen and Wang, Fei and De Sa, Christopher and Kuleshov, Volodymyr},
  booktitle = 	 {Proceedings of the 40th International Conference on Machine Learning},
  pages = 	 {36336--36354},
  year = 	 {2023},
  editor = 	 {Krause, Andreas and Brunskill, Emma and Cho, Kyunghyun and Engelhardt, Barbara and Sabato, Sivan and Scarlett, Jonathan},
  volume = 	 {202},
  series = 	 {Proceedings of Machine Learning Research},
  month = 	 {23--29 Jul},
  publisher =    {PMLR},
  pdf = 	 {https://proceedings.mlr.press/v202/wang23ah/wang23ah.pdf},
  url = 	 {https://proceedings.mlr.press/v202/wang23ah.html},
  abstract = 	 {While diffusion models excel at generating high-quality samples, their latent variables typically lack semantic meaning and are not suitable for representation learning. Here, we propose InfoDiffusion, an algorithm that augments diffusion models with low-dimensional latent variables that capture high-level factors of variation in the data. InfoDiffusion relies on a learning objective regularized with the mutual information between observed and hidden variables, which improves latent space quality and prevents the latents from being ignored by expressive diffusion-based decoders. Empirically, we find that InfoDiffusion learns disentangled and human-interpretable latent representations that are competitive with state-of-the-art generative and contrastive methods, while retaining the high sample quality of diffusion models. Our method enables manipulating the attributes of generated images and has the potential to assist tasks that require exploring a learned latent space to generate quality samples, e.g., generative design.}
}

Metadata:DOI registry

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: https://creativecommons.org/licenses/by/4.0/

InfoDiffusion: Representation Learning Using Information Maximizing Diffusion Models

Abstract

Topics

Keywords

Table of Contents

Knowls

Knowl 1 — Tractable Objective for Information-Maximizing Diffusion Models

Knowl 2 — Global Optimality of the InfoDiffusion Objective

Knowl 3 — Auxiliary-Variable Diffusion Model Generative Formulation

Knowl 4 — Auxiliary Latent Conditioning via Adaptive Group Normalization

Knowl 5 — Two-Phase Prior Sampling and Latent Space Inversion

Knowl 6 — Discrete Representation Learning with Relaxed-Bernoulli Distributions

Knowl 7 — Unification of Autoencoders, VAEs, and Diffusion Models under InfoDiffusion

Knowl 8 — Disentanglement and Representation Quality on 3DShapes and CelebA

Knowl 9 — Latent Representation Quality and Image Generation FID on Benchmarks

Knowl 10 — Representation Learning Comparison to Self-Supervised and Contrastive Methods

References

Citation

BibTeX

Access the Paper