Fisher Flow Matching for Generative Modeling over Discrete Data

TL;DR

Fisher-Flow introduces a geometric flow-matching approach for discrete data generative modeling using the Fisher-Rao metric, enabling principled continuous reparameterization and improved performance over existing methods.

Contribution

The paper proposes Fisher-Flow, a novel flow-matching model for discrete data based on Riemannian geometry, with theoretical optimality and practical improvements demonstrated on real-world benchmarks.

Findings

Fisher-Flow outperforms prior diffusion and flow-matching models on benchmarks.

The model leverages Riemannian geometry for continuous reparameterization of discrete data.

Gradient flow in Fisher-Flow is proven to be optimal in KL divergence reduction.

Abstract

Generative modeling over discrete data has recently seen numerous success stories, with applications spanning language modeling, biological sequence design, and graph-structured molecular data. The predominant generative modeling paradigm for discrete data is still autoregressive, with more recent alternatives based on diffusion or flow-matching falling short of their impressive performance in continuous data settings, such as image or video generation. In this work, we introduce Fisher-Flow, a novel flow-matching model for discrete data. Fisher-Flow takes a manifestly geometric perspective by considering categorical distributions over discrete data as points residing on a statistical manifold equipped with its natural Riemannian metric: the $\textit{Fisher-Rao metric}$. As a result, we demonstrate discrete data itself can be continuously reparameterised to points on the positive orthant of the $d$-hypersphere $\mathbb{S}^d_+$, which allows us to define flows that map any source distribution to target in a principled manner by transporting mass along (closed-form) geodesics of $\mathbb{S}^d_+$. Furthermore, the learned flows in Fisher-Flow can be further bootstrapped by leveraging Riemannian optimal transport leading to improved training dynamics. We prove that the gradient flow induced by Fisher-Flow is optimal in reducing the forward KL divergence. We evaluate Fisher-Flow on an array of synthetic and diverse real-world benchmarks, including designing DNA Promoter, and DNA Enhancer sequences. Empirically, we find that Fisher-Flow improves over prior diffusion and flow-matching models on these benchmarks.