PaperlayerView exact revision on arXiv

Source-linked AI summary

JLT: Clean-Latent Prediction in Latent Diffusion Transformers

Funing Fu, Tenghui Wang, Guanyu Zhou, Junyong Cen, Qichao Zhu

arXiv:2605.27102v2cs.CVcs.LG

TL;DR

The paper asks whether clean-data prediction remains useful after images are mapped into a learned latent space. It introduces JLT, a controlled latent Transformer comparison, and finds that clean-latent prediction outperforms matched velocity prediction, with local Gaussian analysis identifying a geometric mechanism.

  • Problem

    The paper asks whether direct clean prediction still matters after images are compressed into a learned latent space.

  • Method

    JLT compares clean-latent and velocity prediction in a fixed FLUX.2 VAE latent space using matched Transformer and training settings.

  • Results

    Clean-latent prediction consistently outperforms matched velocity prediction, with JLT-B/1 reaching FID-50K 2.56 versus 6.56 for DiT-B/1.

  • Takeaways & Limitations

    Prediction targets in latent diffusion are geometric modeling choices rather than merely interchangeable algebraic parameterizations.

  • Takeaways & Limitations

    The study focuses on ImageNet 256 × 256 and a 130M-parameter JLT-B/1 configuration, rather than characterizing all latent diffusion objectives.

Abstract

from arXiv · show

Flow matching with clean-data prediction has shown that regressing the clean point can exploit low-dimensional structure more effectively than predicting an ambient noised quantity. We ask whether this principle remains useful after images are mapped into a learned latent space, where compression has already removed much of the raw pixel variability. We introduce JLT, a 130M latent diffusion Transformer over frozen FLUX.2 VAE codes, and compare clean-latent prediction with a matched velocity-prediction DiT under the same representation, backbone, and training settings. Although the three variables x, epsilon, and v are linearly convertible for a fixed corruption time, a local Gaussian analysis shows that velocity regression inherits an isotropic target-covariance floor and amplifies low-variance latent directions, while clean prediction damps them. On ImageNet 256 x 256, JLT-B/1 obtains FID-50K 2.50 with classifier-free guidance, with a large matched-target gap over velocity prediction. These results suggest that prediction targets in latent diffusion are representation-dependent geometric choices, rather than interchangeable algebraic parameterizations.

1 Introduction

The paper asks whether clean-data prediction remains useful in a compressed latent space and tests it through a controlled comparison of clean-latent and velocity targets. It finds both an empirical target gap and a local geometric explanation for that gap.

  • The study asks whether direct clean prediction still matters after images are mapped into a compressed latent space.
  • JLT compares clean-latent and velocity targets using the same FLUX.2 VAE representation, Base-scale Transformer configuration, and 250K-step training setting.
  • 2.56 FID-50K is achieved by JLT-B/1 versus 6.56 for DiT-B/1 under the matched latent comparison.
  • 14.81 FID-50K is achieved by JLT-B/2 versus 28.71 for DiT-B/2 under the matched comparison.
  • The controlled study identifies clean-latent prediction as the empirical target choice and explains its advantage through covariance geometry.

2 Related Work

Related work frames prediction targets as geometric choices under finite capacity, while latent diffusion and Transformer studies emphasize representation and architecture. JLT differs by fixing the latent representation and isolating the target parameterization.

  • DDPM popularized noise prediction, while later diffusion and flow formulations introduced velocity parameterization and broader target-design choices.
  • JiT argues that clean images occupy structured low-dimensional subsets, whereas noise and velocity contain ambient components unsupported by the data distribution.
  • Latent diffusion reduces synthesis cost by training in autoencoder latents, and DiT studies Transformer blocks over latent patches.
  • SiT emphasizes controlled flow and diffusion comparisons using the same Transformer backbone, parameter count, and GFLOPs.
  • REPA and RiT vary or augment representation geometry, whereas this work fixes the FLUX.2 VAE representation for its target comparison.

3 Method

The method fixes a latent representation and compares clean-latent with velocity prediction, then analyzes their local Gaussian regression geometry. The analysis links target choice to covariance floors, conditional ambiguity, and amplification of weak latent directions.

  • 3.1 Formulation and prediction targets: The formulation defines clean latents x and Gaussian noise ϵ in a shared coordinate system before applying a linear corruption path.
  • 3.1 Formulation and prediction targets: For fixed corruption time, x-, ϵ-, and v-parameterizations are algebraically convertible, but affine readouts scale their prediction errors differently.
  • 3.1 Formulation and prediction targets: JLT outputs the clean latent x, while matched DiT outputs v = x−ϵ from the same corrupted latent and training setting.
  • 3.2 Target-geometry analysis: The target-geometry analysis uses a local linear-Gaussian approximation rather than claiming a complete theory of generative modeling.
  • 3.2 Target-geometry analysis: Velocity prediction has covariance Σ + I, adding an isotropic unit floor to every latent direction, while clean prediction retains covariance Σ.
  • 3.2 Target-geometry analysis: The local model shows that velocity can have larger conditional ambiguity than clean prediction despite affine convertibility after prediction.
  • 3.2 Target-geometry analysis: As λ_i approaches zero, clean prediction attenuates weak directions while velocity prediction can amplify them through a coefficient tending to −1/(1 −t).
  • 3.3 Architecture and training settings: JLT uses a 130M Base-scale latent Transformer with a fixed FLUX.2 VAE tokenizer and 250K-step training for the /1 and /2 variants.

