Application of the L2HMC algorithm to simulations in lattice QCD.

Overview

l2hmc-qcd CodeFactor

📊 Slides

📒 Example Notebook


Overview

The L2HMC algorithm aims to improve upon HMC by optimizing a carefully chosen loss function which is designed to minimize autocorrelations within the Markov Chain, thereby improving the efficiency of the sampler.

This work is based on the original implementation: brain-research/l2hmc/.

A detailed description of the L2HMC algorithm can be found in the paper:

Generalizing Hamiltonian Monte Carlo with Neural Network

by Daniel Levy, Matt D. Hoffman and Jascha Sohl-Dickstein.

Broadly, given an analytically described target distribution, π(x), L2HMC provides a statistically exact sampler that:

  • Quickly converges to the target distribution (fast burn-in).
  • Quickly produces uncorrelated samples (fast mixing).
  • Is able to efficiently mix between energy levels.
  • Is capable of traversing low-density zones to mix between modes (often difficult for generic HMC).

L2HMC for LatticeQCD

Goal: Use L2HMC to efficiently generate gauge configurations for calculating observables in lattice QCD.

A detailed description of the (ongoing) work to apply this algorithm to simulations in lattice QCD (specifically, a 2D U(1) lattice gauge theory model) can be found in doc/main.pdf.

l2hmc-qcd poster

Organization

Dynamics / Network

The base class for the augmented L2HMC leapfrog integrator is implemented in the BaseDynamics (a tf.keras.Model object).

The GaugeDynamics is a subclass of BaseDynamics containing modifications for the 2D U(1) pure gauge theory.

The network is defined in l2hmc-qcd/network/functional_net.py.

Network Architecture

An illustration of the leapfrog layer updating (x, v) --> (x', v') can be seen below.

leapfrog layer

Lattice

Lattice code can be found in lattice.py, specifically the GaugeLattice object that provides the base structure on which our target distribution exists.

Additionally, the GaugeLattice object implements a variety of methods for calculating physical observables such as the average plaquette, ɸₚ, and the topological charge Q,

Training

The training loop is implemented in l2hmc-qcd/utils/training_utils.py .

To train the sampler on a 2D U(1) gauge model using the parameters specified in bin/train_configs.json:

$ python3 /path/to/l2hmc-qcd/l2hmc-qcd/train.py --json_file=/path/to/l2hmc-qcd/bin/train_configs.json

Or via the bin/train.sh script provided in bin/.

Features

  • Distributed training (via horovod): If horovod is installed, the model can be trained across multiple GPUs (or CPUs) by:

    #!/bin/bash
    
    TRAINER=/path/to/l2hmc-qcd/l2hmc-qcd/train.py
    JSON_FILE=/path/to/l2hmc-qcd/bin/train_configs.json
    
    horovodrun -np ${PROCS} python3 ${TRAINER} --json_file=${JSON_FILE}

Contact


Code author: Sam Foreman

Pull requests and issues should be directed to: saforem2

Citation

If you use this code or found this work interesting, please cite our work along with the original paper:

@misc{foreman2021deep,
      title={Deep Learning Hamiltonian Monte Carlo}, 
      author={Sam Foreman and Xiao-Yong Jin and James C. Osborn},
      year={2021},
      eprint={2105.03418},
      archivePrefix={arXiv},
      primaryClass={hep-lat}
}
@article{levy2017generalizing,
  title={Generalizing Hamiltonian Monte Carlo with Neural Networks},
  author={Levy, Daniel and Hoffman, Matthew D. and Sohl-Dickstein, Jascha},
  journal={arXiv preprint arXiv:1711.09268},
  year={2017}
}

Acknowledgement

This research used resources of the Argonne Leadership Computing Facility, which is a DOE Office of Science User Facility supported under contract DE_AC02-06CH11357. This work describes objective technical results and analysis. Any subjective views or opinions that might be expressed in the work do not necessarily represent the views of the U.S. DOE or the United States Government. Declaration of Interests - None.

Hits

Stargazers over time

