Fisher Flow Matching for Generative Modeling over Discrete Data
Authors: Oscar Davis
Samuel Kessler
Mircea Petrache
Ismail Ilkan Ceylan
Michael M. Bronstein
Joey Bose
Venue: NeurIPS 2024 poster (2024)
Keywords: Flow matching
Generative models
Riemannian manifolds
Discrete data
Primary area: generative_models
TL;DR: We propose a novel flow matching approach for discrete data by continuously reparameterising categorical distributions on the hypersphere.

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 \emph{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.

Fisher Flow Matching for Generative Modeling over Discrete Data
Oscar Davis'
Samuel Kessler'
Mircea Petrache
ismail ilkan Ceylan'
Michael Bronstein
Avishek Joey Bose
University Oxford, "Pontificia Universidad Católica de Chile, 3 *Aithyra
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 contin- uous data settings, such as image or video generation. In this work, we introduce LISHER-FLOWE novel flow-machhing model for discrete data. FISHER- FLOW takes manifestly geometric perspective by considering categorical distributions over discrete data as points residing on statistical manifold equipped with natu- ral Riemannian metric: the Fisher-Rao metric. As result, we demonstrate discrete data itself can be continuously reparameterised to points on the positive orthant of the d-hypersphere $ which allows us to define flows that map any source dis- tribution to target in a principled manner by transporting mass along (closed-fomm) geodesics of $4. Furthermore, the learned flows in FISHER-FLOO 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 KI. divergence. We evaluate FISHER-FLOW on array of synthetic and diverse real-world benchmarks, including designing DNA Promoter, and DNA Enhancer sequences. Empirically, we find that FISHER-FI.OW improves over prior diffusion and flow-matching models these benchmarks. Our code is available at https: / //gi thub. com/
Introduction
The recent success of generative models operating on continuous data such as images has been watershed moment for Al exceeding the wildest expectations just few years ago [69, 22]. A key driver of this progress has come from substantial innovations in simulation-fre generative models, the most popular of which include diffusion [38, 66] and flow matching methods [48, 73], leading to a plethora of advances in image generation [16, 30, 55J, video generation [21, 14], audio generation [59]. and 3D protein structure generation [77, 19], name few.
In contrast, analogous advancements in generative models over discrete data domains, such as language models [1, 72], have been dominated autoregressiive models [79], which attribute simple factorisation of probabilities over sequences. Modern autoregressive models, while impressive, have several key limitations which include the slow sequential sampling of tokens in sequence, the assumption of a specified ordering over discrete objects, and the degradation of performance
*Correspondence oscar, davis@cs. ox. ax. uk.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

without important inference techniques such as nucleus sampling [39]. It is expected that further progress will come from the principled equivalents of diffusion and fow-matching approaches for categorical distributions in the discrete data setting.
While appealing, one central barrier in constructing diffusion and flow matching over discrete spaces lies in designing an appropriate forward process that progressively corrupts discrete data. This often involves the sophisticated design of transition kernels [8. 23, 2, 50], which hits an ideal stationary distribution-itsell remaining an unclear quantity in the discrete setting. An alternative path to design- ing discrete transitions is to instead opt for a continuous relaxation of discrete data over a continuous space, which then enables the simple application of flo-mmacching and diffusion. Consequently, past work has relied on relaxing discrete data to points on the interior of the probability simplex [9. 68].
However, since the probability simplex is not Fuclidean, not possible to utilise Gaussian probability paths-the the stationary distribution of an uninformative prior is uniform rather than Gaussian [26]. One possible remedy is to construct conditional probability paths on the simplex using Dirichlet distributions [68], but this can lead to undesirable properties that include complex parameterisation of the vector field. An even greater limitation is that flows using Dirichlet paths are not general enough to accommodate starting from a non-uniform (source) prior ceteriuutionnmmmprriii generative modeling applications. These limitations motivate the following research question: Can we find a continuous reparameterisation of discrete data allowing us to learn a push-forward map between any source and target distribution?
Categorical
A
q
Figure 1: FR metric on Á and the corresponding path on S7
A geodesic connecting zo and æ1 using the
downstream
Present work. In this paper, we propose FiSHER-FLOW, a new flow matching-based generative model for discrete data. Our key geometric insight is to endow the probability simplex with its natural Riemannian metric- Fisher-Rao Rao metric- which transforms the space into a Riemannian manifold and induces a different geometry compared to past approaches of Avdeyev et al. [9], Stark al. [68]. Moreover, using this Riemannian manifold, we exploit a well-known geometric construction: the probability simplex under the Fisher-Rao metric isometric the positive orthant of the d-dimensional hypersphere $ [6] (see Figure 1). By operating on $ we obtain a more flexible and numerically stable parameterisation of learned vector fields well as the ability to familiar metric-namely, the Euclidean metric & restricted to the sphere, which leads to better training dynamics and improved performance. As a result, FISHER- Riemannian Flow Matching (RFM) [26], and our designed flows enjoy explicit and numerically favorable formulas for the trajectory connecting a pair of sampled points between any source and target distribttin-efficcivvel. generalising previous flow models [68].
becomes an instance of
On a theoretical front, we prove in Proposition I that optimising the dow-matching objective with FISHER-Fi -FLOW is an optimal choice for matching categorical distributions on the probability simplex. More precisely, converges to the Fisher-Rao flow in the space of probabilities. In addition, we show in Proposition how to design straighter flows, leading to improved training dynamics, by solving the Riemannian optimal transport problem on s Empirically, we investigate FISHER-FLOW on sequence modeling over synthetic categorical densities as well biological sequence design tasks in DNA promoter and DNA enhancer design. We observe that our approach obtains improved performance to comparable discrete diffusion and flow matching methods of Austin et al. [8]. Stark et al. [68].
show the direction of the optimal induced gradient flow in the space of probabilities
2 Background
The main task of generative modeling is to approximate the target distribution, Paata E P(M), over a probability space (M, Σ, P), using a parametric model po- The choice of M = Rd appears in the classical setup of generative modeling over continuous domains, e.g., images; while for categorical distributions over discrete data, we identify M = P(A), where A = (0....,d d} represents the categories corresponding to an alphabet with d + 1 elements. In this paper, we consider problem settings where the modeler has access to Paata as an empirical distribution from which the samples are drawn identically and independently. Such an empirical distribution corresponds to the training set

---

used to train a generative model and is denoted by D modeling in these settings is learn parameters 6 of generative model po that minimises the forward KI. divergence, KL (Pdatal (po), in other words, maximises the log-likelihood of data under po-
[ A standard approach to generative
2.1 Information geometry
The space of probability distributions P = P(X) over set X can be endowed with geometric structure. Let w be the parameters of a distribution such that the map w H Pu E P is injective. We note that this map is distinguished from the generative model, 0 + pe, as O corresponds to parameters of the neural network rather than the parameters of the outpur distribution being approximated. For instance, if we seek model multi-variate Gaussian N Σ) in Re, the parameters of the distribution are w = (14,5) 2), while 0 can be the parameters of an arbitrary deep neural network.
Our distributions Pw are taken to be family of distributions parameterised by subset of vectors w = (w',.. ........ w E Ω S Rd, with its usual topology. If the distributions Pw are absolutely continuous w.r.t reference measure j over X. with densities (x), x E X,w E Ω2 then the injective map w E 12 + Pw E ()) defines a statistical manifold (cf. Amari [4], Ay et al. [1]]):
Md := (p)(.) | (w'...., wa) E LRCR..
(1)
Note that M is identified as d-dimensional submanifold in the space of absolutely continuous prob- ability distributions P()). If pa (z)x) is differentiable w, then M inherits a differentiable structure. We can then define a metric that converts Md into a Riemannian manifold. Moreover, the parameters w are the local coordinates and the map w + Pa is a global parameterisation for the manifold.
As for the choice of metric, the minimisation of the forward KL divergence, DyL (Pdatal ( under mild conditions, suggests natural prescription of a Riemannian metric on MM [13]. We can arrive at this result by inspecting the log-likelihood of the generative model, logPu. inhormiiformaaiion matrix whose (i, entry G(w) =[9 (w)]k is defined as (Olog Pw
P.: and constructing the
[wr(w) = = (%
( Plgg Pw Ow"
)
(Dlos aws
Pw due
(2)
for 1 S أ S. d, where is the reference measure on 52, which must satisfy the property that all Pu are absolutely continuous with respect to μ.. In this setting, the manifestation of the Fisher-ifformation matrix is not a mere coincidence: it is the second-order Taylor approximation of DKL (Ppstt. Ilpw) in a local neighborhood of Pur in its local coordinates, w. Furthermore, the matrix is symmetric and positive-definit, consequently defining a Riemannian metric. It is called the Fisher-Rae Rao metric and it equips family of inner products at the tangent space Ty Md x Ty... Md + R. that are continuous on the statistical manifold. MM (they vary smoothly, in case the map E + Pa E (p) is assumed to be smooth). Beyond arising as natural consequence of KL minimisation in the generative modeling setup, the Fisher-Rao metric is the unique metric invariant to reparameterisation of M (see Ay et al. [11, Thm. 1.2])-a fact we later exploit in $3.2 to build more scalable and more numerically stable generative models.
2.2 Flow matching over Riemannian manifolds
probability path on a Riemannian manifold, Md, is a continuous interpolation between two distributions, po. P1 P(M)), indexed by time t. Let pt be a distribution a probability path that connects Po to P1 and consider its associated flow, t/c, and vector field, ut- We can learn a continuous normalising flow (CNF) by directly regressing the vector field, ue with parametric one, E T .Md where Md is the tangent bundle. In effect, the goal of learning is to match the flow-termed flow- matching of the target vector field and can be formulated into a simulation-free free training objective [48, FM], provided pt satisfies the boundary conditions, Pu Pdata and P1 = Pprior- As stated, the vanilla flow matching objective is intractable generates pt- Instead, we can opt to regress Un against a conditional vector field, u, (x<r(z)) generating conditional probability path (x(2x12), and use it to recover the target unconditional path: (x.))
generally do not have access to the closed-form of un that
J. Pe ( (=))d) Similarly, the vector field un can also be recovered by marginalising conditional vector fields, us (x z). This allows us to state the CFM objective for Riemannian manifolds [26]:
Lrctu (f) = E2 (xz. (ws (t, xx)) us (1)
t t~ L((0 1).
(3)
If,. as in case, we take X = A {0..... d). then can fix u ボ4ं Ed 8, and then M = P(X).
3

---

As FM and CFM objectives have the same gradients [73, 48], at inference. we can generate by sampling from Pı, and using Up to propagate the ODE backwards in time. The central question in the Riemannian setting corresponds to then finding Ia and u4 (z)2222). For simple geometries, one can always exploit the geodesic interpolant to construct 2; = expx, (1 logz, (x1)) and uy (xx|2) = dt. Instead of computing the time derivative explicitly, we may also use general closed-form expression for Ug, based on the geometry of the problem: U = lox,, (11 t), cf. [18].
Notation and convention. We use & E [0, I] to indicate the time index of a process such that t = 0 corresponds to Pdata and t = corresponds to the terminal distribution of a (stochastic) process to be defined later. Typically, this will correspond to an easy-to-sane sample from source distribution. Henceforth, we use subscripts to denote the time index index-iee. Pt-and reserve superscripts to designate indices over coordinates in (parameter) vector, e.g., w¹ E (w .......)).
3 Fisher Flow Matching
We now establish new methodology to perform discrete generative models under a flow-matching paradigm which we term as DASHER-FLOWW Intuitively, our approach begins with the realisation that discrete data modeled as categorical distributions over d categories can be parameterised to live on the d-dimensional probability simplex, Ad, whose relative interior, Ád, can be identified Riemannian manifold endowed with the Fisher-Rao metric [4, 11, 56]. Additionally, we leverage the sphere map, which defines a diffeomorphism between the interior of the probability simplex and the positive orthant of a hypersphere, $4. As a result, generative modeling over discrete data is amenable to continuous parameterisation over spherical manifolds and offers the following key advantages:
(A1) Continuous reparameterisation We can now seamlessly define conditional probability paths directly the Riemannian manifold S4 enabling to treat discrete generative modeling as continuous, through Riemannian flow matching on the hypersphere.
(A2) Flexibility of source distribution. In stark contrast with prior work [23, 68), our conditional probability paths can map any source distribution to a desired target distribution by leveraging the explicit analytic expression of the geodesies on Sd
(A3) Riemannian optimal transport. As the sphere map is an isometry of the interior of the
probability simplex, we can perform Riemannian OT using the geodesic cost on St. to construct
a coupling between Po and p1, leading to straighter flows and lower variance training.
In the following subsections, we detail first how to construct the continuous reparameterisatioon used in FISHER-- $3.1. An algorithmic description of the training procedure of FiSHER-FLOW is presented in Algorithm 1. We justify the use of the Fisher-Rao metric in $3.3 by showing that induces gradient flow that minimises the KL divergence. Finally, we discuss the sphere map in $3.2, and conclude by elevating the constructed flows to minimise the Riemannian OT problem in $3.4.
3.1 Reparameterising discrete data on the simplex
We now take our manifold Md = A fa E 1 2 0} as the d-dimensional simplex. We seek to model distributions over this space which we denote as P(Ad). )). We can represent categori- cal distributions, p(x), over K- = d d 1 categories A by placing a Dirac 6, with weight p on each vertex E {0,.. ... .J.3 Thus discrete probability distribution given by categorical can be converted into continuous representation over A by representing the categories p' as mixture of point masses at each vertex of A This allows us to write our data distribution Pdata over discrete objects as:
d
Pdata (x)
Ep'6(=-0 -;), ei),
i=0
where e; are K = d + one-hot vectors representing the vertices of the probability simplex While the vertices of A are still discrete, the relative interior of the probability simplex, denoted as Ad = (x Ad: x 0}. is continuous space, whose geometry can be leveraged to build our
'We denote, with a slight abuse of notation, the probability of category i by p', i.e., Σ,p' = *Note that e, E A" represents Dirae mass 6, E P(A), thus Eq. means that Pdata p'ós, ((P(A)) ~ P(A "). The traditional form E. p 8, P(A)) is recovered via the identification à; +
4

---

method, FISHER-FLOW. Consequently, we may move Dirac masses on the vertices of the probability simplex to its interior- and thereby performing continuous reparameterisation by simply applying any smoothing function a : A Å e.g., label smoothing as in supervised learning [70].
Defining a Riemannian metric. Relaxing categorical distributions to the relative interior, Á enables us to consider a more geometric approach to building generative models. Specifically, this construction necessitates that we treat Á as a statistical Riemannian manifold wherein the geometry of the problem corresponds to classical information geometry [4, 11, 56]. This leads natural choice of Riemannian metric: the Fisher-Rao metric, defined as, on the tangent space at an interior point p E Ad,
Vu, E E TTA", (FR(D)(U, o] = (),
u
(5)
p
2
In the above equation, the inner product normalisation by vp is applied component-wise. After normalising by vp the inner product on the simplex becomes synonymous with the familiar Euclidean inner product 12- However, near the boundary of the simplex, this "tautological" parameterisation of the metric by the components of p is numerically unstable due to division by zero This motivates a search for a more numerically stable solution which we find through the sphere-map in $3.2. As Riemannian metric is choice of inner product that varies smoothly, it can be readily used to define geometric quantities of interest such as distances between points or angles, as well as a metric-induced norm. We refer the interested reader to $B for more details on the geometry of Ád
3.2 Flow Matching from A S via the sphere map
The continuous parameterisation of categorical distributions to the interior of the probability simplex, while theoretically appealing, can be prone to numerical challenges. This is primarily because in practice we do not have the explicit probabilities of the input distribution, but instead, one-hot encoded samples which means that we must flow to a vertex. More concretely, this implies that when t + 1, I₁ + e for some E [d], therefore implying ||u||,, oc, where I1-11z, denotes the norm at point 24- This occurs due to the metric normalisation vp, applied component-wise In addition, the restriction of Uy to be the tangent space imposes architectural constraints on the network. What we instead seek is a flow parameterisation without any architectural restrictions or numerical instability due to the met- ric norm. We achieve this through the sphere map, Á + s4. which is diffeomorphism between the interior of the simplex and an open subset of the positive orthant of a d-dimensional hypersphere.
9 A
sd.
p p 8 (p) vT, 8 8 p (s) 22.
(6)
gd
Ad
In Eq. 6, both the sphere map and its inverse are operations that are applied element-wise. The sphere map reparameterisation identifies the Fisher-Rao geometry of d to the geometry of a hypersphere, whose Riemannian metric is induced by the Euclidean inner product of Rd+1 It is easy show that 2 (the sphere map scaled by 2) preserves the Riemannian metric of A i.e., is an isometry, and that therefore all geometric notions such as distances are also preserved. However, a key benefit we obtain that we can extend the metric the boundary of the manifold without introducing numerical instability as the metric at the boundary does not require us to divide by zero
Building conditional paths and vector fields on Sd. On any Riemannian manifold Md that admits a probability density, it is possible to define a geodesic interpolant that connects two points between samples 2g Po 21 ~ P1. A point traversing this interpolant, indexed by time E [0, 1]. can be expressed as Z22 = cxpx. (t log.., (. On general Riemannian manifolds, it is often not possible to obtain analytic expressions for the manifold exponential and logarithmic maps and result, travers- ing this interpolant requires the costly simulation of the Euler- Guler-Lagrange equations. Conveniently, under the Rish metric Á admits simple analytic expressions for the exponential and logarith- mie and consequently the geodesic interpolant. Moreover, due to the sphere-map & in eq. (6) the geodesic interpolant is also well-defined on Sd. Such a result means that the conditional flow x; on $4. can be derived analytically from the well-known geodesics on hypersphere, iee,, they are the great circles but restricted to the positive orthant. Consequently, we may build all of the conditional flow machinery using well-studied geometric expressions for $ in numerically stable manner.
5

---

The target conditional vector field associated at can also be written in closed-form (z x logz, (z1)/(1 t) and computed exactly on $ Intuitively, ut moves at constant velocity from 24 in the direction of xj and presents simple regression target to learn the vector field up. One practical benefit of learning conditional vector fields on S is that it allows for more flexible parameterisation of the vector field network ri. Specifically, the network to can be unconstrained and output directly in the ambient space gd+1 after which we can orthogonally project them to the tangent space of zr. This is possible since we can take an extrinsic view the geometry and isometrically embed 84 to the higher dimensional ambient space due to the Nash embedding theorem [35]. More formally, we have that oo xT)) = Dx, (Co (t, (6, x4)), where up is the output in IRd and Φ2,- 12d + Tz, sd and is defined as 62, (7) = U- - (Ize 8)22z
In the absence of any knowledge we can choose an uninformative prior on s which is the uniform density over the manifold pi (xI) = v G(z1)/ Jsg v G(zI), where G is the matrix repre- sentation of the Riemannian metric. However, a key asset of our construction, in contrast, to [23, 68], is that p1 can be any source distribution since we operate on the interpolant-level by building geodesics between two points, 2o, 21 E $4. We now state the Riemannian CFM objective for $
Egg (6) o(00..)) Pe (z. la m)) Ilvo .x) log,, (a1) 1)
t ~~ ((0, 1).
In a nutshell, our recipe for learning conditional flow matching for discrete data first maps the input data to $4. Then we learn to regress target conditional vector fields on St by performing Riemannian CFM which can be done easily as the hypersphere is a simple geometric object where geodesics can be stated explicitly. At inference, our flow pushes forward a prior on p1 $4 desired target, Po. which is then finally mapped back to Á using y A discrete category can then be chosen using any decoding strategy such as sampling using the mapped categorical or greedily by simply selecting the closest vertex of the probability simplex A to the final point at the end of inference.
3.3 The Fisher-Rao metric from Natural gradient descent
We now motivate the choice of the Fisher-Rao metric as not only a natural choice but also the optimal one on the probability simplex. We show that gradient descent of the general form 50 + argmin|oojs [( + 50) (for L()) [(pe) as in Eq. 7) converges to the gradient flow (of parameterised probabilities Pg. or of probability paths po.:) with respect to the Wasserstein dis- tance on ((M) induced by Fisher-Rao metric gFR over M. Equivalently, we get the canonical metric over 4 due to the isometry. This presents a further justification for the use of the Fisher-Rao metric.
In order to present the gradient flow of L P(M*) + R in which (M .g) is Riemannian manifold, we recall the basics of geometry over probability spaces [5, 76]. If d, is the geodesic distance associated to g then W2, will be the optimal transport distance over P = P(M)) with cost &y (z,z) y). Then (P(M)), W2,)) is an dinitte-diensionall Riemannian manifold, in which for P E P(M we have the tangent space T,P ~ [V & (M")] (T(TTMM " i.e., the closure of gradient vector fields with respect to :d norm. This norm defined by the Riemannian tensor gd induced by which at v, T,P is given by g (v. w) == fya w(x)); dp(z). In particular, note that a choice of Riemannian metric g over Md specifies unique metric g" over P.
In the following, at the onset, we assume a bounded metric, g, over A which we use only to state our Lipschitz dependence assumptions. If we compare categorical densities (elements of M = ((A) via -i-divergenc,, then it is natural to compare distributions μt, E P(M") via the Wasserstein-like WKL (j,, v)= min (t.)) [TKL (pwllpw))) In the next proposition, we show that the Fisher- Rao metric appears naturally in the continuum limit of our gradient descent over P(M)).
Proposition 1. Assume that there exists bounded Riemannian metric g over Ad such that the p(0) is Lipschitz and differentiable from O to (P(M), Wa.). parumeterisation map 0 + + p=
Then the "natural gradient" descent of the form:
p(On+1) argmin {C(p(0n-1)) WKL,(p(0++1) ee}
(8)

---

approximates, as c 0", the gradient flow of L on manifold (P(M)), W gen.2) with metric FFR induced by Fisher-Rao metric 9FR:
a d. p. V DER ((p(O(s)))
(9)
For the proof, see $C. We distinguish the results of Proposition from those of Natural Gradients used in classical NN optimisation such as KFAC [52]. Note that in regular NN training, Natural Gradients [58] implement second-order optimisation to tame the gradient descent, at a nontrivial computational cost. Thus, the above proposition implies that, just by selecting JFR metric over Ad. we directly get the benefits that are equivalent to the regularisation procedure of Natural Gradient.
3.4 FISHER-FLOW Matching with Riemannian optimal transport
We now demonstrate how to build conditional flows that minimise a Riemannian optimal transport (OT) cost. Flows constructed by following an optimal transport plan enjoy several theoretical and practical benefits: 1. They lead to shorter global paths. No two paths cross which leads to lower variance gradients during training. 3. Paths that follow the transport plan have lower kinetic energy which often leads to improved empirical performance due to improved training dynamics [64].
The Riemannian optimal transport for FISHER-FI.OW can be stated for either Á under the Fisher-Rao metric or S Both instantiations lead to the same optimal plan due to the isometry between the two manifolds. Specifically, we couple q(1)) via the Optimal Transport (OT) plan π (I0, x1) under square-distance cost c(x, y) == dl v).ie. π (zo, 21) will be the minimiser of ((z,,z))<< (x )] amongst all couplings n' of fixed marginals ((zo). Now, recall that Wasserstein distance W2 over P(S4) is defined as W2(nu, v) = min, ((.,)))aa (d (x. v)),
in which the minimisation is amongst transport plans from μ v, defined as probability measures over $4 x 8d whose two marginals are respectively μ and v [75]. Since 3d. is a smooth bounded uniquely geodesic Riemannian manifold with boundary, the metric space (P(S)), W2) is uniquely geodesic and we have the following "informal" proposition (see $D for the full statement):
Proposition 2. For any two Borel probability measures Po.P1 P((4), there exists a unique OT-plan between Po.Pi- If e e2 21) is the constant- speed parameterisation of the unique geodesic of extremes To and and e 84 8 s4 is given by e1(x0,z1)= = exp2o (t logz, (x1)), then there exists a unique Wasserstein geodesic (Pt)ce(o.1) connecting po to PL, given by
Pt (ct) E P(S)), te 10.1].
(10)
The complete statement of Proposition 4 along with its proof is provided in $D. As a consequence of Proposition 2 use the Wasserstein geodesic our target conditional probability path. Operationally, this requires us to sample from marginals 2o ~ and ø ~ Pı and solve for the OT plan using the squared distance on Sd the cost which is done using the Sinkhorn algorithm [32].
3.5 Training FISHER-FLOW
Generalising to sequences. Many problems in generative modeling over discrete data are concerned with handling set or sequence of discrete objects. For complete generality, we extend FISHER- FLOw to a sequence of discrete data by modeling it as a Cartesian product of categorical distributions. Formally, for a sequence of length à we have a distribution over product manifold P(A)
P(A))x x-. x xP((A)) Equipping each manifold in the product with the Fisher- metric allows us to extend the metric in a natural way to A. Moreover, by invoking the diffcomorphism using the sphere- achieve the product of d-hyperspheres restricted to the positive orthant. Stated explicitly, a sequence of categorical distributions is P(S+) == ((())) x ... x (( ($4))) )x)- Due to the factorisation of the metric across the product space, we can build independent flows each manifold ş and couple them in a natural way using the product metric to induce a flow on S..
map y independently
Training. We detail our method for training FiSHER-FI.ON in Algorithm 1 in $F2. Training FISHER-FiE FL.Ow requires two input distributions: a source and a target one. In the case of unconditional generation, one can take po (($)),, ), by default. In some settings, it is possible to incorporate
7

---

(a) Real distribution
(b) FF à
FF-OT A
(d) FF on 8]
(e) FF-OT on S
Figure 2: Synthetic experiments on learning a distribution resembling smiley face on A
10P
301
10
Mitinal XL diergeoce
301
10-2
10-9
LaBic HLEP PM
(00) Splant
10-4
Sphiny
Leass
Fisman Flom (OC, Sdin))
Livisn
Nimrler of categeiss, K
(a) Results of the ablation study. Missing points for Multinomial Flow [41] and D3PM [8] are NaNs.
(b) KL divergence of FISHER NISHER-FLOW against DIRICHLET FM.
Figure 3: Toy experiment from Stark et al. [68]. Minimal KL divergence over 5 seeds is reported.
additional conditional information, This can easily be accommodated in FISHER--FLOW by directly inputting this into the parameterised vector field network with vg(+) becoming 0g (-(c).
Experiments
We now investigate the empirical caliber of FISHER-FLOW on range of synthetic and real-world benchmarks outlined below. Unless stated otherwise, all instantiations of FISHER-FLOWW the optimal transport coupling on minibatches. Exact implementation details are included in $F.
4.1 Synthetic experiments
Density estimation. In this first experiment, we model an empirical categorical distribution visualized on Á Figure we observe that FISHER-FLE FLOW instantiated on S? with is the best modeling the ground truth distribution. Both learning the simplex and the positive orthant benefit from OT.
Density learning in arbitrary dimensions. We also consider the toy experiment of Stark al. [68]. where we seek to model random distribution over (A))a for K N The KI. divergence between the estimated distribution over 512,000 samples and the true generated distribution used as the evaluation metric. Details provided in $F.411. Results in Figure 3b demonstrate that FISHER-FLOW outperforms DIRICHLET FM, while remaining competitive against D3PM [8] and Multinomial Flow [41], especially in high dimensions, in which both exhibit unstable behaviour. We also conduct an ablation Figure 3a and find that using optimal transport helps for both FISHER-Fi FF.OW on Á and $d with the latter leading to the best performance.
4.2 Promoter DNA sequence design
We assess the ability of FISHER- HISHER-FLLW to generate DNA sequences. Promoters are DNA sequences that determine where gene DNA is transcribed into RNA; they contribute to determining how much transcription happens [36]. The goal this task is to generate promoter DNA sequences conditioned a desired transcription signal profile. Solving this problem would enable one to control the expression level of any synthetic gene, e.g., the production of antibodies. For detailed dataset background, see SF. in Avdeyev et al. [9].

