Diffusion bridges vector quantized Variational AutoEncoders

1 Introduction

• VQ-VAE quantization uses a rough gradient approximation, while its prior is initially uniform and is later estimated with a separate autoregressive model.Joint training of the prior and VQ-VAE is described as challenging.Source: §1 · p. 1
• Autoregressive PixelCNN sampling proceeds sequentially over pixels or latent variables, so runtime depends on lattice size and can become large.The fixed ordering can accumulate prediction errors during inference.Source: §1 · p. 2
• The authors describe continuous-space diffusion as a way to learn complex discrete distributions with support of cardinality K^N, using an embedded space and an uninformative target distribution.They state that the framework also extends beyond VQ-VAE to continuous embeddings and sampling from discrete laws.Source: §1 · p. 2
• The proposed prior uses a continuous Markov-chain latent vector and diffusion bridges to produce samples matching the data, with discrete latent states defined as random functions of that chain.The framework is presented as a way to train the prior together with the encoder and decoder, bypassing separate autoregressive-prior training.Source: §1 · p. 2
• The paper presents a toy-dataset study and an efficiency comparison with the original VQ-VAE PixelCNN prior on miniImagenet.The passage specifies the comparison but does not report its outcome.Source: §1 · p. 2

2 Related Works

• Figure 1 separates the image encoding-and-reconstruction pathway from a vertical diffusion denoising pathway between latent states.The caption identifies the top pathway as resembling the original VQ-VAE.Source: Figure 1 · p. 3
• Continuous diffusion corrupts data through Markovian steps and learns iterative denoising, while sampling applies the learned process without a fixed pixel order.The passage notes that sampling time depends on the number of steps T.Source: §2 · pp. 2–3
• Multinomial diffusion learns to denoise categorical samples, but its transition matrix is reported not to scale well with large embedding vocabularies.Absorbing MASK states can distinguish corrupted dimensions from original ones.Source: §2 · p. 3
• The proposed approach uses continuous bridges because VQ-VAE’s discrete distribution depends on a continuous distribution, and the authors characterize it as simpler and more scalable than the discrete denoising alternatives.Source: §2 · p. 3
• Latent score-matching work can train encoder, decoder, and score matching end-to-end, whereas this method is designed for discrete VQ-VAE latent spaces.Source: §2 · p. 4

3.1 Model and loss function

• The framework introduces a diffusion-based generative VQ-VAE for efficient joint training of the prior and variational approximation.The input is modeled as depending on a hidden discrete state, with an encoder providing a variational approximation.Source: §3.1 · p. 4
• Continuous latent states form a Markov chain, while the discrete variables are random functions of those states and are not themselves Markovian.The model considers a sequence of discrete latent states with support in E^N.Source: §3.1 · pp. 4–5
• The ELBO combines a reconstruction term, a diffusion term, and an additional regularization term.The regularization term is described as connecting the diffusion and quantization parts through a gradient pathway and codebook regularization.Source: §3.1 · p. 6, §3.2 · p. 8

3.2 Application to Ornstein-Uhlenbeck processes

• The Ornstein–Uhlenbeck SDE adds noise to normalized inputs toward a target state, and its associated bridge supplies the denoising process.The process is parameterized by positive drift and noise parameters and Brownian motion; its bridge law is explicit for this diffusion.Source: §3.2 · p. 6
• At T = 0, the loss reduces to the standard VQ-VAE loss; for T > 0, the framework permits more flexible discrete-state parameterizations.The Gumbel-Softmax trick is given as an efficient differentiable parameterization, with time-dependent scaling parameters.Source: §3.2 · p. 7
• The regularization term can bring codebooks and continuous latent states together or apart depending on the choices of τ and τ_t.The final training procedure is described as end-to-end.Source: §3.2 · p. 8

4.1 Toy Experiment

• In the toy experiment, denoising brings samples near five codebooks; the Gumbel-Softmax trick also supports end-to-end training with randomly initialized codebooks.The toy setup uses five 2D points and eight codebook centers; codebooks can be trained alongside the bridge prior.Source: §4.1 · pp. 8–9

4.2 Image Synthesis

• Image-prior comparisons use pretrained, fixed encoders, decoders, and codebooks; end-to-end runs first pretrain the VQ-VAE, then train its components with diffusion.Because the VQ-VAE changes in the end-to-end setup, its NLL cannot be directly compared with PixelCNN or the previous diffusion model.Source: §4.2 · pp. 9–10

4.3 Quantitative results

• PixelCNN performs poorly on later pixels in latent-space KL, while the diffusion prior remains consistent across pixels.Figure 4 displays empirical KL distances between the true distribution and both priors; darker squares indicate lower values.Source: Figure 4 · p. 11, §4.3 · p. 11
• Sampling a batch of 32 images is bottlenecked by the 1,000 sequential diffusion steps, while latent-space NLL evaluation estimates masked-patch likelihoods by Monte Carlo.The stated sampling-time evaluation uses a GTX TITAN Xp GPU; the conditional likelihood is intractable and replaced by a Monte Carlo estimate.Source: §4.3 · pp. 10–11

