Merge pull request #69 from atong01/forest_flow

Adding Forest-Flow: Generating and Imputing Tabular Data via Diffusion and Flow-based Gradient-Boosted Trees

Author: kilianFatras

README.md

@@ -23,7 +23,7 @@
 
 ## Description
 
-Conditional Flow Matching (CFM) is a fast way to train continuous normalizing flow (CNF) models. CFM is a simulation-free training objective for continuous normalizing flows that allows conditional generative modeling and speeds up training and inference. CFM's performance closes the gap between CNFs and diffusion models. To spread its use within the machine learning community, we have built a library focused on Flow Matching methods: TorchCFM. TorchCFM is a library showing how Flow Matching methods can be trained and used to deal with image generation, single-cell dynamics and (soon) SO(3) data and tabular data.
+Conditional Flow Matching (CFM) is a fast way to train continuous normalizing flow (CNF) models. CFM is a simulation-free training objective for continuous normalizing flows that allows conditional generative modeling and speeds up training and inference. CFM's performance closes the gap between CNFs and diffusion models. To spread its use within the machine learning community, we have built a library focused on Flow Matching methods: TorchCFM. TorchCFM is a library showing how Flow Matching methods can be trained and used to deal with image generation, single-cell dynamics, tabular data and soon SO(3) data.
 
 <p align="center">
 <img src="assets/169_generated_samples_otcfm.png" width="600"/>
@@ -107,8 +107,8 @@ List of implemented papers:
 - Building Normalizing Flows with Stochastic Interpolants (Albergo et al. 2023a) [Paper](https://openreview.net/forum?id=li7qeBbCR1t)
 - Action Matching: Learning Stochastic Dynamics From Samples (Neklyudov et al. 2022) [Paper](https://arxiv.org/abs/2210.06662) [Code](https://github.com/necludov/jam)
 - Concurrent work to our OT-CFM method: Multisample Flow Matching: Straightening Flows with Minibatch Couplings (Pooladian et al. 2023) [Paper](https://arxiv.org/abs/2304.14772)
-- Soon: SE(3)-Stochastic Flow Matching for Protein Backbone Generation (Bose et al.) [paper](https://arxiv.org/abs/2310.02391)
-- Soon: Generating and Imputing Tabular Data via Diffusion and Flow-based Gradient-Boosted Trees (Jolicoeur-Martineau et al.) [paper](https://arxiv.org/abs/2309.09968)
+- Generating and Imputing Tabular Data via Diffusion and Flow-based Gradient-Boosted Trees (Jolicoeur-Martineau et al.) [Paper](https://arxiv.org/abs/2309.09968) [Code](https://github.com/SamsungSAILMontreal/ForestDiffusion)
+- Soon: SE(3)-Stochastic Flow Matching for Protein Backbone Generation (Bose et al.) [Paper](https://arxiv.org/abs/2310.02391)
 
 ## How to run
 
@@ -155,7 +155,7 @@ python -m ipykernel install --user --name=torchcfm
 
 ## Project Structure
 
-The directory structure of a new project looks like this:
+The directory structure looks like this:
 
 ```
 

examples/notebooks/mnist_example.ipynb

@@ -253,6 +253,8 @@
    "outputs": [],
    "source": [
     "# follows example from https://github.com/google-research/torchsde/blob/master/examples/cont_ddpm.py\n",
+    "\n",
+    "\n",
     "class SDE(torch.nn.Module):\n",
     "    noise_type = \"diagonal\"\n",
     "    sde_type = \"ito\"\n",

examples/notebooks/training-8gaussians-to-moons.ipynb

@@ -843,6 +843,8 @@
    ],
    "source": [
     "# %%time\n",
+    "\n",
+    "\n",
     "class MLP2(torch.nn.Module):\n",
     "    def __init__(self, dim, out_dim=None, w=64, time_varying=False):\n",
     "        super().__init__()\n",

examples/tabular/README.md

@@ -0,0 +1,20 @@
+# Forest-Flow experiment on the Iris dataset using TorchCFM
+
+This notebook is a self-contained example showing how to train the novel Forest-Flow method to generate tabular data [(Jolicoeur-Martineau et al. 2023)](https://arxiv.org/abs/2309.09968). The idea behind Forest-Flow is to **learn Independent Conditional Flow-Matching's vector field with XGBoost models** instead of neural networks. The motivation is that it is known that Forests work currently better than neural networks on Tabular data tasks. This idea comes with some difficulties, for instance how to approximate Flow Matching's loss, and this notebook shows how to do it on a minimal example. The method, its training procedure and the experiments are described in [(Jolicoeur-Martineau et al. 2023)](https://arxiv.org/abs/2309.09968). The full code can be found [here](https://github.com/SamsungSAILMontreal/ForestDiffusion).
+
+To run our jupyter notebooks, installing our package:
+
+```bash
+cd ../../
+
+# install torchcfm
+pip install -e '.[forest-flow]'
+
+# install ipykernel
+conda install -c anaconda ipykernel
+
+# install conda env in jupyter notebook
+python -m ipykernel install --user --name=torchcfm
+
+# launch our notebooks with the torchcfm kernel
+```