README.md

---
language:
- en
license: apache-2.0
tags:
- image-generation
- latent-diffusion
- denoising-diffusion
- computer-vision
- generative-model
paperswithcode_id: akatsuki-neo/JLT
---

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

<div align="center">

[![Paper](https://img.shields.io/badge/arXiv-2605.27102-b31b1b)](https://arxiv.org/abs/2605.27102)
[![Project Page](https://img.shields.io/badge/GitHub%20Pages-JLT-blue)](https://akatsuki-neo.github.io/JLT)
[![Code](https://img.shields.io/badge/Code-GitHub-blue)](https://github.com/akatsuki-neo/JLT)

</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">
  <sub><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.</sub>
</p>

## Authors

[**Funing Fu**](https://github.com/chinoll)<sup>1\*</sup> · [**Tenghui Wang**](https://github.com/spawner1145)<sup>2\*</sup> · [**Guanyu Zhou**](https://the-martyr.github.io/)<sup>3</sup> · Junyong Cen<sup>1</sup> · Qichao Zhu<sup>4</sup>  
<sup>1</sup> Independent Researcher · <sup>2</sup> Wuhan University of Technology · <sup>3</sup> Technology Innovation Institute · <sup>4</sup> 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<sub>t</sub>* and time *t*, predicting clean data *x* and predicting velocity *v = x - ε* are algebraically equivalent under the sample-wise linear path *z<sub>t</sub> = 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)<sup>-1</sup>* 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 *Σ<sub>v</sub> = Σ<sub>x</sub> + 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<sub>v</sub>*, and to **2.38** (IS 256.88) with unweighted clean MSE *L<sub>x</sub>*.
- **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:

| Scale | Model | Target | Loss | FID-50K ↓ | IS ↑ |
|:---|:---|:---:|:---:|---:|---:|
| Base | Baseline | velocity *v* | *L<sub>v</sub>* | 6.56 | 132.12 |
| Base | **JLT-B/1** | **clean latent *x*** | *L<sub>v</sub>* | **2.84** | **204.83** |
| Base | **JLT-B/1** | **clean latent *x*** | *L<sub>x</sub>* | **2.38** | **256.88** |
| Large | Baseline | velocity *v* | *L<sub>v</sub>* | 2.12 | 236.21 |
| Large | **JLT-L/1** | **clean latent *x*** | *L<sub>v</sub>* | **1.83** | **301.07** |
| Huge | Baseline | velocity *v* | *L<sub>v</sub>* | 1.60 | **327.41** |
| Huge | **JLT-H/1** | **clean latent *x*** | *L<sub>v</sub>* | **1.19** | 271.96 |

## Model Family

| Model | Depth | Width | Heads | Params | Tokenizer |
|:---|---:|---:|---:|---:|:---|
| **JLT-B/1** | 12 | 768 | 12 | 130.5M | FLUX.2 VAE (frozen) |
| **JLT-L/1** | 24 | 1024 | 16 | 458.1M | FLUX.2 VAE (frozen) |
| **JLT-H/1** | 32 | 1280 | 16 | 951.3M | FLUX.2 VAE (frozen) |

## Usage

### Download Checkpoints
```bash
# 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`:
```bash
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](https://github.com/akatsuki-neo/JLT).

## Citation

```bibtex
@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](https://github.com/LTH14/JiT) — Base architecture
- [FLUX.2 VAE](https://huggingface.co/black-forest-labs/FLUX.2-klein-4B) — Latent space
- [Li & He. "Back to Basics"](https://arxiv.org/abs/2511.13720) — Clean prediction insight