TL;DR

We introduce ReDi, a novel iterative method that reduces factorization error by rectifying the coupling between source and target distribution.

Overview

Discrete Flow-based Models (DFMs) are powerful generative models for high-quality discrete data but typically suffer from slow sampling speeds due to their reliance on iterative decoding processes. This reliance on a multi-step process originates from the factorization approximation of DFMs, which is necessary for handling high-dimensional data. In this paper, we analyze the factorization approximation error using Conditional Total Correlation (TC), and reveal its dependence on the coupling. To address the challenge of efficient few-step generation, we propose Rectified Discrete Flow (ReDi), a novel iterative method that reduces the underlying factorization error (measured as Conditional TC) by rectifying the coupling between source and target distributions. We theoretically prove that each ReDi step guarantees a monotonic decreasing Conditional TC, ensuring its convergence. Empirically, ReDi significantly reduces Conditional TC and enables few-step generation. Moreover, we demonstrate that the rectified couplings are well-suited for training efficient one-step models on image generation. ReDi offers a simple and theoretically grounded approach for tackling the few-step challenge, providing a new perspective on efficient discrete data synthesis.

Project Structure

  ├ image/                                        <- MaskGIT based ReDi-image code
  |    ├── Metrics/                               <- Evaluation tool
  |    |      ├── inception_metrics.py                  
  |    |      └── sample_and_eval.py
  |    |    
  |    ├── Network/                             
  |    |      ├── Taming/                         <- VQGAN architecture   
  |    |      ├── reweight_mlp.py                 <- Embedding architecture  
  |    |      └── transformer.py                  <- Transformer architecture  
  |    |
  |    ├── Trainer/                               <- Main class for training
  |    |      ├── trainer.py                      <- Abstract trainer     
  |    |      └── vit.py                          <- Trainer of MaskGIT
  |    |
  |    ├── Scripts/                               <- Shell scripts for training/evaluation.
  |    |      ├── create_rectified_dataset.sh     <- Rectify the model 
  |    |      ├── finetune_model.sh               <- Finetune the model (from origin MaskGIT)     
  |    |      ├── test_model.sh                   <- Test the model
  |    |      └── train_model.sh                  <- Train the model
  |    |
  |    ├── compute_tc.py                          <- Compute the TC
  |    ├── download_models.py                     <- Download the pretrained models
  |    ├── LICENSE.txt                            <- MIT license
  |    ├── requirements.txt                       <- Help to install env 
  |    ├── README.md                              
  |    └── main.py                                <- Main
  |
  ├ text/                                         <- Duo based ReDi-text code
  |    ├── config/                                <- Config files for datasets/denoising networks/noise schedules/LR schedules.
  |    |      └── config.yaml                     <- Main config file
  |    |
  |    ├── integral/
  |    |    
  |    ├── models/                                <- Denoising network architectures. Supports [DiT](https://arxiv.org/abs/2212.09748) and AR transformer.
  |    |      ├── dit.py                          <- DiT structure
  |    |      ├── ema.py                          <- EMA model
  |    |      └── unit_test_attention.py          <- Attention module
  |    |
  |    ├── scripts/                               <- Shell scripts for training/evaluation.
  |    |      ├── distil_*                        <- Distillate the model    
  |    |      ├── eval_*                          <- Evaluate the model     
  |    |      ├── gen_ppl_*                       <- Measure the generation perplexity     
  |    |      ├── gen_ppl_tc_*                    <- Measure the generation perplexity and total correlation score
  |    |      ├── rectifi_*                       <- Rectify the model
  |    |      ├── train_*                         <- Train the model 
  |    |      └── zero_shot_*                     
  |    |
  |    ├── algo.py                                <- Main model structures: Algorithms such as DUO, MDLM, AR, SEDD, D3PM, ReDi.
  |    ├── dataloader.py                          <- Dataloader and tokenizer module
  |    ├── LICENSE                                <- Apache License 2.0
  |    ├── main.py                                <- Main
  |    ├── metrics.py                             <- Metrics module
  |    ├── README.md                              
  |    ├── requirements.txt                       <- Help to install env 
  |    ├── trainer_base.py                        <- Boiler plate trainer using pytorch lightning.
  |    └── utils.py                               <- LR scheduler, logging, `fsspec` handling.

Usage

To get started, you can follow the process in "Usage" part of each image and text README.

Experiment score

Image

Step	Model	FID(↓)	IS(↑)	Prec.(↑)	Rec.(↑)	Den.(↑)	Cov.(↑)
1	MaskGIT	95.16	12	0.26	0.12	0.17	0.35
	SDTT	90.40	14	0.31	0.13	0.21	0.34
	Di4C	90.32	13	0.26	0.24	0.17	0.33
	ReDi1	37.43	49	0.63	0.51	0.78	0.86
	ReDi2	21.80	90	0.74	0.52	1.05	0.93
	ReDi3-distill	11.68	182	0.83	0.44	1.25	0.96
4	MaskGIT	10.90	184	0.83	0.46	1.18	0.96
	SDTT	8.97	205	0.88	0.41	1.43	0.97
	Di4C	6.20	216	0.87	0.52	1.33	0.98
	ReDi1	7.58	228	0.87	0.46	1.33	0.98
	ReDi2	7.86	240	0.87	0.44	1.31	0.97
8	MaskGIT	6.51	227	0.89	0.48	1.38	0.98

Text

License

The ReDi-image(Halton-MaskGIT-based ReDi) is licensed under the MIT License, and the ReDi-text(DUO-based ReDi) is licensed under the Apache License 2.0. From this, this project is licensed under the Apache License 2.0.

Acknowledgments

In this repo, we integrate Halton-MaskGIT and DUO as the structural foundation for our implementation, which involves both image and text models to demonstrate the ReDi method.

The pretrained VQGAN ImageNet is from the Halton-MaskGIT and LlamaGen official repository.

Citation

Cite our paper using:

@misc{yoo2025redirectifieddiscreteflow,
      title={ReDi: Rectified Discrete Flow}, 
      author={Yoo, Jaehoon and Kim, Wonjung and Hong, Seunghoon},
      year={2025},
      eprint={2507.15897},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2507.15897}, 
}