This is the official JAX implementation for the paper Bidirectional Normalizing Flow: From Data to Noise and Back:

@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}
}

BiFlow is a new framework for Normalizing Flows that eliminates the requirement for an exact analytic inverse. It enables direct generation, greater modeling flexibility, and computational efficiency, all within a single unified approach.

For training code and checkpoints of iTARFlow, please refer to this branch.

Run install.sh to install the dependencies (JAX+TPUs). Log in to WandB to track your experiments if needed.

bash scripts/install.sh
wandb login YOUR_WANDB_API_KEY

You can quickly verify your setup with our provided checkpoint.

Note that slight differences in FID/IS may arise due to different computation setups. Our results are computed from TPU v5p-64.

1. Download the checkpoint and FID stats:

   • Download the pre-trained checkpoint from the table above.
   • Download the FID stats file from here. Our FID stats is computed on TPU and JAX, which may slightly differ from those computed on GPU and PyTorch. You can also compute FID stats using prepare_dataset.py if needed.

2. Unzip the checkpoint:

   unzip <downloaded_checkpoint.zip> -d <your_ckpt_dir>

   Replace <downloaded_checkpoint.zip> and <your_ckpt_dir> with your actual paths.

3. Set up the config:

   • Set load_from in configs/eval_config.yml to the path of <your_ckpt_dir>.
   • Set fid.cache_ref to the path of the downloaded FID stats file.

4. Launch evaluation:

   bash scripts/eval.sh JOB_NAME

Before training, you need to prepare the ImageNet dataset and compute latent representations:

Download the ImageNet dataset and extract it to your desired location. The dataset should have the following structure:

imagenet/
├── train/
│   ├── n01440764/
│   ├── n01443537/
│   └── ...
└── val/
    ├── n01440764/
    ├── n01443537/
    └── ...

Update the data paths in scripts/prepare_data.sh:

IMAGENET_ROOT="YOUR_IMGNET_ROOT"
OUTPUT_DIR="YOUR_OUTPUT_DIR"
LOG_DIR="YOUR_LOG_DIR"

Run the data preparation script to compute latent representations:

IMAGE_SIZE=256 COMPUTE_LATENT=True bash ./scripts/prepare_data.sh

Parameters:

• IMAGE_SIZE: Image size for processing (256, 512, or 1024). Latent sizes will be 32x32, 64x64, or 128x128 respectively.
• COMPUTE_LATENT: Whether to compute and save the latent dataset (True/False)
• COMPUTE_FID: Whether to compute FID statistics (True/False)

The script will:

• Encode ImageNet images to latent representations using a VAE model
• Save the latent dataset to OUTPUT_DIR/
• Compute FID statistics and save to OUTPUT_DIR/imagenet_{IMAGE_SIZE}_fid_stats.npz
• Log progress to LOG_DIR/$USER/

Run the following commands to launch training:

bash scripts/train.sh JOB_NAME

Note: Update the environment variables in scripts/train.sh before running:

• DATA_ROOT: Path to your prepared data directory
• LOG_DIR: Path where to save training logs

The training system uses two config files:

• configs/default.py - Base configuration with all default hyperparameters
• configs/train_config.yml - Model-specific overrides for BiFlow-B/2 training

The system merges these files, allowing you to customize only the parameters you need.

To create a custom experiment:

1. Create a new config file (e.g., configs/my_exp_config.yml)
2. Update the launch script to use your config:

   # In launch.sh, change the config line to:
   --config=configs/load_config.py:my_exp

Example custom config:

training:
    num_epochs: 80                  # Train for fewer epochs

model:
    model_str: BiFlow_B_2               # Use BiFlow-B/2 model

for more details on configuration options, refer to configs/default.py and configs/train_config.yml.

This repo is under the MIT license. See LICENSE for details.

This repository is a collaborative effort by Kaiming He, Hanhong Zhao, Zhicheng Jiang, Xianbang Wang, Qiao Sun and Yiyang Lu, developed in support of several research projects, including MeanFlow, improved MeanFlow.

We gratefully acknowledge the Google TPU Research Cloud (TRC) for granting TPU access. We hope this work will serve as a useful resource for the open-source community.