On the flow matching interpretability

TL;DR

The paper tackles the interpretability gap in flow matching by constraining each generative step to pass through equilibrium states of a known physical process. It implements a physics-informed diffusion framework on the 2D Ising model with a cooling schedule, using an encoder $\phi$, a vector-field predictor $v_\theta$, and a projector $P_\theta$ to map between discrete spins and a continuous latent space. The approach preserves thermodynamic fidelity while delivering speedups over traditional Monte Carlo sampling, especially as lattice size grows, with each latent-step corresponding to a concrete thermodynamic transition at inverse temperatures $\beta_j$. This work advances interpretable diffusion models for scientific applications and opens avenues for applying physics-informed constraints to other diffusion-based generative systems.

Abstract

Generative models based on flow matching have demonstrated remarkable success in various domains, yet they suffer from a fundamental limitation: the lack of interpretability in their intermediate generation steps. In fact these models learn to transform noise into data through a series of vector field updates, however the meaning of each step remains opaque. We address this problem by proposing a general framework constraining each flow step to be sampled from a known physical distribution. Flow trajectories are mapped to (and constrained to traverse) the equilibrium states of the simulated physical process. We implement this approach through the 2D Ising model in such a way that flow steps become thermal equilibrium points along a parametric cooling schedule. Our proposed architecture includes an encoder that maps discrete Ising configurations into a continuous latent space, a flow-matching network that performs temperature-driven diffusion, and a projector that returns to discrete Ising states while preserving physical constraints. We validate this framework across multiple lattice sizes, showing that it preserves physical fidelity while outperforming Monte Carlo generation in speed as the lattice size increases. In contrast with standard flow matching, each vector field represents a meaningful stepwise transition in the 2D Ising model's latent space. This demonstrates that embedding physical semantics into generative flows transforms opaque neural trajectories into interpretable physical processes.

Figures (4)

• Figure 1: Pipeline architecture showing two consecutive cooling steps. On the left, the spin state, on the right, the latent one. Starting from spin configuration $\boldsymbol{x}^i_{\beta_j}$, the encoder $\phi$ maps it to latent space $\boldsymbol{z}^i_{\beta_j}$, the vector field $v_\theta$ predicts the next latent state $\boldsymbol{z}^i_{\beta_{j+1}}$, and the projector $P_\theta$ decodes back to spin configuration $\boldsymbol{x}^i_{\beta_{j+1}}$. This process is iterated autoregressively through the cooling schedule.
• Figure 2: Thermodynamic observables for the $N=48$ lattice as a function of temperature. The vertical dotted line indicates the critical temperature.
• Figure 3: Comparison of decoding methods in the encoder-vector field pipeline across increasing temperature values (left to right) for grid size $N = 48$. Further comparisons for grid sizes $N = 32$ and $N = 64$ are provided in the Technical Appendix.
• Figure 4: Comparison of physical observables distributions defined in \ref{['eq:quantities']} over the cooling schedule (grid size $N = 48$). $C_v$ and $\chi$ on a log scale. The observables for grid size $N = 32$ and $N = 64$ are provided in the Technical Appendix.