Comments
  • Remove upper bound on python_requires

    Remove upper bound on python_requires

    (I'm moving between meetings so can iterate on this more later, so excuse the very brief Issue for now).

    At the moment the project has an upper bound on python_requires

    https://github.com/saforem2/l2hmc-qcd/blob/2eb6ee63cc0c53b187e6d716f4c12f418c8b8515/setup.py#L165

    Assuming that you're intending l2hmc to be a library and not an application, then I would highly recommend removing this for the reasons summarized in Henry's detailed blog post on the subject.

    Congrats on getting l2hmc up on PyPI though! :snake: :rocket:

    opened by matthewfeickert 2
  • Alpha

    Alpha

    Pull upstream alpha branch into main

    Major changes

    • new src/ hierarchical module organization
    • Contains skeleton implementation of 4D SU(3) lattice gauge model
    • Framework independent configuration
      • Unified configuration system simplifies logic, same configs used for both tensorflow and pytorch experiments
      • Plan to be able to specify which backend to use through config option
    • Unified (and framework independent) configurations between tensorflow and pytorch implementations

    Note: This is still very much a WIP. Many existing features still need to be re-implemented / updated into new code in src/.

    Todo

    • [ ] Write unit tests
    • [ ] Use simple configs for end-to-end workflow test + integrate into CI
    • [ ] dynamic learning rate scheduling
    • [ ] Test 4D SU(3) numpy code
    • [ ] Write tensorflow and pytorch implementations of LatticeSU3 objects
    • [ ] Improved / simplified ( / trainable?) annealing schedule
    • [ ] Distributed training support
      • [ ] horovod
      • [ ] DDP for pytorch implementation
      • [ ] DeepSpeed from Microsoft??
    • [ ] Testing / inference logic
    • [ ] Automatic checkpointing
    • [ ] Metric logging
      • [ ] Tensorboard?
      • [ ] Sacred?
      • [ ] build custom dashboard? plot.ly?
    • [ ] Setup packaging / distribution through pip
    • [ ] Resolve issue
    opened by saforem2 1
  • Alpha

    Alpha

    opened by saforem2 1
  • Rich

    Rich

    General improvements, rewrote logging methods to use Rich for better formatting.

    • Adds dynamic (trainable) step size eps for each separate x and v updates, seems to generally increase the total energy towards the middle of the trajectory but it remains unclear if this corresponds to an improvement in the tunneling rate
    • Adds methods for calculating autocorrelations of the topological charge, as well as notebooks for generating the plots
    • Updates to the writeup in doc/main.pdf
    • Will likely be last changes to writeup before public release of official draft
    opened by saforem2 1
  • Dev

    Dev

    • Updates to README

    • Ability to load network with new training instance

    • Updates to doc/, removes old sections related to debugging the bias in the plaquette

    opened by saforem2 1
  • Saveable model

    Saveable model

    Complete rewrite of dynamics.xnet and dynamics.vnet models to use tf.keras.functional Models.

    Additional changes include:

    • Non-Compact Projection update for gauge fields
    • Ability to specify convolution structure to be prepended at beginning of gauge network
    opened by saforem2 1
  • Dev

    Dev

    Removes models/gauge_model.py entirely.

    Instead, a base dynamics class is implemented in dynamics/dynamics.py, and an example subclass is provided in dynamics/gauge_dynamics.py.

    opened by saforem2 1
  • Split networks

    Split networks

    Major rewrite of existing codebase.

    This pull request updates everything to be compatible with tensorflow >= 2.2 and removes a bunch of redundant legacy code.

    opened by saforem2 1
  • Dev

    Dev

    • Dynamics object is now compatible with tf >= 2.0
    • Running inference on trained model with tensorflow now creates identical graphs and summary files to numpy inference code
    • Inference with numpy now uses object oriented structure
    • Adds LaTeX + PDF documentation in doc/
    opened by saforem2 1
  • Cooley dev

    Cooley dev

    Adds new GaugeNetwork architecture as the default for training GaugeModel

    Additionally, replaces pickle with joblib for saving data as .z compressed files (as opposed to .pkl files).

    opened by saforem2 1
  • Testing

    Testing

    Implemented nnehmc_loss calculation for an alternative loss function using the approach suggested in https://infoscience.epfl.ch/record/264887/files/robust_parameter_estimation.pdf.

    This modified loss function can be chosen (instead of the standard loss described in the original paper) by passing --use_nnehmc_loss as a command line argument.

    opened by saforem2 1
  • Packaging and PyPI distribution?

    Packaging and PyPI distribution?

    As you've made a library and are using it as such:

    # snippet from toy_distributions.ipynb
    
    # append parent directory to `sys.path`
    # to load from modules in `../l2hmc-qcd/`
    module_path = os.path.join('..')
    if module_path not in sys.path:
        sys.path.append(module_path)
    
    # Local imports
    from utils.attr_dict import AttrDict
    from utils.training_utils import train_dynamics
    from dynamics.config import DynamicsConfig
    from dynamics.base_dynamics import BaseDynamics
    from dynamics.generic_dynamics import GenericDynamics
    from network.config import LearningRateConfig
    from config import (State, NetWeights, MonteCarloStates,
                        BASE_DIR, BIN_DIR, TF_FLOAT)
    
    from utils.distributions import (plot_samples2D, contour_potential,
                                     two_moons_potential, sin_potential,
                                     sin_potential1, sin_potential2)
    

    do you have any plans and/or interest in packaging it as a Python library so it can either be pip installed from GitHub or be distributed on PyPI?

    opened by matthewfeickert 5
Releases(0.12.0)
Owner
Sam Foreman
Computational science Postdoc at Argonne National Laboratory working on applying machine learning to simulations in lattice QCD.
Sam Foreman
Source code of the paper Meta-learning with an Adaptive Task Scheduler.

ATS About Source code of the paper Meta-learning with an Adaptive Task Scheduler. If you find this repository useful in your research, please cite the

Huaxiu Yao 16 Dec 26, 2022
Optimus: the first large-scale pre-trained VAE language model

Optimus: the first pre-trained Big VAE language model This repository contains source code necessary to reproduce the results presented in the EMNLP 2

314 Dec 19, 2022
Wileless-PDGNet Implementation

Wileless-PDGNet Implementation This repo is related to the following paper: Boning Li, Ananthram Swami, and Santiago Segarra, "Power allocation for wi

6 Oct 04, 2022
Revitalizing CNN Attention via Transformers in Self-Supervised Visual Representation Learning

Revitalizing CNN Attention via Transformers in Self-Supervised Visual Representation Learning This repository is the official implementation of CARE.

ChongjianGE 89 Dec 02, 2022
This repository contains the code used for the implementation of the paper "Probabilistic Regression with HuberDistributions"

Public_prob_regression_with_huber_distributions This repository contains the code used for the implementation of the paper "Probabilistic Regression w

David Mohlin 1 Dec 04, 2021
Compact Bilinear Pooling for PyTorch

Compact Bilinear Pooling for PyTorch. This repository has a pure Python implementation of Compact Bilinear Pooling and Count Sketch for PyTorch. This

Grégoire Payen de La Garanderie 234 Dec 07, 2022
[NeurIPS 2019] Learning Imbalanced Datasets with Label-Distribution-Aware Margin Loss

Learning Imbalanced Datasets with Label-Distribution-Aware Margin Loss Kaidi Cao, Colin Wei, Adrien Gaidon, Nikos Arechiga, Tengyu Ma This is the offi

Kaidi Cao 528 Jan 01, 2023
Code release for NeuS

NeuS We present a novel neural surface reconstruction method, called NeuS, for reconstructing objects and scenes with high fidelity from 2D image inpu

Peng Wang 813 Jan 04, 2023
Offical implementation for "Trash or Treasure? An Interactive Dual-Stream Strategy for Single Image Reflection Separation".

Trash or Treasure? An Interactive Dual-Stream Strategy for Single Image Reflection Separation (NeurIPS 2021) by Qiming Hu, Xiaojie Guo. Dependencies P

Qiming Hu 31 Dec 20, 2022
The audio-video synchronization of MKV Container Format is exploited to achieve data hiding

The audio-video synchronization of MKV Container Format is exploited to achieve data hiding, where the hidden data can be utilized for various management purposes, including hyper-linking, annotation

Maxim Zaika 1 Nov 17, 2021
Coursera - Quiz & Assignment of Coursera

Coursera Assignments This repository is aimed to help Coursera learners who have difficulties in their learning process. The quiz and programming home

浅梦 828 Jan 04, 2023
DeepLearning Anomalies Detection with Bluetooth Sensor Data

Final Year Project. Constructing models to create offline anomalies detection using Travel Time Data collected from Bluetooth sensors along the route.

1 Jan 10, 2022
Sentiment analysis translations of the Bhagavad Gita

Sentiment and Semantic Analysis of Bhagavad Gita Translations It is well known that translations of songs and poems not only breaks rhythm and rhyming

Machine learning and Bayesian inference @ UNSW Sydney 3 Aug 01, 2022
Learning Compatible Embeddings, ICCV 2021

LCE Learning Compatible Embeddings, ICCV 2021 by Qiang Meng, Chixiang Zhang, Xiaoqiang Xu and Feng Zhou Paper: Arxiv We cannot release source codes pu

Qiang Meng 25 Dec 17, 2022
Experiments on continual learning from a stream of pretrained models.

Ex-model CL Ex-model continual learning is a setting where a stream of experts (i.e. model's parameters) is available and a CL model learns from them

Antonio Carta 6 Dec 04, 2022
Unofficial implementation of Pix2SEQ

Unofficial-Pix2seq: A Language Modeling Framework for Object Detection Unofficial implementation of Pix2SEQ. Please use this code with causion. Many i

159 Dec 12, 2022
NAS-FCOS: Fast Neural Architecture Search for Object Detection (CVPR 2020)

NAS-FCOS: Fast Neural Architecture Search for Object Detection This project hosts the train and inference code with pretrained model for implementing

Ning Wang 180 Dec 06, 2022
Converts geometry node attributes to built-in attributes

Attribute Converter Simplifies converting attributes created by geometry nodes to built-in attributes like UVs or vertex colors, as a single click ope

Ivan Notaros 12 Dec 22, 2022
The fastest way to visualize GradCAM with your Keras models.

VizGradCAM VizGradCam is the fastest way to visualize GradCAM in Keras models. GradCAM helps with providing visual explainability of trained models an

58 Nov 19, 2022
ColossalAI-Examples - Examples of training models with hybrid parallelism using ColossalAI

ColossalAI-Examples This repository contains examples of training models with Co

HPC-AI Tech 185 Jan 09, 2023