Equivalent Flows, Unequal Learning: Clean-Latent Prediction in Transformers

arXiv:2605.27102 · cs.CV, cs.LG · Submitted 2026-05-26 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "Equivalent Flows, Unequal Learning".

Jane: 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.

Tom: First, who's behind it and why it matters.

Title and authors: Tom: So, let's start by looking at the title of "Equivalent Flows, Unequal Learning: Clean-Latent Prediction in Transformers." It immediately tells us that while mathematically you can convert the two prediction tasks into each other after some processing, their learning behaviors are actually quite different.

Jane: That’s a crucial distinction to make; it’s not just an algebraic trick to make one task look like the other, but something deeper about how the model learns from the data distribution itself. The authors are essentially showing that this difference in learning behavior isn't just some artifact of how we set up the equations.

Lu: They focus on comparing clean-latent prediction against a matched velocity-prediction DiT under identical representation and training settings, which is a very controlled experimental setup to isolate the target effect they are investigating. This comparison is what makes their study so focused on this specific mechanism.

Meng: It’s good that they explicitly mention keeping the backbone and training settings fixed; that level of control tells me they are trying to ensure whatever difference we see isn't just because one architecture is inherently better than the other in a vacuum.

Lalam: What excites me about this title is the "Unequal Learning" part; it suggests that there’s an inherent asymmetry in how these two targets interact with the latent space structure, which hints at a richer understanding of diffusion dynamics.

The paper's summary: Tom: Now moving into what they actually found, the core finding of "Equivalent Flows, Unequal Learning: Clean-Latent Prediction in Transformers" is that clean-latent prediction consistently outperforms matched velocity prediction when working within a fixed latent space.

Jane: That’s a big empirical result; the paper shows that when we regress for the actual clean latent point, the model achieves significantly better results than when we try to predict just the noisy velocity. They found this separation isn't due to how compressed the latent space is, but rather because of geometric differences in the target distributions themselves.

Lu: They use a local Gaussian analysis to explain this empirically; they show that velocity prediction introduces an isotropic covariance floor that gets added to every clean-latent direction, which essentially amplifies those low-variance directions in the latent data.

Meng: So, it’s not just about noise reduction; it’s about how the model handles different levels of structural information within that compressed representation. The paper highlights a specific mathematical structure where velocity prediction adds a uniform floor to every direction, whereas clean prediction dampens those variations.

Lalam: It really helps us understand that the target choice matters immensely in diffusion models, suggesting that simply following standard procedures might not always lead to the best results if we don't consider this geometric relationship between the targets.

The paper's improvements: Tom: The authors suggest a couple of key ways to think about this, moving beyond just running one target or the other. They highlight that velocity prediction can have larger conditional ambiguity than the clean target even after they are related by an affine transformation.

Jane: That’s fascinating because it means that even though mathematically you can convert x, epsilon, and v into each other, their behavior in a trained model is not identical because of how the training process handles uncertainty differently for those two targets.

Lu: They also point out something specific about low-variance directions; when the eigenvalue lambda i tends to zero, the clean-target coefficient goes toward zero, but the velocity-target coefficient trends towards negative one over one - t. This is a concrete mathematical detail explaining why clean prediction attenuates those weak directions while velocity prediction can amplify them.

Meng: From an implementation view, this means we might need a mechanism that accounts for this local covariance spectrum when choosing our loss functions, rather than just picking one target blindly. It suggests that the "geometric modeling choice" they mention in their conclusion is something engineers should pay attention to.

Lalam: I see the implication for culture here; it pushes us to be more deliberate about our model design choices, recognizing that a seemingly simple algebraic rewrite can hide fundamentally different learning dynamics depending on the target we select.

Conclusion: Tom: So, wrapping up this discussion on "Equivalent Flows, Unequal Learning: Clean-Latent Prediction in Transformers," the central message is that target parameterization in latent diffusion models represents a geometric modeling choice, not just an algebraic rewrite. The clean-latent prediction substantially lowers the difficulty of denoising by attenuating directions weakly supported by the latent data distribution.

Jane: Exactly; it’s about how clean-latent prediction specifically manages low-variance directions better, which is a subtle but important mechanism that leads to better synthesis quality in practice, especially when using classifier-free guidance.

Lu: This research provides a very concrete mechanism for why target choice remains important even after the images are mapped into a fixed latent space; it’s about how the local Gaussian approximation of the distribution behaves near data regions with little variation.

Meng: For practical application, this suggests that we need to be careful not to assume universality across all diffusion objectives; real latent distributions aren't always perfectly Gaussian, and their local covariance can change.

Lalam: I think this study serves as a valuable piece of evidence showing that targeting the clean point offers a specific advantage in exploiting low-dimensional structure, and we should keep looking into these target-geometry effects in future work.

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