---

Table 1: MSE the transcription profile condi- tioned on generated promoter DNA sequences over the test set. The last MSE and PPL values are from 5 independent experiments. The remaining numbers are taken directly from Stark et al. [68].
Table 2: Perplexities (PPL) values for different
methods for enhancer DNA generation. Lower PPL is better. Values are an average and standard error over 5 seeds.
Method
Melanoma PPL. (J)
Fly Brain PPL (1)
Medel
MSE (1)
PPL
Rundom Sequence
885.5 88
885.88
DIRESION [MI-ENCOOING
0.011
Language Model
2.19 0.10
DImISION CONE-SIOr EXCODING)
0,900
DIRICHLET FM
2.25 2: 0.01
2:25 0.00
DSPM-INNITOMM
0,08
FISHER-FLe FLOW (ours)
1.4 0.1
à 0.66
0,033
LANGUAGE MODEL
0.001
DIRICHLET
0,004
0.006
PISHEx-
1.7
Table 3: Results on QM9. Higher better. The baseline numbers are taken from the cited papers. The numbers reported for FlowMol are those for the uniform distribution and end-point parameteri- sation. Our numbers for 1,000 molecules.
Figure 4: Generated molecules using FISHER- FLOW QM9.
Method
Atoms (S)
Muls Val (%)
Mols, S ((3)
Fistru--F Pi ow (ours)
98.6
35.3
58.2
30DO [44]
90.1
98.9
98.7
IquiFM
99.4
FlowMal [20]
Results. Our experimental evaluation closely follows prior work [9, 68). We report the MSE between the signal our conditionally generated sequence and the target one, a human genome promoter sequence (MSE in Table 1), both given by the same pre-trained Sei model [25]. We train model on 88,470 promoter sequences, each of length 1,024, from a database of human promoters [40], each sequence having an associated expression level indicating the likelihood of transcription at each DNA position. As shown in Table 1, FISHER--FLON outperforms baseline methods DDSM [9] and DIRICHLET FM [68] on the MSE evaluation. Perplexities (PPL) from FISHER-FLOW on the test set are also better than the baselines and, on average, improve over DIRICHLET FM.
4.3 Enhancer DNA design
Enhancers DNA sequences that regulate the transcription of DNA in specific cell types (e.g,, melanoma cells). Prior work has made use of generative models for designing enhancer DNA sequences in specific cells [71]. Following Stark al. [68], report the perplexity over sequences as the main measure of performance. We also include results on Fréchet Biological Distance (FBD) with pre-trained classifiers provided in DIRICHLET FM [68], ef. $F..3. Nevertheless, those classifiers perform poorly on cell-type classification of Enhancer sequences, with test set accuracies of 11.5% and 11.2% on the Melanoma and Fly Brain datasets, respectively; thus, metries derived from these are not representative of model quality, which we still included in $F.433 for transparency.
Results. We report results in Table 2. with FBD reported in Table 5. We observe that DUSHER-FLOW obtains significantly better performance than DIRICHLET FM, which highlights its ability to fit the distribution of Melanoma and FlyBrain DNA enhancer sequences. Moreover, also note that our method improves over the language model baseline on both datasets, which bolsters the belief that FISHER-Fie -FOOW be used in similar settings to those of autoregressive models.
4.4 De novo molecule generation
In this experiment, we evaluate FISHER-FLOW's ability to generate molecules unconditionally. a.k.a. "de novo". The difficulty in this task is that we are interested in generating the positions of the molecules, their atom types, their charges, and the bonds between these, resulting in a high dimensional space with both discrete and continuous data (Rd)) x (Λ²)" x (Ae)" x (At)r², where N' is the number of atoms, possible atom types, e charges, and bonds. We train model over the QM9 dataset [61, 60]. We report the percentage of stable atoms within molecules, valid molecules, and stable molecules. Our implementation mostly based on that of [29].

---

Results. We report our results in Table 3. We also provide some qualitative examples in Figure 4. As we can see, FISHER-FLOW compares well on all metries to SIMPLEX-FLON on all metries. Nonethe- less, it must be reported that the latter, trained with Gaussian prior, endpoint parameterisation and cosine time schedule performed substantially better than both flow-based methods, closing the gap with the other baselines. It is likely that more extensive exploration of priors, time parameterisations and other hyperparameters would increase FISHER-FLOW performance.
4.5 Language modelling
Finally, we test the language modelling capa- bilities of FISHER-FI.OW. To do so, train the model on the LMIB dataset [24], a large language modelling dataset containing about 800,000 words. For this experiment, we extend FISHER-FLOW
Table 4: Test perplexities on the LM1B dataset. All baselines are taken from concurrent work MDLM by Sahoo et al. [62]. Best diffusion or flow- matching method is in bold font.
masked path is done by [62, 65]: we define the probability path as Pr KyPar + (1-Ak (1-Kk)Punnis where k : [0, 1] - [0, 1] is noise scheduler. Here, PM is the Fisher- Rao geodesic between the target, To, and the designated mask token M, while Punit is also a Fisher-Rao geodesic between a sample from uniform distribution and zo- It is thus con- vex combination of probability paths. Using denoising architecture enables us to rewrite the original loss as a weighted negative log likelihood - EEllog xt)]. This allows us to calculate an upper bound on the test perplexity, natural evaluation metric for language modelling [62, 65].
Method
Parameters
PPL
BERT BERT-MOUTH MOUTH
110M
142,,9
Dilusson
D3PM (ABSORB)
70M
77.50
DIFFUSION-LM
DIPUUSIOBBRRE
110M
63.78
SEDD (33B TOKENS)
110M
2.3279
DM/FM AR
TRANSFORMER (33R TOKENS)
110M
22.32
TRANSFORSER (327B TOKENS)
110M
20.86
MDLM (33B TUKENS)
110M
27.04
FISHER- FLOW (33B TOKENS) fours)
MDLM (327B TOKENS)
110M
S 2..00
FISHER-FLOW (3278 TOKENS) (ours)
110M
5:2-42
Results. The results are given in Table As one can observe, using the Fisher-Raa metric enables better performance than MDLM. Yet, the gap with auto- nuto-regressive methods is still significant.
5 Related work
Geometric generative models. There are several methods for defining generative models over Riemannian manifolds, the most pertinent to this work include diffusion models [43, 28], normalising flows [20, 53, 15, 26]. For molecular tasks that require generating nodes and edges, equivariant variants of diffusion and flow-based models are a natural choice [42, 78].
Discrete diffusion and flow models. Discrete generative models diffusion and flow models can be categorised into either relaxations continuous spaces [47, 27]. methods that use continuous-tin Markov chains with sophisticated transition kernels [8, 80, 23, 50], with some matching autoregressive models [34]. Defining discrete data on the simplex has also been explored in the context of generative models [37, 51, 68]. FISHER--
is fundamentally different from existing works [8, 23, 2, 50] in that we consider a continuous relaxation of the discrete space and construct vector fields on S Finally, concurrent to our work Dunn and Koes [29] propose simplex flow matching, and Boll et al. [17] introduced e-geodesic flows that leverage the Fisher-Rao Rao metric on the assignment manifold [17]. Simplex flow-matching differs from FISHER-FU FLOW in that it does not make use of the Fisher-Rao metrie. We include detailed comparison between FISHER-FLOW in relation to DFM and e-Geodesic Flow Matching [17] in I.
6 Conclusion
In this paper, we introduce FISHER-FLOW novel generative model for discrete data. Our approach offers a novel perspective and reparameterises discrete data to live on the positive orthant of d-hypersphere, which allows us to learn categorical densities by performing Riemannian flow matching. Empirically, FISHER-FLOW improves performance on synthetic and biological sequence design tasks over comparative discrete diffusion and flow matching models while being more general as a framework. While FISHER- FFOO enjoys favorable theoretical properties with strong empirical performance, our method is not fully developed for language modeling domains. Consequently, natural direction for future work to design variations of FISHER- FF.OW capable of handling larger sequence lengths and discrete categories found in language domains.
10

