Technical Summary: Flow Matching with Arbitrary Auxiliary Paths (AuxPath-FM)

Problem Statement

Continuous-time generative models, including score-based diffusion models and flow matching (FM) frameworks, have established themselves as powerful tools for modeling complex data distributions. In conditional generative modeling, incorporating side information (e.g., class labels) is essential for controllable synthesis. However, existing approaches typically introduce such information exclusively through the conditioning of neural networks, leaving the underlying probability paths unchanged. While recent theoretical work suggests flow matching adapts to intrinsic manifold structures, the influence of the path's stochastic component on trajectory-level control remains under-explored.

Furthermore, prior path-based formulations, such as stochastic interpolants and Schrödinger Bridge methods, generally rely on Gaussian noise or Brownian dynamics for stochastic interpolation. This restriction limits the flexibility of designing probability paths for structured guidance. Although recent works have explored non-Gaussian perturbations or discrete state-space flows, these methods often remain specialized to specific data types or fail to unify continuous and discrete priors under a single, general framework. There is a need for a principled method that allows the generative trajectory to be directly shaped by diverse auxiliary signals, rather than relying solely on network-level modulation.

Methodology: AuxPath-FM

The authors introduce Flow Matching with Arbitrary Auxiliary Paths (AuxPath-FM), a framework that generalizes conditional flow matching by incorporating an auxiliary random variable $\eta$ drawn from an arbitrary distribution directly into the probability path.

1. Arbitrary Auxiliary Probability Paths

Unlike standard interpolation schemes that restrict auxiliary terms to Gaussian noise, AuxPath-FM constructs trajectories of the form:
$$X_t = a(t)X_1 + b(t)X_0 + c(t)\eta$$
where:

• $X_0 \sim p_0$ (base distribution) and $X_1 \sim p_1$ (target data distribution).
• $\eta \sim p_\eta$ is an auxiliary variable drawn from an arbitrary distribution (e.g., Gaussian, Uniform, Laplace, Rademacher, or discrete distributions).
• $a(t), b(t), c(t)$ are time-dependent scalar functions satisfying boundary conditions ($a(0)=0, b(0)=1, c(0)=0$ and $a(1)=1, b(1)=0, c(1)=0$) to ensure the path transports samples from $p_0$ to $p_1$.

2. Theoretical Consistency

The paper establishes that this construction preserves the continuity equation. Theorem 1 demonstrates that if conditional vector fields $u_t(x | x_0, \eta)$ satisfy the continuity equation, the marginal vector field $u_t(x)$, obtained by aggregating conditional fields, generates the marginal probability path $p_t(x)$. This provides a principled justification for designing complex paths via auxiliary variables.

3. Training Objective and Equivalence

Directly optimizing the marginal flow matching objective is intractable due to integrals over $(X_0, \eta)$. Instead, the authors optimize a conditional objective defined on tractable paths:
$$L_{\text{AuxPath}} = \mathbb{E}_{t, X_0, X_1, \eta} \left[ \| v_\theta(X_t, t) - \dot{X}_t \|^2 \right]$$
where $\dot{X}_t = \dot{a}(t)X_1 + \dot{b}(t)X_0 + \dot{c}(t)\eta$. Theorem 2 proves that this conditional objective is equivalent to the marginal flow matching objective up to a constant independent of the network parameters $\theta$. Thus, optimizing the conditional loss yields the same optimal solution as the intractable marginal objective.

4. Semantic Auxiliary Variables and Guidance

To enable conditional generation, the framework defines $\eta = F_\phi(Y)$, where $Y$ is a class label and $F_\phi$ is a lightweight neural network trained via prototype learning to approximate the class-wise centroid ($F_\phi(Y) \approx \mathbb{E}[X_1 | Y]$).

This formulation enables Trajectory-Level Classifier-Free Guidance (CFG). The conditional and unconditional velocity fields are expressed as:
$$v(x, t, y) = v_\theta(x, t) + \dot{c}(t)F_\phi(y)$$
$$v(x, t, \emptyset) = v_\theta(x, t) + \dot{c}(t)F_\phi(\emptyset)$$
Applying standard CFG yields a guidance term that acts as a global drift. Crucially, this approach decouples guidance from the backbone model: unlike conventional CFG which requires two evaluations of the backbone $v_\theta$, AuxPath-FM requires only one evaluation of $v_\theta$ and applies guidance through the lightweight module $F_\phi$.

Key Contributions

1. Generalization of Probability Paths: The paper generalizes flow matching by allowing the auxiliary component of the probability path to follow any arbitrary distribution, moving beyond the restriction to Gaussian noise found in stochastic interpolants and Schrödinger Bridge methods.
2. Theoretical Foundation: It provides a rigorous theoretical proof that the marginal flow remains consistent with the continuity equation and that the conditional training objective is equivalent to the marginal objective.
3. Trajectory-Level Control: The framework enables the direct encoding of structured semantic information (e.g., class labels) into the generative trajectory itself, rather than relying solely on network conditioning.
4. Efficient Classifier-Free Guidance: The authors propose a trajectory-level formulation of CFG that reduces inference computational cost by approximately half, as it avoids double evaluation of the backbone network.
5. Unified Perspective: The method unifies various interpolation schemes (including discrete and heavy-tailed priors) under a single framework, offering a flexible perspective for structured continuous-time generative modeling.

Experimental Results

The authors evaluated AuxPath-FM on MNIST, CIFAR-10, and ImageNet-1k datasets.

• Conditional Generation Performance: On MNIST and CIFAR-10, combining standard conditional flow matching (CFM) with auxiliary path guidance (using learned semantic prototypes) consistently improved semantic accuracy (classification accuracy) and generative quality (FID, sFID) compared to baseline CFM or auxiliary guidance alone.
• Robustness to Distribution Choice: Experiments on the Ring-64 dataset and ImageNet-1k demonstrated that the framework maintains stable generative performance across diverse auxiliary distributions (Gaussian, Uniform, Laplace, Rademacher). Label-guided auxiliary variables achieved the best performance, indicating improved controllable multi-modal generation.
• Guidance Scale Effects: Increasing the guidance scale $w$ in trajectory-level CFG improved semantic alignment while maintaining competitive generative quality.
• Efficiency: The trajectory-level CFG formulation successfully reduced inference costs (measured in GFLOPs) by approximately 50% compared to standard CFG while preserving the effectiveness of guidance.
• Model Adaptation: The framework successfully converted pre-trained unconditional flow models into conditional generators with minimal additional cost and without architectural changes.

Significance and Claims

The paper claims that AuxPath-FM provides a principled and general foundation for probability path design. By allowing the auxiliary variable to follow arbitrary distributions, the framework offers both theoretical generality and practical flexibility.

The authors emphasize that their method shifts the paradigm of conditional generation from purely network-level modulation to trajectory-level design. This allows for the direct embedding of guidance signals into the generative dynamics. The significance of this work lies in its ability to:

• Unify existing interpolation schemes as special cases.
• Enable diverse geometric properties for generative flows through different priors.
• Provide a computationally efficient mechanism for classifier-free guidance.
• Offer a simple yet effective method for upgrading existing unconditional models to conditional ones.

The paper concludes that this approach offers a unified perspective for structured continuous-time generative modeling, validated by extensive experiments showing that semantic-auxiliary CFG maintains superior generative quality while significantly enhancing computational efficiency.