4 Experiments

The experiments isolate prediction targets in matched latent diffusion models and show that clean-latent prediction consistently outperforms velocity prediction across patch sizes and training.

  • FID improves from 6.56 to 2.56 at VAE-grid patch /1 when replacing velocity prediction with clean-latent prediction.
  • FID improves from 28.71 to 14.81 at VAE-grid patch /2, showing the target effect persists under more aggressive tokenization.
  • Clean-latent variants retain lower FID and higher Inception Score than velocity counterparts throughout the matched-target training curves.
  • The /1 clean-latent model reaches the low-FID regime by roughly 100K steps and maintains a margin through the final checkpoint.
  • Table 2 contextualizes guided JLT against representative ImageNet 256 × 256 baselines rather than forming an unrestricted leaderboard.

5 Conclusion and Discussion

The study concludes that clean-state prediction improves matched latent diffusion results and that target parameterization reflects latent-space geometry rather than mere algebraic equivalence. The supporting theory is intentionally local and the empirical evidence remains bounded to a controlled ImageNet setting.

  • Under matched B-scale conditions, replacing velocity regression with clean-latent prediction substantially lowers denoising difficulty and improves ImageNet synthesis quality.
  • Velocity prediction adds an isotropic covariance floor and amplifies weakly supported low-variance latent directions, whereas clean prediction attenuates them.
  • Holding representation, Transformer scale, optimizer, batch size, and sampling settings fixed identifies target geometry—not latent compression alone—as the matched-ablation difference.
  • JLT applies the clean-prediction principle to fixed VAE latents, so its comparison with raw-pixel JiT should not be read as interchangeability between latent and pixel models.
  • The theory does not establish global optimality across tokenizers, noise schedules, loss weightings, or samplers, and real latent covariances may vary by class and spatial position.
  • The empirical scope is ImageNet 256 × 256 with a 130M-parameter JLT-B/1 configuration, not a complete characterization of latent diffusion objectives.

A Target Conversions and Error Scaling

For fixed corruption time, any target can be converted to the others, but the resulting direct regression losses differ because readout rescales prediction errors. The implementation therefore compares output parameterizations while preserving the training and sampling convention.

  • For fixed t, any one target determines the other two endpoint variables through an affine readout from the predicted target and known mixture z_t.
  • Algebraic target convertibility does not imply identical finite-model training behavior because conversion reweights direct prediction errors across noise levels.
  • Clean-latent prediction uses a clean-to-velocity readout, v_hat^(x)_theta = (x_hat_theta − z_t)/(1 − t), for its training loss.
  • The clean-to-velocity readout weights clean-prediction errors by (1 − t)^−2, whereas the matched baseline directly predicts v = x − ϵ.

B Residual-Variance Derivation

The local Gaussian derivation expresses residual risks for clean, noise, and velocity targets through the latent-coordinate variance term D_i. It concludes that velocity risk is clean-target risk scaled by (1 − t)^−2, within its stated local scope.

  • For one latent coordinate, D_i = t^2λ_i + (1 − t)^2 enters the joint Gaussian conditioning calculation.
  • The conditioning calculation uses Var(z_i) = D_i and target-specific covariances with z_i to derive the coordinatewise residual risks.
  • Summing coordinatewise conditional errors in the eigenbasis yields the local squared-error residual risks for the three prediction targets.
  • For fixed t in [0, 1), R_v(t) = R_x(t)/(1 − t)^2 under local Gaussian squared-error regression.The passage states that this is a mechanism for the controlled target gap, not a universal optimality theorem.

C Implementation Details

The matched comparison keeps optimization and evaluation conditions fixed across target choices. Only the direct prediction target changes between matched rows.

  • Matched rows use the same representation, Transformer scale, optimizer, batch size, time-sampling setting, and evaluation protocol.The optimizer uses AdamW with effective batch size 1024 and an actual learning rate of 2 × 10^−4 after batch-size scaling.
  • The ablation changes only the direct prediction target while holding the remaining training and evaluation settings constant.

D Additional Geometry Diagnostics

The analysis proposes empirical diagnostics for whether velocity targets have greater effective complexity and conditional uncertainty than clean-latent targets. These checks are presented as future validation tools, not main-result evidence.

  • The effective rank of y_v should exceed that of y_x when the clean latent spectrum is anisotropic.
  • kNN covariance estimates around corrupted latents should show larger conditional uncertainty for velocity targets over the effective training range.
  • Finite-capacity probes trained on identical corrupted inputs should fit y_x more easily than y_v.
  • These diagnostics are future validation tools rather than evidence used in the main result.
Loading 2605.27102v2…