{"id":13478491,"url":"https://github.com/lollcat/fab-torch","last_synced_at":"2026-01-14T17:36:05.225Z","repository":{"id":37779405,"uuid":"430081955","full_name":"lollcat/fab-torch","owner":"lollcat","description":"Flow Annealed Importance Sampling Bootstrap (FAB). ICLR 2023. ","archived":false,"fork":false,"pushed_at":"2024-03-11T19:02:51.000Z","size":279039,"stargazers_count":66,"open_issues_count":3,"forks_count":8,"subscribers_count":2,"default_branch":"master","last_synced_at":"2025-12-22T11:16:25.573Z","etag":null,"topics":["annealed-importance-sampling","boltzmann-distribution","boltzmann-generator","normalizing-flow"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/lollcat.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null}},"created_at":"2021-11-20T11:16:47.000Z","updated_at":"2025-11-19T16:22:44.000Z","dependencies_parsed_at":"2024-01-16T07:22:07.000Z","dependency_job_id":"a0903d05-e90c-452b-88b9-ddcdb8117f47","html_url":"https://github.com/lollcat/fab-torch","commit_stats":null,"previous_names":[],"tags_count":3,"template":false,"template_full_name":null,"purl":"pkg:github/lollcat/fab-torch","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lollcat%2Ffab-torch","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lollcat%2Ffab-torch/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lollcat%2Ffab-torch/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lollcat%2Ffab-torch/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/lollcat","download_url":"https://codeload.github.com/lollcat/fab-torch/tar.gz/refs/heads/master","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lollcat%2Ffab-torch/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":28428766,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-01-14T16:38:47.836Z","status":"ssl_error","status_checked_at":"2026-01-14T16:34:59.695Z","response_time":107,"last_error":"SSL_connect returned=1 errno=0 peeraddr=140.82.121.5:443 state=error: unexpected eof while reading","robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":false,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["annealed-importance-sampling","boltzmann-distribution","boltzmann-generator","normalizing-flow"],"created_at":"2024-07-31T16:01:57.795Z","updated_at":"2026-01-14T17:36:05.211Z","avatar_url":"https://github.com/lollcat.png","language":"Python","funding_links":[],"categories":["Python"],"sub_categories":[],"readme":"# Flow Annealed Importance Sampling Bootstrap (FAB)\n\n## Overview\nCode for the paper [Flow Annealed Importance Sampling Bootstrap (FAB)](https://arxiv.org/abs/2208.01893).\n\n**FAB in JAX**: See the JAX implementation of the FAB algorithm in the [fab-jax](https://github.com/lollcat/fab-jax) repo. The `fab-jax` code is cleaner, faster and easier to use - hence we recommend it over the `fab-torch` code. Additionally, the `fab-jax` code applies FAB to some new problems, including the commonly used, challenging, 1600 dimensional log Gaussian\nCox process [[Møller et al., 1998](https://api.semanticscholar.org/CorpusID:120543073), [Arbel et al., 2021](https://arxiv.org/abs/2102.07501), [Mathews et al., 2022](https://arxiv.org/abs/2201.13117), [Zhang et at., 2023](https://arxiv.org/abs/2310.02679)]. \n\n\nSee [About the code](#about-the-code) for further details on how to use the FAB codebase on new problems.\nPlease contact us if you need any help running the code and replicating our experiments. \n\n## Methods of Installation\n\nThe  package can be installed via pip by navigating in the repository directory and running\n\n```\npip install --upgrade .\n```\n\nTo run the alanine dipeptide experiments, you will need to install the [OpenMM Library](http://openmm.org/)\nas well as [`openmmtools`](https://openmmtools.readthedocs.io/en/stable/). This can be done via conda.\n\n```\nconda install -c conda-forge openmm openmmtools\n```\n\n## Experiments\n\n**NB**: See README within experiments/{problem-name} for further details on training and evaluation for each problem. \n\n**NB**: Quickstart notebooks are simply to get up and running with the code with some visualisation of results after \na little bit of training. To replicate the results from the paper run the python commands described below.\n\n### Gaussian Mixture Model\n\nQuickstart (**NB** just for getting started, to replicate results from paper see python command below)\n\n\u003ca href=\"https://colab.research.google.com/github/lollcat/fab-torch/blob/master/experiments/gmm/fab_gmm.ipynb\" target=\"_parent\"\u003e\u003cimg src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/\u003e\u003c/a\u003e\n\nFor this problem we use a mixture of 40 two dimensional Gaussian distributions. \nThis allows for easy visualisation of the various methods for training the flow.\nWe provide a [colab notebook](experiments/gmm/fab_gmm.ipynb) with an example of training a flow on the GMM problem, \ncomparing FAB to training a flow with KL divergence minimisation.\nThis can be run in a short period of time (10 min) and provides a clear visualisation of how FAB is \nable to discover new modes and fit them. \n\nTo run the experiment for the FAB with a prioritised replay buffer (for the first seed), use the following command:\n```\npython experiments/gmm/run.py training.use_buffer=True training.prioritised_buffer=True\n```\nTo run the full set of experiments see the [README](experiments/gmm/README.md) for the GMM experiments. \n\nThe below plot shows samples from various trained models, with the GMM problem target contours in the background.\n![Gaussian Mixture Model samples vs contours](experiments/gmm/plots/MoG.png)\n\n### Many Well distribution\nQuickstart (**NB** just for getting started, to replicate results from paper see python command below)\n\n\u003ca href=\"https://colab.research.google.com/github/lollcat/fab-torch/blob/master/experiments/many_well/fab_many_well.ipynb\" target=\"_parent\"\u003e\u003cimg src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/\u003e\u003c/a\u003e\n\n\nThe Many Well distribution is made up of multiple repeats of the Double Well distribution, \nfrom the [original Boltzmann generators paper](https://www.science.org/doi/10.1126/science.aaw1147).\n\nWe provide a [colab notebook](experiments/many_well/fab_many_well.ipynb) comparing FAB to training a flow via KL divergence minimisation, on the \n6 dimensional Many Well problem, where the difference between the two methods is apparent after a \nshort (\u003c5 min) training period. This experiment can be run locally on a laptop using just CPU. \n\nAdditionally, we provide the colab notebook \n\u003ca href=\"https://colab.research.google.com/github/lollcat/fab-torch/blob/master/demo/many_well.ipynb\" target=\"_parent\"\u003e\u003cimg src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/\u003e\u003c/a\u003e\nwhich demos inference with the flow trained with FAB (+prioritised buffer) on the 32 dim Many Well problem. \n\nTo run the experiment for the FAB with a prioritised replay buffer (for the first seed) on the \n32 dimensional Many Well problem, use the following command:\n```\npython experiments/many_well/run.py training.use_buffer=True training.prioritised_buffer=True\n```\nTo run the full set of experiments see the [README](experiments/many_well/README.md) for the Many Well experiments. \n\nThe below plot shows samples for our model (FAB) vs training a flow by reverse KL divergence \nminimisation, with the Many Well problem target contours in the background. \nThis visualisation is for the marginal pairs of the distributions for the first four elements of the x.\n![Many Well distribution FAB vs training by KL divergence minimisation](experiments/many_well/plots/many_well.png)\n\n### Alanine dipeptide\n\nIn our final experiment, we approximate the Boltzmann distribution of alanine dipeptide in an \nimplicit solvent, which is a molecule with 22 atoms and a popular model system. The molecule\nis visualized in the figure below. The right figure shows the probability density of for the\ndihedral angle $\\phi$ comparing the ground truth, which was obtrained with a molecular dynamics\n(MD) simulation, the models trained with our method as well as maximum likelihood on MD samples.\n\n![Alanine dipeptide and its dihedral angles; Comparison of probability densities](experiments/aldp/plots/aldp_phi.png)\n\nFurthermore, we compared the Ramachandran plots of the different methods in the following figure.\n\n![Ramachandran plot of alanine dipeptide](experiments/aldp/plots/ramachandran.png)\n\nThe weights for the flow model trained with FAB are available on [huggingface](https://huggingface.co/VincentStimper/fab). \nAdditionally, we provide the colab notebook \n\u003ca href=\"https://colab.research.google.com/github/lollcat/fab-torch/blob/master/demo/aldp.ipynb\" target=\"_parent\"\u003e\u003cimg src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/\u003e\u003c/a\u003e\nwhich demos inference with this trained model.\n\nTo reproduce our experiment, use the [`experiments/aldp/train.py`](experiments/aldp/train.py) script.\nThe respective configuration files are located in [`experiments/aldp/config`](experiments/aldp/config).\nWe used the seeds 0, 1, and 2 in our runs.\n\nThe data used to evaluate our models and to train the flow model with maximum likelihood is provided \non [Zenodo](https://zenodo.org/record/6993124#.YvpugVpBy5M). If you want to use the configuration files\nin [`experiments/aldp/config`](experiments/aldp/config) as is, you should put the data in the \n[`experiment/aldp/data`](experiments/aldp/data) folder.\n\n[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.6993124.svg)](https://doi.org/10.5281/zenodo.6993124)\n\n\n### About the code \nThe main FAB loss can be found in [core.py](fab/core.py), and we provide a simple training loop to \ntrain a flow with this loss (or other flow - loss combinations that meet the spec) in [train.py](fab/train.py) \nThe FAB training algorithm **with** the prioritised buffer can be found in [train_with_prioritised_buffer.py](fab/train_with_prioritised_buffer.py). \nAdditionally, we provide the code for running the SNR/dimensionality analysis with p and q set to independent Gaussians.\nin the [fab-jax-old](https://github.com/lollcat/fab-jax-old) repository.\nFor training the CRAFT model on the GMM problem we forked the \n[Annealed Flow Transport repository](https://github.com/deepmind/annealed_flow_transport). \nThis fork may be found [here](https://github.com/lollcat/annealed_flow_transport), and may be used for training the CRAFT model.\n\nAs we are still adding improvements to the efficiency and stability of the code, make sure you use the latest version.\nAdditionally, if you spot any areas of the code that could be improved then make an issue and we will be more \nthan happy to fix it.\nFor the version of the code that was used in the paper see our [releases](https://github.com/lollcat/fab-torch/releases).\n\n#### Applying FAB to a new problem:\nThe most important thing to get right when applying FAB to a given problem is to make sure that AIS is returning reasonable samples,\nwhere by reasonable we just mean that the samples from AIS are closer to the target than the flow. \nSimply visualising the samples from the flow and AIS provides a good check for whether this is the case.\nMaking sure that the transition kernel (e.g. HMC) is working well (e.g. has well tuned step size) is key for AIS to work well.\n\nAn additional source of instability can be if the target energy function gives spurious values to points that have extreme values.\nFor example, evaluating the density of a zero-mean unit variance Gaussian on a point that has a value of 100 will give a spurious values. \nOne can fix this by manually setting the log prob of the target to be -inf for regions that are \nknown to be far outside of where samples from the target lie. \n\nFeel free to contact us if you would like any help getting FAB to work nicely!\n\n\n### Normalizing Flow Libraries\nWe offer a simple wrapper that allows for various normalising flow libraries to be plugged into \nthis repository. The main library we rely on is \n[normflows](https://github.com/VincentStimper/normalizing-flows). \n\n\n## Citation\n\nIf you use this code in your research, please cite it as:\n\n\u003e Laurence I. Midgley, Vincent Stimper, Gregor N. C. Simm, Bernhard Schölkopf, José Miguel Hernández-Lobato.\n\u003e Flow Annealed Importance Sampling Bootstrap. The Eleventh International Conference on Learning Representations. 2023.\n\n**Bibtex**\n\n```\n@inproceedings{\nmidgley2023flow,\ntitle={Flow Annealed Importance Sampling Bootstrap},\nauthor={Laurence Illing Midgley and Vincent Stimper and Gregor N. C. Simm and Bernhard Sch{\\\"o}lkopf and Jos{\\'e} Miguel Hern{\\'a}ndez-Lobato},\nbooktitle={The Eleventh International Conference on Learning Representations },\nyear={2023},\nurl={https://openreview.net/forum?id=XCTVFJwS9LJ}\n}\n```\n\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flollcat%2Ffab-torch","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Flollcat%2Ffab-torch","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flollcat%2Ffab-torch/lists"}