ReCAP: Plug-and-Play Context Feature Reuse for Efficient Masked Generation

Official PyTorch Implementation of Plug-and-Play Context Feature Reuse for Efficient Masked Generation (ReCAP).

🚀 Overview

ReCAP provides an efficient KV-caching mechanism for accelerating the inference of Masked Generative Models (MGMs) with bidirectional attention, achieving up to 3.2× speedup with negligible quality loss. It is compatible with popular MGM framework such as MaskGIT and MAR.

🛠️ Environment Setup

conda create -n recap python=3.10
conda activate recap
pip install -r requirements.txt

📘 Example Usage

📌 MaskGIT-r + ReCAP

1. Download the pre-trained VQGAN tokenizer to assets/vqgan_jax_strongaug.ckpt.
2. Download the pre-trained MaskGIT to examples/MaskGIT+ReCAP/ckpts/nnet_ema.pth.
3. Run the following scripts to evaluate the pre-trained MaskGIT model with ReCAP:

cd examples/MaskGIT+ReCAP
bash run_maskgit_recap.sh

📌 MAR + ReCAP

1. Download the pre-trained VAE to examples/MAR+ReCAP/pretrained_models/vae/kl16.ckpt.
2. Download the pre-trained MAR-L to examples/MAR+ReCAP/pretrained_models/mar_large/checkpoint-last.pth and MAR-H to examples/MAR+ReCAP/pretrained_models/mar_huge/checkpoint-last.pth.
3. Run the following scripts to evaluate the pre-trained MAR models with ReCAP:

cd examples/MAR+ReCAP
bash run_mar_recap.sh

📊 Calculate FID/IS

1. Download the FID-stats to assets/fid_stats/imagenet256_guided_diffusion.npz.
2. Run the script eval_metrics.py

📄 Citation

If you find this work useful, please consider citing our paper:

@misc{liu2025recap,
      title={Plug-and-Play Context Feature Reuse for Efficient Masked Generation}, 
      author={Xuejie Liu and Anji Liu and Guy Van den Broeck and Yitao Liang},
      year={2025},
      eprint={2505.19089},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2505.19089}, 
}

📮 Contact

For any questions, please feel free to contact us via email: [liebenxj@gmail.com]

🎉 Acknowledgements

This code is built upon the following repositories:

• AutoNAT
• MAR