---

Acknowledgements
We thank Alexander Tong for his generous time, help, and guidance in helping with the language modeling experiments. OD is supported by both Project CETI and Intel. MP is supported by CenIA and by Chilean Fondecyt grant n. 1210426. AJB is partially supported by an NSERC Post-doc fellowship. This research is partially supported by EPSRC Turing AI World-Leading Research Fellowship No. EP/X040062/1 and EPSRC AI Hub on Mathematical Foundations of Intelligence: An "Erlangen Programme" for AI No. EP/Y028872/1

---

References
[1] Achiam, Adler, S. Agarwal, L. Ahmad, Akkaya, F. L. Aleman, D. Almeida, J. Altenschmidt, S. Altman, S. Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023. (Cited on page 1)
[2] S. Alamdari, N. Thakkar, R. van den Berg, A. X. Lu, N. Fusi, A. P. Amini, and K. K. Yang. Protein generation with evolutionary diffusion: sequence all you need. bioRxiv, pages 2023-09, 2023, (Cited on pages 2 and 10)
[3] S.-I. Amari. u-divergence is unique, belonging to both f-divergence and bregman divergence classes. IEEE Transactions on Information Theory, 55(11))4925-4931 2009. (Cited on page 21)
[4] S.-I. Amari. Information geometry and its applications, volume 194. Springer, 2016. (Cited on pages 3, 4, 17, and 21)
[5] L. Ambrosio, N. Gigli, and G. Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2005. (Cited on pages and 18)
[6] Åström, S. Petra, B. Schmitzer, and C. Schnörr. Image labeling by assignment. Journal of Mathematical Imaging and Vision, 58:21 1-238, 2017. (Cited on page 2)
[7] Z. K. Atak, I. Taskiran, J. Demeulemeester, C. Flerin, D. Mauduit, L. Minnoye, G. Hulselmans, V. Christiaens, G.-E. Ghanem, J. Wouters, et al. Interpretation of allele-specifie chromatin accessibility using cell state-aware aware deep learning, Genome research, 31(6):0082-1098 2021. (Cited on page 23)
[8] J. Austin, D. D. Johnson, Ho, D. Tarlow, and R. Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 34: 17981-17993, 2021. (Cited on pages 2, 8, and 10)
[9] P. Avdeyev, C. Shi, Y. Tan, K. Dudnyk, and J. Zhou. Dirichlet diffusion score model for biological sequence generation. In International Conference Machine Learning, pages 1276-1301, PMLR, 2023. (Cited pages 2, 8, and 9)
[10] D. Axen, M. Baran, R. Bergmann, and K. Rzecki. Manifolds.jl: An extensible julia framework for data analysis on manifolds. ACM Transactions on Mathematical Sofwware 49(4), dec 2023. doi: 10.1 145/361 1255/661296... (Cited pages 17 and 22)
[11] N. Ay, J. Jost, H. Vân Lê, and L. Schwachhöfer. Information geomelry, volume 64. Springer, 2017. (Cited on pages 3, 4, 5, 17, 18, 21, and 22)
[12] J. L. Ba, J. R. Kiros, and G. E. Hinton. Layer normalization. arXiv preprint arXiv: 2016. (Cited on page 23)
[13] V. Balasubramanian. geometric formulation of occam's razor for inference of parametric distributions. In Maximum Entropy and Bayesian Methods. Springer Netherlands, 1996. (Cited
on page 3)
[14] O. Bar-Tal, H. Chefer, O. Tov, Herrmann, Paiss, S. Zada, Ephrat, J. Hur, G. Liu, A. Raj, Y. Li, M. Rubinstein, T. Michaeli, O. Wang, D. Sun, T. Dekel, and I. Mosseri. Lumiere: A space-time diffusion model for video generation, 2024. (Cited on page 1)
[15] H. Ben-Hamu, S. Cohen, J. Bose, B. Amos, A. Grover, M. Nickel, R. T. Chen, and Y. Lip- man. Matching normalizing flows and probability paths manifolds. arXiv preprint arXiv:2207.0m 0471 2022. (Cited on page 10)
[16] J. Betker, G. Goh, L. Jing, T. Brooks, J. Wang, L. Li, L. Ouyang, J. Zhuang, J. Lee, Y. Guo, al. Improving image generation with better captions. Computer Science. https://edn. openai. com/papers/dall-e-3 pdf, 2(3):8, 2023. (Cited on page 1)
[17] B. Boll, D. Gonzale--Alvaran and C. Schnörr. Generative modeling of discrete joint distribu- tions by e-geodesic flow matching on assignment manifolds, arXiv preprint arXiv:2402.0784. 077466 2024. (Cited on pages 10, 21, and 22)
12

---

[18] A. J. Bose, T. Akhound-Sadepp K. Fatras, G. Huguet, J. Rector- Brotor-Brooos,, C.-H. Liu, A. C. Nica, M. Korablyov, M. Bronstein, and A. Tong. Se (3)-stochastic flow matching for protein backbone generation. arXiv preprint 2rXiv:2310.02.990 2023. (Cited on pages 4 and 22)
[19] A. J. Bose, T. Akhound-Sadegh. G. Huguet, K. Fatras, J. Rector- Brooks, C.-H. Liu, A. C. Nica, M. Korablyov, M. Bronstein, and A. Tong. Se(3)-stochastic flow matching for protein backbone generation. In The International Conference on Learning Representations (ICLR), 2024. (Cited on page 1)
[20] J. Bose, A. Smofsky, R. Liao, P. Panangaden, and W. Hamilton. Latent variable modelling with hyperbolic normalizing flows. International Conference on Machine Learning, pages 1045-1055. PMLR, 2020. (Cited page 10)
[21] T. Brooks, B. Peebles, C. Holmes, W. DePue, Y. Guo, L. Jing, D. Schnurr, J. Tay- lor, T. Luhman, E. Luhman, C. Ng. R. Wang, and A. Ramesh. Video genera- tion models world simulators. 2024. URL https:// video- inmeerationmoddell as- world- as-woldd
com/research/
minldd-smulators. (Cited on page 1)
[22] S. Bubeck, V. Chandrasekaran, R. Eldan, J. Gehrke, E. Horvitz, E. Kamar, P. Lee, Y. T. Lee, Y. Li, S. Lundberg, et al. Sparks of artificial general intelligence: Early experiments with gpt-4. arXiv preprint arXiv: 2023. (Cited page 1)
[23] A. Campbell, J. Yim, R. Barzilay, T. Rainforth, and T. Jaakkola. Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design. arXiv preprint arXiv:2402.04499. 2024. (Cited on pages 2, 4, 6. 10, and 20)
[24] C. Chelba, T. Mikolov, M. Schuster, Q. Ge, T. Brants, P. Koehn, and T. Robinson. One billion word benchmark for measuring progress in statistical language modeling, 2014. (Cited on page 10)
[25] K. M. Chen, A. K. Wong, O. G. Troyanskaya, and J. Zhou. A sequence-based global map of regulatory activity for deciphering human genetics. Nature genetics, 54(7):940 949, 2022. (Cited on page 9)
[26] R. T. Chen and Y. Lipman. Riemannian flow matching on general geometries. arXiv preprint arXiv:2302.03666. 2023. (Cited on pages 2, 3, and 10)
[27] T. Chen, R. Zhang, and G. Hinton. Analog bits: Generating discrete data using diffusion models with selficooniiioning. arXiv preprint arXiv:2208.04200 2022. (Cited on page 10)
[28] V. De Bortoli, E. Mathieu, M. Hutchinson, J. Thornton, Y. W. Teh, and A. Doucet. Riemannian score-based generative modelling. Advances in Neural Information Processing Systems, 35: 2406-2422, 2022. (Cited on page 10)
[29] I. Dunn and D. R. Koes. Mixed continuous and categorical flow matching for 3d de novo molecule generation. arXiv preprint arXiv:2404. 19739, 2024. (Cited on pages 9, 10, and 24)
[30] Esser, S. Kulal, A. Blattmann, R. Entezari, J. Müller, H. Saini, Y. Levi, D. Lorenz, A. Sauer, F. Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. arXiv preprint arXiv:2403.03206, 2024. (Cited on page 1)
[31] R. Flamary. N. Courty, A. Gramfort, M. Z. Alaya, A. Boisbunon, S. Chambon, L. Chapel, A. Corenflos, K. Fatras, N. Fournier, I.. Gautheron, N. T. Gayraud, Janati, A. Rakotomamoniy, I. Redko, A. Rolet, A. Schutz, V. Seguy, D. J. Sutherland, R. Tavenard, A. Tong, and T. Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1-8, 2021. URL http://jmir wwtg/appers//2222-666
html. (Cited on page 22)
[32] A. Genevay, G. Peyré, and M. Cuturi. Learning generative models with sinkhorn divergences. In International Conference on Artificial Intelligence and Statistics, pages 1 1608-1617. PMIR, 2018. (Cited on page 7)
[33] I. Granieri. On action minimizing measures for the monge-kantorovich problem. Nonlinear Differential Equations and Applications NoDEA, 14:125-152 152, 2007. (Cited on page 20)
13

