Bidirectional Normalizing Flow: From Data to Noise and Back

This repo contains code that accompanies the research paper, Bidirectional Normalizing Flow: From Data to Noise and Back.

Setup

The code is tested on Python3.10, and install dependencies with:

pip install -r requirements.txt

Preparing datasets

Download the datasets you want to experiment with:

• Imagenet
• Imagenet64
• AFHQ

Save the training files only in data/<dataset>/<category>/<filename>, the code does not use the validation/test files.

Compute and save stats for the true data distribution

# Files are saved in ./data
torchrun --standalone --nproc_per_node=8 prepare_fid_stats.py --dataset=imagenet64 --img_size=64  # Unconditional
torchrun --standalone --nproc_per_node=8 prepare_fid_stats.py --dataset=imagenet --img_size=64    # Conditional
torchrun --standalone --nproc_per_node=8 prepare_fid_stats.py --dataset=imagenet --img_size=128   # Conditional
torchrun --standalone --nproc_per_node=8 prepare_fid_stats.py --dataset=afhq --img_size=256       # Conditional

Note: To run on a single GPU, replace torchrun with python like this:

python prepare_fid_stats.py --dataset=imagenet --img_size=64    # Conditional

Training

Toy experiments on MNIST, this can be run locally with MPS (Macbooks) or CPU.

jupyter notebook train_local.ipynb

Reproducing results from the paper

# Unconditional ImageNet64 density modelling (16 GPUs, fp32)
torchrun --standalone --nproc_per_node=8 train.py --dataset=imagenet64 --img_size=64 --channel_size=3\
  --patch_size=2 --channels=768 --blocks=8 --layers_per_block=8\
  --noise_type=uniform --batch_size=256 --epochs=100 --lr=1e-4 --nvp\
  --sample_freq=1000 --logdir=runs/imagenet64-uncond-bpd

# Unconditional ImageNet64 generation(8 GPUs)
torchrun --standalone --nproc_per_node=8 train.py --dataset=imagenet64 --img_size=64 --channel_size=3\
  --patch_size=2 --channels=768 --blocks=8 --layers_per_block=8\
  --noise_std=0.05 --batch_size=256 --epochs=200 --lr=1e-4 --nvp\
  --sample_freq=5 --logdir=runs/imagenet64-uncond

# Conditional ImageNet64 (8 GPUs)
torchrun --standalone --nproc_per_node=8 train.py --dataset=imagenet --img_size=64 --channel_size=3\
  --patch_size=2 --channels=768 --blocks=8 --layers_per_block=8\
  --noise_std=0.05 --batch_size=256 --epochs=200 --lr=1e-4 --nvp --cfg=0 --drop_label=0.1\
  --sample_freq=5 --logdir=runs/imagenet64-cond

# Conditional ImageNet128 (need to run on 4 nodes, 32 GPUs total)
torchrun --standalone --nproc_per_node=8 train.py --dataset=imagenet --img_size=128 --channel_size=3\
  --patch_size=4 --channels=1024 --blocks=8 --layers_per_block=8\
  --noise_std=0.15 --batch_size=768 --epochs=320 --lr=1e-4 --nvp --cfg=0 --drop_label=0.1\
  --sample_freq=20 --logdir=runs/imagenet128-cond

# AFHQ (8 GPUs)
torchrun --standalone --nproc_per_node=8 train.py --dataset=afhq --img_size=256 --channel_size=3\
  --patch_size=8 --channels=768 --blocks=8 --layers_per_block=8\
  --noise_std=0.07 --batch_size=256 --epochs=4000 --lr=1e-4 --nvp --cfg=0 --drop_label=0.1\
  --sample_freq=200 --logdir=runs/afhq256

For single-GPU

python train.py --dataset=imagenet64 --img_size=64 --channel_size=3\
  --patch_size=2 --channels=768 --blocks=8 --layers_per_block=8\
  --noise_std=0.05 --batch_size=32 --epochs=200 --lr=1e-4 --nvp\
  --sample_freq=5 --logdir=runs/imagenet64-uncond
# etc...

Sampling

Use the notebook to generate samples from a model checkpoint. Inside the notebook is an option to download a pretrained checkpoint on AFHQ.

jupyter notebook sample.ipynb

Evaluating BPD

python evaluate_bpd.py --dataset=imagenet64 --img_size=64 --channel_size=3\
  --patch_size=2 --channels=768 --blocks=8 --layers_per_block=8\
  --ckpt_file=runs/imagenet64-uncond-bpd/imagenet64_model_2_768_8_8_uniform.pth

Evaluating FID

Multi-GPU (8 GPUs)

# Conditional ImageNet64, samples saved in runs/imagenet64-cond/eval
torchrun --standalone --nproc_per_node=8 evaluate_fid.py --dataset=imagenet --img_size=64 --channel_size=3\
  --patch_size=2 --channels=768 --blocks=8 --layers_per_block=8\
  --noise_std=0.05 --cfg=2.3 --nvp --batch_size=1024\
  --ckpt_file=runs/imagenet64-cond/imagenet_model_2_768_8_8_0.05.pth\
  --logdir=runs/imagenet64-cond/eval

For single-GPU

# Conditional ImageNet64, samples saved in runs/imagenet64-cond/eval
python evaluate_fid.py --dataset=imagenet --img_size=64 --channel_size=3\
  --patch_size=2 --channels=768 --blocks=8 --layers_per_block=8\
  --noise_std=0.05 --cfg=2.3 --nvp --batch_size=32\
  --ckpt_file=runs/imagenet64-cond/imagenet_model_2_768_8_8_0.05.pth\
  --logdir=runs/imagenet64-cond/eval

BibTeX

@article{lu2025bidirectional,
  title={Bidirectional Normalizing Flow: From Data to Noise and Back},
  author={Lu, Yiyang and Sun, Qiao and Wang, Xianbang and Jiang, Zhicheng and Zhao, Hanhong and He, Kaiming},
  journal={arXiv preprint arXiv:2512.10953},
  year={2025}
}