Description

There is much interest in building ultra-fast samplers that map a density that is easy to sample from, typically, an n-dimensional normal to a desired n-dimensional density. One way to compute this mapping is to solve the reverse-time diffusion equation [1], which is an integro-differential equation. In Ref. [2], the integral in this equation is approximated using Monte Carlo integration where the integrand is averaged over N (~5K – 10K) points sampled from the desired distribution. Solving this equation is relatively slow, therefore, typically a neural network is trained to model the mapping from the normal to the desired density using training data generated by repeatedly solving the differential equation.

In this project, an alternative approach is investigated: modeling the solution to the differential equation using a physics-informed neural network (PINN) [3]. There is a large upfront cost in training the PINN, but this is subsequently amortized over the fast sampling using the PINN. Various neural network architectures for the PINN will be investigated.

Duration

Total project length: 175/350 hours.

Difficulty Level

• Intermediate/Advanced

Task ideas

• Map a 3D zero mean, unit variance, diagonal normal to a 3D non-Gaussian density using a PINN. The inputs to the PINN are t, x, y, z — that is, the reverse time t ∈ [1, 0] and a point sampled from the 3D normal. The output of the PINN is the vector solution u(t, x, y, z). Since the PINN is conditioned on x, y, z, during training the points can be sampled from any convenient distribution, including quasi-random sampling such as Sobol sampling. (Of course, when used we must sample from a 3D normal.)
• Repeat with increasingly complex 3D non-Gaussian densities.
• Optional: Apply what has been learned from 1 and 2 to build a fast calorimeter simulator. Use Dataset 1 from the Fast Calorimeter Simulation Challenge 2022 [4].
• Publish the results in an ML paper.

Expected results

• Trained graph-based jet classifier
• Benchmarks on selected datasets

Test

• Using PyTorch, solve the damped harmonic oscillator [5] using a PINN. Choose fixed initial conditions:
  x(0) = x₀, dx/dz(0) = v₀, with x₀ = 0.7 and v₀ = 1.2.
  Condition the PINN on damping ratios in the range ξ = 0.1 to 0.4.
  Solve on the domain z ∈ [0, 20]:
  d²x/dz² + 2ξ·dx/dz + x = 0

Requirements

• Experience with numerical solution of ordinary differential equations.
• Familiarity with PyTorch.

Difficultly Level

Advanced

Mentors

• Harrison B. Prosper (Florida State University)
• Pushpalatha Bhat (Fermilab)
• Sergei Gleyzer (University of Alabama)

Links

1. Cheng Lu†, Yuhao Zhou†, Fan Bao†, Jianfei Chen†, Chongxuan Li‡, Jun Zhu, DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps, arXiv:2206.00927v3, 13 Oct 2022.

2. Yanfang Lui, Minglei Yang, Zezhong Zhang, Feng Bao, Yanzhao Cao, and Guannan Zhang, Diffusion-Model-Assisted Supervised Learning of Generative Models for Density Estimation, arXiv:2310.14458v1, 22 Oct 2023.

3. S. Cuomo et al., Scientific Machine Learning through Physics-Informed Neural Networks: Where we are and What’s next, https://doi.org/10.48550/arXiv.2201.05624.

4. https://calochallenge.github.io/homepage/

5. https://en.wikipedia.org/wiki/Harmonic_oscillator

Please DO NOT contact mentors directly by email. Instead, please email ml4-sci@cern.ch with Project Title and include your CV and test results. The mentors will then get in touch with you.

Corresponding Project

• GENIE

Participating Organizations

• Alabama
• FSU
• Fermilab

---