---

[34] I. Gulrajani and T. B. Hashimoto. Likelihood based diffusion language models. Advances in Neural Information Processing Systems, 36, 2024. (Cited on page 10)
[35] M. Gunther. Isometric embeddings of riemannian manifolds, kyoto, 1990. In Proc. Intern. Congr. Math., pages 1137-1143. Math. Soc. Japan, 1991. (Cited on page
[36] V. Haberle and A. Stark. Eukaryotic core promoters and the functional basis of transcription initiation. Nature reviews Molecular cell biology, 19(1))621-637, 2018. (Cited on page 8)
[37] X. Han, S. Kumar, and Y. Tsvetkov. Ssd-Im; Semi-autoregressive simplex-based diffusion language model for text generation and modular control. arXiv preprint arXiv:22 17432, 2022. (Cited on page 10)
[38] J. Ho, Jain, and P. Abbeel. Denoising diffusion probabilistic models, 2020. (Cited on page
[39] A. Holtzman, J. Buys, L. Du, M. Forbes, and Y. Choi The curious case of neural text degeneration. arXiv preprint arXiv:194.0097 2019. (Cited on page 2)
[40] C.-C. Hon, J. A. Ramilowski, J. Harshbarger, N. Bertin, J. Rackham, J. Gough, E. Denisenko, Schmeier, T. M. Poulsen, J. Severin, et al. An atlas human long non-coding ras with accurate 5' ends. Nature, 543(7644):199 204, 2017. (Cited page
[41] E. Hoogeboom, D. Nielsen, P. Jaini, P. Forré, and M. Welling, Argmax flows and multinomial diffusion: Learning categorical distributions, 2021. URL https://arxx:: http:/awwwv.aaaaaaa/a///a/ httg/ab//200 2202. 05379. (Cited on page 8)
[42] E. Hoogeboom, V. G. Satorras, C. Vignac, and M. Welling. Equivariant diffusion for molecule generation in 3d. In International Conference on Machine Learning, pages 8867-8887. PMLR,
2022. (Cited on page 10)
[43] C.-W. Huang, M. Aghajohari, J. Bose, P. Panangaden, and A. C. Courville. Riemannian diffusion models. Advances in Neural Information Processing Systems, 35:2750-2766 2022. (Cited on page 10)
[44] H. Huang, L. Sun, B. Du, and W. Lv. Learning joint 2d & 3d diffusion models for complete molecule generation, 2023. URL httph:::///arttti org/ abs absaaaas///// /2005 12347. (Cited on page 9)
[45] J. Janssens, S. Aibar, I. Taskiran, J. N. Ismail, A. E. Gomez, G. Aughey, K. Spanier, F. V. De Rop. C. B. Konzalez-BBlas M. Dionne, et al. Decoding gene regulation in the fly brain. Nature, 601(7894):630-633 2022. (Cited on page 23)
[46] M. Kochurov, R. Karimov, and S. Kozlukov. Geoopt: Riemannian optimization in pytorch, 2020. (Cited on page 22)
[47] X. Li, J. Thickstun, I. Gulrajani, P. S. Liang, and T. B. Hashimoto. Diffusion-Im improves controllable text generation. Advances in Neural Information Processing Systems, 35:4328- 4343, 2022. (Cited on page 10)
[48] Y. Lipman, R. Q. Chen, H. Ben-Hamu, M. Nickel, and M. Le. Flow matching for generative modeling. International Conference on Learning Representations (ICLR), 2023. (Cited on pages 1, 4, and 21)
[49] I. Loshchilov and F. Hutter. Decoupled weight decay regularization, 2019. (Cited on page 23)
[50] A. Lou, C. Meng, and S. Ermon. Discrete diffusion language modeling by estimating the ratios of the data distribution. arXiv preprint arXiv:2310.168.. 2023. (Cited on pages 2 and 10)
[51] R. K. Mahabadi, H. Ivison, J. Tae, J. Henderson, I. Beltagy, M. E. Peters, and A. Cohan. Tess: Text-to-text celf-conditioned simplex diffusion. arXiv preprint arXiv:2305.08373 2023, (Cited on page 10)
[52] J. Martens and R. Grosse. Optimizing neural networks with kronecker-factoree approximate curvature. In International conference on machine learning, pages 2408-2417. PMLR, 2015. (Cited page 7)
14

---

[53] E. Mathieu and M. Nickel. Riemannian continuous normalizing flows. Advances in Neural Information Processing Systems,
2020. (Cited on page 10)
[54] R. J. McCann. Polar factorization of maps on riemannian manifolds. Geometric & Fumctional Analysis GAFA, 11(3):589-608, 2001. (Cited on page 20)
[55] on page 1)
Midjourney, https: http:///wwwwmm////j/ooom.. ourney htt.ey.ccmmomm//.//.e 2023. Accessed: 2023-09-09, (Cited
[56] F. Nielsen. An elementary introduction to information geometry. Entropy, 22(10): 1 100, 2020. (Cited on pages 4 and 5)
[57] A. d. Oord, S. Dieleman, H. Zen, K. Simonyan, O. Vinyals, A. Graves, N. Kalchbrenner,
A. Senior, and K. Kavukcuoglu. Wavenet: A generative model for raw audio. arXiv preprini arXiv: 03499, 2016. (Cited page 23)
[58] R. Pascanu and Y. Bengio. Revisiting natural gradient for deep networks. arXiv preprint arXiv:1301.3584, 2013. (Cited on page
[59] A. Radford, J. W. Kim, T. Xu, G. Brockman, C. McLeavey, and I. Sutskever. Robust speech recognition via large-scale weak supervision. In International Conference Machine Learning, pages 28492-28518, PMLR, 2023. (Cited on page 1)
[60] R. Ramakrishnan, P. O. Dral, M. Rupp, and O. A. von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific Data, 1(1):140022, 08 2014. ISSN 2052-4463. doi: 00.1038/datt..2144.22 URL https: //do:::///d...e org/ 10. 10w.///at/aaaa.../ 2014. 22. (Cited on page 9)
[61] L.. Ruddigkeit, R. Deursen, L. C. Blum, and J.-L.. Reymond. Enumeration of 166 billion organic small molecules the chemical universe database gdb-17. Journal of Chemical Information and Modeling, 52(1 1)2664 227552-2877, 11 2012. ISSN 1549-9596. doi: 10.1021/ ci300415d. URL https::///ooo orgoo////// 1tt///....10//////1 ci300415.. (Cited on page 9)
[62] S. S. Sahoo, M. Arriola, Y. Schiff, A. Gokaslan, E. Marroquin, J. T. Chiu, A. Rush, and V. Kuleshov. Simple and effective masked diffusion language models. arXiv preprint 24X0::2406.07524, 2024. (Cited on page 10)
[63] F. Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY, 55(58-33):94, 2015. (Cited on pages 19 and 20)
[64] N. Shaul, R. T. Chen, M. Nickel, M. Le, and Y. Lipman. On kinetic optimal probability paths for generative models. In International Conference on Machine Learning, pages 30883-30907. PMLR, 2023, (Cited on page 7)
[65] Shi, K. Han, Z. Wang, A. Doucet, and M. Titsias. Simplified and generalized masked diffusion for discrete data, 2024. URL https://arxx org/ abs/a/s///200 04329. (Cited on page 10)
[66] Y. Song, J. Sohk-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole. Score- generative modeling through stochastic differential equations. International Conference on Learning Representations (ICLR). 2021. (Cited on page I)
[67] Y. Song, J. Gong.
M. Xu, Z. Cao, Y. Lan, S. Ermon, Zhou, and W.-Y. Ma. Equivariant flow matching with hybrid probability transport, 2023. URL https://arxivv wtw.aawwww..rra//////.a... 07168. (Cited on page 9)
[68] H. Stark, B. Jing, C. Wang, G. Corso, B. Berger, R. Barzilay, and T. Jaakkola. Dirichlet flow matching with applications to dna sequence design. arXiv preprint arXiv:2402.0584), 2024. (Cited on pages 2, 4, 6, 8, 9, 10, 20, 21, 22, 23, and 24)
[69] Steinhardt. AI Forecasting: One Year In
www.eet..ggoo..tt. io/
ai si-forecasting- one- oner in/, an/,20 2022. Accessed: 2023-09-09, (Cited on page 1)
15

---

