Normalizing flow neural networks by JKO scheme
Authors: Chen Xu, Xiuyuan Cheng, Yao Xie
NeurIPS 2023 | Conference PDF | Archive PDF | Plain Text | LLM Run Details
| Reproducibility Variable | Result | LLM Response |
|---|---|---|
| Research Type | Experimental | Experiments with synthetic and real data show that the proposed JKO-i Flow network achieves competitive performance compared with existing flow and diffusion models at a significantly reduced computational and memory cost. |
| Researcher Affiliation | Academia | Chen Xu School of Industrial and Systems Engineering Georgia Tech Xiuyuan Cheng Department of Mathematics Duke University Yao Xie School of Industrial and Systems Engineering Georgia Tech |
| Pseudocode | Yes | Algorithm 1 Block-wise JKO-i Flow training |
| Open Source Code | Yes | Codes are available at https://github.com/hamrel-cxu/JKO-i Flow. |
| Open Datasets | Yes | The four high-dimensional real datasets (POWER, GAS, MINIBOONE, BSDS300) come from the University of California Irvine (UCI) machine learning data repository, and we follow the pre-processing procedures of [Papamakarios et al., 2017]. |
| Dataset Splits | No | POWER: 1.85M training sample and 205K test sample. GAS: 1M training sample, 100K test sample. MINIBOONE: 32K training sample, 3.6K test sample. BSDS300: 1.05M training sample, 250K test sample. (No explicit mention of validation set size or split.) |
| Hardware Specification | Yes | Specifically, on a single A100 GPU, our experiments took 24 hours on CIFAR10 and 30 hours on Imagenet-32. |
| Software Dependencies | No | All experiments are conducted using Py Torch [Paszke et al., 2019] and Py Torch Geometric [Fey and Lenssen, 2019]. (Specific version numbers for PyTorch and PyTorch Geometric are not stated.) |
| Experiment Setup | Yes | The initial schedule is by setting h0 = 0.75, ρ = 1.2. We set hmax = 5 for rose, checkerboard, and Olympic rings, and hmax = 3 for fractal tree. For reparametrization and refinement, we use η = 0.5 in the reparameterization iterations for all examples. The learning rate is 5e-3. On MNIST, we fix the batch size to be 2000 during training, and we train 15K batches per JKO-i Flow block. We fix the learning rate to be 1e-3. |