Hamiltonian Score Matching and Generative Flows

TL;DR

This work explores the potential of deliberately designing force fields for Hamiltonian ODEs, introducing Hamiltonian velocity predictors (HVPs) as a tool for score matching and generative models, and presents two innovations constructed with HVPs: Hamiltonian Score Matching (HSM) and Hamiltonian Generative Flows (HGFs).

Abstract

Classical Hamiltonian mechanics has been widely used in machine learning in the form of Hamiltonian Monte Carlo for applications with predetermined force fields. In this work, we explore the potential of deliberately designing force fields for Hamiltonian ODEs, introducing Hamiltonian velocity predictors (HVPs) as a tool for score matching and generative models. We present two innovations constructed with HVPs: Hamiltonian Score Matching (HSM), which estimates score functions by augmenting data via Hamiltonian trajectories, and Hamiltonian Generative Flows (HGFs), a novel generative model that encompasses diffusion models and flow matching as HGFs with zero force fields. We showcase the extended design space of force fields by introducing Oscillation HGFs, a generative model inspired by harmonic oscillators. Our experiments validate our theoretical insights about HSM as a novel score matching metric and demonstrate that HGFs rival leading generative modeling techniques.

Figures (6)

• Figure 1: Results of training HSM on a Gaussian mixture. The score vector field faithfully recovers gradients of the density. The optimal velocity predictor is zero everywhere.
• Figure 2: Evolution of various HGFs in joint coordinate-velocity space from $t=0$ (blue) to $t=T$ (red) with trajectories (black). Data distribution $\pi(x)=0.4 * \mathcal{N}(-2,1)+ 0.6* \mathcal{N}(2,1)$. Diffusion models and flow matching have zero force fields, i.e. the velocity does not change. Diffusion models do not converge in finite time (here, $T=3$). The coupled distribution in FM allow for a convergence for $T=1$. Both distort the joint distribution. Oscillation HGFs only rotate the distribution.
• Figure 3: Empirical investigation of Hamiltonian score discrepancy (HSD). (a) The Taylor approximation is a good approximation. (b) Hamiltonian score discrepancy is strongly correlated with explicit score matching loss. (c) Signal-to-noise ratio is significantly better for HSM vs DSM for low $\sigma$.
• Figure 4: Image generation examples based on Oscillation HGFs for FFHQ.
• Figure 5: Data distribution (left) and velocity distribution (right) used for Reflection HGFs as initial distribution. With the above starting conditions, a reflection (=”infinite force”) at the boundaries of the domain is used to simulate trajectories forward (this can be computed in closed form in a simulation-free manner).

Theorems & Definitions (11)

• Theorem 1
• Proposition 1
• Theorem 2
• Proposition 2: Taylor approximation of HSM loss
• Proposition 3
• Proposition 4
• proof
• Proposition 5: Symplectic property
• proof
• Lemma 1