[70] C. Szegedy, V. Vanhoucke, S. loffe, J. Shlens, and Z. Wojna. Rethinking the inception archi- tecture for computer vision. In Proceedings of the IEEE conference on computer vision and paltern recognition, pages 281 8-2866,, 2016. (Cited on page 5)
[71] I. Taskiran, K. I. Spanier, H. Dickmänken, N. Kempynek, A. Pančíková, E. C. Ekşi, G. Hulsel- mans, J. N. Ismail, K. Theunis, R. Vandepoel, et al. Cell-type-directer design of synthetic enhancers. Nature, 626(7997):212-255) 220122222222222 2024. (Cited page 9)
[72] G. Team, R. Anil, S. Borgeaud, Y. Wu, J.-B. Alayrac, J. Yu, R. Soricut, J. Schalkwyk, A. M. Dai, A. Hauth, et al. Gemini: family of highly capable multimodal models. arXiv preprint arXiv:2312. 1805, 2023. (Cited on page 1)
[73] A. Tong, N. Malkin, G. Huguet, Y. Zhang. J. Rector-Brooks, K. Fatras, G. Wolf, and Y. Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. arXiv preprint .332.00482. 2023. (Cited on pages I, 4, and 21)
[74] N. N. Čencov. Statistical decision rules and optimal inference. Number 53, American Mathe- matical Soc., 2000. (Cited on page 18)
[75] C. Villani. Topics in Optimal Transportation. Graduate studies in mathematics. American Mathematical Society, 2003. ISBN 9781470418045 URL https /tttpss//books http://ooooo.oogll..ee/ooo co..gee.......... co. uk/ books?id=Myy Jooker?iAKyjJjEEAAAAJ (Cited on page 7)
[76] C. Villani. Optimal transport: old and new, volume 338. Springer, 2009. (Cited on pages 6, 18, 19, and 20)
[77] L. Watson, D. Juergens, N. R. Bennett, B. L. Trippe, J. Yim, H. E. Eisenach, W. Ahern, A. J. Borst, R. Ragotte, L. F. Milles, et al. De novo design of protein structure and function with rfdiffusion. Nature, pages 1-3, 2023. (Cited on page
[78] M. Xu, Yu, Y. Song, Shi. S. Ermon, and J. Tang. Geodiff: A geometric diffusion model for molecular conformation generation. arXiv preprint arXiv:2203.02923, 2022. (Cited page 10)
[79] G. U. Yule. On method of investigating periodicities in disturbed series with special reference to wolfer's s sunspot numbers. Statistical Papers of George Udny Yule, pages 389-420, 1971. (Cited on page 1)
[80] L. Zhao,
X. Ding, L. Yu, and L. Akoglu. Improving and unifying discrete&ccomtiuuus--iiie discrete denoising diffusion. arXiv preprint arXiv:2402.03701. 2024. (Cited on page 10)
[81] F. Aström, S. Petra, B. Schmitzer, and C. Schnörr. Image labeling by assignment. Journal of Mathematical Imaging and Vision, 58(2):211-238, Jan. 2017. ISSN 1573-7683, doi: (Cited on page 17)
URL http: //tx doi. waai.org/1...0
016 0702-4.
16

---

A Broader Impacts
We would like to emphasise that our paper is mainly theoretical and establishes generative modeling of discrete data using flow matching by continuously reparameterising points onto statistical manifold equipped with the Fisher-Rao metric. However, more broadly discrete generative modeling based on diffusion models and flow matching has important implications in various fields. In biology, these models enable the generation of novel biological sequences, facilitating the development of new therapeutics. However, the same technology poses risks if exploited maliciously. as could be used to design harmful substances or biological weapons. In language modeling, the capability to generate coherent and contextually relevant text can significantly enhance productivity, creativity, and communication. Nevertheless, the advent of superhuman intelligence through advanced language models raises concerns about potential misuse, loss of human control, and ethical dilemmas, highlighting the need for robust oversight and ethical guidelines.
B Geometry of the Simplex
We introduce here very briefly properties of geometry on the simplex that we use in this paper. Our main reference for these results is Aström al. [81]. Note that our implementation for most of these properties relies on that of Axen et al. [10], which we port to Python. Recall that d-simplex, for No, is defined as A == {-r E Rd+1|1 x 1, 2 Oj. When equipped with the Fisher-Rao metric, it becomes Riemannian manifold that is isometric to the positive orthant of the d-sphere of in gd+1 That is to say, 1/ : A - $4, (Zo..... za) + (2. 2//a)) is a diffeomorphism, where $ = {I E IR ||||| = 2, x 2 0}; we call &b the "sphere-map".
In the following, Á denotes the interior of the simplex, and T,A = fo E Rd+1, :1x = 0} the tangent space at point p. The exp map the simplex is given by, for all p Å TAA
2
2
vP |
(x))(v) = 2
cos cos(llu)))
+
sin (12 Il),
(11)
2
where up == vDr and squares, square roots and quotients of vectors are meant element- Similarly, the log map given by, for p.q E Åd
ds-(p.9) (VPq (VP. /g)p), log.o (xI) = VI (/P, vg)
(12)
where the product is meant element-wise, and the distance is
d 2 aarcos(((// vq)).
(13)
The Riemannian metric at point p Å for vectors u, TA is given by
u
(u, o)pp =
(14)
vp
Finally, for parallel transport, use the sphere-map, perform invert the mphere-map.
the sphere, and
The relevance of the Fisher-Rao metric stems from the following two characterisations:
The Fisher-Rao metric the leading-order approximation of the Kullback Leibler divergence [4, 11]. Recall the general setting: if a d-dimensional manifold of probability densities Md parameterised by a differentiable map 0 + pe from a submanifold O FD (note that the requirement D = d is not necessary for the following computations to make sense), then for fixed fu we may Taylor-expand
@p(On) 0g)
p(0) p(0o)
00
203 + o(10 0o1).
17

---

and straightforward computation gives
Ölog llgg
(rI(o((0)ll(())
Σ j,k=
(0 02) -00) Bp(o)
+010 -0ol3
303
a0k
18=80
gik (66)) - 8g 0t 00 + o -0ol)).
2
j,
Thus the matrix g(0)) = (a, (8o)) defines the quadratic form on the tangent space Te, θ which best approximates ( (())) in the limit 0 - Bu- In the coordinates @ A R when we parameterise probabilities over K = d+1 classes numbered 0, O.... via the "tautological" parameterisation 0 for p E Ad, (explicitly. in this parameterisation class has probability (i(i) 8i p¹), then we obtain Blog 881 pi 5(i f) and
(p) Ep(o)
Ologg p Ô llogp
0li -
1 5)
1 -k) = p
1 8(i k))
000
00h
Thus u(p)(u, grn E n'v" before.
The Fisher-Rao metric up rescaling, the only metric that is preserved under sufficient statistics. First, for 2 I K' < K, define a map M. P((K)) - P[[K]) to be a Markov map if there exist probability measures 91 .1..... 9K' P((]) such that for p P((K')) we have M(p) = Ek= p(k)q). In other words, representing probability spaces as simplices and denoting d d= k= 1, K' 1, we have that M is Markov map if the simplex A is affinely mapped under M to a d' -imeennnnional simplex in A (the vertices of the image simplex have been denoted above by Q1..... ......... q.²).
Then a restatement of Chentsov's 's theorem [74, Thm. 11.1], [11, Thm. 1.2] is that if a sequence of Riemannian metries Ja over A defined for d 2 2 satisfies the property that for any d' S any Markov morphism M : A - A is an isometry with respect to metries ga, Ja. then there exists C' > O such that each of the Dd is C' times the Fisher-Rao metric on A common reformulation, interpreting the Markov map reparameterisations M P((K))
P[[K] of P(K"") as sufficient statistics, is to say that Fisher-Rao metrics are (up to common rescaling for all d) the only metrics that are invariant under mapping probability measures to sufficient statistics.
C Details and proofs for Section 3.3
We here recall the setup: we are considering a loss function C P(M") R. in which M is Riemannian manifold, specifically it will be the simplex A endowed with Riemannian metric g. Points Pw E Md represent categorical distributions, as M was obtained from ((A) by parametrising it with the simplex Ad, thus inducing a differentiable structure.
The space ((M") is then endowed with the Wasserstein distance W2,, induced by the Riemannian geodesic distance of (Md, .). Then P = (P(M)), W2.9) can be given Riemannian metric structure too, defined as follows [5, 76]. For p E P the tangent space T,P is identified with the 2 g)- closure of the space of vector fields t: 'dd - TMd which are the gradient a C!-function
ab : Md + I.. Here we have for = Ve E T,P
| ( S [|u(p)) d()
and the corresponding Riemannian tensor induced by g over u, TpP is given by (0(,, w) / MM (u(pw), w(pw)) d(p=)))
For further details see Villani [76, Ch. 13].
In order to find well-behaved metric over Md, we start by considering Ma (which in our case the statistical manifold parameterising the space of categorical probabilities P(A)) with the KI.- divergence comparison tool for its clements. We will use this divergence in order regularise
18

---

the gradient descent of a loss function L: P(M") + R, and to do so we introduce the KL-optimum coupling which for jr, v E P(M4)) takes the value
[DKI. (Pull=))) a E P(M x Md) has marginals f4, "} In words, WKL. determines the smallest average displacement required for moving to u, in which displacements between elements of M = (PA)) are quantified by diverrivege
WKLY(U, v) == min (p
We then use this distance to regularise the gradient descent of C, and show that then the gradient descent converges to the Wasserstein gradient flow on L, for precisely the Wasserstein distance W2.9en induced by the 9FR -etric over M
Here we consider a Riemannian metric structure g on Md, which we assume be bounded on the interior Ád, i.e. to have bounded coefficients when expressed in the parametrisation, which is only used in order to give rough Lipschitz hypothesis on the underlying parametrisations.
Proposition 3 (extended version of Proposition 1). Assume that g is a bounded Riemannian metric over A such that the parametrisation map 0 p = p(0)= A - (P(M)), W2,g) is
Lipschitz and differentiable. Then the "natural gradient" descent of the form:
p(0n+1) E argmin (L(p(On-1)) WKI. approximates, as E + 0*, the gradient flow of C on manifold (P(M)), W 2F with metric PPn induced by Fisher-Raa Rao metric gER:
[e}
(15)
d ds
p(((s)) = Vo L(p(B(s)))
(16)
Proof. We restrict the discussion to the case that p(0) is supported in the region A : {o E Rd
If r = 1, x; 2 c, 1 S S d), and the general result can be recovered by taking 0+. Restricted to this set, it is easy to verify that DxL is bounded.
Step 1. Note that by a small modification of the proof, we can apply Villani [76, Thm. 10.42) to A with cost equal to DKL, and obtain that the WKL.dditance between an admissible competitor p(0 50) in Eq. 15 and p(0..) is realised by a transport plan 760, such that we have p(0 + 50) = =TTpp(0) By definition of WKL and due to Chebyshev's inequality, for all C > O, the set of points Sc that 7 moves by more than C' in Dit-distance has p(0) -eessure not larger than L/C. Furthermore, 76 is uniformly bounded over A \ Sc by our initial hypothesis. By approximating this transport plan by a flow (one can adapt the ideas from e.g., Santambrogio [63, Thm. 4.4] for this contstruction) over Sc, can find vector field 90 such that 2d (pw)
logp.. (Tw()))) + oz (168)) for Per E Sc, with error uniformly bounded in Par E ... We then extend 59 arbitrarily outside Sc- This procedure associates to each small enough change 50 vector field Usa E po) which whose time-c flow, denoted Oesa (t = -,- pushes measure p(0) to measure approaching p(0 + 50) in the limit c + O, C + D0.
Step 2. We approximate the optimisation problem Eq. 15. For the constraint, we recall that as noted in Appendix B, we have Taylor expansion (pw [ 1 ||l w'l + O[[ w w1.)). For approximating L we use its differentiability and get L(p(8)) L(p(O)) + +d[(p(()) for T.P. Thus minimisation problem Eq. 15 is well approximated, (in the limits mentioned in the previous step) by
p(0n+1)= = (dhr (1, p(Om), Uso E argmin, (c (cd(((())))) : (0, s) (), =1).
(17) in which we used rescaling compared to previous step, given by ev. This means that we used the associated flow up to time rather than time e, and thus the minimisation has to be taken amongst elements E Tp and we approximate the constraint by (v, u) = 1, which replaces the correct constraint WKL (p(0 + 80), p(0))
Step 3. In the optimisation Eq. 17, we have a quadratic constraint over the vector space Tp(..)) P. and thus we can use Lagrange multipliers, and for the optimiser we need to look for critical points of + edC(p(0)) + 2/2 (v, 1) in which X is the Lagrange multiplier, to be fixed at the end using the constraint. This gives the following characterisation of the optimiser ugg=
Vu P, (Ugo. w) dd(p()))
V
DFn L(p(8)),
(18)
e
19

---

in which we just use the classical definition of the gradient on manifold.
This means that in the approximation of the step p(0) p(0 + 50) must move in the negative-grrs -rradiiinnnn direction L at p(0), as desired.
D Optimal Transport proofs
Proposition 4 (extended version of Proposition 2). For any two Borel probability measures po. P1 E P(S+), the following hold:
1. There exists unique OT-plan n between po.P1-
2. For E [0, 1] ler (u( 21) be the constant-sper masstant-spedd parameterisation of the unique geodesic of extremes zo and æ1, defining the map
6 : x S + S+, ee(2o, ) == exPao (1log (r1)).
(19)
Then there exists unique Wasserstein geodesic (pr)se(o,1) connecting po to pt, and it is given by
Pt = (e1) E P(S_), t E 10.
(20)
3. For every point 21 in the support of pt, there exists unique pair (Zo, 21) in the support of the optimal transport plan π such that I, = ex(z,, 21). Furthermore, the assignment 14. (z0, .1) is continuous in It-
4. The probability path (pe)te(o.1] has velocity field ua logz, (I1) - -logz, (No), which is uniquely determined over the support of pt-
5. The above probability measure path and associated velocity fields (pe, us ue)eelo.1 are minimisers of the following kinetic energy minimisation problem
{/1 p.p [I O.p. + div(per)) =0, 0, Po Po: PI p1 (pe,ee)ceoo.)
min
(21)
Proof. For point 1, we can use Villani [76, Thm. 10.28] (the simpler Villani [76, Thm. 10.41] also applies, with the minor modification that we work on manifold with boundary). To verify its conditions, note that Md C S is a subset of Riemannian manifold and has (d-11)dimension) measure, and that cost c(x, y) = dl (xx y) is convex, thus it has unique superdifferential and Vzc(z,-) injective, as required.
For points 2 and 3, we note that by Villani [76, Cor. 7.22] (see also McCann [54]), in general Polish spaces displacement interpolants as given by Eq. 19 and Eq. 20, coincide with Wasserstein geodesics.
simplified version of the proof of 4. is present in Santambrogio [63, Prop. 5.30J. For the general case, we can Villani [76, Thm. 10.28), in particular eq. (10.20) therein. Note that for c(x, y) = d y), as indicated in Example 10.36 this equation corresponds to the equation of geodesics in the underlying manifold. Then we just note that u4 is the velocity field of constant speed geodesic.
Point 5 is special case of Villani [76, Thm. 7.21], see also Granieri [33].
E Relation to prior work on the simplex
E.1 Dirichlet Flow matching
In this appendix, we discuss how flow matching can be done on the simplex using Dirichlet conditional
probability paths. This recovers the simplex flows designed in Stark et al. [68], Campbell et al. [23].
The equivalent of uniform density over A is given by a Dirichlet distribution with parameter vector 1, i.e. Pi (x1) Dir(z1;a a 1). This is the starting point for defining flow between our data
20

---

distribution, Po. and the Dirichlet prior As proven in Stark et al. [68] we can reformulate Eq. 3 using a cross-entropy objective,
(o) Br, (lo m4) un (xrl=))
(22)
(x)l=)) log (2
(23)
Here, we parameterise a denoising classifier which predicts a denoised sample zo from It. which is built using the conditioner Such a parameterisation naturally restricts the vector field to move tangentially to the simplex and also training is simplified as we do not need to explicitly construct
the conditional vector field ur (x)(2) during training. At inference, we can recover (t, 24)
Σ "t (z IzD = (z) = e1(z²) and follow the vy by integrating time to
Designing conditional paths. There are two primary points of attack when designing flow-matching model. We can either define an interpolant (x)ea)) with initial conditions tu Zo. which we can differentiate to obtain ut, i.e., (>))//)) = (x(z))=); or we can operate on the distributional level and specify conditionals (zzl22) (z) from which suitable vector field can be recovered.
If we take the interpolant perspective. one can easily implement the linear interpolant [48, 73), which gives the following conditional vector field:
Lxp + (1 ue ue(xe/20. x1) Ft
( v1)
(24)
20
20
21
(25)
t
Unfortunately, in the case of flow matching on the simplex, the linear interpolant has undesirable properties in that the intermediate distribution induced by the flow must quickly reduce support over A by dropping vertices progressively for > 0 [68].
Operating directly on the distribution level, we can define pt themselves being Dirichlet distribu- tions indexed by t such that, O, we have uniform mass over A and that, at 1, we reach a vertex. One choice of parameterisation that fulfills these desiderata is
p=(2 e,) a
(26)
where t' = f(t) is monotonic function of the original time variable t such that f(0) = 0 and lim,-1- f(1) = co. Clearly, /' = 0 recovers the uniform prior as a = 1 while + oo increases the mass of e, while other vertices remain constant. Given the conditional Eq. 26, one corresponding vector field that satisfies the continuity equation is
B(t + 1,
u.( = =;) = C(zs, (z; - ei), C(zx,i,,t) = - (t 1. 1)
(27)
(1 x, 12)
dg I (a. b) is the derivative of the regularised incomplete beta function [68, Appendix A. and C 1/t as regular linear flow matching.
where (a, b) =
F.2 e-geodesics on the Assignment manifold
In this appendix, we survey other common geometries implied by the theory of a-divergences on statistical manifolds, described in more detail in Amari [4, Ch. 4] or Ay et al. [11, Ch. 2], of which the case of e-connections was proposed in relation to flow-matching in Boll et al. [17].
In what has become a fundamental paper for the field Information Geometry, Amari [3] unified several commonly used parameterisations of statistical manifolds, in the theory of so-called a- connections. Without entering full details (which can be found in the mentioned references), on statistical manifold, endowed with Fisher-Rao metric, one can introduce a 1-parameter family of affine connections, so-called a-connections with a E [-1, 1]. where a = 0 corresponds to Fisher- Rao Levi-Civita connection, and other notable values are the m-connection for = -1 and the e-connection for a = Furthermore, specific classes of a-divergences which for a = 0 recover KL divergence have been introduced as adapted to the corresponding a-comnections.
21

---

Algorithm 1 WISHER-FLLOW training $
1: Input: Source and target distributions, Pr. Po. flow network up.
2: while Training do
3:
t, To; 21
(0, 1). Pu, P1 = Pdata
4:
gy (1))
Since æ is one-hot encoded, it is on $4.
5:
I0,
6:
Ie exPa am log (I1))
Geodesic interpolant between To. St & Calculated either explicitly or with numerical approximation. CFISHER. + [la((, ux xi)
7:
ue(ze/20 21) Lt
8:
9:
Update(0, VoLFisHEk-FLOW)
10:
return Up
In general, choice of differential geometric connection allows to define ad-hoc covariant derivatives, and corresponds to explicit formula for associated geodesics (curves whose tangent vector has zero covariant derivative).
For the case of a-connections categorical probabilities P(A),, explicit formulas be given (see Ay et al. [11, Ch. 2)), recovering, for m-cconecctoons, interpretations as mixtures, with geodesics equal straight lines in A parameteriistion,, and for e-connections geodesics can be interpreted as exponential mixtures, as elucidated in Ay et al. [11, Ch. 2] and illustrated in Boll et al. [17].
For the case of e-connections, concurrent work [17] has proposed to use the corresponding explicit parameterisation geodesics flow-matching, leaving as an open question the adaptation of Optimal Transport ideas to the framework.
Implementation Details
F.1
General Remarks
All of our code is implemented in Python, using PyTorch. For the implementation of the manifold functions (such as log, exp, geodesic distance, etc.), we have tried two different versions. The first direct port of Manifolds JL [10], originally written in Julia; the second one used the geoopt library [46] as a back-end. The latter performed noticeably theter-the underlying reason being probably better numerical stability of the provided functions.
As the optimal transport part, it is essentially an adaptation of that of FoldFlow [18), which itself relies on the POT library [31].
F.2 FISHER-FLOW Algorithm
We provide pseudo-code for training FISHER-FI.OW Algorithm 1.
F.3
Compute Resources
All experiments are run on single Nvidia A10 or RTX A6000 GPUs.
F.4
Experiments
E.4.1 Toy Experiment
We reproduce most hyper-parameters, except for the number of epochs trained for 500 instead of 540,000. Nonetheless, major modification from the original setting is the size of the dataset. Indeed, in the original dataset code of Stark al. [68]². one can observe that the points are generated each retrieval, and the defined length of the dataset is of 10", thus amounting to 540,000 10" training
hhttps://githin cow. wwwpsesttrrrrrrrkiraeeee -llor wwwowwaacciing/eeeee btapaiingg/ooo//mnn//ttttee dattasee..
Py#L53, retrieved October 30, 2024.
22

---

points by the end of the training process. This results in an unrealistic learning setup. To slightly toughen the experiment, we limit the training set size to 100,000 points.
Note that the model with which we train our method is a much simpler architecture than that of DIRICHLET FM (which was the one used in Stark et al. [68]), ours consisting exclusively of (residual) MLPs. For lower dimensions, it has less parameters, and slightly more in higher dimensions. The other baselines were run with our MLP too.
Table 5: Fréchet Biological Distance (FBD) and perplexities (PPL) values for different methods for enhancer
DNA generation. Lower FBD and PPL are better. Values are an average and standard error over 5 different runs.
Method
Melanoma FBD ()) Melanoma PPL. (1)
Fly Brain FBD (1)
Fly Brain PPL. ())
Random Sequence
619.0 + 0.8
895.88
832.4 + 0.3
895.88 2.19 + 0.10
Language Model
35.4 + 0.5
2.22 +: 0.09
25.7 + 1.0
DIRICILET FM
7.3 + 1.2
2.25 + 0.01
6.8 + 1.8
2.25 +: 0.02
TISHER-FLOW (ours)
27.5 + 2.6
1.4
3.8 + 0.3
1.4 +: 0.66
F.4.2
Promoter DNA
We train our generative models for 200,000 steps with a batch size of 256. We cache the best checkpoint over the course of training according to the validation MSE between the true promoter signal and the signal from the Sei model conditioned on the generated promoter DNA sequences. We use the same train/val/test splits as Stark et al. [68] of size 88,470/3.9333 9337774933/7
The generative model used for FISHER-FI.OR with an initial embedding for the DNA. Each block consists of a LayerNorm [12] followed by convolutional layer with kernel size and ReLU activation and a residual connection. As we stack
and DFM Stark et al. [68] is 20 layer 1-d CNN
the layers we increase the dilation and padding of the convolutional allowing the receptive field to grow [57]. In general, we use the Adam W optimiser [49].
Our Language Model implementation is identical to Stark et al. [68] and we use the pre-trained checkpoint provided by the authors and evaluated on the test set.
E.4.3 Enhancer DNA
We consider two DNA enhancer datasets, the fly brain enhancer dataset with 81 classes [45], the classes are different cell types, and the melanoma enhancer dataset with 47 classes [7]. Both datasets are comprised of DNA sequences of length 500. We use the same train/val/test splits as Stark et al. [68] of size 70, 00,892/8,900 022/8/9666900111 0,1 for the human melanoma and 83,726/ 33,726/10,500 10,66///////// /05/00,,4// 10,444 for the fly brain enhancer DNA dataset.
The generative model for our experiments for FISHER-FLOW and DFM [68] is the same as used for our promoter DNA experiments. Specifically, we use a 20 layer 1-d CNN with an initial embedding for the DNA. Each block consists of a LayerNorm [12] followed by convolutional layer with kernel size 9 and ReLU activation followed by a residual connection. As we stack the layers we increase the dilation and padding of the convolutional allowing the receptive field to grow [57].
We train for a total of 450,000 steps with a batch size of 256 where we cache the best checkpoint according to the validation FBD. The test set results in Table 2 are using the best checkpoint according to the validation FBD.
To calculate the FBD we compare embeddings from FISHER--LOO with a shallow 5 layer classifier embeddings originally trained classify the cell type given the enhancer DNA sequences. Our Language Model implementation is identical to Stark et al. [68] and we use the pre-trained checkpoint provided by the authors and evaluated on the test set.
E,5 Additional Metrics
For complete transparency, we also report the Fréchet Biological Distance (FBD) for the DNA enhancer generation experiments as initially reported in Stark et al. [68].
The FBD computes Wasserstein distance between Gaussians of embeddings generated and training samples under a pre-trained classifier trained to predict cell-type of enhancer sequences. Versus those embeddings from the generative model under consideration. So crucially there is dependence on classifier features.
23

---

On FlyBrain we find that FISHER-FLOW also improves over DFM in FBD being roughly = 2x better while DFM better on Melanoma. However, we caveat both FBD results by noting the trained classifiers provided in DFM [68] obtain a test set accuracy of 11.5% and 11.2% on the Melanoma dataset and FlyBrain dataset respectively. Moreover, switching out the pre-trained classifier for another trained from scratch caused large variations in FBD metrics. As result, the low-accuracy classifiers do not provide reliable representation spaces needed to compute FBD metrics. Consequently, FBD in this setting is noisy metric that is loosely correlated with model performance, so we opt to report perplexities in Table
F.6 De Novo Molecule Generation
Following the setup of Dunn and Koes [29], we report the following metrics for our model on de novo molecule generation over the QM9 and GeomDrugs datasets: percentage of stable atoms, percentage of stable molecules, percentage of valid molecules. Note that the following inference scheme is used, when training a model J on endpoint prediction:
(a'11)
A
loga, (ān))
F+41 exPx,
(28)
1 a-a(t) a(t)
where A 1/N, and > 0 the number of integration steps.
24

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The claims made in the abstract and introduction are backed up by our theoretical contributions and experiments results.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discuss limitations in the conclusion. In particular that we have not developed our method for language tasks.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
25

---

Answer: [Yes]
Justification: The theory underpinning our method is backed up by theoretical results outlined in the main paper and elaborated in the appendix. We have stated all assumptions and referenced relevant prior work.
Guidelines:
The answer NA means that the paper does not include theoretical results, All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: We provide all the details of the experimental setup to reproduce our method in our experiments section and in the appendix. We include our code as .zip file as supplementary material with instructions to reproduce our results. Our code will be made public upon publication.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurlPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the partieular way they provide for reproducibilitty In the case of closed-source models, it may be that access to the model is limited in
26

---

some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We include a .zip file with our code base and the necessary commands to reproduce our experiments. All the datasets use open-access.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.c cc/
Perides/CodeSuumis
ammollcyy for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / //nips wttpccccppccppppttcccccccccccttt http//guiee///oooeee /oaeeeee///oSSeeeiiiiiiiiee ssionPolicy) for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/eetails
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification:
We outline our experimental settings in detail in our Appendix F.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We make significant effort to produce results with means and standard errors over different runs with different random seeds. For our method and the main baselines we consider. This is in stark contrast to prior work.
Guidelines:
The answer NA means that the paper does not include experiments.
27

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We outline the compute resources required in Appendix F.3.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: /ttps/////eu.....eee cc
Answer: [Yes]
Justification:
We reviewed the code of ethics and our research line with the code.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We consider the broader impact of work in Appendix A.
Guidelines:
28

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g.. gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper does not put forward models and data which pose such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: All existing assets which are used are open-source and properly cited.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
29

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasttssss has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is available online, the authors encouraged to reach out the asset's s creators.
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We do not release any new assets. We will release documented code base upon publication.
Guidelines:
The answer NA means that the paper does not release assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Our paper does not involve human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our paper does not involve human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
30

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
31


Reviews:
[{"review_id": "HhLYdFNh3K", "reviewer": "Reviewer_11qA", "summary": "The paper proposed a novel flow-based generative model called *Fisher-Flow* for discrete data. The model uses the Fisher metric to deduce a Riemmanian geometric structure of the statistical manifold. The authors also demonstrated the connections to natural gradient descent and optimal transport. Experiments on DNA datasets demonstrated better performance.", "questions": "1. Previous existing diffusion- or flow-based discrete generative models have achieved considerable success in image and text domains (see Weakness 1). As a general discrete generation model, will the proposed model also achieve comparable performance on these image or text generation tasks? Will it have scalability issues regarding datasets with a large number of categories? More experiments on other generative domains will be needed to demonstrate the model's effectiveness and scalability.\n2. Can you provide more baselines in the visualization of the density estimation toy dataset and some quantitative evaluation metrics for this task to better demonstrate the generation quality? With the wide range of existing discrete diffusion or flow based models, more should also be compared as baselines for both toy examples and the enhancer generation task. Specifically, will the proposed model outperform simple linear flow matching? See Weakness 2 for some existing diffusion or flow baselines.\n3. Why the original numbers from the Dirichlet FM paper were not used? If you reran the experiments in the Dirichlet FM paper, can you provide justifications for why you only reran the ones for Dirichlet FM but not for other baselines like linear flow matching? It is also unclear why the results of linear flow matching in the original Dirichlet FM paper were omitted -- it has better performance than FisherFlow according to the Dirichlet FM paper. The reported numbers in this paper were **significantly worse** than those in the original paper. If original numbers were used, the proposed model cannot outperform the Dirichlet FM model. See Weakness 3 for details.\n4. How was perplexity calculated for flow-based models? Is it calculated in an autoregressive fashion as the language models? (See Weakness 4) Flow-based model is not an autoregressive model. During generation, all tokens are simultaneously denoised into meaningful sequences. If you instead modeled the joint distribution $\\log p_\\theta(x_{1:N})$, how was this marginal log-likelihood calculated for arbitrary input sequence $x_{1:N}$? How can you deal with the randomness in the initial samples $p_0$ if a deterministic cross-entropy loss needs to be computed? Furthermore, how was perplexity calculated for random sequences? The log probability $\\log p(x)$ for one-hot random sequences should either be 0 or infinite.\n5. OT seems to lead to better performance. Will it significantly add to the computation time during training? \n6. There is a notation inconsistency in the paper about whether to refer to the Dirichlet flow matching model as *Dirichlet FM* (e.g., Table 1 & 2) or *DFM* (main text).", "limitations": "Limitations were mentioned in the conclusion section, and the potential negative societal impact was adequately addressed.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The geometric perspective of the manifold of categorical distributions, to the best of my knowledge, is a novel extension of flow matching models for discrete data. The use of the Fisher-Rao metric can induce a Riemannian structure for continuous parameterization.\n\n2. The authors proposed a diffeomorphism to the unit sphere to avoid numerical instability during training and demonstrated better performance in experiments.\n\n\n3. The authors provided connections between the proposed model and natural gradient descent and Riemannian optimal transport with mathematical proof. The proposed model can potentially share such theoretical benefits. Ablation studies also supported such claims.", "weaknesses": "1. Although being claimed as a general discrete generation model, **the proposed model was only tested on the specific task of DNA generation without further justification of its performance on more complex data such as image and text generation**. The two DNA datasets used in the paper have a small cardinality of 4 categories. Existing diffusion- or flow-based discrete generation models (e.g. BFN [1], D3PM [2], SEDD [3]) demonstrated good results on image generation (discretized pixel values or VQVAE tokens) and/or text generation (tokens) tasks which have higher cardinalities to show their scalability. It remains unclear whether the proposed model can be adapted for more general real-world tasks with a larger number of categories.\n\n2. The choice of baselines is **too limited to be convincing enough** to demonstrate the proposed model's superior performance.\n   1. For the density estimation task, **no baseline was compared** (results were more of an ablation study). For the toy data in higher dimensions, **only Dirichlet FM was compared**. For these two tasks on simplex, models like multinomial flow [4] should be tested to provide a comparison.\n   2. For the enhancer DNA design task, **only Dirichlet FM was compared as a flow-/diffusion-based baseline**, despite that many other discrete diffusion models have been proposed, e.g., BFN [1], D3PM [2], SEDD [3]. Some of them were included in the related work section but none was tested as baselines in this task. \n   3. Moreover, one simple but effective and commonly used baseline is missing across the board, which is linear flow matching on simplex (which was used in DirichletFM [5] and other multimodal flow matching papers as well).\n\n3. Another critical concern lies in the **inconsistency of the reported performance of the same baseline models in this paper and the original paper**. Numbers for the Dirichlet FM reported in this paper for the promoter design task are **significantly worse** than those in the original paper. The original paper reported an MSE of 0.0269 for their best model (Table 1 in [5]) which is far better than the number indicated in this paper (0.034, Table 1 in this paper). Moreover, the linear flow matching model that achieved 0.0281 in Table 1 of [5] (outperforming the MSE of FisherFlow) was dropped in this paper. Similar large discrepancies in the baseline model's performance can be also noted in the FBD scores for enhancer design. Comparing Table 2 in [5] and Table 3 in this paper, the performance for the baseline Dirichlet FM was also significantly worse. As the numbers for other baselines were directly copied without change, it is unclear why the authors reported DirichletFM's performance differently. **Using the original MSEs or scores, the results indicated that the proposed model was not as good as Dirichlet FM or even Linear FM**. \n\n4. In terms of evaluation metrics, **it is unclear whether perplexity is a valid metric** to compare flow/diffusion based models against language models. Perplexity was initially proposed for autoregressive language models, in which it was calculated directly over conditional probabilities of discrete tokens:\n$$\nPPL(X)=\\exp\\left(-\\frac{1}{N}\\sum_{k=1}^N\\log p_\\theta(x_k|x_{<k})\\right)\n$$ \nHowever, in the continuous flow/diffusion setting, such probabilities may **not be well-defined** as they are neither autoregressive nor defined over discrete space. Therefore, their PPL cannot be calculated in this way. If the authors instead tried to calculate the joint distribution $\\log p_\\theta(x_{1:N})$, it remains unclear how to leverage the unconditional flow model to calculate this log-likelihood for arbitrary given input sequence $x_{1:N}$, as the flow is defined over the continuous simplex space with infinite possible initial probability samples $p_0$. The authors did not explain how they were able to calculate the PPL in the paper, and it is unclear whether the derived PPL is comparable to the ones computed for autoregressive models. \n\n\n[1] Graves, Alex, et al. \"Bayesian flow networks.\" *arXiv preprint arXiv:2308.07037* (2023).\n\n[2] Austin, Jacob, et al. \"Structured denoising diffusion models in discrete state-spaces.\" *Advances in Neural Information Processing Systems* 34 (2021): 17981-17993.\n\n[3] Lou, Aaron, Chenlin Meng, and Stefano Ermon. \"Discrete diffusion language modeling by estimating the ratios of the data distribution.\" arXiv preprint arXiv:2310.16834 (2023).\n\n[4] Hoogeboom, Emiel, et al. \"Argmax flows and multinomial diffusion: Learning categorical distributions.\" *Advances in Neural Information Processing Systems* 34 (2021): 12454-12465.\n\n[5] Stark, Hannes, et al. \"Dirichlet flow matching with applications to dna sequence design.\" *arXiv preprint arXiv:2402.05841* (2024).\ndirectly", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the time and effort they spent on reviewing our work. We are appreciative of the fact that the reviewer found our geometric perspective a “novel extension” to flow matching and that this leads to numerically stable training which is supported by better empirical performance. We now address the concerns raised by the reviewer.\n## Experiments beyond DNA\nWe acknowledge the reviewer's valid concern. In this rebuttal we included two more standard experimental settings for discrete data in molecular generation (QM9) and language modelling over text (LM1B). We kindly invite the reviewer to read our global response, and the 1pg rebuttal PDF, which gives more detail on these experiments. In summary, we found Fisher-Flows to be scalable to over 800k categories in text and achieve the best upper bound to text PPL amongst discrete diffusion models. For QM9 our results outperform flow matching and are competitive, almost saturating the task, with molecular diffusion models. We hope that our new experiments sufficiently demonstrates the scalability of Fisher-Flows to much higher dimensions and across a greater breadth of domains.\n## Choice of baselines\nAs suggested by the reviewer we’ve now included multiple baselines as part of our 1pg rebuttal PDF. Specifically, we added Linear Flow Matching (Euclidean metric) and Multinomial Flows to the synthetic task, where we found the latter to be numerically unstable as we increased the number of categories. We also added Linear Flow Matching to our DNA Promoter and Enhancer experiments where we found Fisher-Flows to be on par. For our new language modelling experiment on LM1B we compared against discrete diffusion models such as SEDD as well D3PM, as well as SOTA concurrent work in masked discrete diffusion in MDLM [1]. We found that Fisher-Flows still outperforms previous methods like D3PM and SEDD and is marginally better than MDLM.\n\n## Reporting of Dirichlet Flow Matching results\n\n>“Why the original numbers from the Dirichlet FM paper were not used? \n\nUsing the reference code implementation of DFM we noticed severe methodological issues that made it clear that the original numbers in their paper were an inaccurate reflection of true performance. We found the following flaws in their experimental setup:\nWe found for the DNA datasets the provided command used `--subset_train_as_val` argument which sets the validation set as the training set. This meant that their reported metrics were measured on the training set.\nWe found their released DNA classifier to have only $11%$ test accuracy which makes it a poor choice as a representation space to compute FBD.\nTheir synthetic data experiment did not standardise a training and validation split and meant that their training set size was a function of run time—effectively infinite. Operationally, DFM was trained for $10^{10}$ points per epoch which is a tremendous amount for a synthetic task. We standardised the experiment to train all models with a more modest 10^6 points.\n\nAfter fixing these issues we reran DFM and found that it performed worse than what was initially reported. We are appreciative of the DFM authors for releasing their code that enabled us to find these experimental flaws, some which we outlined in Appendix F.4.2 and F.4.3 but will discuss further in our updated manuscript. \n\n## Reporting PPL as a metric\n\nWe agree with the reviewer that calling our reported result perplexity on DNA is inaccurate. To compute our metric for DNA sequences we used the continuous time change of variable of the flow along the ODE trajectory over the manifold to compute the log-likelihood of the observed sequence (as done in Riemannian Flow Matching). Since the output space is discrete this likelihood on the manifold is an upper bound to the discrete log likelihood. More accurately this is a valid Evidence Lower Bound (ELBO) which was proven in concurrent masked discrete diffusion models [1, 2]. We also note that for autoregressive models the used PPL metric is a specific factorisation over the joint $\\log p(x_{1:N})$ and any valid ELBO is comparable. Nevertheless, the reviewer is correct and we will amend our reported results by indicating this upperbound when reporting PPL for diffusion/flow matching models. \n\nFinally, we have also included a new Gen PPL eval metric for DNA sequences. This is the PPL supplied a larger pre-trained autoregressive model on the outputs samples of a generative model. We found this metric could not distinguish the performance of DFM and Fisher-Flows, but we report it for completeness.\n## Q1 \nPlease see our global response for new larger scale experiments, including on text.\n## Q2\nPlease see above about additional baselines as well as our global response.\n## Q3\nPlease see our response above which finds flaws in the empirical setup of DFM.\n## Q4\nWe computed an upper bound to the PPL using the continuous time change of variable formula associated with an ODE. This is a valid ELBO see [1-2] for proof.\n## Q5\nOT in Fisher-Flows is not computationally expensive as both the sphere and the simplex have closed form distance functions which allows us to efficiently calculate the OT plan over mini-batches. This adds minimal overhead to training and is not affected by scaling. \n## Q6\nThank you for pointing this out. We will fix it in the updated paper.\n## Conclusion\nWe thank the reviewer again for their time. We believe we have answered to the best of our ability all the great questions raised by the reviewer. We hope our answers allow the reviewer to consider potentially upgrading their score if they see fit. We are also more than happy to answer any further questions. \n\n## References\n[1] Sahoo et al. \"Simple and Effective Masked Diffusion Language Models.\" arXiv preprint arXiv:2406.07524 (2024).\n\n[2] Shi et al. \"Simplified and Generalized Masked Diffusion for Discrete Data.\" arXiv preprint arXiv:2406.04329 (2024).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "J3ihNnI0FB", "reviewer": "Reviewer_cNL3", "summary": "This paper proposes a framework that enables flow matching over a d-dimensional simplex, by instantiating a riemannian flow matching algorithm using the Fisher-Rao metric. Some motivation in connecting the Fisher-Rao metric to natural gradient descent and Riemannian optimal transport is used to justify the choice of Fisher-Rao metric. Experimental validation is carried out on DNA sequences.", "questions": "### Regarding the mathematical justifications\n\nI find that Sections 3.3 and 3.4 are circular in their reasoning. \n\nIn Section 3.3, it is stated \"the choice of Fisher-Rao metric [is] the optimal one on the probability simplex\". However, here \"optimality\" is defined in terms of the expected KL, W_KL. This is an arbitrary choice of optimality. There is obviously a correspondence between a distance function and a Riemannian metric, and any metric can be stated as the optimal one for some distance function. So:\n- Why is considering the W_KL ball in Eq 8 useful? Are there computational reasons to use this? Are there theoretical justifications over other distance functions?\n\nIn Section 3.4, a similar problem arises. Here it is stated that the flows (1) lead to shorter global paths and (3) have lower kinetic energy. However, this is circular because the distance of a path (and consequently, the kinetic energy) is an implication of the chosen Riemannian metric g. Here the correspondence between distance function <==> Riemannian metric shows up again. \n- Why do we care about the lengths of paths defined by the Fisher-Rao metric?\n- Does the choice of Fisher-Rao metric give velocity fields that are faster to simulate?\n\nI'm not completely sure what the point of Proposition 2 is. Equation 10 looks to me like a definition of p_t given a coupling pi. That is, given that we've already solved the optimal transport coupling, we can then take any probability path p_t that transports between the marginals of pi, so Proposition 2 is just stating a definition of a particular choice of p_t. \n- Here do you mean to imply something about this choice of p_t? E.g., that it allows expressing W_2 in terms of p_t through a dynamical formulation, where the kinetic energy (as defined by the Fisher-Rao metric) shows up?\n\n### Regarding empirical validation\n\nThe paper makes the following claims:\n(1) flexibility of source distribution\n(2) the sphere map has better numerical stability properties\n(3) the choice of Fisher-Rao metric enables continuous reparameterisation and OT\n\nThe experiments do not yet completely justify these points, from what I can tell. \n\n(1) There is no experiment that uses a flexible source distribution. \n- Is there a case for using something that isn't a Dirichlet distribution which can be handled by Dirichlet FM? Or at least, a justification over the uniform noise distribution?\n\n(2) It's unclear why the sphere map helps. Concretely, there is a division by p in the computation of the inner product (Eq 5). However, this inner product is unneeded for the geodesic computation, since we can easily solve the geodesic on the sphere and then map back to the simplex. So this inner product seems to show up only in Eq 7, the training objective. However, given that this metric only depends on x_t, and that the model conditions on x_t, the optimal solution of CFM (Eq 7) is unaffected. So one can easily use any other metric for CFM without affecting the optimal solution. \n- Have the authors tried just using another metric for CFM and training on the simplex representation?\n\n(3) Similarly, I think a good ablation to do is just to use the Euclidean metric here. One can easily use the interpolants defined by the Fisher-Rao metric instead of linear, but otherwise just use regular Euclidean flow matching. Otherwise, related to the first concern, it is unclear why the connections to kinetic energy (as in Riemannian optimal transport) is important.\n\n### Clarify about experiment setup\n\nAs someone not familiar with DNA sequences, the experiment section does not provide enough information for me to gauge the usefulness of these experiments. For instance, it should at least be stated what d is (I imagine d=4?). My subjective opinion is that many methods could work reasonable well with such a small d. It would be even better if the authors could comment on the real world implications of their improved metrics (how does PPL affect a scientist who would use this generative model?)", "limitations": "The paper discusses a bit the limitations regarding higher d, but does not provide details on why the current approach cannot be applied. I am also curious why the authors did not consider relatively toy datasets (such as variants of MNIST) that would be more understandable / appealing for the general machine learning community. For instance, it is unclear to me if non-DNA experiments would be too hard, or even too easy, compared to the chosen experiments.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The proposed method is quite a natural application of riemmanian flow matching to distributions over simplices.", "weaknesses": "- I felt the writing is a bit too mathematically dense, and the justifications for the Fisher-Rao metric feel a bit circular (see below for questions on Propositions 1 & 2).\n- Additional ablation experiments on non-toy data would be good to have to verify some of the claims in the paper.\n- It's unclear to me (due to lack of expertise in the application area) to judge how well this approach performs.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and effort assessing our paper. We now address the key points in the review, while additional experiments are included in the global response and rebuttal PDF.\n## Circular justifications for the Fisher Rao\nWe understand that certain aspects of the theory could be explained more clearly. We now explain our reasoning step by step. First, note that a distance function on a Riemannian manifold is fully determined by the choice of metric, and that choosing a different distance function amounts to choosing a different metric.\n\nWe require the following key desiderata from a metric to employ flow matching:\n- Analytic expressions of manifold operations like the exponential and logarithmic map.\n- Easy parameterisation of the tangent space in order to define vector fields.\n- Interpretability and naturality, i.e. connection to established statistical and physical theory.\n- Numerically stability of the metric over the entire manifold.\n\nThe use of the FR metric satisfies both points 1 and 2 above and more and point 4 is achieved via the sphere map. In addition, we argue that the pervasiveness of the use of the forward KL divergence in generative models such as diffusion, flows, and VAE, testifies to its naturalness. As to point 3, the KL divergence is natural due to its strong link to fundamental concepts from Information Theory, being interpretable in terms of entropy. In this sense, KL is perhaps the simplest natural divergence over categorical probabilities.\n\nWe now explain why our choice does not lead to circular reasoning using Prop 1.\nNote that Prop 1 starts off with hypotheses over a generic metric “g” in the $W_{2,g}$ distance from the hypothesis. The thesis of the proposition shows that natural gradient with respect to KL-divergence-based $W_{KL}$ distance produces naturally the $W_{2, g_{FR}}$ gradient flow, with now “$g_{FR}$” being FR metric, which was not assumed in the hypotheses. Thus this is not a circular statement, as it deduces FR as the limit of KL-divergence, at the level of probability spaces. \n\n>Are there computational reasons to use this?\n\nThe FR metric leads to theoretical benefits as it is also the Wasserstein geodesic over probability paths on the simplex. The Euclidean metric, as explored in Dirichlet Flow Matching, does not enjoy such benefits and arguably is one reason that is empirically less performant than Fisher-Flows. \n\n>In Section 3.4, a similar problem arises…\n\nWe believe there may be a slight misunderstanding as the logic of our setup is to first pick a metric and then show OT is computationally feasible. Indeed other metrics are possible and will lead to different OT maps; but, crucially, they may not be easy to compute for generative models. Fortunately, for the FR metric the distance is simply the distance on a sphere, which is **easy to compute** allowing us to efficiently compute the OT cost $c(x_0, x_1)$.\n\n>Why do we care about the lengths of paths defined by the Fisher-Rao metric? … velocity fields faster to simulate?\n\nShorter paths lead to lower variance training: on the sphere and the simplex. The paths do not cross at intermediate points $x_t$, which would have added to noisier training—in line with the literature see [1,2]. We note OT often leads to vector fields that are faster to simulate using numerical solvers as they incur smaller per step error, as studied for $SE(3)$ in [3].  \n\n>Here do you mean to imply something about this choice of p_t? … \nProp. 2 shows that, by using FR geodesics as our flow, we automatically use the geodesic path between $p_0$ and $p_1$ with respect to the natural Wasserstein metric over probability measures. A consequence of this is that the links to kinetic energy and dynamical formulations apply (see Prop. 4). \n## Empirical validation\n>There is no experiment that uses a flexible source distribution\nWe invite the reviewer to kindly read our global response which includes new experiments over text that uses a new source distribution which is a convex combination of a specialised mask token and the uniform distribution. \n\n>It's unclear why the sphere map helps ….\n\nThe metric appears when we compute $x_t$ due to exp and log map on the simplex (see Appendix B). Looking at equation 11, for the exp map, we notice we need to compute the norm which is induced by the Riemannian metric so $<v^2_p, v^2_p>_{FR}$. Similarly, the log map on the simplex in equation 12 explicitly requires the inner product. Exploiting the sphere map, we obtain $x_t$ and avoid numerical instability near the boundary of the manifold. Other metrics beyond the FR are indeed possible on the simplex but it may not have analytic formulae for $x_t$.\n\n>Similarly, I think a good ablation to do is just to use the Euclidean metric here …\n\nPlease see our global response PDF where we have included the Euclidean metric in Linear Flow Matching on the toy datasets as well as our DNA experiments. \n## Clarifying experimental setup\nIn addition to DNA sequences we have now also included more standard generative modelling experiments on molecules and larger scale language modelling that demonstrate the scalability of Fisher-Flows on a host of different domains. \n\nWe thank the reviewer again for their valuable feedback. We hope that our rebuttal addresses their questions and concerns, and we kindly ask the reviewer to consider fresher evaluation of our paper if the reviewer is satisfied with our responses. We are also more than happy to answer any further questions that arise.\n## References\n[1] Bose et. al (2024). SE(3)-Stochastic Flow Matching for Protein Backbone Generation. arXiv preprint arXiv:2310.02391.\n\n[2] Tong et. al. Improving and generalizing flow-based generative models with minibatch optimal transport. arXiv preprint 2302.00482, 2023.\n\n[3] Klein et. al  \"Equivariant flow matching.\" Advances in Neural Information Processing Systems 36 (2024).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WARkfxWuQ2", "reviewer": "Reviewer_YHTA", "summary": "The paper aims to build a Flow Matching [1] based generative model for discrete data. The approach models the discrete data as a categorical distribution that resides on the simplex, thereby translating the problem into continuous flows. Equipped with the Fisher-Rao metric, the simplex is identified as a Riemannian manifold, and the task is reduced to Riemannian Flow Matching [2]. To further simplify the geometry of the space, the sphere map is used, which identifies the simplex with the positive orthant of a hypersphere. Finally, the conditional path between every pair $(x_0, x_1)$ is defined via the geodesic interpolant induced by the FIsher-Rao metric. For marginalization, two joint distributions are considered: 1. $x_0$ and $x_1$ are independent, i.e., $\\pi(x_0,x_1) = p_0(x_0)p_1(x_1)$ as in [2]. 2. $\\pi(x_0,x_1)$ is the optimal transport (OT) plan between $p_0$ and $p_1$, approximated with minibatch [3].", "questions": "1. Can the author expand on the difference of his use in Riemannian Optimal Transport with a mini batch compared to previous works?\n2. The author claims that [5] is restricted in the choice of possible source distributions. Can the author please elaborate on that?\n3. Could the author possibly provide more experimental evidence to demonstrate that the proposed change in geometry indeed contributes to a considerable improvement and that the gain in results is not mainly from the OT plan?", "limitations": "It is suspected that the proposed method might not be as effective in large dimensions (i.e., large vocabulary size ) compared to other methods.\n\n[1] Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., & Le, M. (2023). Flow Matching for Generative Modeling. arXiv preprint arXiv:2210.02747.\n\n[2] Chen, R. T. Q., & Lipman, Y. (2023). Flow Matching on General Geometries. arXiv preprint arXiv:2302.03660.\n\n[3] Bose, A. J., Akhound-Sadegh, T., Huguet, G., Fatras, K., Rector-Brooks, J., Liu, C.-H., Nica, A. C., Korablyov, M., Bronstein, M., & Tong, A. (2024). SE(3)-Stochastic Flow Matching for Protein Backbone Generation. arXiv preprint arXiv:2310.02391.\n\n[4] Stark, Hannes, et al. \"Dirichlet flow matching with applications to dna sequence design.\" arXiv preprint arXiv:2402.05841 (2024).\n\n[5] Campbell, Andrew, et al. \"Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design.\" arXiv preprint arXiv:2402.04997 (2024).\n\n[6] Tong, Alexander, et al. \"Improving and generalizing flow-based generative models with minibatch optimal transport.\" arXiv preprint arXiv:2302.00482 (2023).\n\n[7] Pooladian, Aram-Alexandre, et al. \"Multisample flow matching: Straightening flows with minibatch couplings.\" arXiv preprint arXiv:2304.14772 (2023).", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "1. The paper is well-written, and the method is rigorously presented.\n\n2. The author successfully frames the proposed method for generative models of discrete data in terms of well-established existing methods, which eases the understanding of the approach and potentially allows for a relatively low cost of implementation.\n\n3. Aside from the computational cost of the OT plan, the method does not introduce additional computational costs compared to previous works.", "weaknesses": "1. The use of Riemannian Optimal Transport with a mini batch is of low novelty. This has also been done in [3] and on euclidean space in  [6,7].\n\n2. Two key contributions of the paper seems to be the change in geometry, that is the introduction of the Fisher-Rao metric and the sphere map. However the paper lacks a clear experimental  demonstration of the portion these two components contribute to results. Specifically:\n\n    2.1. Regarding Figure 2, allegedly,  the author claim that Fisher-Flow (FF)-OT on $\\mathbb{S}_+^2$ (e) performs best, however it is not clear, one may even say that FF-OT on the $\\Delta^2$ (c) perform better. Furthermore, even without FF-OT it is not clear that using the sphere map helps ( (b) vs (d)).\n\n    2.2. In Figure 3.a the author compares FF simplex, FF-OT simplex, FF sphere, and FF-OT sphere, that shows that FF-OT sphere preforms the best. It is also apparent that with increasing dimension size the advantage of the use of the OT plan decreases. And Figure 3.b shows that FF-OT sphere out performs [4]. However, comparing Figure 3.a and 3.b it seems that the other three, i.e., FF simplex, FF-OT simplex, FF sphere are actually outperformed by [4] which uses flows on the simplex but no OT plan. To conclude, this experiment suggests that the apparent advantage of FF-OT sphere on [4] may exist only for relatively low dimensions.\n\n    2.3. The other two experiments only show FF-OT sphere.\n\n    2.4. No ablation that focus on the use of Fisher-Rao metric compared to other potential metrics (though the author does provide a theoretical justification).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and detailed feedback. We focus here on addressing the key clarification points raised in this review while the global response contains new experiments on molecules and language modelling with additional baselines.\n## Riemannian OT is of low-novelty\nWe acknowledge the reviewer’s concern that OT for Riemannian manifolds has previously been considered in [1]. Unfortunately, the general theory of optimal transport on manifolds Villani 2003 [2] (Theorem 2.47) does not give a prescription for a fast computation implementation of OT on manifolds—i.e., analytic expressions for the geodesic cost $c(x_0, x_1)$.\nWhile [1] found a closed form solution to this cost for $SE(3)$, it cannot be generally extended to all Riemannian manifolds, since, for arbitrary manifolds, one must measure the geodesic by simulating the expensive Euler-Lagrange equations, which are 2nd order PDEs. \n\nOur work with Fisher-Flows pushes the advantage of Riemannian OT within generative modelling by showing that: \n1.) The probability simplex and the sphere both have closed form expressions for the OT cost $c(x_0, x_1)$.\n2.) We demonstrate that OT via Fisher-Flows can be used for discrete data. We highlight that this is the first use of OT for generative modelling over discrete data. \n3.) As proved in Appendix D, OT over the probability simplex elicits new benefits, such as the fact that $p_t$ becomes the Wasserstein geodesic between $p_0$ and $p_1$, and the associated velocity field $v_t$ minimises the kinetic energy. These two points were not rigorously proven in [1].\n## Clearer experimental demonstration over the benefits of the Fisher-Rao metric and Sphere map \n### OT in Fig 2 \nWith regards to Fig. 2, we find that both c) and e) to be of comparable visual quality, and, more importantly, of a higher one than that of their non-OT counterparts. As this is a small toy synthetic density estimation task, the benefits of the Fisher-Rao metric and the added numerical benefits of the sphere map are less pronounced. \n### Comparison of FR on Fig 3\nSince the sphere is more numerically stable and, for practical problems—where the number of dimensions can greatly increase—we expect Fisher-Flows on the sphere to perform the best—regardless of OT. We see this in Fig 1 of the 1pg rebuttal PDF, Dirichlet Flow Matching is often worse than FF-no OT on the sphere as we increase the number of categories. With regards to OT, our experimental findings are in line with previous use of OT in generative models, where lowering variance in training and shorter paths during inference (fewer ODE integration steps [3,4]) are observed. \n### The other two experiments only show FF-OT sphere. \nOur new experiments on molecule generation and language modelling (see 1pg rebuttal PDF) used Fisher-Flows on the sphere and without OT. \nIn our language modelling results, we found that Fisher-Flow slightly outperforms discrete diffusion approaches (e.g., D3PM, SEDD) while marginally outperforming concurrent SOTA work on Masked Discrete Diffusion (MDLM [5]) released after the submission deadline. This suggests that geometry plays a key role in improvement as we compare against conventional discrete diffusion methods on higher dimensional problem setting ~800k categories (text). \n### No ablation on other potential metrics.\nThank you for this suggestion! We have included new ablations in our 1 pg rebuttal PDF that use the Euclidean metric (Linear Flow Matching) on the simplex for both the synthetic data and DNA Promoter and Enhancer experiments. We observe that the Fisher-Rao metric is still on par and better than Dirichlet FM, in terms of empirical performance. Finally, considering other metrics beyond the Fisher-Rao and Euclidean is an interesting thought; however, it is unclear whether it is possible to easily obtain closed-form geodesics, which are needed for simulation-free training as mandated in the flow matching framework.\n## Q1. Use of Riemannian OT vs. Prior Work \nPlease see our detailed response under the heading “Riemannian OT is of low-novelty” for a deeper discussion. \n## Q2. Flexibility of source distribution in Campbell et. al 2024?\nWe note that Campbell et. al 2024 explicitly state that the “conditional flows we use in this paper linearly interpolate towards $x_1$ from a uniform prior or an artificially introduced mask state, M.” They do not make any theoretical claims about the generality of their starting source distribution unlike we do for Fisher-Flows. While it is potentially possible to accommodate for a more flexible source distribution in their CTMC framework, it is not immediately clear how to do this easily. We will refine our discussion over Campbell et. al 2024 in our paper to highlight this aspect.\n## Q3. Ablation on geometry vs. OT?\nPlease see our previous response on Fisher-Flows without OT on language modelling. \n## Closing comment\nWe hope that our responses were sufficient in clarifying all the great questions asked by the reviewer. We thank the reviewer again for their time, and we politely encourage the reviewer to consider updating their score if they deem that our responses in this rebuttal along with the new experiments merit it.\n## References\n[1] Bose et. al (2024). SE(3)-Stochastic Flow Matching for Protein Backbone Generation. arXiv preprint arXiv:2310.02391.\n\n[2] Villani. Optimal Transport: Old and New. Grundlehren der mathematischen Wissenschaften. Springer Berlin Heidelberg, 2008. ISBN 9783540710509.\n\n[3] Tong et. al. Improving and generalizing flow-based generative models with minibatch optimal transport. arXiv preprint 2302.00482, 2023.\n\n[4] Klein et. al . \"Equivariant flow matching.\" Advances in Neural Information Processing Systems 36 (2024).\n\n[5] Sahoo et al. \"Simple and Effective Masked Diffusion Language Models.\" arXiv preprint arXiv:2406.07524 (2024).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all the reviewers for their time and constructive comments. We are grateful that the reviewers appreciated the novelty of our geometric perspective to discrete generative modelling (R 11qA) and that it is a natural application of Riemannian flow matching over simplices (R cNL3). We are also heartened to hear that the reviewers agree that our parametrization of flows over the sphere is beneficial for numerical stability and leading to demonstrated empirical benefits with little computation overhead (R 11qA, YHTA).  We now address shared concerns raised by the reviewers, and summarise our new experiments and ablations included in the supplementary PDF.
## New experiments on complex data (R 11qA, R cNL3)
The reviewers expressed concerns regarding the empirical verification of Fisher-Flows against other domains. We included, in our rebuttal PDF, 2 additional standard discrete generative modelling tasks: molecule generation (QM9), and language modelling on the One Billion Words Dataset (LM1B). These new domains enable us to test the scalability of Fisher-Flows against SOTA discrete generative models for molecules and text.
### Experimental setup and Flow parameterization
To enable scalable and effective training of our flows, we parametrize the loss using the target prediction framework by learning a denoising network to predict the target discrete target token given $x_t$. We note that the vector field can easily be recovered for inference and used to predict the next Euler step—thus no extra cost is added to sampling.
### Language Modelling
We train Fisher-Flows on the LM1B dataset, which has a vocabulary size of ~800k words. This experiment seeks to address the scalability concerns as well as enable direct comparison to a more standard discrete generative modelling setting. 

We define Fisher-Flows by using the following probability path $p_t = \kappa_t * p_{M} + (1 - kappa_t) * p_{unif}$, where $\kappa_t \in [0,1]$. Here $p_{M}$ is the Fisher probability path between the target $x_0$ and designated mask token $M$, while $p_{unif}$ is also a Fisher probability path from a sample from a uniform distribution and $x_0$. Thus $p_t$ is a convex combination of probability paths using a noise scheduler $\kappa_t$. Using a denoising architecture enables us to rewrite the original loss as a weighted negative log likelihood $L = \mathbb{E}[-\log p(x_0 | x_t)]]$. Crucially, this allows us to calculate an upper bound to the test perplexity, which is a natural evaluation metric for language modelling and employed by concurrent SOTA works on discrete diffusion MDLM and MD4 [1, 2], who prove that this is a valid Evidence Lower Bound. 

We report our results in Table 2 of the 1pg PDF and note that Fisher-Flows is the best performing discrete diffusion/flow-matching model achieving a test perplexity of $\leq 26.51$ and $\leq 22.42$ after training on $33B$ and $327B$ tokens respectively. We note that our method outperforms the baselines suggested by the reviewer of D3PM, SEDD, as well as marginally outperforming concurrent masked diffusion MDLM [1]. 
### Molecule Generation
We instantiate Fisher-Flows on QM9 where four different data types must be produced: positions, charges, atom types, and the possible bonds between them. Furthermore, the positions are continuous rather than discrete which allows us to showcase the feasibility of using Fisher-Flows in a mixed setting of discrete and continuous data. 

We report our results in Table 3 of the rebuttal PDF. We find that Fisher Flows outperforms a comparable flow matching baseline in EquiFM [3] and is competitive with a diffusion baseline in Jodo [4], all saturating the benchmark. We also include examples of generated samples in Fig. 2 from Fisher Flows to qualitatively evaluate our approach.
## New ablations and evaluation metrics (R 11qA, RcNL3, YHTA)
To make our comparisons more robust, we extended the synthetic experiment with a variety of new baselines; the results are shown in Fig 1 of the rebuttal PDF. We included Linear Flow matching, D3PM, and Multinomial Flow, as suggested by the reviewers. Fisher-Flows clearly outperform Linear FM across all dimensions. Although between dimensions 40 and 100/120 Fisher-Flow performs marginally worse than both D3PM and Multinomial Flows, past dimensions 100 and 120 these two’s KL divergences diverge to infinity (hence their absence from the figure) indicating that some categories have been sampled essentially zero times. Our Fisher-Flow parametrization does not incur this numerical failure mode.
### Generative Perplexity
We updated our DNA sequence results in Table 1 of the rebuttal PDF by first amending the reported Perplexity as an upper bound to the true test perplexity, as our flows yield a valid ELBO. In addition, we ålso included a common evaluation method (Gen-PPL) in the perplexity induced by autoregressive language models that we train on the same dataset as the evaluated generative model. We train small GPT2 DNA language models on respective datasets and assess the flow model’s DNA generations under these. Random sequences obtain gen-ppls of over 4 for Promoter DNA, fly-brain and melanoma DNA LMs. The gen-ppl values for Fisher-Flows, Dirichlet Flow Matching and Linear Flow matching for all DNA datasets are all around 1, exhibiting no clear winner under this metric.

We also include a Linear Flow Matching baseline and find that Fisher Flows is on par and both are better than Dirichlet Flow Matching.
## References
[1] Sahoo et al. "Simple and Effective Masked Diffusion Language Models." arXiv preprint arXiv:2406.07524 (2024).

[2] Shi et al. "Simplified and Generalized Masked Diffusion for Discrete Data." arXiv preprint arXiv:2406.04329 (2024).

[3] Song  et al. "Equivariant flow matching with hybrid probability transport." arXiv preprint arXiv:2312.07168 (2023).

[4] Huang et al. "Learning joint 2d & 3d diffusion models for complete molecule generation." arXiv preprint arXiv:2305.12347 (2023).

Meta-review: This paper proposes a flow matching method for generative modeling of discrete data. The authors proceed by operating on the (continuous) simplex of discrete probability distributions, where they use Riemannian flow matching along with the Fisher-Rao metric. The authors also propose to map the simplex onto the positive orthant of a hypersphere, which they argue improves performance.

The authors provided a thorough rebuttal addressing concerns by the reviewers, who now all unanimously agree that this paper meets the bar for publication at NeurIPS. I also agree and thus recommend acceptance. For the camera-ready version of the paper, I encourage the authors to include the additional experiments and ablations carried out during the rebuttal, and to more carefully articulate the empirical benefits of using the sphere map.
Decision: Accept (poster)