FlowTrunc: Flow-Based Truncated Diffusion for Efficient Text-to-Image Generation

⭐ Key Results

• FID Score: 33.47 (vs. baseline TDPM: 75.48)
• Truncation Level ($T_{trunc}$): Only 4 diffusion steps (vs. standard 1000 steps)

📝 Abstract

Diffusion probabilistic models have emerged as powerful generative models capable of producing high-fidelity images. However, their slow sampling process—often requiring hundreds of iterative denoising steps—makes them computationally expensive. FlowTrunc addresses this by combining Truncated Diffusion Probabilistic Models (TDPM) with a Transformer-based normalizing flow prior (TARFlow).

By replacing the adversarially-trained GAN prior used in standard TDPM with an explicit flow-based prior, FlowTrunc achieves:

• Better mode coverage
• More accurate modeling of the truncated latent distribution
• Significantly improved image quality (42-point FID improvement)
• Efficient generation with only 4 diffusion steps

📊 Experiments

Dataset

• CIFAR-10: 60,000 color images (32×32) across 10 object classes
• Training: 50,000 images
• Test: 10,000 images
• Preprocessing: Linear scaling to [0, 1] range
• Evaluation: FID computed on 50,000 generated images

Training Configuration

TDPM (GAN Prior)

Parameter	Value
Truncated timestep (T_trunc)	4
Initial LR (Discriminator)	10^-4
Initial LR (Generator)	10^-5
Adam β1	0.5
Adam β2	0.9
EMA decay	0.9999
Training epochs	471 (184,000 steps)

FlowTrunc (Flow Prior)

Parameter	Value
TDPM settings	Same as above
Patch size	4
Channel dimension	512
Number of flows	4
Layers per flow	8
Training epochs	215 (84,000 steps)

Hardware:

• Training: Dual NVIDIA L40S GPUs (~2 days, 6 hours)
• Sampling: Single NVIDIA RTX 4070 Laptop GPU (8GB VRAM)

📈 Results

Model	FID ↓	Sampling Time (s/img) ↓
TDPM (GAN prior)	75.48	0.01286
FlowTrunc (Flow prior)	33.47	0.02662

FlowTrunc achieves a 42-point FID improvement over baseline TDPM while maintaining competitive generation speed. The flow-based prior provides better mode coverage and more accurate modeling of the truncated latent distribution at $T_{trunc} = 4$.

🚀 Installation

Prerequisites

• Python 3.10.11
• CUDA 11.8 (for GPU support)

Setup

1. Clone the repository:

git clone https://github.com/aysanorhun/flow-diffusion.git
cd flow-diffusion

2. Install dependencies using uv:

uv sync

💻 Usage

Training

Train FlowTrunc on CIFAR-10:

python main.py --config configs/cifar10.yml --doc flowtrunc_experiment --exp ./experiments

Key arguments:

• --config: Path to configuration file
• --doc: Experiment name (creates log folder)
• --exp: Path for saving experiments
• --dataroot: Path for downloaded data (default: ./data)
• --seed: Random seed (default: 1234)

Sampling

Generate samples from a trained model:

python main.py --config configs/cifar10.yml --doc flowtrunc_experiment --sample --exp ./experiments

Testing

Evaluate model and compute FID:

python main.py --config configs/cifar10.yml --doc flowtrunc_experiment --test --exp ./experiments

📁 Project Structure

flow-diffusion/
├── README.md                 # This file
├── main.py                   # Main training/sampling script
├── configs/
│   └── cifar10.yml          # CIFAR-10 configuration
├── models/
│   ├── diffusion.py         # Diffusion model (U-Net)
│   └── __init__.py
├── runners/
│   ├── diffusion.py         # Training/sampling runner
│   └── __init__.py
├── functions/
│   ├── losses.py            # Loss functions
│   ├── denoising.py         # Denoising functions
│   └── __init__.py
├── datasets/
│   └── __init__.py          # Dataset loaders
├── analysis/
│   ├── tensorboard_reader.py  # TensorBoard log analysis
│   └── flow_plots/            # Training plots
├── docs/
│   └── flowtrunc.pdf        # Full technical report
└── transformer_flow.py      # TARFlow implementation

📖 Full Report

For complete technical details, mathematical derivations, and additional experimental results, please refer to the full technical report.