Causality--Δ: Jacobian-Based Dependency Analysis in Flow Matching Models

TL;DR

The paper investigates how tiny latent perturbations propagate through flow matching models by leveraging Jacobian-vector products (JVPs) to estimate dependency structures in generated features. It derives closed-form drift and Jacobian expressions for Gaussian and mixture targets, revealing that locally affine dynamics emerge even in globally nonlinear flows, and validates these insights on synthetic data, MNIST, and CelebA. By composing flows with a classifier, it demonstrates attribute-level dependencies and shows that conditioning on small Jacobian norms reduces correlations consistent with a common-cause structure, while clearly distinguishing this from do-interventions. The results provide a practical perturbation-based lens for interpreting flow-based generative models and bridge causal reasoning with local linear approximations in continuous normalizing flows.

Abstract

Flow matching learns a velocity field that transports a base distribution to data. We study how small latent perturbations propagate through these flows and show that Jacobian-vector products (JVPs) provide a practical lens on dependency structure in the generated features. We derive closed-form expressions for the optimal drift and its Jacobian in Gaussian and mixture-of-Gaussian settings, revealing that even globally nonlinear flows admit local affine structure. In low-dimensional synthetic benchmarks, numerical JVPs recover the analytical Jacobians. In image domains, composing the flow with an attribute classifier yields an attribute-level JVP estimator that recovers empirical correlations on MNIST and CelebA. Conditioning on small classifier-Jacobian norms reduces correlations in a way consistent with a hypothesized common-cause structure, while we emphasize that this conditioning is not a formal do intervention.

Figures (6)

• Figure 1: Flow Matching between two standard Gaussian distributions.
• Figure 2: Time Evolution of Mixture of Gaussians Example.
• Figure 3: (Left) is the approximate "derivative image" of an MNIST five; (Middle) is the finite difference of pixel pairs over MNIST fives; (Right) is the true join distribution of the pixel pairs.
• Figure 4: Estimated Conditional Dependence Structure of the CelebA dataset.
• Figure 5: The empirical ground truth correlation matrix, the estimated correlation matrix under observation, and the estimated correlation matrix after conditioning on small $\|\mathbf{J}^{19}\|$.