# One Pass Is Not Enough: Recursive Latent Refinement for Generative Models [**Project page**](https://mehdie79.github.io/Portfolio_Mehdi_Esmaeilzadeh/project/rtm-imle/)  |  [**Paper**](https://arxiv.org/abs/2605.15309) Official implementation of the **Recursive Token Mapper (RTM)**, which replaces the single-pass MLP mapper of style-based generators with a small recursive block that refines the latent across H weight-shared cycles. ![AFHQ baseline vs RTM](assets/mainpage_afhq_baseline_vs_rtm.jpg) *Left: baseline single-pass mapper. Right: RTM (ours). Recursive refinement produces more diverse coat patterns and better sample quality (FID 4.79 vs 4.99, Recall 0.565 vs 0.507 on AFHQ-v1 512x512).* ## Abstract Despite remarkable progress, image generation is far from solved. The dominant metric, FID, conflates sample fidelity with mode coverage and is close to being saturated. Yet a model can still exhibit mode collapse while achieving a low FID, since a handful of sharp, near-duplicate images can outscore a model that faithfully covers the full data distribution. We argue that precision and recall are essential complements to FID, and that because FID is already saturated, the more meaningful goal is to improve diversity and coverage. Achieving high recall requires a model that explicitly prioritizes mode coverage, unlike most generative models, which optimize sample fidelity. We introduce **RTM**, which replaces the single-pass latent mapping in style-based generators with an iterative refinement process, and show that this consistently improves both quality and diversity. Integrated with Implicit Maximum Likelihood Estimation (IMLE), which optimizes mode coverage by design, RTM achieves the highest precision and recall among current state-of-the-art approaches while maintaining competitive FID, with improvements across CIFAR-10, CelebA-HQ at 256x256, and nine few-shot benchmarks. RTM also improves StyleGAN2 and StyleGAN2-ADA on CIFAR-10 and AFHQ-v1 at 512x512, demonstrating that the benefit is not specific to IMLE. Unlike flow-matching baselines that achieve competitive FID at the expense of coverage, recursive refinement improves both quality and diversity simultaneously. ## Architecture ![Mapper architecture](assets/architecture.png) The mapper M projects noise z into a small set of latent tokens that are refined through H weight-shared cycles of token-mixing and channel-mixing MLPs, then projected back into the style vector w. ## Repository layout ``` . ├── assets/ # Hero image and architecture diagram ├── rtm_core.py # RTM building blocks ├── models.py hps.py sampler.py train.py evaluate.py ├── prepare_cifar10.py prepare_celeba_hq.py ├── helpers/ visual/ lpips/ ├── scripts/ # CIFAR-10 + CelebA-HQ-256 train/eval ├── fewshot/ # Few-shot pipeline (Obama, AnimalFace-cat, ...) │ ├── train.py evaluate.py models.py hps.py sampler.py │ ├── README.md │ └── scripts/ └── studiogan/ # StyleGAN2 / StyleGAN2-ADA + RTM (AFHQ, CIFAR-10) ├── src/ # Vendored PyTorch-StudioGAN with RTM patched in ├── scripts/ # AFHQ-v1 + CIFAR-10 train/eval └── README.md ``` ## Setup ```bash virtualenv -p python venv && source venv/bin/activate pip install -r requirements.txt pip install -i https://test.pypi.org/simple/ dciknn-cuda==0.1.15 ``` The StudioGAN pipeline shares the same environment, plus a few extra dependencies; see `studiogan/README.md`. ## Datasets ```bash python prepare_cifar10.py --out_dir ./datasets/cifar10 python prepare_celeba_hq.py --src /path/to/celeba-hq --out_dir ./datasets/celeba-hq-256 ``` For the few-shot benchmarks drop each image folder under `./datasets//` and point `--data_root` at it. For AFHQ-v1 (StudioGAN pipeline) follow the standard StudioGAN data layout: `/{train,valid}/{cat,dog,wild}/*.png`. ## Training CIFAR-10 (RS-IMLE + RTM): ```bash bash scripts/train_cifar10.sh ``` CelebA-HQ-256 (RS-IMLE + RTM): ```bash bash scripts/train_celebahq256.sh ``` Few-shot (RS-IMLE + RTM): ```bash cd fewshot && bash scripts/train.sh /path/to/dataset ``` StyleGAN2-ADA + RTM: ```bash DATA_DIR=/path/to/AFHQ NUM_GPUS=4 bash studiogan/scripts/train_afhq.sh DATA_DIR=/path/to/cifar-10 bash studiogan/scripts/train_cifar10.sh ``` ## Pretrained checkpoints Checkpoints will be released soon! ## Evaluation RS-IMLE on CIFAR-10 (FID with 50,000 samples, Precision/Recall with 50,000): ```bash bash scripts/eval_cifar10.sh /path/to/model.th ``` RS-IMLE on CelebA-HQ-256 (FID with 30,000 samples, Precision/Recall with 30,000): ```bash bash scripts/eval_celebahq256.sh /path/to/model.th ``` Few-shot (FID with 5000 samples, Precision/Recall with 1000): ```bash cd fewshot && bash scripts/evaluate.sh /path/to/model.th /path/to/dataset ``` StyleGAN2-ADA + RTM on AFHQ-v1 / CIFAR-10 (FID and improved Precision/Recall via StudioGAN's PRDC): ```bash CKPT_DIR=./runs/afhq_rtm/checkpoints/ DATA_DIR=/path/to/AFHQ bash studiogan/scripts/eval_afhq.sh CKPT_DIR=./runs/cifar10_rtm/checkpoints/ DATA_DIR=/path/to/cifar-10 bash studiogan/scripts/eval_cifar10.sh ``` All evaluation scripts use the same configuration that produced the headline numbers in the paper. ## Citation If you find this work useful, please cite: ```bibtex @misc{esmaeilzadeh2026onepass, title = {One Pass Is Not Enough: Recursive Latent Refinement for Generative Models}, author = {Mehdi Esmaeilzadeh and Alexia Jolicoeur-Martineau and Chirag Vashist and Ke Li}, year = {2026}, eprint = {2605.15309}, archivePrefix = {arXiv}, primaryClass = {cs.CV}, url = {https://arxiv.org/abs/2605.15309}, doi = {10.48550/arXiv.2605.15309} } ``` ## Acknowledgments This work has been built on top of these codebases: - [RS-IMLE (Vashist et al., ECCV 2024)](https://github.com/serchirag/rs-imle) - [Tiny Recursive Model (Jolicoeur-Martineau, 2025)](https://github.com/SamsungSAILMontreal/TinyRecursiveModels) - [PyTorch-StudioGAN (POSTECH-CVLab)](https://github.com/POSTECH-CVLab/PyTorch-StudioGAN) — see `studiogan/LICENSE-StudioGAN` and `studiogan/LICENSE-NVIDIA` for upstream licenses.