4.4 Qualitative results

• Diffusion samples show wide variety and greater large-scale spatial coherence than PixelCNN samples, although the diffusion prior is not class-conditioned.Figure 5 places diffusion samples above PixelCNN samples.Source: §4.4 · p. 11, Figure 5 · p. 12
• For 9 × 9 centered and top-left masks on the 21 × 21 feature map, the diffusion model is much less sensitive to the selected masked region than PixelCNN.Source: §4.4 · pp. 12–13

5 Conclusion

• The framework uses a diffusion probabilistic model to learn dependencies between continuous latent variables alongside VQ-VAE encoding and decoding.The authors describe their numerical experiments as preliminary and identify scaling and hyperparameter study as future directions.Source: §5 · p. 13
• The authors report conceptual improvements over the VQ-VAE prior and first numerical results on middle-scale image generation.They identify scaling to larger models, optimizing hyperparameters, adding diffusion-method tricks, and studying end-to-end regularization as future research avenues.Source: §5 · p. 13
• Figure 7 shows conditional samples from diffusion and PixelCNN for three images, alongside masked and full VQ-VAE reconstructions.The samples are conditioned on a top-left mask.Source: Figure 7 · p. 14
• Figures 8 and 9 show the last 500 denoising steps at regular intervals for conditional and unconditional sampling, respectively.The conditional chain is conditioned on the picture's outer part; the captions say the first 500 steps are not visually informative.Source: Figure 8 · p. 14, Figure 9 · p. 15

6 Details on the loss function

• The proof of Lemma 3.1 begins from its definition.Source: §6 · p. 15
• The proof then decomposes the last term.Source: §6 · p. 15
• The excerpt states that this concludes the proof.Source: §6 · p. 15

7 Inpainting diffusion sampling

• Given a known sub-part of picture X, the procedure predicts its complementary pixels.Source: §7 · p. 16
• The known region's n encoded latent vectors are retained, while N − n terminal latent vectors are sampled from an uninformative Gaussian distribution.Source: §7 · p. 16
• At each step, a neural network predicts the preceding latent state, which is then sampled using the forward bridge noising process.Source: §7 · p. 16

8 Additional regularisation considerations

• The Gumbel-Softmax formulation provides an efficient and differentiable parameterization for computing Lreg.Source: §8 · p. 16
• The regularizer derivation considers the first term depending on zt that produces non-zero gradients.Source: §8 · p. 16
• For large t, negative γt pushes codebooks away from zt to prevent too-early specialization or matching codebooks with noise.Source: §8 · p. 16
• For small t, positive γt helps match codebooks with ze when corruption is small; the reported schedule uses τ = 1 and τt from 10 to 0.1.Source: §8 · p. 16

• The diffusion network is a 57M-parameter U-net-like model, slightly smaller than the 95.8M-parameter PixelCNN; the VQ-VAE encoder and decoder total 65M parameters.Source: §9 · p. 16

• The network in Figure 10 uses a time embedding and a few linear or 1D-convolutional layers, totaling about 5,000 parameters.Source: §10 · p. 16, Figure 10 · p. 17
• Discrete states during diffusion are represented by nearest-centroid sequences; at t = 0, a small-covariance centered Gaussian yields a uniform sequence.The table shows discrete samples from the denoising process.Source: Table 4 · p. 17, §10 · p. 17
• End-to-end training uses Gumbel-Softmax to make the reconstruction-likelihood path differentiable and update the codebooks.Source: §10 · p. 17
• In the toy example, the third loss term is unnecessary for good results with Lreg^e, but Lreg_t can improve codebook learning.With decreasing γ_t, the regularizer pulls the encoded and selected vectors together when denoising is nearly complete.Source: §10 · p. 17
• Joint denoising and codebook training gives well-positioned codebooks and naturally addresses dead codebooks.At high noise, Gumbel-Softmax gives all codebooks non-zero gradient; Figure 11 shows their learned positions and notes that codebook indexes need not match Gaussian indexes.Source: §10 · p. 17, Figure 11 · p. 18
• With q1 = 0 and q5 = 4 fixed, the bridge samples the intermediate sequence elements; three shown samples are valid sequences.Source: §10 · p. 18, Figure 12 · p. 18

11.1 Cifar

• CIFAR visuals show VQ-VAE reconstructions and samples from both the PixelCNN prior and the diffusion prior.Source: Figure 14 · p. 19, Figure 13 · p. 19

11.2 MiniImageNet

• MiniImageNet visuals show reconstructions of encoded, discretized images and samples from the model.Source: Figure 15 · p. 19, Figure 16 · p. 20
• The miniImageNet conditional-sampling figure contrasts samples conditioned on the left image region with PixelCNN samples in the same context.It also shows VQ-VAE reconstructions of the original images.Source: Figure 17 · p. 20
• A miniImageNet figure shows the sampling denoising chain at regular intervals, conditioned on the left part of the picture.Source: Figure 18 · p. 21
• The PixelCNN conditional-sampling visual pairs original miniImageNet images with samples conditioned on the left side; each row represents a validation-set class.Source: Figure 19 · p. 21