Summary
A method called Markovian flow matching (MFM) for training neural ODEs (continuous normalizing flows) by flow matching to sample from distributions given as unnormalized density functions is proposed. The method obtains samples on which to perform FM training by an MCMC procedure that alternates two kinds of kernels: Metropolis-adjusted Langevin in the target space (as has been done in past work) and Gaussian proposals in the latent space, i.e., the pullback of the target by the learned ODE. An annealing scheme is additionally used. MFM is evaluated on toy low-dimensional tasks, a synthetic density from statistical physics, and a high-dimensional data-derived density (log-Gaussian Cox process).
Strengths
- The idea to do MCMC in the latent space of a CNF by pullback along the flow map is original, as far as I know.
- There are two general families methods for training continuous-time models to sample Boltzmann densities: using maximum likelihood on (possibly reweighted) approximate samples (e.g., FAB and the proposed MFM) and using distribution-matching objectives (e.g., DDS, PIS), as well as hybrid approaches.
- It's interesting to see progress on making methods of the first kind efficient, and to see them combined with CNFs. In the past, e.g., in [Tong et al. "Improving and generalizing flow-based generative models", D.2], FM has been used with MCMC, but only in the target space.
- No serious problems with the writing, including the global structure of the paper and the exposition of preliminaries.
Weaknesses
- Some math bugs (imprecise exposition):
- L72: $v_t$ being **any** time-dependent vector field does not imply existence of a diffeomorphic flow map (even if it is continuous in $t$, if that is how we interpret "runs continuously in the unit interval"). One also needs integrability conditions.
- L75: $p_0$ has to be strictly positive for the path to take values in $\mathbb{R}^+$ and for the discussion the discussion that follows to make sense.
- Comparison with SDE models: Approximate (MCMC) samples can also be used to train diffusion models (neural SDEs):
- First, this could be done either by minimising a variational bound on the log-likelihood, similar to what is done here for CNFs -- is it possible to do such a comparison? At present the SDE baseline (DDS) is not using MCMC, only "on-policy" forward exploration.
- Second, methods such as that in [Sendera et al. "Improved off-policy training of diffusion samplers"] use MCMC on the target density to obtain samples for training with a distribution-matching, not maximum-likelihood, objective.
- The main weakness is that the experiments are not very comprehensive or convincing:
- The 2D GMMs are toy problems; all methods, including DDS, should find the modes when appropriately tuned on such problems, and so it is hard to draw any conclusions from the results.
- On the field system, MFM, although much faster than FAB, is performing **much** worse in the metrics used.
- On LGCP:
- The use of this task for benchmarking samplers of Boltzmann densities is questionable in the first place, since there is no known ground truth (even the exact normalising constant is diagreed upon) -- see, for example, the aforementioned [Sendera et al., B.1] for a discussion of inconsistent or irreproducible evaluation in past work on this dataset.
- Why use the nonstandard metric (KSD)? Most past work has used estimates of the log-partition function, which here is reported in the appendix.
- On that more standard metric, MFM is significantly underperforming (Table 8).
- On both the field system and LGCP, how are we to know if the proposed method is useful on this problem without a point of comparison? I would strongly suggest a long-run HMC, MALA, and SMC (for example, run for the same wall time as the shortest and longest of the ML methods evaluated) as baselines.
- I would also strongly suggest evaluations on a more comprehensive set of benchmarks that are common in past work on samplers of Boltzmann densities (as in the FAB and DDS papers, among others).
Questions
Minor:
- L25: You probably want to specify "a new sample $y$" and in the subsequent lines clarify that the "mild conditions" are also needed for the target to be the *only* stationary distribution.
- L77: Dash between "change" and "of" instead of hyphen.
- L86: Wrong placement of parentheses.
- L246: I do not believe [78] introduced this parametrisation -- it was in [80] (and possibly in earlier work as well).
- I don't understand equation (18), which involves $x_{65}$ in the first sum although the dimensionality is 64. Is there is mistake, or should the indices be interpreted modulo $d$?
- L308: Typo in "comparable".