ReDi: Rectified Discrete Flow

2025-07-20

AI Summary

Discrete flow models (DFMs) suffer from low sampling efficiency due to multi-step iterative decoding, necessitated by factorized approximations for high-dimensional discrete data. This approximation introduces error that we rigorously characterize via Conditional Total Correlation (TC), whose magnitude depends on the coupling strength between source and target distributions. To address this, we propose ReDi—a novel iterative correction framework that, for the first time, treats Conditional TC as a differentiable optimization objective. ReDi employs a theoretically grounded, monotonically decreasing coupling correction mechanism to suppress redundant inter-distribution dependencies. Crucially, it operates post-hoc without model retraining and enables few-step or even single-step generation. Experiments demonstrate that ReDi significantly reduces Conditional TC, yielding substantial improvements in both generation quality and speed. Furthermore, it successfully transfers to training highly efficient single-step image generators.

Technology Category

• Machine Learning: Deep Generative Models & Autoencoders
• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Mixed Discrete/Continuous Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

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 rigorously characterize the approximation error from factorization using Conditional Total Correlation (TC), which depends on the coupling. To reduce the Conditional TC and enable efficient few-step generation, we propose Rectified Discrete Flow (ReDi), a novel iterative method that reduces factorization error 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. Code is available at https://github.com/Ugness/ReDi_discrete

Problem

• Reduces slow sampling in discrete flow models
• Minimizes factorization error via rectified coupling
• Enables efficient few-step discrete data generation

Innovation

• Rectified Discrete Flow reduces factorization error
• Monotonic decrease in Conditional Total Correlation
• Enables efficient few-step image generation