Optimal Flow Matching: Learning Straight Trajectories in Just One Step

Authors: Nikita Kornilov, Petr Mokrov, Alexander Gasnikov, Aleksandr Korotin

NeurIPS 2024 | Conference PDF | Archive PDF | Plain Text | LLM Run Details

Reproducibility Variable Result LLM Response
Research Type Experimental We demonstrate the potential of OFM in the series of experiments and benchmarks ( 4).
Researcher Affiliation Academia Nikita Kornilov Skolkovo Institute of Science and Technology R.Center for AI, Innopolis University Moscow Institute of Physics and Technology kornilov.nm@phystech.edu Petr Mokrov Skolkovo Institute of Science and Technology petr.mokrov@skoltech.ru Alexander Gasnikov Innopolis University Moscow Institute of Physics and Technology Steklov Mathematical Institute of RAS gasnikov@yandex.ru Alexander Korotin Skolkovo Institute of Science and Technology Artificial Intelligence Research Institute a.korotin@skoltech.ru
Pseudocode Yes Algorithm. The Optimal Flow Matching pseudocode is presented in listing 1.
Open Source Code Yes The code of our OFM implementation and the conducted experiments is available at https://github.com/Jhomanik/Optimal-Flow-Matching.
Open Datasets Yes We run our OFM, FM based methods and OT solvers on OT Benchmark [33]. (...) FFHQ dataset [29].
Dataset Splits No The paper mentions 'train FFHQ sample' and 'test FFHQ sample' but does not specify clear train/validation/test splits with percentages, sample counts, or explicit use of a validation set.
Hardware Specification Yes In the Illustrative 2D experiment, the training takes 1.5 hours on a single 1080 ti GPU. (...) Totally, all the benchmark experiments (both with Ind and MB plan π) take 3 days on three A100 GPUs. In the ALAE experiment, the training stage lasts for 5 hours on a single 1080 ti GPU.
Software Dependencies No The paper mentions 'Py Torch implementation' and various libraries/optimizers (e.g., LBFGS, Adam, W2GN_ICNN, CPF_ICNN) but does not provide specific version numbers for any of them.
Experiment Setup Yes We aggregate the hyper-parameters of our Algorithm 1 and utilized ICNNs for different experiments in Table 3. In all our experiments as the Sub Opt optimizer we use LBFGS (torch.optim.LBFGS) with Ksub optimization steps and early stopping criteria based on gradient norm. (...) As the Opt optimizer we adopt Adam with learning rate lr and other hyperparameters set to be default.