Zero Flux: Flow-Based Comparison of High-Dimensional Discrete Distributions

Motivated by the success of modern flow-based generative models in modeling complex data, we study two-sample testing through the lens of flow-based methods. We propose the Zero-Flow Two-Sample Test (ZF2ST), built on the zero-flow criterion, which characterizes distributional equality through a time-reversal antisymmetry of a learnable velocity field. We extend this criterion and further develop the Zero-Flow Discrepancy, an identifying discrepancy that controls the Wasserstein distance, and derive a variational representation in terms of a witness function. This representation naturally leads to a witness-based test whose power is governed by the signal-to-noise ratio (SNR), allowing direct power maximization for witness learning. ZF2ST learns the witness on one data split and performs testing on held-out samples, thereby maintaining Type-I error control and admitting a simple asymptotic null distribution. Experimentally, ZF2ST performs competitively across various synthetic and real-world benchmarks, while showing particularly strong performance in distinguishing image distributions from different sources.

Flow Matching has recently emerged as a popular class of generative models for simulating a target distribution $μ_1$ from samples drawn from a source distribution $μ_0$. This framework relies on a fixed coupling between $μ_0$ and $μ_1$, and on a deterministic or stochastic bridge to define an interpolating process between the two distributions. The time marginals of this process can then be approximately sampled by estimating the transition rates, or more generally the generator, of its Markovian projection. This framework has recently been extended to the case of discrete source and target distributions, under the name Discrete Flow Matching (DFM). However, theoretical guarantees for such models remain scarce. In this paper, we study two DFM models on $\mathbb{Z}_m^d = \{0,\ldots,m-1\}^d$, sampled through time discretization, and derive non-asymptotic associated bounds for both of them. In contrast to previous work, we establish non-asymptotic bounds in Kullback--Leibler divergence for the early-stopped version of the target distribution. We also derive explicit convergence guarantees in total variation distance with respect to the true target distribution. Importantly, these bounds rely only on an approximation error assumption, relaxing standard score assumptions used in earlier works, while also yielding improved dependence on the vocabulary size $m$ and the dimension $d$.

Fair representation learning with a continuous sensitive attribute $S$ requires a representation $Z$ that is statistically independent of $S$. Existing criteria, including generalized demographic parity, the expectation of integral probability metrics (EIPM), and mutual information, enforce this independence by averaging a per-value discrepancy between the conditional law $P_{Z \mid S=s}$ and the marginal $P_Z$ over the law of $S$. This approach requires a nonparametric surrogate for the conditional law at each sensitive value. We propose evaluating independence through a single joint discrepancy $d\left(P_{Z, S}, P_Z \otimes P_S\right)$ between the joint law and the product of its marginals. We establish a disintegration identity; on decomposable witness classes it equals the conditional-integral functional that EIPM and generalized demographic parity instantiate. By reaching the same target without the conditional law, this discrepancy can be estimated directly from samples via a dependence statistic rather than conditional smoothing. We take the Hilbert-Schmidt independence criterion (HSIC) as an instance of the joint discrepancy $d$ to investigate the statistical efficiency of replacing the conditional formulation. The HSIC estimator is a closed-form $O\left(n^2\right)$ statistic that converges at the $O\left(n^{-1 / 2}\right)$ rate, in contrast to the nonparametric $O\left(n^{-2 / 5}\right)$ rate of the conditional-route estimators. We prove this instance is equivalent to the conditional maximum mean discrepancy (MMD) integral up to an explicit spectral tail. The corresponding algorithmic implementation, i.e., FRHSIC, attains fairness-accuracy tradeoffs comparable to conditional-route basel es while reducing per-epoch training time.