Table of Contents
Fetching ...

Diversified Flow Matching with Translation Identifiability

Sagar Shrestha, Xiao Fu

TL;DR

This work addresses content misalignment in unpaired domain translation by enforcing translation identifiability via diversified distribution matching (DDM). It introduces diversified flow matching (DFM), an ODE-based flow matching method that uses a bilevel learning loss and private nonlinear interpolants to guarantee DDM identifiability, while also producing explicit transport trajectories. A non-overlapping support assumption enables a practical, two-stage implementation that avoids the computational burden of a full bilevel optimization. Experiments on synthetic data, unpaired image translation, and swarm navigation demonstrate that DFM achieves translation identifiability, superior trajectory consistency, and robust performance compared to GAN- and diffusion-based baselines, with potential for trajectory-aware domain transfer in real-world applications.

Abstract

Diversified distribution matching (DDM) finds a unified translation function mapping a diverse collection of conditional source distributions to their target counterparts. DDM was proposed to resolve content misalignment issues in unpaired domain translation, achieving translation identifiability. However, DDM has only been implemented using GANs due to its constraints on the translation function. GANs are often unstable to train and do not provide the transport trajectory information -- yet such trajectories are useful in applications such as single-cell evolution analysis and robot route planning. This work introduces diversified flow matching (DFM), an ODE-based framework for DDM. Adapting flow matching (FM) to enforce a unified translation function as in DDM is challenging, as FM learns the translation function's velocity rather than the translation function itself. A custom bilevel optimization-based training loss, a nonlinear interpolant, and a structural reformulation are proposed to address these challenges, offering a tangible implementation. To our knowledge, DFM is the first ODE-based approach guaranteeing translation identifiability. Experiments on synthetic and real-world datasets validate the proposed method.

Diversified Flow Matching with Translation Identifiability

TL;DR

This work addresses content misalignment in unpaired domain translation by enforcing translation identifiability via diversified distribution matching (DDM). It introduces diversified flow matching (DFM), an ODE-based flow matching method that uses a bilevel learning loss and private nonlinear interpolants to guarantee DDM identifiability, while also producing explicit transport trajectories. A non-overlapping support assumption enables a practical, two-stage implementation that avoids the computational burden of a full bilevel optimization. Experiments on synthetic data, unpaired image translation, and swarm navigation demonstrate that DFM achieves translation identifiability, superior trajectory consistency, and robust performance compared to GAN- and diffusion-based baselines, with potential for trajectory-aware domain transfer in real-world applications.

Abstract

Diversified distribution matching (DDM) finds a unified translation function mapping a diverse collection of conditional source distributions to their target counterparts. DDM was proposed to resolve content misalignment issues in unpaired domain translation, achieving translation identifiability. However, DDM has only been implemented using GANs due to its constraints on the translation function. GANs are often unstable to train and do not provide the transport trajectory information -- yet such trajectories are useful in applications such as single-cell evolution analysis and robot route planning. This work introduces diversified flow matching (DFM), an ODE-based framework for DDM. Adapting flow matching (FM) to enforce a unified translation function as in DDM is challenging, as FM learns the translation function's velocity rather than the translation function itself. A custom bilevel optimization-based training loss, a nonlinear interpolant, and a structural reformulation are proposed to address these challenges, offering a tangible implementation. To our knowledge, DFM is the first ODE-based approach guaranteeing translation identifiability. Experiments on synthetic and real-world datasets validate the proposed method.

Paper Structure

This paper contains 27 sections, 4 theorems, 35 equations, 10 figures, 4 tables, 2 algorithms.

Key Result

Theorem 2.2

shrestha2024towards Suppose that $\{p_{\boldsymbol{x} | u^{(q)}}\}_{q=1}^Q$ satisfies the SDC. Let $\widehat{\boldsymbol{g}}$ be any optimal solution of the DDM criterion eq:diverse_dist_matching. Then, we have $\widehat{\boldsymbol{g}} = \boldsymbol{g}^\star,~a.e.$

Figures (10)

  • Figure 1: [Columns 2-4] Content misalignment issues in both GAN and FM based UDT (CycleGANzhu2017unpaired, FMlipman2022flow, FM-OTtong2023improving) [Column 5] Result by DDM-GAN shrestha2024towards.
  • Figure 2: The idea of DDM. The variable $u^{(q)}$ can often be defined as attributes that are not supposed to change across domains. In shrestha2024towards, it was shown $Q \geq 2$ suffices to underpin the translation identifiability.
  • Figure 3: (a) Samples of two pairs of conditional distributions. (b) Linear interpolant at $t \in [0,1]$ for interpolant trajectories. (c) Actural trajectories learned by solving \ref{['eq:naive_fm_ddm']}. (d) $\widehat{\boldsymbol{v}}_{\frac{1}{2}}$ points towards the $-\mathbb{E}[\boldsymbol{x}]$.
  • Figure 4: [Left] Success case of solving \ref{['eq:fm_ddm_sum']}. [Right] Failure case of solving \ref{['eq:fm_ddm_sum']}.
  • Figure 5: Trajectories returned by all methods for the 3D synthetic data.
  • ...and 5 more figures

Theorems & Definitions (10)

  • Definition 2.1: Sufficiently Diverse Condition (SDC)
  • Theorem 2.2: Translation Identifiability
  • Definition 3.1: Transport of Measures
  • Definition 3.2: DDM Satisfaction
  • Proposition 3.4
  • Definition 3.7: Non-intersecting Interpolants
  • Theorem 2.1
  • Definition 5.1: Interpolable density, albergo2023stochastic Def. D.1
  • proof
  • Proposition 5.2: albergo2023stochastic Proposition D.1