Reproduce “Fast Flow Matching based Conditional Independence Tests for Causal Discovery” on OpenResearch

About this paper

Shunyu Zhao, Yanfeng Yang, Shuai Li, Kenji Fukumizu

This work introduces the Flow Matching-based Conditional Independence Test (FMCIT), which leverages a single, pre-trained flow matching model to accelerate conditional independence testing by rapidly generating samples through imputation. Integrated into a guided PC-style algorithm (GPC-FMCIT), the approach provides improved accuracy in causal skeleton recovery for non-linear, high-dimensional data, while controlling Type-I error rates.

Key insights

• Training a single generative model on the joint distribution and leveraging it for conditional imputation across all CI queries significantly reduces computational overhead for causal discovery.
• Flow Matching's ODE-based sampling, combined with Picard/RePaint imputation, provides a faster and more stable alternative to SDE-based diffusion models for conditional sample generation.
• A two-stage guided skeleton learning approach, with a fast initial screening and a targeted, budgeted refinement using a more robust CI test, effectively manages the combinatorial complexity of causal discovery in high dimensions.

Summary by alphaXiv.