InfoDiffusion: Representation Learning Using Information Maximizing Diffusion Models

Yingheng Wang; Yair Schiff; Aaron Gokaslan; Weishen Pan; Fei Wang; Christopher De Sa; Volodymyr Kuleshov

cs.LGJun 14, 2023

InfoDiffusion: Representation Learning Using Information Maximizing Diffusion Models

Authors: Yingheng Wang, Yair Schiff, Aaron Gokaslan, Weishen Pan, Fei Wang, Christopher De Sa, Volodymyr Kuleshov

Organizations: Department of Computer Science, Cornell University, Ithaca, NY, USA · Department of Computer Science, Cornell Tech, New York City, NY, USA · Department of Population Health Sciences, Weill Cornell Medicine, New York City, NY, USA

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 is designed to improve latent space quality and to prevent 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.

Figures & tables

Figure 1: InfoDiffusion produces semantically meaningful latent space for a diffusion model. ( Top ) Smooth latent space. ( Bottom ) Disentangled, human-interpretable factors of variation.

Figure 2: Flow chart demonstrating auxiliary-variable diffusion model with mutual information and prior regularization.

	Semantic latents	Discrete latents	Custom prior	Max MI	HQ samples
AE	✖	✖	✖	✖	✖
VAE	✔	✔	✔	✖	✖
$\beta$ -VAE	✔	✔	✔	✖	✖
AAE	✔	✔	✔	✔	✖
InfoVAE	✔	✔	✔	✔	✖
DDPM	✖	✖	✖	✖	✔

Table 1: Comparison of InfoDiffusion model to other auto-encoder ( top ) and diffusion ( bottom ) frameworks in terms of enabling semantic latents, discrete latents, custom priors, mutual information maximization (Max MI), and high-quality sample generation (HQ samples).

\beta

-VAE with

\beta_{\text{VAE}}>1

penalizes rather than maximizes mutual information; AAEs support categorical latents through adversarial matching of a categorical prior ( Makhzani et al., 2015 ) .

Panel of Figure 3 — Figure 3: $\mathbf{z}$ captures high-level semantic detail. Varying $\mathbf{x}_{T}\sim$ $\mathcal{N}(\mathbf{0},\mathbf{I})$ (across the columns in each row) changes lower level detail in the image. Red box indicates original image.

Panel of Figure 4 — Figure 4: Finding disentangled dimensions in InfoDiffusion’s auxiliary latent variable $\mathbf{z}.$ Images are produced through a linear traversal along a particular dimension, spanning values from -1.5 to 1.5.

Panel of Figure 5 — Figure 5: Latent space interpolation. InfoDiffusion has a smooth latent space and maintains high image generation quality. Reconstructions of the two original images are at the left and right ends of each row and are marked by red boxes.

	FashionMNIST		CIFAR10		FFHQ
	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.532\pm 0.002$	$316.7\pm 0.4$
VAE	$0.796\pm 0.002$	$63.4\pm 1.6$	$0.342\pm 0.004$	$177.2\pm 3.2$	$0.488\pm 0.002$	$307.7\pm 0.1$
beta-VAE	$0.779\pm 0.004$	$66.9\pm 1.8$	$0.253\pm 0.003$	$183.3\pm 3.1$	$0.510\pm 0.001$	$222.6\pm 0.2$
InfoVAE	$0.807\pm 0.003$	$55.0\pm 1.7$	$0.357\pm 0.005$	$160.7\pm 2.5$	$0.545\pm 0.001$	$195.2\pm 0.5$
DiffAE	$0.822\pm 0.006$	$5.6\pm 0.1$	$0.442\pm 0.004$	$34.1\pm 0.1$	$0.565\pm 0.001$	$12.1\pm 0.1$

Table 2: Latent quality, as measured by classification accuracies for logistic regression classifiers trained on the auxiliary latent vector

\mathbf{z},

and FID. We report mean

\pm

one standard deviation. Darkly shaded cells indicate the best while lightly shaded cells indicate the second best. See Table 11 for the performance of varying hyperparameters. InfoDiffusion rows without “w/Learned Latent” are generated with

\mathbf{z}\sim p(\mathbf{z})

( Section D.1 ). Latent quality and FID are measured with

d_{z}=32

and

d_{z}=256

models, respectively ( Table 7 ). For InfoDiffusion,

\pm

