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
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
✔
✔
✔
✖
✖
β -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). β -VAE with βVAE>1 penalizes rather than maximizes mutual information; AAEs support categorical latents through adversarial matching of a categorical prior ( Makhzani et al., 2015 ) .
Figure 3: z captures high-level semantic detail. Varying xT∼N(0,I) (across the columns in each row) changes lower level detail in the image. Red box indicates original image.
Figure 4: Finding disentangled dimensions in InfoDiffusion’s auxiliary latent variable z. Images are produced through a linear traversal along a particular dimension, spanning values from -1.5 to 1.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. ↑
FID ↓
Latent Qual. ↑
FID ↓
Latent Qual. ↑
FID ↓
AE
0.819±0.003
62.9±2.1
0.336±0.005
169.4±2.4
0.532±0.002
316.7±0.4
VAE
0.796±0.002
63.4±1.6
0.342±0.004
177.2±3.2
0.488±0.002
307.7±0.1
beta-VAE
0.779±0.004
66.9±1.8
0.253±0.003
183.3±3.1
0.510±0.001
222.6±0.2
InfoVAE
0.807±0.003
55.0±1.7
0.357±0.005
160.7±2.5
0.545±0.001
195.2±0.5
DiffAE
0.822±0.006
5.6±0.1
0.442±0.004
34.1±0.1
0.565±0.001
12.1±0.1
Table 2: Latent quality, as measured by classification accuracies for logistic regression classifiers trained on the auxiliary latent vector z, and FID. We report mean ± 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 z∼p(z) ( Section D.1 ). Latent quality and FID are measured with dz=32 and dz=256 models, respectively ( Table 7 ). For InfoDiffusion, ± is the standard deviation over three training seeds, and w is the guidance weight on 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 ↑
Image Qual.
TAD ↑
Attrs ↑
Latent Qual. ↑
FID ↓
AE
0.219±0.001
✖
0.042±0.004
1.0±0.0
0.759±0.003
90.4±1.8
VAE
0.276±0.001
✖
0.000±0.000
0.0±0.0
0.770±0.002
94.3±2.8
beta-VAE
0.281±0.001
✖
0.088±0.051
1.6±0.8
0.699±0.001
99.8±2.4
InfoVAE
0.134±0.001
✖
0.000±0.000
0.0±0.0
0.757±0.003
77.8±1.6
DiffAE
0.177±0.001
✔
0.114±0.033
2.8±0.4
0.817±0.001
8.4±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 z . We report means ± one standard deviation for quantitative metrics. Darkly shaded cells indicate the best while lightly shaded cells indicate the second best. For InfoDiffusion ( λ=0.1 , ζ=0.9 ), ± is the standard deviation over three training seeds; FID is measured with dz=256 models, w is the guidance weight on z , and image quality on 3DShapes is judged from 16 samples drawn with z∼p(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 ↑
Latent Qual. ↑
SIMCLR † 0 0 footnotetext: The weights are taken from the respective PyTorch repositories of these models.
✖
2048
0.192±0.015
0.812±0.003
MOCO-v2 †
✖
2048
0.279±0.025
0.846±0.001
DINO †
✖
384
0.000±0.000
0.592±0.003
InfoDiffusion
✔
32
0.439±0.134
0.844±0.005
Table 4: Representation learning comparison to contrastive methods (InfoDiffusion: λ=0.1 , ζ=0.9 ; mean ± 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 ± 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 ↑
ATTRS ↑
Latent Qual. ↑
SIMCLR ∗
✖
\pagecolorgray!200.062±0.005
2.6±0.5
0.757±0.002
PDAE ‡
✔
0.009±0.001
1.0±0.0
0.767±0.003
InfoDiff.
✔
0.439±0.134
5.7±0.6
0.844±0.005
Table 5: Representation learning comparison to SIMCLR and PDAE with 32-dimensional latents (InfoDiffusion: λ=0.1 , ζ=0.9 ; mean ± 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 z . We report means ± 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 z
Channels
Pool
Crop
t
Steps (epochs)
Eval.
Gen.
dec. / enc.
3DShapes
3×64×64
10
–
40 / 48
flatten
✗
uniform
16,875 (5)
FashionMNIST
1×32×32
32
256
40 / 48
flatten
✓
logit-normal
23,400 (50)
CIFAR10
3×32×32
32
256
80 / 96
mean
✓
logit-normal
19,500 (50)
FFHQ
3×64×64
32
256
80 / 96
flatten
✓
logit-normal
24,600 (50)
CelebA
3×64×64
32
256
80 / 96
mean / flatten
✓
logit-normal
63,600 (50)
Appendix
Table 7: Dataset-specific settings for InfoDiffusion. The two dimensionalities of 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.
Figure 6: The encoder of InfoDiffusion: the downsampling half of a U-Net (residual blocks without time conditioning; self-attention at 16×16 and 8×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 x0 . For ζ<1 , the head outputs the mean and log-variance of a Gaussian qϕ(z∣x0) ; for discrete latents, its output is quantized with FSQ.
Figure 7: The U-Net decoder of InfoDiffusion, conditioned on the time t and the auxiliary variable z . The embeddings of t and z are passed to every residual block. With latent dropout, the embedding of z is replaced by a learned null embedding ∅ . C is the base width and k the channel multiplier of a level.
Figure 8: Left : the auxiliary residual block. The normalized features are modulated first by the time embedding (AGN) and then by the embedding of 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. ↑
FID ↓
Latent Qual. ↑
FID ↓
Latent Qual. ↑
FID ↓
Latent Qual. ↑
FID ↓
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 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 z. Both variants use λ=0.1 , ζ=0.9 and a single training seed; latent quality is measured with dz=32 models and FID with dz=256 models, using the learned latent prior without guidance. Darkly shaded cells indicate best results.
CelebA
TAD ↑
ATTRS ↑
Ours (with bott. only)
0.654
6.0
Ours
0.555
6.0
Appendix
Table 9: Comparison between two modeling choices: conditioning on 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 ( λ=0.1 , ζ=0.9 , dz=32 ). Darkly shaded cells indicate best result.
Latent
Seeds
Latent Qual. ↑
TAD ↑
Attrs ↑
z -only MSE ↓
Continuous, (λ,ζ)=(0.1,0.9)
3
0.844±0.005
0.439±0.134
5.7±0.6
0.221
FSQ ( L=5 )
3
0.849±0.003
0.370±0.240
5.6±0.9
0.231
Binary sign(h) , MMD 0.01
3
0.712±0.014
0.030±0.026
0.5±0.4
0.337
Binary sign(h) , no regularizer
1
0.687
0.000
0.0
0.365
Relaxed Bernoulli, MMD 0.01
3
0.574±0.011
0.000
0.0
0.588
Relaxed Bernoulli, KL 0.1
3
0.613±0.008
0.000
0.0
0.437
Appendix
Table 10: Discrete latents on CelebA ( dz=32 ). Latent quality (average AUROC), TAD, and captured attributes follow the protocol of Table 3 ; ‘ z -only MSE’ is the mean squared error (images in [−1,1] ) of decoding 1,000 test images from their latents and a random xT . Means ± standard deviations are over three seeds; other rows use one seed.
FashionMNIST
CIFAR10
FFHQ
Latent Qual. ↑
FID ↓
Latent Qual. ↑
FID ↓
Latent Qual. ↑
FID ↓
AE
0.819±0.003
62.9±2.1
0.336±0.005
169.4±2.4
0.532±0.002
316.7±0.4
VAE
0.796±0.002
63.4±1.6
0.342±0.004
177.2±3.2
0.488±0.002
307.7±0.1
beta-VAE
0.779±0.004
66.9±1.8
0.253±0.003
183.3±3.1
0.510±0.001
222.6±0.2
InfoVAE
0.807±0.003
55.0±1.7
0.357±0.005
160.7±2.5
0.545±0.001
195.2±0.5
DiffAE
0.822±0.006
5.6±0.1
0.442±0.004
34.1±0.1
0.565±0.001
12.1±0.1
Appendix
Table 11: Latent quality, as measured by classification accuracies for logistic regression classifiers trained on the auxiliary latent vector z, and FID. We report mean ± 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 ↑
Image Qual.
TAD ↑
Attrs ↑
Latent Qual. ↑
FID ↓
AE
0.219±0.001
✖
0.042±0.004
1.0±0.0
0.759±0.003
90.4±1.8
VAE
0.276±0.001
✖
0.000±0.000
0.0±0.0
0.770±0.002
94.3±2.8
beta-VAE
0.281±0.001
✖
0.088±0.051
1.6±0.8
0.699±0.001
99.8±2.4
InfoVAE
0.134±0.001
✖
0.000±0.000
0.0±0.0
0.757±0.003
77.8±1.6
DiffAE
0.177±0.001
✔
0.114±0.033
2.8±0.4
0.817±0.001
8.4±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 z∼p(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 z . We report means ± 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).
Figure 10: Uncurated samples of InfoDiffusion ( λ=0.1 , ζ=0.9 , dz=256 ) with z drawn from the learned latent prior and guidance w=2 .
Figure 11: Uncurated samples of InfoDiffusion ( λ=0.1 , ζ=0.9 , dz=256 ) with z∼p(z) and guidance w=2 .
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-13× faster, indicating that jointly learning the latent space is a key step toward making latent diffusion competitive for text generation.
Viacheslav Meshchaninov, Alexander Shabalin, Egor Chimbulatov +4
Constructor University, Bremen, Germany · HSE University, Moscow, Russia · Applied AI Institute, Moscow, Russia +2
Joint image-feature generative modeling has recently emerged as an effective strategy for improving diffusion training by coupling low-level VAE latents with high-level semantic features extracted from pre-trained visual encoders. However, existing approaches rely on a fixed representation space, constructed independently of the generative objective and kept unchanged during training. We argue that the representation space guiding diffusion should itself adapt to the generative task. To this end, we propose Coevolving Representation Diffusion (CoReDi), a framework in which the semantic representation space evolves during training by learning a lightweight linear projection jointly with the diffusion model. While naively optimizing this projection leads to degenerate solutions, we show that stable coevolution can be achieved through a combination of stop-gradient targets, normalization, and targeted regularization that prevents feature collapse. This formulation enables the semantic space to progressively specialize to the needs of image synthesis, improving its complementarity with image latents. We apply CoReDi to both VAE latent diffusion and pixel-space diffusion, demonstrating that adaptive semantic representations improve generative modeling across both settings. Experiments show that CoReDi achieves faster convergence and higher sample quality compared to joint diffusion models operating in fixed representation spaces.
Diffusion language models intrinsically fail to capture correlations between decoded tokens, which leads to a harsh trade-off between sampling quality and throughput. To solve this issue, we propose DiLaDiff, a variant of masked diffusion language models with three components: (1) a continuous latent space with semantic capabilities, learned by an auto-encoder fine-tuned from an existing masked diffusion language model; (2) a latent diffusion model learning the prior over the encoder distribution; (3) a consistency model distilling the learned prior into a few-step latent generative model. We show that, even without distillation, our latent-guided diffusion model outperforms the masked diffusion baseline while significantly accelerating inference. Consistency distillation further lowers the computational overhead of continuous diffusion, such that the latent is generated in negligible time compared to discrete decoding.
Jean-Marie Lemercier, Tomas Geffner, Karsten Kreis +3