Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Towards Fast Coarse-graining and Equation Discovery with Foundation Inference Models

Manuel Hinz, Maximilian Mauel, Patrick Seifner, David Berghaus, Kostadin Cvejoski, Ramses J. Sanchez

TL;DR

The paper tackles extracting low-dimensional latent SDEs from high-dimensional dynamics by decoupling coarse-grain variable discovery from dynamics inference via Foundation Inference Models (FIMs). FIMs provide zero-shot drift and diffusion estimators, f_{FIM} and G_{FIM}, to fix the dynamical class, enabling a dynamics-constrained autoencoder that trains only the encoder/decoder. A Dynamics-constrained Autoencoder Loss uses these estimates to predict the next observation from the encoded state, reinforcing simulation-consistent latent dynamics. A proof-of-concept on a stochastic double-well system embedded in synthetic video demonstrates rapid latent-variable discovery and robust reconstructions, suggesting fast, reusable coarse-graining pipelines for complex, high-dimensional systems.

Abstract

High-dimensional recordings of dynamical processes are often characterized by a much smaller set of effective variables, evolving on low-dimensional manifolds. Identifying these latent dynamics requires solving two intertwined problems: discovering appropriate coarse-grained variables and simultaneously fitting the governing equations. Most machine learning approaches tackle these tasks jointly by training autoencoders together with models that enforce dynamical consistency. We propose to decouple the two problems by leveraging the recently introduced Foundation Inference Models (FIMs). FIMs are pretrained models that estimate the infinitesimal generators of dynamical systems (e.g., the drift and diffusion of a stochastic differential equation) in zero-shot mode. By amortizing the inference of the dynamics through a FIM with frozen weights, and training only the encoder-decoder map, we define a simple, simulation-consistent loss that stabilizes representation learning. A proof of concept on a stochastic double-well system with semicircle diffusion, embedded into synthetic video data, illustrates the potential of this approach for fast and reusable coarse-graining pipelines.

Towards Fast Coarse-graining and Equation Discovery with Foundation Inference Models

TL;DR

The paper tackles extracting low-dimensional latent SDEs from high-dimensional dynamics by decoupling coarse-grain variable discovery from dynamics inference via Foundation Inference Models (FIMs). FIMs provide zero-shot drift and diffusion estimators, f_{FIM} and G_{FIM}, to fix the dynamical class, enabling a dynamics-constrained autoencoder that trains only the encoder/decoder. A Dynamics-constrained Autoencoder Loss uses these estimates to predict the next observation from the encoded state, reinforcing simulation-consistent latent dynamics. A proof-of-concept on a stochastic double-well system embedded in synthetic video demonstrates rapid latent-variable discovery and robust reconstructions, suggesting fast, reusable coarse-graining pipelines for complex, high-dimensional systems.

Abstract

High-dimensional recordings of dynamical processes are often characterized by a much smaller set of effective variables, evolving on low-dimensional manifolds. Identifying these latent dynamics requires solving two intertwined problems: discovering appropriate coarse-grained variables and simultaneously fitting the governing equations. Most machine learning approaches tackle these tasks jointly by training autoencoders together with models that enforce dynamical consistency. We propose to decouple the two problems by leveraging the recently introduced Foundation Inference Models (FIMs). FIMs are pretrained models that estimate the infinitesimal generators of dynamical systems (e.g., the drift and diffusion of a stochastic differential equation) in zero-shot mode. By amortizing the inference of the dynamics through a FIM with frozen weights, and training only the encoder-decoder map, we define a simple, simulation-consistent loss that stabilizes representation learning. A proof of concept on a stochastic double-well system with semicircle diffusion, embedded into synthetic video data, illustrates the potential of this approach for fast and reusable coarse-graining pipelines.
Paper Structure (3 sections, 3 equations, 2 figures)

This paper contains 3 sections, 3 equations, 2 figures.

Table of Contents

  1. Introduction
  2. Discovering Latent SDEs with FIM
  3. Conclusions

Figures (2)

  • Figure 1: Left: Stochastic double-well system embedded in a $51\times51$ pixel space. The frame shown is averaged over 100 time steps. Center: Drift of the ground-truth latent dynamics (dashed black) compared with the drift inferred by FIM-SDE using the true coordinates (blue) and using the learned coordinates (orange). Right: Diffusion function, shown analogously to the center panel.
  • Figure 2: Left: Training trajectories of 10 independently trained autoencoders with different random initializations, illustrating the stability of the optimization. Center: Drift of the ground-truth latent dynamics (dashed black), compared with the drift inferred by FIM-SDE using the true coordinates (blue), and using the 10 learned coordinates from the independently trained autoencoders (orange). Shaded regions denote confidence intervals. Right: Diffusion function, shown analogously to the center panel.