is the standard deviation over three training seeds, and

w

is the guidance weight on

\mathbf{z}

( Section 4.4 ). DiffAE and the FFHQ entries of the other baselines are re-trained; the remaining baseline entries come from earlier runs ( Appendix C ).

	3DShapes		CelebA
	DCI $\uparrow$	Image Qual.	TAD $\uparrow$	Attrs $\uparrow$	Latent Qual. $\uparrow$	FID $\downarrow$
AE	$0.219\pm 0.001$	✖	$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$	✖	$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}$	✖	$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$	✖	$0.000\pm 0.000$	$0.0\pm 0.0$	$0.757\pm 0.003$	$77.8\pm 1.6$
DiffAE	$0.177\pm 0.001$	✔	$0.114\pm 0.033$	$2.8\pm 0.4$	$0.817\pm 0.001$	$8.4\pm 0.1$

Table 3: Disentanglement and latent quality metrics and FID. For 3DShapes, we check the image quality manually and label the models which generate high-quality images with check marks (‘Image Qual.’). The visualization of the samples is shown in Figure 9 in Appendix I . For CelebA, ‘Attrs.’ counts the number of “captured” attributes when calculating the TAD score. ‘Latent Quality’ is measured as AUROC scores averaged across attributes for logistic regression classifiers trained on the auxiliary latent vector

\mathbf{z}

. We report means

\pm

