🔥❤️🔥🥳🍾
New paper alert!
❤️🔥🔥🍾🥳
Quite stoked about this recent collaboration with
@msalbergo
and
@evdende2
! We introduce a new class of generative models that we call **flow map matching models**, which learn the flow map of a probability flow ODE.
Very excited to announce that in September 2024 I will be joining the faculty at
@CarnegieMellon
as a tenure-track assistant professor of computational mathematics and as an affiliated faculty in machine learning at
@mldcmu
!
Very excited to announce new work () with Eric Vanden-Eijnden from
@NYUCourant
. We leverage a correspondence between the Fokker-Planck equation and the score-based diffusion method of
@DrYangSong
,
@StefanoErmon
to compute the entropy of active systems. 1/n
Our results confirm previous theoretical and computational predictions that the entropy production for systems undergoing motility-induced phase separation is concentrated on interfaces, e.g. and . More results on arxiv! 7/n
*Very* excited to broadcast some recently-published work with Eric Vanden-Eijnden at
@NYU_Courant
, which just appeared in
@MLSTjournal
a few weeks ago (). Some eye candy from the publication below.
Remarkably, a single network trained with 4096 particles also generalizes to new regions of the phase diagram, as evidenced by varying the packing fraction (phi). We trained with phi=0.5. 4/n
@sp_monte_carlo
Following up on
@msalbergo
and refining the suggestion by
@mufan_li
, you can derive an *equality* for the KL between any two time marginals by appealing to the corresponding Fokker-Planck equations. This also works for transport equations, but the result is weaker. See the paper!
Our first paper with the new eLife model is out:
The work studies, primarily computationally, a model for microbial evolution, where the strength of epistasis can be tuned, and shows that this leads to a slow-down of the fitness trajectory dynamics.
We develop a new, flow-based method for solving the time-dependent Fokker-Planck equation. Motivated by recent advances in generative modeling -- in particular, score-based diffusion -- we recast the Fokker-Planck equation as an equivalent, score-based transport equation.
Fundamentally, the key is view the system from a *deterministic* perspective along the probability flow, rather than a stochastic perspective. This leads to simpler, more interpretable dynamics, along with hard-to-estimate quantities like the entropy production rate. 6/n
In addition, the method can be used to generate other interesting outputs, such as effective deterministic phase portraits for stochastic systems, as shown in the first tweet.
Code is written in jax and available at .
The key to scale to such high-dimensionality is to introduce a spatially-local transformer neural network that operates directly on the particles. We found that when trained with 4096 particles it makes excellent predictions for much higher-dimensional systems. 3/n
The movie above contains 32,768 active swimmers undergoing motility-induced phase separation, corresponding to a 131,072-dimensional Fokker-Planck equation. The particles are colored by two different definitions of their contribution to the entropy production rate. 2/n
We show that by performing a kind of sequential, self-consistent score matching procedure, one can control the KL divergence from the model to the ideal solution of the equation. This leads to a loss that can be efficiently minimized over neural nets.
Anyone know if there is support for any of the recent memory efficient attention implementations such as FlashAttention or FlashAttention 2 in jax? Tagging some jax maestros --
@SingularMattrix
@PatrickKidger
-- any network library okay with me.
And here's the method applied to a low-dimensional system describing the motion of an active swimmer, where we can visualize 10^5 samples simultaneously.
Practically, this gives access to several interesting quantities that are not available from trajectories of the underlying SDE, such as evaluation of the density and estimation of the entropy production rate, a quantity of interest in the active matter community.
@QuanquanGu
Flow matching models (equivalently stochastic interpolants) are more general and include diffusions as a special case. They can bridge two arbitrary densities exactly in finite time -- diffusions use Gaussians and need infinite time. See for more details!
@atalwalkar
@JunhongShen1
@__tm__157
This is really cool stuff. Coming from a background in computational math and numerical PDEs, it’s surprising to me that LLMs would work here, because the data is so fundamentally different. I’ll have a closer look at the paper, but any idea why the LLM has a useful inductive
Practically, this gives access to several interesting quantities that are not available from trajectories of the underlying SDE, such as evaluation of the density and the entropy production rate, a quantity of interest in the active matter community.
Solving the resulting equation via the method of characteristics leads to the probability flow ODE, which dynamically builds a transport map from the initial condition to the solution at any later time.
We show that by performing a kind of sequential, self-consistent score matching procedure, one can control the KL divergence from the model to the ideal solution of the equation. This leads to a loss that can be efficiently minimized over neural nets.
We apply the method to several high-dimensional systems of interacting particles, for which standard PDE solvers do not apply. Here's the method applied to a 100-dimensional system of interacting particles chasing a moving trap, where we visualize the evolution of a single sample
We develop a new, flow-based method for solving the time-dependent Fokker-Planck equation. Motivated by recent advances in generative modeling -- in particular, score-based diffusion -- we recast the Fokker-Planck equation as an equivalent, score-based transport equation.