Wuhan University of Technology · Hangzhou Jiyi Artificial Intelligence Co., Ltd.

cs.CV, cs.LG

Submitted: 2026-05-26

Updated: 2026-09-28

Code: https://github.com/akatsuki-neo/JLT

Importance score: 77/100

The gist: 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.

Key concepts

Clean-Latent Prediction (JLT): This method trains a model to directly predict the clean latent representation 'x' from noise. It exploits the low-dimensional structure of the data distribution more effectively than predicting noisy quantities, leading to better performance in image generation tasks.
Matched Velocity Prediction (DiT): This baseline predicts the velocity vector 'v' (the difference between a clean point and noise). The analysis shows that this approach introduces different geometric constraints compared to predicting the clean point, resulting in a less optimal target for denoising.
Target-Geometry Effect: This refers to how the mathematical structure of the prediction target influences model performance. The paper argues that the gap between clean prediction and velocity prediction is due to these inherent geometric differences in their respective distributions, not just simple compression effects.
Low-Variance Directions: These are directions in the latent space where data variance is very low (eigenvalues approaching zero). Clean prediction tends to attenuate errors in these directions, while velocity prediction can amplify them, providing a concrete reason for the observed performance gap.

Terminology

Transformer — 30× in this paper · explained in 2 episodes
Latent Space — 15× in this paper · explained in 7 episodes
Ablation — 12× in this paper · explained in 1 episode
Diffusion Model — 10× in this paper · explained in 8 episodes
Diffusion Models — 9× in this paper · explained in 6 episodes
Flow Matching — 6× in this paper · explained in 7 episodes
Alignment — 4× in this paper · explained in 1 episode
Core — 4× in this paper · explained in 2 episodes
Generative Model — 4× in this paper · explained in 2 episodes
Latent Diffusion Model — 4× in this paper · explained in 2 episodes
Latent Diffusion Models — 4× in this paper · explained in 2 episodes
Distillation — 3× in this paper · explained in 2 episodes
Generative Models — 3× in this paper · explained in 2 episodes
Latent Representation — 3× in this paper · explained in 2 episodes
Generative Modeling — 2× in this paper · explained in 3 episodes
Rectified Flow — 2× in this paper · explained in 3 episodes
Aggregation — 1× in this paper · explained in 1 episode
Intrinsic Dimension — 1× in this paper · explained in 2 episodes
Loss Function — 1× in this paper · explained in 2 episodes
Nash Equilibrium — 1× in this paper · explained in 2 episodes
Tokenization — 1× in this paper · explained in 2 episodes

Summary

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. This study investigates whether this principle remains useful after images are mapped into a learned latent space by comparing clean-latent prediction with a matched velocity-prediction DiT under fixed representation and training settings, finding that clean prediction consistently outperforms velocity regression due to geometric differences in the target distributions.

How it works

The research instantiates a comparison between two types of direct targets within a controlled latent diffusion setting: clean-latent variants (JLT) and matched velocity variants (DiT). The clean-latent model is parameterized to predict the clean latent, denoted as "x, while the matched velocity baseline is parameterized to predict the velocity, denoted as v = x - ϵ." Both models operate within a fixed FLUX.2 VAE latent space, using the same Base-scale Transformer configuration and training settings (250K steps/200 epochs).

The core finding stems from a local Gaussian analysis of the regression problem. The marginal target covariances are defined as: Cov(yx) = Σ, Cov(yϵ) = I, Cov(yv) = Σ + I. This mathematical structure reveals that velocity prediction adds the same isotropic unit floor to every clean-latent direction, whereas clean prediction damps them.

Key Results and Mechanisms

The empirical results demonstrate a significant performance gap. On ImageNet 256 × 256, JLT-B/1 obtains an FID-50K of 2.50 with classifier-free guidance, which is substantially better than the matched velocity prediction (DiT-B/1) FID of 6.56 under the same settings. This separation is interpreted as a target-geometry effect rather than a consequence of latent compression alone.

The theoretical explanation for this gap lies in conditional ambiguity and Bayes estimators. The analysis shows that velocity prediction can have larger conditional ambiguity than the clean target even though both are affinely related after prediction. Furthermore, when considering low-variance directions (where the eigenvalue λi → 0), the clean-target coefficient tends to 0, while the velocity-target coefficient tends to-1/(1 − t). This mechanism explains why clean prediction attenuates low-variance directions, whereas velocity prediction can amplify them, providing a concrete reason for the empirical gap.

Experimental Setup and Control

The study maintains strict control over confounding variables to isolate the target effect. The comparison is held fixed across:

