Generator Matching: Generative modeling with arbitrary Markov processes

We want to thank the reviewer for considering our paper as "original and significant overall" and that it "could very well become a must-read for those interested in $generative modeling$ ". We understand that the main concern of the reviewer is notation and understandability of our work that we are happy to address.

Densities vs. measures

A. Why do we use probability measures and not restrict ourselves to densities? It was a deliberate choice for us to build a generative modeling framework that is not restricted to probability measures that admit a density with respect to a common reference measure. Let us give several examples where this is relevant:

• Discrete support: The probability path $p\_t=\\delta\_{tz}$ on $S=\\mathbb{R}$ is a path of point measures corresponding to a line connecting $0$ and $z\\in\\mathbb{R}$. This probability path does not admit a density with respect to the Lebesgue measure. However, this is not an artificially constructed example but such paths actually appear in cases where one conditions probability paths on start and end points (see e.g. Tong et al (2023)). There are other examples of data distributions with discrete support, e.g. when using latent space encodings with a VQ-VAE.
• Complex state spaces: Generator Matching (GM) would also work for generative modeling over infinite-dimensional spaces, e.g. path space $C($ 0,1 $)$ for trajectory modelling. However, measures over paths of stochastic processes in general do not have a density with respect to a common reference measure (they often only have relative densities, e.g. see Girsanov's theorem). Another example is trans-dimensional generative modeling (see e.g. section H.7).
• Universal characterization: The universal characterization in theorem 1 only holds for general probability measures and jump measures (if we restrict ourselves, we would lose this result).

B. Notational changes to make density vs measure clearer to the reader Thank you for pointing out a potential confusion regarding the notation of probability measures vs probability densities. We made several changes:

• Different notation for density and measure: As you suggested, we introduce now formally the reference measure $\\nu$. Further, we now distinguish in the notation between a probability measure $p\_t$ and its Radon-Nikodym derivative $\\frac{dp\_t}{d\\nu}$ with respect to a reference measure $\\nu$ throughout our entire work. When operating on Euclidean space with the Lebesgue measure, we suggest to keep the old notation $p\_t(x)$ to denote a density - this is to avoid confusion with readers who are familiar with the notation of densities in Euclidean space but not with general Radon-Nikodym derivatives. However, we always state this clearly then (see e.g. A.5.1 or description of table 1).
• Assumptions stated explicitly: Thank you for pointing out that it should be highlighted more clearly when we assume that a density exists. Whenever we assume that a density exists, we state this assumption now more explicitly. For example, in table 1, we state explicitly which rows assume a density exists. Similarly, we modified the appendix to make it clear.

This should dissolve the confusion between the density and a measure.

Dot product

We are adapting your notation and writing the action of $p\_t$ on a test function in the dot product notation: \\langle p\_t,f\\rangle :=\\mathbb{E}\_{x\\sim p\_t} $ f(x) $ . We note that if $p\_t$ is a probability measure without a density, then the above is a "symbolic" dot product (i.e. it is not a dot product defined on a Hilbert space) and is simply a definition. However, it always is a bounded linear functional motivating this notation.

F-parameterization of generator

Thank you for suggesting to rewrite the generator and the $F$-parameterization in the form $\\langle \\mathcal{K}f,F\\rangle$ for a linear operator $\\mathcal{K}$ independent of the parameters. We call this now a linear parameterization and added a section in A.6. that explains this. We further show that all propositions of our work also hold for these linear parameterizations.

Discrete difference operator: We note that for discrete state spaces, it is standard parameterizes the jump kernel as a rate matrix $Q$ for which it holds $Q\_t(x;x)=-\\sum\\limits\_{y\\neq x}Q(y;x),\\quad Q(y;x)\\geq\\text{for all }y\\neq x$ (see section A.5.4). In particular, columns sum to $0$. Therefore, in your suggested example, it would hold that $\\sum\\limits\_{y\\in S}(f(y)-f(x))Q(y;x)= \\sum\\limits\_{y\\in S}f(y)Q(y;x)=f^TQ(\\cdot;x)$ Therefore, the generator application can be expressed as a simple matrix multiplication. This is the notation we use in table 1.

We have adapted all changes in our updated revised manuscript that you can view on OpenReview. Would you consider this notation as clear? We are very happy to integrate further feedback. Thank you again for your helpful response.

Practical benefits of GM.

It is unclear to me, and difficult to gauge, in how far the GM framework has benefits over simple, naive combinations of different modeling paradigms. I.e. how will the GM framework help future researchers in practice to build better models, other than only providing a theoretical unificiation?

Indeed, one contribution of GM lies in a theoretical unification of many existing generative models and also a universal characterization of the design space of Markovian generative models. However, there is a whole variety of practical benefits for future machine learning introduced via GM. We are happy to provide several examples:

1. Markov superpositions: GM enables the construction of Markov superpositions of different classes of Markov processes generating the same probability path (section 7.1). Without GM, there would be no basis for combining for example a jump model and a flow model at the same time and achieving the improvement that we report in our experiments would not be possible.
2. Multimodal constructions: GM enables the construction of multimodal models. Such a rigorous construction for arbitrary modalities does not exist to the best of our knowledge and naturally requires the GM abstraction. To address your comment, we added an explicit "recipe" to build a multimodal model out of unimodal models in section C.6. in detail. This will help future researchers to construct models for arbitrary combinations of modalities. We also included new experiments for multimodal protein generation improving state-of-the-art and highlighting the practical utility of the GM framework.
3. Jump models: In our work, we introduced a scalable way of training jump models on Euclidean spaces in this work. As we show, this model is a powerful model achieving performance that rivals first versions of GANs or diffusion models. We believe that with more engineering and optimization efforts of the jump kernel and sampling, this model could well become equally or more powerful than existing state-of-the-art models, while having potentially other benefits such as increased diversity of samples (see e.g. the protein generation results).
4. Modality-agnostic generative models: The GM framework also allows to build generative models agnostic of the modality. In section 5.1., we derive a general solution of a jump process on arbitrary state spaces (that we then end up using for proteins). Therefore, any modality could be modelled with that process.
5. Loss functions: The GM enables a whole diversity of loss functions (i.e. any Bregman divergence works). This gives future researchers a whole variety of loss functions to use. To illustrate the utility of this, we did additional experiments to show one can improve existing flow and diffusion models with a minimal change in the code. We explain this in the general response.

Questions:

Is the linearity of the operators a limiting factor in the modeling capacity?

If a generator of a Markov process exists, then it must necessarily be a linear operator - in the same way as the "common" derivative of a function is always a linear operation if it exists. While generators are linear as operators on function spaces, the neural network we use is still highly non-linear. Therefore, the linearity of the generator is rather a feature that allows us to develop the GM framework.

What are the limitations of GM?

GM considers Markovian generative models and ones that are in continuous time. For this reason, GANs or VAEs are not instantiations of GM. Further, sampling with a model from GM requires iterative simulation of a Markov process. This can require a number of steps for sampling (i.e. one might require distillation as done for diffusion models).

Line 128. What exactly is meant with "simple". If I understand correctly, it needs to be linear, why not use that word?

We did not use the word "linear" because if one considers $p\_{t+h|t}$ as an operator (as explained in section 4 in the paragraph on "test functions"), it is always a linear operator - even for "large" h. We choose the word "simple" to motivate the decomposition in equation (1) that is then rigorously defined later in the section.

Figure 2. What is ME in the legend? This figure could benefit some extra attention. Larger text size, match the naming in the legends with the vertical namings of the rows.

Thank you for pointing this out. "ME" was a typo and should stand for "MS" (Markov superposition). We also increased the font size.

Minor comments. Thank you for the suggested edits. We integrated them into the new version.

Together with our general response, we hope that the above addresses all of your comments. We welcome any additional discussions. Thank you again for your helpful feedback and response.

We thank the reviewer for insightful comments and for taking the time to review our work.

Time-reversals.

Why does focusing on matching of probability paths remove the need for Anderson's result?

Indeed, the idea of time-reversal is not needed in our framework. We added a detailed explanation in section A.4. that we summarize here. In diffusion models, one "noises" data via a "forward process". Every such "forward process" also defines a unique probability path but not vice versa. Our framework aims to train a model to approximate a solution to the KFE following a desired probability path $p\_t$. If we have defined $p\_t$ via a noising Markov process ("forward process"), then this is equivalent to finding a Markov process in reverse time that has the same marginals. This is a weaker condition than to find a true time-versal process in the sense used by Anderson (1982) - the true time-reversal does not only concern marginals $X\_t\\sim p\_t$ but the joint distribution across time points (see A.4). However, for generative modeling, we often only use the final point $X\_1$ of the Markov process as a sample. Therefore, it only matters whether $X\_t$ has the desired marginals. For example, consider the probability flow ODE in diffusion models (Song et al, 2020) which does not constitute a time-reversal of a diffusion process in the sense of Anderson (1982) but it fulfils the KFE in our sense. This shows that restricting oneself to proper time-reversals imposes unnecessary mathematical complexity and can even lead to suboptimal solutions. Still, a "true" time-reversal can still be used as one solution to the KFE for the probability path $p\_t$. See section A.4. for details.

Related work on jump models. We have added an extended related work section including a subsection dedicated to jump processes (section H) - here discussing the 2 mentioned references:

• Blackout diffusion: Thank you for pointing out this reference. We reference it now in the related work section and added an extended discussion in H.4. of how "blackout diffusion" is an interesting instantiation of GM on discrete spaces. We note that every continuous-time Markov process on discrete spaces is a jump model (see section A.5.4) and that we had discussed a variety of such models in the related work section of the submission. Hence, one can find a detailed discussion of blackout diffusion in section H.4.
• Levy processes: Yoon et al (2023) construct SDEs driven by Lévy processes as a forward process. The fractional score serves as a linear parameterization of the generator (see appendix A.6) and the fractional denoising score matching loss (i.e. Theorem 4.3. in Yoon et al (2023)) is a special case of the CGM loss (proposition 2 in our work). While Levy processes in principle include jump processes, we note that the work by Yoon et al (2023) only considers learning the flow component of the Markov processes (the rest are hyperparameters), while here we show how to also learn the jump kernel or the diffusion coefficient. See section H.7. for details.

Generators and KFE.

It remains unclear why using generators and adjoint KFE leads to “correct” generated samples.

For all state spaces and Markov processes of interest, the KFE is both a sufficient and necessary criterion for a Markov process to generate a probability path. We made this clearer adding a comment on page 5 (below equation (6)) and adding an additional section A.2.

Regularity assumptions. We made the regularity assumptions again clear in A.2. and also listed our motivation for choosing them.

Do the authors expect that [...] the whole framework can be applied to any process with just Markovian property?

In order for a generator of a Markov process to be defined, property 2 in the definition of a Feller process must hold (see section A.1.2). Property 1 in the definition of a Feller process, might not be strictly necessary. However, we use it because (i) most of the mathematical literature (incl. the one we cite) uses it and (ii) any function on a computer is defined on a compact set. We stress this in the main paper now.

Why does regularity assumption 2 (finite number of discontinuities) is needed?

While our framework would not break, simulating an infinite number of discontinuities would necessarily induce uncontrollable sampling error. We therefore came to the conclusion that such a model would therefore most likely be not interesting for the purposes of generative modeling.

Other questions

• Page 3, line 120: This refers to sampling the initial state of the Markov process.
• $\\delta\_z$ refers to a Dirac delta measure.
• Line 129-130: “Informal” refers to the fact that a Taylor approximation is only defined if a derivative of a function is defined - this definition is then introduced formally later in the section.

Thank you again for your helpful feedback and response. We welcome any additional discussions.

Dear authors,

Thank you for addressing my concerns in your revised manuscript and providing answers on my questions. I read revised version and the new version is easier to understand. Overall, the exposition of the framework has been improved. I agree with other reviewer that this paper will become a must-read for those interested in flow/sde generative models. I increased my score.

We thank the reviewer for insightful comments including the extensive feedback on the supplementary material.

Generator abstraction.

GM primarily applies to only 4 specific examples: flow, diffusion, jump in [Euclidean space], and CTMC in finite state spaces.

We stress that GM applies much more broadly (table 1 just lists examples). For example, as shown in our experiments, we also include Markov processes on manifolds (see also section H.5 on related work). Taking pairwise multimodal combinations gives us $\~80$ different classes of Markov processes. Further, there are different linear parameterizations of a generator and we can take Markov superpositions of all mentioned models. This leads to an explosive number of modelling options and design choices for GM.

Loss abstraction.

Although the theoretical concepts like target-affine (TA) loss and Bregman divergences are elegant, they lack practical utility [...].

We have done additional experiments showing that certain Bregman divergences improve training stability. New theoretical contributions allow us to drop the definition of TA-losses altogether such that we can simplify the manuscript. See general response for details.

Multimodal data.

It would certainly be more compelling to apply GM to multimodal tasks in protein design [...].

We have added multimodal protein experiments (see general response and the updated version of the paper) showcasing that GM improves the state-of-the-art.

Performance gain vs computational cost of training network.

The performance gain seems to be minimal at the cost of training a separate network [...].

In all of our protein experiments, we did not need to re-train any model for this result but repurposed an existing model (also for the new multimodal experiments). We made this clearer now. We note that even for re-training, the computational cost is essentially the same (only the final layer of the network has more channels).

Details for protein experiments. We have added the requested details and more in Appendix G.4.

The target task is restricted to protein with length less than 128, a relatively unchallenging setting.

We have used a standard benchmark which goes up to proteins with length 300.

Designability, as an important metric for evaluating protein backbone, is not presented.

We have included a comprehensive table with designability, diversity, and novelty in Appendix G.4. We don't discuss designability in the main paper because as noted in (Yim et al., 2024), "designability can be artificially high if the generative model samples the same designable protein repeatedly." Diversity measures the number of unique proteins that pass the designability criteria.

[...] More up-to-date work like FoldFlow (Bose et al., 2023) and FoldFlow-2 (Huguet et al., 2024) are missing from the benchmarks.

We have included FoldFlow in our new results (see general response ). FoldFlow2 code is not open sourced at time of writing, so we are unable to compare to them. We note that these newly reported FoldFlow results are worse than reported in their manuscript. We have reached out to the authors to verify we have ran the code from their github correctly.

Literature review for jump models: We included a very extended discussion of all related work in the supplementary material (section H) including a subsection on jump models in section H.7 discussing the suggested references.

Reflected SDEs.

I don't fully understand the argument in appendix D.2, starting from line 1473.

We added references and a brief explanation of reflected SDEs to make it more understandable. Adding the mentioned boundary conditions is necessary for a uniform distribution because of the following: To use the Fokker-Planck equation, it is necessary that the density is differentiable. Here, the density of a uniform distribution is not differentiable (see section A.3.). It might be helpful to look at the derivation in A.5.2 of the adjoint Fokker-Planck and see that the partial integration (equation (72)-(74)) has to be adapted.

Are there alternative choices of lambda_t that also satisfy the constraint?

Indeed, there are alternative choices for the jump intensity $\\lambda\_t$. Our derivation shows that the $\\lambda\_t$ is the minimal jump intensity with a state-independent jump kernel $J\_t$ that fulfils the jump continuity equation. However, we can artificially increase the number of jumps. For example, choosing

$\\begin{align\*} \\lambda\_t(x)=\\frac{1}{(1-t)^3}\\left(z-\\frac{(t+1)z}{2}\\right)^2,\\quad J\_t(x;x')=p\_t(x|z) \\end{align\*}$

also gives a solution. We tested this option but the version with minimal jump intensity gave better results.

Questions/Comments: Thank you for the suggested edits. We integrated them into the new version and we also simplified the notation adapting the notation requested by RJkg.

Thank you again for your helpful feedback.

   Developing improved numerical schemes for sampling are valuable directions for future work

     The task of image-text joint generation [...] seems highly promising for future exploration.

Thank you for your interest in this example. All information you requested is in the paper and we happily point you to the places where the information can be found that you request:

1. Use of uniform distribution: This example illustrates the flexibility of the generator matching framework, also with respect to the prior distribution. We could also use a Gaussian distribution in principle (the derivations in section E.3) are general. However, the function $G_0(x)$ cannot analytically be computed for a Gaussian (there is no closed formula) but for a uniform it can be (see section E.3).

2. Relation between two paths: One is a special case of the other. So the relation is to plugin the uniform distribution for $p_0$.

3. Arbitrary dimensions: We state in line 261 that we do everything for $d=1$ as proposition 4 easily allows to extend it to arbitrary dimensions.

4. Diagonal matrix: This is a common convention in the notation of SDEs. The term $\sigma_t(x|z)$ can be considered a diagonal matrix.

5. Diffusion generators: In fact, in general SDE generators include a drift (not necessarily a score). However, as stated in the paper, we make the choice of setting the drift to be zero in this example (line 274, 275). Further, as stated in the description of table 1, we most of the time make diffusion processes have zero drift in the work because we can use a Markov ensemble of a flow and a diffusion to include arbitrary drifts (i.e. this is no way a restriction). However, you can also just assume that there is a drift term throughout our work. The theory checks out the same way. This was done for brevity.

6. Reflection process: As stated in line 2177, we point to the literature for a definition of reflection SDEs. Explaining this mathematical background is out of scope for our work.

7. Cumulative distribution functions: They are correct as written. We integrate twice (the arguments appear as limits in the integrand).

Let us know if you have any further questions.