JKO Scheme for Learning Diffusion Processes — Poster

论文上下文

Paper title: Learning diffusion at lightspeed: Efficient JKO Scheme Optimization Abstract: A conference poster presenting a method for learning diffusion processes using the Jordan-Kinderlehrer-Otto (JKO) scheme. It details the mathematical formulation, training efficiency, scaling to high dimensions via heatmaps, and applications to single-cell dynamics. Paper body (method & results): Learning diffusion at lightspeed Antonio Terpin ETH Zürich aterpin@ethz.ch Nicolas Lanzetti ETH Zürich lnicolas@ethz.ch Martín Gadea ETH Zürich mgadea@ethz.ch Florian Dörfler ETH Zürich dorfler@ethz.ch Abstract Diffusion regulates numerous natural processes and the dynamics of many suc- cessful generative models. Existing models to learn the diffusion terms from observational data rely on complex bilevel optimization problems and model only the drift of the system. We propose a new simple model, JKOnet∗, which bypasses the complexity of existing architectures while presenting significantly enhanced representational capabilities: JKOnet∗recovers the potential, interaction, and internal energy components of the underlying diffusion process. JKOnet∗mini- mizes a simple quadratic loss and outperforms other baselines in terms of sample efficiency, computational complexity, and accuracy. Additionally, JKOnet∗pro- vides a closed-form optimal solution for linearly parametrized functionals, and, when applied to predict the evolution of cellular processes from real-world data, it achieves state-of-the-art accuracy at a fraction of the computational cost of all existing methods. Our methodology is based on the interpretation of diffusion pro- cesses as energy-minimizing trajectories in the probability space via the so-called JKO scheme, which we study via its first-order optimality conditions. Source code: https://github.com/antonioterpin/jkonet-star 1 Introduction Diffusion processes govern the homeostasis of biological systems [40], stem cells reprogramming [20, 36], and the learning dynamics of diffusion models [16, 22, 54] and transformers [19, 52]. The diffusion process of interest often originates from three quantities: a drift term due to a potential field, the interaction with other particles, and a stochastic term. If these three components are known, predictions follow from simple forward sampling [27] or the recent work in optimization in the probability space [1, 3, 11, 25, 34, 38, 41]. In this paper, we consider the case when the diffusion process is unknown, and we seek to learn its representation from observational data. The problem has been addressed when the trajectories of the individual particles are known [5, 32], but it is often the case that we only have “population data”. For instance, single-cell RNA sequencing techniques enabled the collection of large quantities of data on biological systems [35], but the observer cannot access the trajectories of individual cells since measurements are destructive [20, 36]. The most promising avenue to circumvent the lack of particle trajectories is the Jordan-Kinderlehrer-Otto (JKO) scheme [24] which predicates that the particles as a whole move to decrease an aggregate energy, while not deviating too much from the current configuration. However, the JKO scheme entails an optimization problem in the probability space. Thus, the problem of finding the energy functional that minimizes a prediction error (w.r.t. observational data) takes the form of a computationally- challenging infinite-dimensional bilevel optimization problem, whereby the upper-level problem is the minimization of the prediction error and the lower-level problem is the JKO scheme. Recent work [1, 9] exploits the theory of optimal transport and in particular Brenier’s theorem [7] to attack this bilevel optimization problem, a model henceforth referred to as JKOnet. Despite promising initial results in [9], this complexity undermines scalability, stability, and generality of the model. 38th Conference on Neural Information Processing Systems (NeurIPS 2024). arXiv:2406.12616v2 [cs.LG] 18 Oct 2024 effects mismatch = W2(µ1, ˆµ1)2 causes mismatch : P i ∇Vθ(xt+1 i ) + 1 τ (xt+1 i −xt i) 2 µt+1 = argmin µ∈P(Rd) Jθ(µ) + 1 2τ W2(µ, µt)2 µ0 µ1 µ2 µT ˆµ1 ˆµ2 ˆµT Figure 1: Given a sequence of snapshots (µ0, . . . , µT ) of a population of particles undergoing diffusion, we want to find the parameters θ of the parametrized energy function Jθ that best explains the particles evolution. Given θ, the effects mismatch is the Wasserstein distance between the observed trajectory and the predicted trajectory obtained iteratively solving the JKO step with Jθ. The first-order optimality condition in [30] applied to the JKO step suggests that the “gradient” of Jθ with respect to each ˆµt vanishes at optimality, i.e., for ˆµt = µt. For Jθ(µ) = R Rd Vθ(x)dµ(x), this condition is depicted on the right. The gradient (dashed blue arrows) of the true V (level curves in dashed blue) at each observed particle xt+1 i (blue circles) in the next snapshot µt+1 opposes the displacement (dotted red arrows) from a particle xt i (red triangles) in the previous snapshot µt. Instead, the gradient (solid green arrows) of the estimated Vθ (level curves in solid green) at each observed particle xt+1 i (square) does not oppose the displacement from a particle xt i in the previous snapshot µt. This mismatch in the causes of the diffusion process is what JKOnet∗minimizes. Furthermore, to be practical, it is limited to learning only potential energies, modelling the underlying physics only partially. Alternatively, [10, 43] learn directly the transport map describing the evolution of the population (i.e., the effects), bypassing the representation of the underlying energy functional (i.e., the causes). Motivated by robustness, interpretability, and generalization, here we seek a method to learn the causes. In [23, 42], the authors try to learn a geometry that explains the observed transport maps. Unfortunately, the cost between two configurations along a cost-minimizing trajectory is often not a metric [47]. Other attempts include recurrent neural networks [21], neural ODEs [15], and Schrödinger bridges [12, 28]. Contributions. We study the first-order necessary optimality conditions for the JKO scheme, an optimization problem in the probability space, and show that these conditions can be exploited to learn the energy functional governing the underlying diffusion process from population data, effectively bypassing the complexity of the infinite-dimensional bilevel optimization problem. We provide a closed-form solution in the case of linearly parametrized energy functionals and a simple, interpretable, and efficient algorithm for non-linear parametrizations. Via exhaustive nu- merical experiments, we show that, in the case of potential energies only, JKOnet∗outperforms the state-of-the-art in terms of solution quality, scalability, and computational efficiency and, in the until now unsolved case of general energy functionals, allows us to also learn interaction and internal energies that explain the observed population trajectories. When applied to predict the evolution of cellular processes, it achieves state-of-the-art accuracy at a fraction of the computational cost. Figure 1 shows an overview of our method, detailed in Section 3. 2 Diffusion processes via optimal transport 2.1 Preliminaries The gradient of ρ : Rd →R is ∇ρ ∈Rd and the Jacobian of ϕ : Rd →Rn is ∇ϕ ∈Rn×d. We say that f : Rd →R has bounded Hessian if ∇2f(x) ≤C for some C > 0 (and some matrix 2 norm ∥·∥). The divergence of F : Rd →Rn is ∇· F and its laplacian is ∇2F. The identity function is Id : Rd →Rd, Id(x) = x. We denote by P(Rd) the space of (Borel) probability measures over Rd with finite second moment. For µ ∈P(Rd), supp(µ) is its support. The Dirac’s delta measure at x ∈Rd, is δx. All the functions are assumed to be Borel, and for f : Rd →R, R Rd f(x)dµ(x) is the (Lebesgue) integral of f w.r.t. µ. If µ is absolutely continuous w.r.t. the Lebesgue measure, µ ≪dx, then it admits a density ρ : Rd →R≥0, and the integral becomes R Rd f(x)ρ(x)dx. The pushforward of µ via a (Borel) map f : Rd →Rd is the probability measure f#µ defined by (f#µ)(B) = µ(f −1(B)); when µ is empirical with N, µ = 1 N PN i=1 δxi, then f#µ = 1 N PN i=1 δf(xi). Given µ, ν ∈P(Rd), we say that a probability measure γ ∈P(Rd × Rd) is a transport plan (or coupling) between µ and ν if its marginals are µ and ν. We denote the set of transport plans between µ and ν by Γ(µ, ν). The Wasserstein distance between µ and ν is W2(µ, ν) :=  min γ∈Γ(µ,ν) Z Rd×Rd ∥x −y∥2dγ(x, y)  1 2 . (1) When µ and ν are discrete, (1) is a linear program. If, additionally, they have the same number of particles, the optimal transport plan is γ = (Id, T)#µ for some (transport) map T : Rd →Rd [39]. When µ is absolutely continuous, γ = (Id, ∇ψ)#µ for some convex function ψ [7]. 2.2 The JKO scheme Many continuous-time diffusion processes can be modeled by partial differential equations (PDEs) or stochastic differential equations (SDEs): Example 2.1 (Fokker-Planck). The Fokker-Planck equation, ∂ρ(t, x) ∂t = ∇· (∇V (x)ρ(t, x)) + β∇2ρ(t, x), (2) describes the time evolution of the distribution ρ of a set of particles undergoing drift and diffusion, dX(t) = −∇V (X(t))dt + p 2βdW(t), where X(t) is the state of the particle, V (x) the driving potential, and W(t) the Wiener process. The pioneering work of Jordan, Kinderlehrer, and Otto [24], related diffusion processes to energy- minimizing trajectories in the Wasserstein space (i.e., probability space endowed with the Wasserstein distance), providing a discrete-time counterpart of the diffusion process, the JKO scheme, µt+1 = argmin µ∈P(Rd) J(µ) + 1 2τ W2(µ, µt)2, (3) where J : P(Rd) →R ∪{+∞} is an energy functional and τ > 0 is the time discretization. Example 2.2 (Fokker-Plank as a Wasserstein gradient flow). The Fokker-Plank equation (2) results from the continuous-time limit (i.e., τ →0) of the JKO scheme (3) for the energy functional J(µ) = Z Rd V (x)dµ(x) + β Z Rd ρ(x) log(ρ(x))dx with dµ(x) = ρ(x)dx. 2.3 Challenges Section 2.2 suggests that we can interpret the problem of learning diffusion processes as the problem of learning the energy functional J in (3). Specifically, the setting is as follows: We have access to sample populations µ0, µ1, . . . , µT , and we want to learn the energy functional governing their dynamics. A direct approach to tackle the inverse problem is a bilevel optimization, used, among others, for the model JKOnet in [9]. This approach bases on the following two facts. First, by Brenier’s theorem, the solution of (3), µt+1, can be expressed1 as the pushforward of µt via the gradient of a convex function ψt : Rd →R and, thus, W2(µt, µt+1)2 = Z Rd ∥x −∇ψt(x)∥2dµt(x). 1Under an absolute continuity assumption. 3 Second, the optimization problem (3) is equivalently written as argmin ψt∈C J(∇ψt#µt) + 1 2τ Z Rd ∥x −∇ψt(x)∥2dµt(x), where C is the class of continuously differentiable convex functions from Rd to R. Therefore, the learning task can be cast into the following bilevel optimization problem, which minimizes the discrepancy between the observations (µt) and the predictions of the model (ˆµt): min J T X t=1 W2(ˆµt, µt)2 s.t. ˆµ0 = µ0, ˆµt+1 = ∇ψ∗ t ˆµt, ψ∗ t := argmin ψ∈C J(∇ψt#ˆµt) + 1 2τ Z Rd ∥x −∇ψt(x)∥2dˆµt(x). (4) A pra