one standard deviation for quantitative metrics. Darkly shaded cells indicate the best while lightly shaded cells indicate the second best. For InfoDiffusion (

\lambda=0.1

\zeta=0.9

\pm

is the standard deviation over three training seeds; FID is measured with

d_{z}=256

models,

w

is the guidance weight on

\mathbf{z}

, and image quality on 3DShapes is judged from 16 samples drawn with

\mathbf{z}\sim p(\mathbf{z})

(first row) or from the learned latent prior ( Appendix C ). DiffAE is re-trained; the other baseline entries come from earlier runs ( Appendix C ).

CelebA	Gen.	Dim.	TAD $\uparrow$	Latent Qual. $\uparrow$
SIMCLR † 0 0 footnotetext: The weights are taken from the respective PyTorch repositories of these models.	✖	2048	${0.192\pm 0.015}$	$0.812\pm 0.003$
MOCO-v2 †	✖	2048	${0.279\pm 0.025}$	${0.846\pm 0.001}$
DINO †	✖	384	${0.000\pm 0.000}$	$0.592\pm 0.003$
InfoDiffusion	✔	32	$0.439\pm 0.134$	$0.844\pm 0.005$

Table 4: Representation learning comparison to contrastive methods (InfoDiffusion:

\lambda=0.1

\zeta=0.9

; mean

\pm

standard deviation over three training seeds). ‘Gen.’ indicates whether the model has generative capabilities. ‘Dim.’ denotes the latent dimension. Disentanglement is measured by TAD. ‘Latent Quality’ is measured as AUROC scores averaged across CelebA attributes for logistic regression classifiers trained on latent representations. We report means

\pm

one standard deviation for quantitative metrics. Darkly shaded cells indicate the best while lightly shaded cells indicate the second best. † denotes that the weights are taken from the PyTorch repository for the method.

CelebA	Gen.	TAD $\uparrow$	ATTRS $\uparrow$	Latent Qual. $\uparrow$
SIMCLR ∗	✖	$\pagecolor{gray!20}{0.062\pm 0.005}$	$2.6\pm 0.5$	$0.757\pm 0.002$
PDAE ‡	✔	${0.009\pm 0.001}$	$1.0\pm 0.0$	$0.767\pm 0.003$
InfoDiff.	✔	$0.439\pm 0.134$	$5.7\pm 0.6$	$0.844\pm 0.005$

Table 5: Representation learning comparison to SIMCLR and PDAE with 32-dimensional latents (InfoDiffusion:

\lambda=0.1

\zeta=0.9

; mean

\pm

standard deviation over three training seeds). ‘Gen.’ indicates whether the model has generative capabilities. ‘Attrs.’ counts the number of “captured” attributes when calculating the TAD score. ‘Latent Quality’ is measured as AUROC scores averaged across attributes for logistic regression classifiers trained on

\mathbf{z}

. We report means

\pm

one standard deviation for quantitative metrics. Darkly shaded cells indicate the best while lightly shaded cells indicate the second best. ∗ denotes that the weights are taken from the PyTorch repository and fine-tuned with an added dense layer. ‡ denotes that the model is trained by us with the codebase provided by this baseline, in earlier runs ( Appendix C ).

Method	Non-Probabilistic	Probabilistic Extension	Regularized Extension
Auto-encoders	AE ( LeCun, 1987 )	VAE ( Kingma & Welling, 2013 )	InfoVAE ( Zhao et al., 2017 )
Diffusion models	DiffAE ( Preechakul et al., 2022 )	Variational Auxiliary-Variable Diffusion Sec. 4	InfoDiffusion Sec. 5

Table 6: Analogy between progress in the space of auto-encoders and similar progress for diffusion models.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

	Input size	Dim. of $\mathbf{z}$		Channels	Pool	Crop	$t$	Steps (epochs)
		Eval.	Gen.	dec. / enc.
3DShapes	$3\times 64\times 64$	10	–	40 / 48	flatten	✗	uniform	16,875 (5)
FashionMNIST	$1\times 32\times 32$	32	256	40 / 48	flatten	✓	logit-normal	23,400 (50)
CIFAR10	$3\times 32\times 32$	32	256	80 / 96	mean	✓	logit-normal	19,500 (50)
FFHQ	$3\times 64\times 64$	32	256	80 / 96	flatten	✓	logit-normal	24,600 (50)
CelebA	$3\times 64\times 64$	32	256	80 / 96	mean / flatten	✓	logit-normal	63,600 (50)

Appendix

Table 7: Dataset-specific settings for InfoDiffusion. The two dimensionalities of

\mathbf{z}

correspond to latent evaluation (‘Eval.’) and generation (‘Gen.’) models. Channels are the base widths of the decoder and the encoder. ‘Pool’ is the operation applied to the final encoder feature map before the linear head; ‘Crop’ indicates the random resized crops of the encoder input; ‘

t

’ is the distribution of the training time step.

Panel of Figure 6 — Figure 6: The encoder of InfoDiffusion: the downsampling half of a U-Net (residual blocks without time conditioning; self-attention at $16\times 16$ and $8\times 8$ ) followed by a linear head on either the spatially averaged or the flattened final feature map. During training, the input is a random resized crop of $\mathbf{x}_{0}$ . For $\zeta<1$ , the head outputs the mean and log-variance of a Gaussian $q_{\phi}(\mathbf{z}\mid\mathbf{x}_{0})$ ; for discrete latents, its output is quantized with FSQ.

Panel of Figure 7 — Figure 7: The U-Net decoder of InfoDiffusion, conditioned on the time $t$ and the auxiliary variable $\mathbf{z}$ . The embeddings of $t$ and $\mathbf{z}$ are passed to every residual block. With latent dropout, the embedding of $\mathbf{z}$ is replaced by a learned null embedding $\varnothing$ . $C$ is the base width and $k$ the channel multiplier of a level.

Panel of Figure 8 — Figure 8: Left : the auxiliary residual block. The normalized features are modulated first by the time embedding (AGN) and then by the embedding of $\mathbf{z}$ ; the second convolution is zero-initialized. Right : the MLP velocity network of the learned latent prior, with adaptive layer normalization (scale, shift, and gate predicted from the time embedding) in each of its six blocks.

	FashionMNIST		CIFAR10		FFHQ		CelebA
	Latent Qual. $\uparrow$	FID $\downarrow$	Latent Qual. $\uparrow$	FID $\downarrow$	Latent Qual. $\uparrow$	FID $\downarrow$	Latent Qual. $\uparrow$	FID $\downarrow$
Ours (with bott. only)	$0.816$	$8.0$	$0.355$	$26.6$	$0.625$	$12.3$	$0.817$	$7.6$
Ours	$0.852$	$7.7$	$0.457$	$25.7$	$0.629$	$12.5$	$0.844$	$7.7$

Appendix

Table 8: Comparison between two modeling choices: conditioning on

\mathbf{z}

at just the bottleneck layer of the UNet (denoted as ‘Ours w/ bott. only’) versus at all layers. ‘Latent quality’ is measured as classification accuracy/AUROC for logistic regression classifiers trained on the auxiliary latent vector

\mathbf{z}.

Both variants use

\lambda=0.1

\zeta=0.9

and a single training seed; latent quality is measured with

d_{z}=32

models and FID with

d_{z}=256

models, using the learned latent prior without guidance. Darkly shaded cells indicate best results.

CelebA	TAD $\uparrow$	ATTRS $\uparrow$
Ours (with bott. only)	$0.654$	$6.0$
Ours	$0.555$	$6.0$

Appendix

Table 9: Comparison between two modeling choices: conditioning on

\mathbf{z}

at just the bottleneck layer of the UNet (denoted as ‘Ours w/ bott. only’) versus at all layers. ‘Attrs.’ counts the number of “captured” attributes when calculating the TAD score. Single training seed (

\lambda=0.1

\zeta=0.9

d_{z}=32

). Darkly shaded cells indicate best result.

Latent	Seeds	Latent Qual. $\uparrow$	TAD $\uparrow$	Attrs $\uparrow$	$\mathbf{z}$ -only MSE $\downarrow$
Continuous, $(\lambda,\zeta)=(0.1,0.9)$	3	$0.844\pm 0.005$	$0.439\pm 0.134$	$5.7\pm 0.6$	$0.221$
FSQ ( $L=5$ )	3	$0.849\pm 0.003$	$0.370\pm 0.240$	$5.6\pm 0.9$	$0.231$
Binary $\mathrm{sign}(h)$ , MMD $0.01$	3	$0.712\pm 0.014$	$0.030\pm 0.026$	$0.5\pm 0.4$	$0.337$
Binary $\mathrm{sign}(h)$ , no regularizer	1	$0.687$	$0.000$	$0.0$	$0.365$
Relaxed Bernoulli, MMD $0.01$	3	$0.574\pm 0.011$	$0.000$	$0.0$	$0.588$
Relaxed Bernoulli, KL $0.1$	3	$0.613\pm 0.008$	$0.000$	$0.0$	$0.437$

Appendix

Table 10: Discrete latents on CelebA (

d_{z}=32

). Latent quality (average AUROC), TAD, and captured attributes follow the protocol of Table 3 ; ‘

\mathbf{z}

-only MSE’ is the mean squared error (images in

[-1,1]

) of decoding

1{,}000

test images from their latents and a random

\mathbf{x}_{T}

. Means

\pm

standard deviations are over three seeds; other rows use one seed.

	FashionMNIST		CIFAR10		FFHQ
	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.532\pm 0.002$	$316.7\pm 0.4$
VAE	$0.796\pm 0.002$	$63.4\pm 1.6$	$0.342\pm 0.004$	$177.2\pm 3.2$	$0.488\pm 0.002$	$307.7\pm 0.1$
beta-VAE	$0.779\pm 0.004$	$66.9\pm 1.8$	$0.253\pm 0.003$	$183.3\pm 3.1$	$0.510\pm 0.001$	$222.6\pm 0.2$
InfoVAE	$0.807\pm 0.003$	$55.0\pm 1.7$	$0.357\pm 0.005$	$160.7\pm 2.5$	$0.545\pm 0.001$	$195.2\pm 0.5$
DiffAE	$0.822\pm 0.006$	$5.6\pm 0.1$	$0.442\pm 0.004$	$34.1\pm 0.1$	$0.565\pm 0.001$	$12.1\pm 0.1$

Appendix

Table 11: Latent quality, as measured by classification accuracies for logistic regression classifiers trained on the auxiliary latent vector

\mathbf{z},

and FID. We report mean

\pm

one standard deviation ( Appendix H ). Darkly shaded cells indicate the best while lightly shaded cells indicate the second best. InfoDiffusion FIDs use

w=1

	3DShapes		CelebA
	DCI $\uparrow$	Image Qual.	TAD $\uparrow$	Attrs $\uparrow$	Latent Qual. $\uparrow$	FID $\downarrow$
AE	$0.219\pm 0.001$	✖	$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$	✖	$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$	✖	$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$	✖	$0.000\pm 0.000$	$0.0\pm 0.0$	$0.757\pm 0.003$	$77.8\pm 1.6$
DiffAE	$0.177\pm 0.001$	✔	$0.114\pm 0.033$	$2.8\pm 0.4$	$0.817\pm 0.001$	$8.4\pm 0.1$

Appendix

Table 12: Disentanglement and latent quality metrics and FID. For 3DShapes, we check the image quality manually and label the models which generate high-quality images with check marks (‘Image Qual.’). The visualization of the samples is shown in Figure 9 in the Appendix I ; for InfoDiffusion, image quality is judged from 16 samples drawn with

\mathbf{z}\sim p(\mathbf{z})

(first row of each setting) or from the learned latent prior (second row). For CelebA, ‘Attrs.’ counts the number of “captured” attributes when calculating the TAD score. ‘Latent Quality’ is measured as AUROC scores averaged across attributes for logistic regression classifiers trained on the auxiliary latent vector

\mathbf{z}

. We report means

\pm

one standard deviation for quantitative metrics ( Appendix H ). Darkly shaded cells indicate the best while lightly shaded cells indicate the second best. InfoDiffusion FIDs use

w=1


(a) VAE

(b) InfoVAE

(c) DiffAE

Appendix

Figure 9: Visualization of image samples from unconditional generation by VAE-based models (a-b) and diffusion models (c-d).

Panel of Figure 10 — Figure 10: Uncurated samples of InfoDiffusion ( $\lambda=0.1$ , $\zeta=0.9$ , $d_{z}=256$ ) with $\mathbf{z}$ drawn from the learned latent prior and guidance $w=2$ .

Panel of Figure 11 — Figure 11: Uncurated samples of InfoDiffusion ( $\lambda=0.1$ , $\zeta=0.9$ , $d_{z}=256$ ) with $\mathbf{z}\sim p(\mathbf{z})$ and guidance $w=2$ .

Diffusion Models

Diffusion models learn to reverse gradual data corruption, enabling flexible generation, editing, forecasting, and inference.

See all 1,169 papers

How to Train Your Latent Diffusion Language Model Jointly With the Latent Space
May 8, 2026Viacheslav Meshchaninov, Alexander Shabalin, Egor Chimbulatov +4Non-Autoregressive Text Generation Latent Diffusion Models
Constructor University, Bremen, Germany · HSE University, Moscow, Russia · Applied AI Institute, Moscow, Russia +2
Coevolving Representations in Joint Image-Feature Diffusion
Apr 19, 2026Theodoros Kouzelis, Spyros Gidaris, Nikos KomodakisRepresentation Learning Latent Diffusion Models
Archimedes, Athena RC · National Technical University of Athens · University of Crete +1
DiLaDiff: Distilled Latent-Augmented Diffusion for Language Modeling
May 22, 2026Jean-Marie Lemercier, Tomas Geffner, Karsten Kreis +3Diffusion Language Models Diffusion Language Model Distillation
NVIDIA

May 8, 2026J/K move · Enter open · S save

How to Train Your Latent Diffusion Language Model Jointly With the Latent Space

Viacheslav Meshchaninov, Alexander Shabalin, Egor Chimbulatov, Nikita Gushchin+3

Constructor University, Bremen, Germany · HSE University, Moscow, Russia · Applied AI Institute, Moscow, Russia+2

Non-Autoregressive Text Generation Latent Diffusion Models Diffusion Language Models Continuous Diffusion Language Models Diffusion Language Model Training

Open paper AskPDF arXiv

Latent diffusion models offer an attractive alternative to discrete diffusion for non-autoregressive text generation by operating on continuous text representations and denoising entire sequences in parallel. The major challenge in latent diffusion modeling is constructing a suitable latent space. In this work, we present the Latent Diffusion Language Model (LDLM), in which the latent encoder, diffusion model, and decoder are trained jointly. LDLM builds its latent space by reshaping the representations of a pre-trained language model with a trainable encoder, yielding latents that are easy to both denoise and decode into tokens. We show that naive joint training produces a low-quality diffusion model, and propose a simple training recipe consisting of an MSE decoder loss, diffusion-to-encoder warmup, adaptive timestep sampling, and decoder-input noise. Ablations show that each component substantially impacts generation performance. On OpenWebText and LM1B, LDLM achieves better generation performance than existing discrete and continuous diffusion language models while being

2{\text -}13\times

faster, indicating that jointly learning the latent space is a key step toward making latent diffusion competitive for text generation.

Abstract

Diffusion Models

Explore similar work

How to Train Your Latent Diffusion Language Model Jointly With the Latent Space

Coevolving Representations in Joint Image-Feature Diffusion

DiLaDiff: Distilled Latent-Augmented Diffusion for Language Modeling