Upload folder using huggingface_hub

3ce19a2 verified 4 months ago

6.36 kB

# One Pass Is Not Enough: Recursive Latent Refinement for Generative Models

[**Project page**](https://mehdie79.github.io/Portfolio_Mehdi_Esmaeilzadeh/project/rtm-imle/) &nbsp;|&nbsp;
[**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/<name>/` and point `--data_root` at
it.

For AFHQ-v1 (StudioGAN pipeline) follow the standard StudioGAN data
layout: `<AFHQ>/{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/<exp>     DATA_DIR=/path/to/AFHQ      bash studiogan/scripts/eval_afhq.sh
CKPT_DIR=./runs/cifar10_rtm/checkpoints/<exp>  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.

	# 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/<name>/` and point `--data_root` at
	it.

	For AFHQ-v1 (StudioGAN pipeline) follow the standard StudioGAN data
	layout: `<AFHQ>/{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/<exp> DATA_DIR=/path/to/AFHQ bash studiogan/scripts/eval_afhq.sh
	CKPT_DIR=./runs/cifar10_rtm/checkpoints/<exp> 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.