JLT

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

<div align="center">

</div> <p align="center"> <img src="https://raw.githubusercontent.com/akatsuki-neo/JLT/main/images/qualitative_main_3x6.png" width="100%"> </p> <p align="center"> _{<b>Figure 1</b>: Matched Base-scale qualitative comparison on class-conditional ImageNet 256×256. Generations from matched B/1 models: clean prediction (x̂) preserves global coherence, boundary crispness, and fine details, whereas direct velocity prediction (v̂) frequently distorts, blurs, or introduces semantic artifacts under identical sampling.} </p>

Authors

Funing Fu^1* · Tenghui Wang^2* · Guanyu Zhou³ · Junyong Cen¹ · Qichao Zhu⁴
¹ Independent Researcher · ² Wuhan University of Technology · ³ Technology Innovation Institute · ⁴ Hangzhou Jiyi AI (*Equal contribution)

Overview

JLT (Just Latent Transformer) addresses a fundamental question in flow-based generative modeling: which quantity should the neural network directly predict?

Given corrupted state z_t and time t, predicting clean data x and predicting velocity v = x - ε are algebraically equivalent under the sample-wise linear path z_t = t·x + (1 - t)·ε, via the exact affine readout: $$\hat{v} = \frac{\hat{x} - z_t}{1 - t}$$

While an ideal oracle can express either form, a finite-capacity Transformer behaves very differently. Clean prediction delegates the known linear residual and time gain (1 - t)^-1 to the analytic readout, sparing the network from internalizing this known response. Moreover, in anisotropic latent spaces (such as the FLUX.2 VAE), clean prediction naturally attenuates low-variance directions, whereas velocity prediction adds an isotropic unit floor Σ_v = Σ_x + I that spreads variance across weakly supported directions.

Across identical architectures, training setups, and frozen FLUX.2 VAE representations:

• Base (B/1): FID-50K improves from 6.56 to 2.84 (IS 204.83) under L_v, and to 2.38 (IS 256.88) with unweighted clean MSE L_x.
• Large (L/1): FID-50K improves from 2.12 to 1.83 (IS 301.07).
• Huge (H/1, 951.3M): Reaches FID-50K 1.19 (IS 271.96).

<p align="center"> <img src="https://raw.githubusercontent.com/akatsuki-neo/JLT/main/images/jlt_thesis_figure.png" width="95%"> </p>

Results

Matched Prediction Target (Class-Conditional ImageNet 256×256)
Within each scale, the VAE, architecture, velocity loss, schedule, and sampler are strictly matched:

Model Family

Usage

Download Checkpoints

# Default Base checkpoint
huggingface-cli download dawn-neo/JLT checkpoint-last.pth

# Or download specific model checkpoints: JLT-L-1.pth, JLT-H-1.pth

Evaluation

Run evaluation using pre-encoded ImageNet latents with main_jlt.py:

python main_jlt.py \
    --model JLT-B/1 \
    --vae_type flux2 \
    --img_size 256 \
    --data_path /path/to/imagenet_latents_256 \
    --use_latent_cache \
    --online_eval \
    --eval_freq 1 \
    --gen_bsz 128 \
    --num_images 50000 \
    --cfg 2.9 \
    --num_sampling_steps 50 \
    --resume /path/to/checkpoint-last.pth \
    --output_dir ./eval_output

For full training recipes and data preparation, see the GitHub repository.

Citation

@article{fu2026jlt,
  title={Equivalent Flows, Unequal Learning: Clean-Latent Prediction in Transformers},
  author={Fu, Funing and Wang, Tenghui and Zhou, Guanyu and Cen, Junyong and Zhu, Qichao},
  journal={arXiv preprint arXiv:2605.27102},
  year={2026}
}

Acknowledgements

• JiT — Base architecture
• FLUX.2 VAE — Latent space
• Li & He. "Back to Basics" — Clean prediction insight