Purrception: Variational Flow Matching for Vector-Quantized Image Generation

1. Problem Statement

The paper addresses a fundamental modeling tension in Vector-Quantized (VQ) image generation. VQ latents represent images as grids of discrete indices (from a codebook) associated with continuous embedding vectors. Existing generative models typically force a trade-off between two degenerate extremes:

• Continuous Methods (e.g., Latent Diffusion, Flow Matching): Operate in the continuous embedding space ($\mathbb{R}^D$). While they preserve geometric smoothness and allow for efficient transport, they ignore the discrete categorical structure. They cannot express uncertainty over multiple plausible codes, lack categorical logits for temperature scaling, and must eventually discretize vectors, potentially losing structural fidelity.
• Discrete Methods (e.g., VQ-Diffusion, Discrete Flow Matching): Operate directly on discrete indices. While they respect the categorical nature of the data, they discard the geometric relationships between embeddings. Predictions become discrete "teleports" between unrelated tokens, making interpolation, uncertainty modeling, and temperature-controlled generation impossible or ineffective.

The core challenge is to create a generative model that leverages the smooth geometry of continuous embeddings while simultaneously utilizing explicit categorical supervision over the codebook indices.

2. Methodology: Purrception

The authors propose Purrception, a novel adaptation of Variational Flow Matching (VFM) specifically designed for VQ latents.

Core Concept

Purrception reframes the flow matching problem as a variational inference task. Instead of directly regressing a velocity field, the model learns a categorical posterior distribution over the codebook indices given the current interpolation point in the continuous space.

Key Components

1. Categorical Posterior over Indices:
   Given a noisy latent $z_t$ at time $t$, the model predicts a probability distribution $\pi_t^\theta(z_t)$ over the $K$ codebook vectors $\{e_k\}$. This distribution is categorical:
   $$q_t^\theta(c | z_t) = \text{Cat}(c | \pi_t^\theta(z_t))$$
   where $c$ represents the discrete code index.

2. Continuous Velocity Field:
   The velocity field $v_t^\theta(z_t)$ is computed as the expectation of the conditional velocity (linear interpolation towards the endpoint) weighted by the learned categorical posterior:
   $$v_t^\theta(z_t) = \sum_{k=1}^K \pi_{t,k}^\theta(z_t) \left( \frac{e_k - z_t}{1-t} \right) = \frac{\mu_t(z_t) - z_t}{1-t}$$
   Here, $\mu_t(z_t)$ is the barycenter (weighted average) of the codebook vectors based on the predicted probabilities. This ensures that uncertainty over multiple plausible codes translates into smooth, geometry-aware transport rather than discrete jumps.

3. Training Objective:
   The training objective reduces to minimizing the Cross-Entropy loss between the predicted categorical posterior and the ground-truth code indices of the target image. This provides direct categorical supervision while the transport dynamics remain continuous.

4. Temperature Scaling ($\tau$):
   The model outputs logits which are converted to probabilities via a softmax with a temperature parameter $\tau$:
   $$\pi_{t,k}^\theta(z_t) = \frac{\exp(\tilde{\pi}_{t,k}^\theta(z_t)/\tau)}{\sum_i \exp(\tilde{\pi}_{t,i}^\theta(z_t)/\tau)}$$

   • Low $\tau$: Sharpens the distribution, enforcing early commitment to a single code (higher fidelity, less diversity).
   • High $\tau$: Flattens the distribution, spreading probability across neighboring codes (more detail/variety, potentially lower fidelity).
     This allows for inference-time control over the trade-off between image sharpness and diversity, a feature absent in purely continuous or purely discrete baselines.

3. Key Contributions

• Hybrid Formulation: Purrception is the first method to successfully combine continuous transport dynamics with categorical supervision in VQ latent spaces, resolving the geometry-vs-discreteness trade-off.
• Uncertainty Quantification: By maintaining a probability distribution over codes, the model can express uncertainty across plausible indices, translating it into smooth motion in the embedding space.
• Temperature-Controlled Generation: Introduces a principled mechanism to control generation quality and diversity at inference time via softmax temperature scaling, which is unique to this hybrid approach.
• Efficiency: Demonstrates that learning discrete structure via categorical supervision accelerates convergence compared to purely continuous methods.

4. Experimental Results

The method was evaluated on ImageNet-1k (256×256) using DiT-L/2 and DiT-XL/2 backbones, comparing against Continuous Flow Matching (CFM), Discrete Flow Matching (DFM), and various state-of-the-art autoregressive and diffusion models.

• Convergence Speed: Purrception converges significantly faster than baselines.
  • With DiT-L/2, it matches CFM/DFM performance in 1.65× to 3.0× fewer iterations.
  • With DiT-XL/2, it is 2.3× faster than CFM and 3.5× faster than DFM.
• Sample Quality (FID):
  • Achieves a competitive FID of 3.88 (using LlamaGen tokenizer).
  • Outperforms all discrete diffusion and masked generative models.
  • Surpasses most autoregressive methods while using fewer parameters or benefiting from faster decoding.
  • While slightly behind top-tier continuous diffusion models (e.g., SiT-XL/2 with FID ~2.06), the gap is attributed to the use of VQ tokenizers (vs. high-quality VAEs) and shorter training schedules.
• Temperature Ablation: Experiments show a U-shaped relationship between temperature ($\tau$) and FID. An optimal $\tau \approx 0.8\text{--}0.9$ yields the best balance, confirming that temperature scaling is an effective, training-free knob for improving sample quality.

5. Significance

Purrception represents a significant step forward in generative modeling by bridging the gap between continuous and discrete paradigms.

• Theoretical Impact: It demonstrates that Variational Flow Matching can be effectively adapted to discrete-continuous hybrid spaces, providing a unified framework that leverages the geometric benefits of continuous flows and the structural benefits of discrete supervision.
• Practical Impact: The method offers a more training-efficient alternative to existing VQ-based generators, reducing computational costs. The ability to control generation via temperature scaling adds a new dimension of flexibility for practitioners.
• Future Directions: The paper suggests extending this hybrid perspective to other modalities (audio, video, 3D) and exploring joint training of the autoencoder and flow model to further close the performance gap with top-tier continuous diffusion models.

In summary, Purrception effectively "has its cake and eats it too" by learning what to choose (via categorical supervision) and where to go (via continuous transport), resulting in faster training and high-quality, controllable image generation.