JLT: Clean-Latent Prediction in Latent Diffusion Transformers

Technical Summary: JLT: Clean-Latent Prediction in Latent Diffusion Transformers

Problem Statement

Denoising diffusion models traditionally reverse a corruption process, but the specific target the neural network is trained to predict varies across architectures. Common targets include the noise ($\epsilon$), the velocity ($v = x - \epsilon$), or the clean data ($x$). While these targets are algebraically convertible for a fixed corruption time, recent research (e.g., JiT) suggests that under finite network capacity and finite data, the choice of prediction target significantly alters the difficulty of the regression problem. Specifically, predicting clean data allows a model to focus on structured, low-dimensional manifold variations rather than reconstructing ambient, off-manifold noise.

The central question addressed by this paper is whether this "clean-prediction" advantage persists when the model operates in a compressed latent space rather than raw pixel space. In latent diffusion, a Variational Autoencoder (VAE) has already removed much of the raw pixel variability. The authors investigate if the geometric properties of the prediction target ($x$ vs. $v$) remain a critical factor even when the representation space is fixed and compressed.

Methodology

The authors introduce JLT (Joint Latent Transformer), a controlled experimental framework designed to isolate the effect of the prediction target within a fixed latent space.

1. Architecture and Setup:

   • Backbone: A Base-scale Transformer (130M parameters) with 12 blocks, hidden dimension 768, and 12 attention heads, mirroring the configuration of JiT-B/16.
   • Representation: The model operates on frozen FLUX.2 VAE latent codes. This ensures the representation space is identical across all compared variants.
   • Training: Models are trained for 250K steps (200 epochs) on ImageNet 256 $\times$ 256.
   • Controlled Variables: The study strictly holds constant the representation (FLUX.2 VAE), architecture scale, optimizer settings, batch size, and sampling protocols.

2. Experimental Comparison:
   The authors compare two primary variants under the same settings:

   • JLT (Clean-Latent Prediction): The network is directly supervised to predict the clean latent vector $x$.
   • DiT (Velocity Prediction): A matched baseline where the network is supervised to predict the velocity vector $v = x - \epsilon$.
   • Variants: Experiments are conducted at two VAE-grid patch sizes, denoted as /1 and /2 (e.g., JLT-B/1 vs. DiT-B/1).

3. Theoretical Analysis:
   The paper provides a local Gaussian analysis to explain the empirical gap. Assuming a local linear-Gaussian approximation of the latent data distribution ($x \sim \mathcal{N}(0, \Sigma)$) and independent noise ($\epsilon \sim \mathcal{N}(0, I)$):

   • Covariance Structure: The covariance of the velocity target is $Cov(y_v) = \Sigma + I$, whereas the clean target is $Cov(y_x) = \Sigma$. This implies velocity prediction adds an isotropic unit floor to every direction. In directions where the clean data has low variance (weakly supported by the manifold), velocity prediction forces the model to learn a unit-variance target, effectively amplifying noise in those directions.
   • Conditional Ambiguity: The analysis shows that the Bayes residual variance for velocity prediction is scaled by a factor of $1/(1-t)^2$ relative to clean prediction. As the eigenvalue $\lambda_i$ of the data covariance approaches zero, the coefficient for clean prediction tends to 0 (attenuating the direction), while the velocity coefficient tends to $-1/(1-t)$ (amplifying the direction).

Key Results

The experiments on ImageNet 256 $\times$ 256 demonstrate a consistent and significant performance gap favoring clean-latent prediction:

• Quantitative Performance:
  • JLT-B/1 (Clean) achieves an FID-50K of 2.50 (with classifier-free guidance), compared to 6.56 for the matched DiT-B/1 (Velocity).
  • JLT-B/2 (Clean) achieves an FID-50K of 14.81, compared to 28.71 for DiT-B/2.
  • The Inception Score (IS) also shows a substantial improvement (e.g., 232.51 vs. 132.12 for the /1 variants).
• Training Dynamics: Training curves indicate that clean-latent models enter the low-FID regime earlier (around 100K steps) and maintain a clear margin over velocity models throughout training.
• Robustness: The advantage holds across different tokenization granularities (/1 and /2), suggesting the benefit is not an artifact of a specific patch size.

Significance and Claims

The paper claims that the choice of prediction target in latent diffusion is a geometric modeling choice rather than a mere algebraic parameterization.

1. Target Geometry over Compression: The results refute the notion that latent compression alone explains the efficiency of latent diffusion. The performance gap exists within the same fixed latent space, proving that the "target geometry" (the distribution of the supervised signal) is a distinct factor from the "representation geometry" (the VAE tokenizer).
2. Mechanism of Improvement: The authors argue that clean prediction is superior because it damps low-variance latent directions (which are often noise or weakly supported manifold features), whereas velocity prediction amplifies these directions by inheriting an isotropic covariance floor.
3. Scope of Claims: The authors are modest about their theoretical claims. They do not assert that clean prediction is globally optimal for all tokenizers, noise schedules, or samplers. Instead, they provide a local Gaussian mechanism that is consistent with the measured empirical gap. The study is presented as a controlled ablation to isolate the target effect, rather than a proposal for a new backbone architecture or a universal replacement for all diffusion objectives.

In conclusion, the paper establishes that even in compressed latent spaces, directly predicting the clean latent state offers a geometric advantage over velocity prediction, leading to significantly higher generation quality on ImageNet.