Representation: A fixed FLUX.2 VAE latent representation.
Architecture Scale: Base-scale Transformer configuration (comparable to JiT-B/16).
Training Setup: 250K steps (200 epochs) with the same optimizer and learning rate schedule as the matched baseline.

To ensure the comparison focuses solely on target geometry, two components of JiT are excluded: repeated in-context class-token concatenation is not used, and the auxiliary ImageNet classification loss explored in JiT is omitted. This ensures that the key comparison is made after those factors have been held constant.

Conclusion

The central conclusion is that target parameterization in latent diffusion models represents a geometric modeling choice, not merely an algebraic rewrite. While algebraically convertible after prediction, the direct regression losses induce different supervised problems. The results suggest that clean-latent prediction substantially lowers the difficulty of denoising by attenuating directions weakly supported by the latent data distribution. The analysis is noted as being deliberately conservative, identifying a mechanism consistent with measured gaps rather than claiming global optimality across all possible diffusion objectives.

Limitations

The study's findings are interpreted within the context of ImageNet 256 × 256 and a specific 130M-parameter JLT-B/1 configuration. The analysis does not prove that clean prediction is globally optimal for every tokenizer, noise schedule, loss weighting, or sampler, as real latent distributions are non-Gaussian and their local covariance can vary. Future validation tools suggested include checking the effective rank of yv and using nonparametric local posterior estimates to confirm the larger conditional uncertainty assigned to the velocity target.

Appendix Details

The paper provides detailed algebraic conversions in Appendix A, showing how targets are linearly convertible after prediction, but emphasizes that this equivalence does not imply identical finite-model training behavior because the readout reweights direct prediction errors across noise levels. The derivation of residual risks in Appendix B formalizes the local squared-error regression at a fixed corruption level. Detailed implementation settings for the optimizer (AdamW with specific hyperparameters) are provided in Appendix C. The study concludes by noting that these results serve as evidence for a target-geometry effect, not as a complete characterization of all latent diffusion objectives.

Improvements for AI systems

Based on the scientific paper JLT: Clean-Latent Prediction in Latent Diffusion Transformers, here are specific improvements for AI systems derived from its findings, and what those improved systems can achieve:

)1. System Improvement: Implement Target-Geometry Aware Denoising Objectives (TGA-DO).

The core finding is that the choice of prediction target (clean latent vs. velocity regression) significantly impacts the model's ability to learn low-variance, high-frequency structural details in compressed latent spaces.

The improved system will replace standard velocity regression targets with a Clean Latent Prediction target for generative tasks operating within fixed VAE latents (like those from FLUX.2).

This TGA-DO objective explicitly leverages the geometric structure of the clean data manifold, allowing the Transformer backbone to focus its capacity on reconstructing structured variations rather than compensating for ambient noise components inherent in velocity fields.

The improved AI system can achieve a substantial reduction in FID (measured at 50K samples) and higher Inception Score (IS) compared to models trained with standard velocity regression targets, particularly when operating under classifier-free guidance, leading to significantly sharper, more coherent high-resolution image synthesis.

System Improvement: Integrate Dynamic Target Switching based on Latent Distribution Analysis.

The paper demonstrates that the performance gap is a function of the local covariance spectrum (anisotropy) of the latent data.

The system will incorporate a module that performs an initial, low-cost analysis of the local covariance structure within the VAE latent space (e.g., via kNN or spectral estimation). Based on this analysis, it will dynamically switch between clean-latent prediction and velocity regression targets during training for specific regions of the latent manifold or specific data classes.

This allows the model to adapt its regression strategy on a per-region basis, mitigating the amplification of low-variance directions observed in velocity targets when they are highly prevalent in that local subspace.

The improved AI system can exhibit superior robustness across diverse datasets and class distributions, maintaining high synthesis quality even when the underlying latent data distribution exhibits strong anisotropy or localized low-variance features.

System Improvement: Develop Robust Latent-Space Feature Extraction (LSAFE).

The paper validates that fixing the representation space (using FLUX.2) isolates a target-geometry effect orthogonal to tokenizer or backbone changes, suggesting that specific latent representations are critical for target choice effectiveness.

The system will utilize the fixed, high-quality FLUX.2 VAE latent space as its primary input manifold rather than relying on raw pixel space or generic feature embeddings.

By training models exclusively in this constrained geometric setting, the system can achieve superior efficiency and quality in generation tasks (e.g., image synthesis) because the prediction targets are optimized for the specific geometry of that fixed latent representation, rather than being generalized across all possible representations.

The improved AI system will demonstrate high efficiency with a controlled parameter count (like JLT-B/130M) while achieving state-of-the-art FID scores on ImageNet, proving that representation alignment and fixed latent spaces are powerful levers for performance enhancement.