Volumetric Optimal Transportation by Fast Fourier Transform
Authors: Na Lei
DONGSHENG An
Min Zhang
Xiaoyin Xu
David Gu
Venue: ICLR 2023 poster (2023)
Keywords: Optimal transport
Monge-Ampere equation
Elliptic PDE
Fast Fourier transform

Abstract: The optimal transportation map finds the most economical way to transport one probability measure to another, and it has been applied in a broad range of applications in machine learning and computer vision. By the Brenier theory, computing the optimal transport map is equivalent to solving a Monge-Amp\`ere equation, which is highly non-linear. Therefore, the computation of optimal transportation maps is intrinsically challenging.

In this work, we propose a novel and powerful method, the FFT-OT (fast Fourier transform-optimal transport), to compute the 3-dimensional OT problems. The method is based on several key ideas: first, the Monge-Amp\`ere equation is linearized to a sequence of linear elliptic PDEs with spacial and temporal variant coefficients; second, the obliqueness property of optimal transportation maps is reformulated as a Neumann boundary condition; and third, the variant coefficient elliptic PDEs are approximated by constant coefficient elliptic PDEs and solved by FFT on GPUs. We also prove that the algorithm converges linearly, namely the approximation error decreases exponentially fast. Experimental results show that the FFT-OT algorithm is more than a hundred times faster than the conventional methods based on the convex geometry. Furthermore, the method can be directly applied for sampling from complex 3D density functions in machine learning and magnifying the volumetric data in medical imaging. 

Published as a conference paper at ICLR 2023
VOLUMETRIC OPTIMAL TRANSPORTATION BY FAST FOURIER TRANSFORM
Na Lei
Dongsheng An
Dalian University of Technology na nale @edlut.edu.c.
Stony Brook University doan@cs. www.ssttoooooooooo edu
Min Zhang
Xiaoyin Xu
Xianfeng Gu
Zhejiang University
Harvard Medical School xxu@bwh. hawvard. edu
Stony Brook University quêcs. stonybookk.ddu
min zhange z ju. edu. en
ABSTRACT
The optimal transportation map finds the most economical way to transport one probability measure another, and it has been applied in a broad range of applica- tions in machine learning and computer vision. By the Brenier theory, computing the optimal transport map is equivalent to solving a Monge-Ampère equation. which is highly non-linear. Therefore, the computation of optimal transportation maps is intrinsically challenging. this work, we propose a novel and powerful method, the FFT-OT (fast Fourier wransformopptimll transport), to compute the 3-dimensional OT problems. The method is based on several key ideas: first, the Mopge-Ampère equation is linearized to sequence of linear elliptic PDEs with spacial and temporal variant coefficients; second, the obliqueness property of op- timal transportation maps is reformulated third, the variant coefficient elliptic PDEs are approximated by constant coeffi- cient elliptic PDEs and solved by FFT on GPUs. We also prove that the algorithm converges linearly. Experimental results show that the FFT-OT algorithm is more than hundred times faster than the conventional methods based on the convex geometry. Furthermore, the method can be directly applied for sampling from complex 3D density functions in machine learning and magnifying the volumetric data in medical imaging.
Neumann boundary condition; and
INTRODUCTION
Optimal transportation (OT) transports one probability measure to another in the most economical
way, and it plays a fundamental role in areas like machine learning Courty et al. (2017); ) Altschuler al. 2020). and computer graphics Solomon et al. (2015) Nader & Guennebaud (2018). Given a Riemannian man- ifold X, all the probability distributions on X form an infinite dimensional space P(X). Given any two distributions ju, E P(X), the optimal transportation map defines a distance between them, and the McCann interpolation McCann (1997) defines the geodesic connecting them. Hence opti- mal transportation equips P(X) with Riemannian metric and defines its covariant differentiation, which provides a variational calculus framework for optimization in it.
al. 2019), computer vision Arjovsky al. 2017): Tolstikhin et al. 2018);
As the optimal transportation problem is highly non-linear, it is quite challenging to compute the OT maps. Recently, researchers have developed many algorithms. The geometric variational approach Aurenhammer al. (1998); Gu et al. (2016); Levy (2013) based on the Brenier theorem Brenier (1991) is capable of achieving high accuracy for low dimensional problems, but it requires com- plicated geometric data structure and the storage complexity grows exponentially as the dimension increases. The Sinkhom method Cuturi (2013) based on the Kantorovich theorem adds an entropic regularizer to the primal problem and can handle high dimensional tasks, but it suffers from the intrinsic approximation error.
indicates equal contribution

---

Published as a conference paper at ICLR 2023
We propose novel method to tackle this challenging problem through Fast Fourier Transforma- tion (FFT). According to the Brenier theorem Brenier (199I), under the quadratic distance cost, the optimal transportation map is the gradient of the Brenier potential, which satisfies the Monge- Ampère equation. With the continuity method Delanoë (1991). J. the Monge-Ampère equation can be linearized sequence of elliptic partial differential equations (PDEs) with spacial and temporal variant coefficients. By iteratively solving the linearized Monge-Ampère equations, we can obtain the OT map. Specifically, we propose to approximate the linearized Monge- Ampèee equation by constant coefficient elliptic PDEs and solve them using the FFT on GPUs.
Our proposed FFT-OT method has many merits: (i) is generalizable for arbitrary dimension; (ii) has a linear convergence rate, namely the approximation error decays exponentially fast; (iii) in each iteration, the computational complexity of FFT is O(n log n), thus our algorithm can solve large scale OT problems; and (iv) it is highly parallelable and can be efficiently implemented on GPUs. We demonstrate the efficiency of the FFT-OT algorithm by solving the volumetric OT problems for machine learning and medical imaging applications including sampling from given 3D density functions and volumetric magnifier. The algorithm also has its own limitations: (i) although it can be generalized to any dimensions, the storage complexity increase exponentially with respect to the dimension, so its power is limited by the memory size of the GPUs; (ii) Since the algorithm uses FFT, the current version of the method only works well for continuous density functions. (iii) In this work, we mainly focus on the computation of the OT map from the uniform distribution to another arbitrary continuous distribution. To extend the method to find the OT map between any two continuous measures, we can compute two OT maps from the uniform distribution to the both continuous measures, then combine them together. The combination will give a reasonable approximation of the OT map Nader & Guennebaud (018)).
Though Lei and Gu Lei & Gu (202T) also uses FFT solve the 2-dimensional OT problem, our method differs their works in the following two aspects: (i) Lei and Gu's 's method uses the fixed point method to compute the 2D OT problems, ours is based on the linearization of the Monge-Ampère operator to solve the 3D OT problems, these are two different methodologies in PDE theory; (ii) In our paper, we also provide the theoretical convergence analysis of the proposed method. For more detailed analysis and related work, please refer to the Appendix A
2 OPTIMAL TRANSPORTATION THEORY
In this section, we review the fundamental concepts and theorems of the OT problem and the Monge- Amperè equation, more details can be found in Villani (2008).
Optimal Transportation Map and the Monge-Ampère equation Suppose the source domain Q is an open set in Rd with the probability measure J4, the target domain Σ is with the probability measure v. Both μ and v have density functions du(x) = f(r)dx and dv(y) g(y)dy, respectively, with the equal total mass: Sn fx)]x = fsz (yyddy, which is called the balance condition.
Suppose T : Ω + Σ is a measurable map. The mapping T is called measure preserving and denoted
as Tap = v if the following relation
M(T v(A)
for every Borel subset A Σ. A cost function e Ω x Σ R measures the transportation cost for transporting the unit mass from 2 E Ω to >. Problem 1 (Monge). The optimal transportation problem finds the measure preserving map with the minimal total transportation cost,
min / clx, (()) TwK=K
The solution to the Monge's problem is called the optimal transport map between a and v. The existence, uniqueness and regularity of OT maps depend on the boundedness and the continuity of the density functions, the convexity of the supporting domains, the continuity of their boundaries, and the cost function. In our current work, we focus on the similar situation in Saumier et al. (013).
The cost function is quadratic Fuclidean distance c(x. y) 1| 1- 1/2
2

---

Published as a conference paper at ICLR 2023
The supports of the source and the target measures the canonical cube Ω [-1, 1|3, which is uniformly convex;
The source and the target measures μ,v are absolutely continuous with respect to the Lebesgue measure, their densities J' are positive and bounded away from zero;
0 < m < f.g < M,
and f, are of class (9)),,
The boundary condition is second boundary condition (OT boundary condition), T(?))
f2
Then according (Villani 2003) Theorem 14.4, Saumier (0133) Theorem 2.1), the OT maps T: : .t - Ω exists and is unique and invertible (μ a.e), and the Brenier potential is of class (22) form some 0 < B a.
Theorem 2. Assume that Q2 p,v, f and are defined as above. Then there exists a convex functiinnnt + R, E (2)) for some O < 3 a, such that Vu pushes p forward to v, (Vu)# v. Moreover, Vu is unique and invertible (μ a.e), and its inverse Ve satisfies (Vv) p..
We call such convex function u the Brenier potential, it satisfies the Monge-Ampére equation,
detD) = f(x) go Vu(x)
with the boundary condition Vu( = E. Then finding the optimal transportation map is equivalent to solving the corresponding Monge-Amper Ampére equation. In the current work, the target measure is always the Lebesgue measure, and the source density f is of class (1)).
Linearized Monge-Ampère Operator
The Monge-Ampère operator is defined as
MA(u] detD2u,
which is highly non-linear. It can be linearized following:
MA(u ev] = det(D +D)x) = deuD2u + TrraceAddj(D'u D³u),
(3)
where Adj(A) is the adjoint (co-factor) matrix of A, Adj(A) = det(A)A-+ Therefore the lin- earized Monge-Ampère operator is defined as
DMA, [v] == Tracc(Adi(D*u)) D²v)= = Σ (="),),,,,
(4)
p.g-1
where (uPT) = Dd) is the adjoint matrix of the Hessian of u, and 0,0, == Dzpdu,
o²
Continuity Method For simplicity, we assume the source domain coincides with the target do- main, that Q E, and the target density g(x) = The Monge-Ampère equation Eqn. (2) is simplified as detl)²u(x) = f(x). Define flow of density as
p(z, t) = (1 - (z), E [0, 1]. The corresponding flow of the Brenier potentials is u(x, t): [0, 1]
(5)
detD", t) = p(x,t), s.t. Vau(r, ())) = 2
where D² t) is the Hessian of u(xx t) with respect to x, and u(x, 1) is the solution to the initial both sides of the linearized Monge-Ampère operator Eqn. 4 we obtain an elliptic PDE with the spacial and temporal variant
Monge-Ampère equation Eqn. 2 Take the derivative w.r.t. time
coefficients of the unknown v( t) :- u(x f), namely the "velocity" of the Brenier potential,
DMA,,
Σ
θ "," ot p(x, f(x)-11
Pig=1
At time = 0, the initial Brenier potential is known as u(æ, 0) = 1111112 Suppose at time t, we have obtained uxx t) already, then we can compute the adjoint matrix w t) of the Hessian (z t), and solve Eqn. (6) to get the velocity v(z, t) t(z,t). t). In turn, move forward to time t St, and update u(z, t + ot) by u(x,t) t) + ü(x, t)8t. By repeating this procedure, eventually we reach time t = and obtain the solution u(x) u(x; 1) to the initial Ampée-Ampère Eqn. (
3

---

Published as a conference paper at ICLR 2023
Obliqueness Boundary Condition Suppose the boundary of Ω is C' almost everywhere, there- fore at C' point z E a9, the outer normal n(x) well defined. For almost every boundary point E 0Ω, the obliqueness condition is represented
(((=),((V((())) 20.
(7)
Suppose Ω is a cuboid and has 6 faces, if boundary point x 0Q is a face, by the cyclic monotonicity of the map and the strict convexity of Villani 22008) its image Vu(z) must be on the same face of z, namely.
(Vu(æ) - zz (x)) =0
(8)
We can rewrite the Brenier potential as u(æ1, 22 ..... ma) Vu(x) - x = Vu(x). By Eqn. (8) v(x) satisfies the Neumann boundary condition,
Σ 2j + v(z1, za), then
dv
(x) 0, E an.
99)
On
Similarly, the velocity the (modified) Brenier potential v in Eqn. (6 also satisfies the Neumann boundary condition. The analysis about the existence and regularity of the solutions Equ. O with boundary condition Eqn. (9) can be found in the supplementary material.
3 COMPUTATIONAL ALGORITHM
Here we introduce the 3-dimensional FFT-OT algorithm, which can be generalized to any dimen- sions. We approximate the Monge-Ampèrt equation by sequence of constant coefficient elliptic PDEs, and solve them by FFT on GPUs. More detailed analysis about the solution of the discretized Monge-Ampère equation, and the proofs of the lemmas and theorems are given by Appendix B
3.1 CONTINUITY METHOD FOR SOLVING THE MONGE-AMPÈRE EQUATION
By using the continuity method, we can solve the Monge-Ampère equation iteratively. For simplic- ity, we assume the target measure is the Lebesgue's measure with g = 1. At the n-th iteration, the Brenier potential is represented as 1 ||a=||2 + M, (I), its Hessian matrix is Hz, == + the corresponding density function is defined as the determinant of the Hessian Pr. = de)(H))) and the velocity of the Brenier potential is u, (z). In the beginning, the Brenier potential Mu(x)) is zero, the Hessian Ho = I and the density is Po = 1. At the n-th step, we compute the adjoint matrix (H)((x)) of the Hessian matrix (x) for any satisfies the variant coefficient elliptic PDE induced by the linearized Monge-Ampère operator,
Ω. According to Eqn. (3), the velocity 2 (x)
DMA.., [vn]
(z)00,, (z) -=(((=)-pp)))
(10)
P:2=0
Note that the right hand side of Eqn. (6) is the difference between the initial and the target densities, whereas here is replaced by the difference between the initial and the current densities. The step length parameter 7 2 1 can be chosen to guarantee the convergence Loeper & Rapetti (2005).
The elliptic PDE Eqn. TTP is with spatially variant coefficients. Although the traditional finite element method (FEM) can solve it using the GMRES algorithm[Saan (2003). this algorithm can not be directly accelerated by GPUs To overcome this difficulty, we approximate Eqn. (10) by a much simpler elliptic PDE with constant coefficients, which can be directly solved using the following FFT-OT algorithm pipeline Alg. [ GPUs in Appendix C
At the n-th iteration, after obtaining the adjoint matrix [x]] x Ω, we compute the mean adjoint matrix [HT
1799 Sg. HRR (x)p))))))
P. = 0,1
(11)
Ja Pm (x)dr
and replace the elliptic PDE Eqn. (T)) with variant coefficients by the elliptic PDE with constant coefficients,
DMA.., [va]= = Σ H [ -f(x)) (x)))),
(12)
p.gma
4

---

Published as a conference paper at ICLR 2023
where DMA is called the mean linearized Monge-Ampére operator.
Then we solve the constant coefficient clliptic PDE Eqn. 12) by FFT Algorithm Alg. 2 in Appendix C Although the original variant coefficient PDE Eqn. 10) is replaced by its constant coefficient approximation Eqn. (12), the algorithm still converges to the solution with a linear convergence rate. This replacement allows the whole algorithm to be solved by FFT GPUs, which greatly improves the computational efficiency.
Theorem 3 (main). Given domain Q C TRd, which is canonical cuboid f2 [-1,1] and positive density function f -2 IR with the balance condition Jn (x))dr = -I dx, suppose the mirror reflection extension Eqn.
14 of the flat torus f T" + R C E (0. 1). then the
Monge-Ampère equation,
ReTD2u()) = f(x), Vu(Ω) 5
can be solved using the FFT-OT Algorithm Alg. D Appendix C In particular, one can choose the step length parameter T, such that there is constant 0 < that the approximation error satisfies
1 Pn+1 < Cy namely the algorithm has a linear convergence rate.
(13)
3.2 FFT SOLVER FOR CONSTANT COFFFICIENT ELLIPTIC PDES
To solve the constant coefficient elliptic PDE Eqn. [ we first extend the PDE the flat torus by mirror reflection, then discretize the domain and compute the differential operators by central difference scheme. Finally the PDE converted to algebraic equations in the frequency domain by FFT and can be efficiently solved on GPUs.
Extension by Mirror Reflection Suppose Ω [0, 1]3 and f IR are given, we extend Ω to = 1, and f to f + R by mirror reflection
f(r. v) = f(1=), Izl, V(z, E Ω.
(14)
By definition, f satisfies the periodic boundary condition and can be treated function defined on the flat torus T.. Ω is one of the fundamental domain of T3, The constant coefficients aP-9 keep unchanged. Then we solve the following constant coefficient elliptic PDE Eqn. (18) [ f with
the periodic boundary condition. Finally, the restriction of ū on Ω gives the initial solution w to L[u] with Neumann boundary condition.
In the following, to avoid using overly complicated symbols, use (u, f, to represent (ü, f. for simplicity.
Tessellation Suppose 2= [-1, 1]3 is the canonical cube (a fundamental domain of flat torus),
we tessellate to the regular cells, and the centers of the cells form a grid M x N x L. The Brenier potential Ω + R. is discretized
tensor U..j. with {i, k) {0..... ... M 11 {O....,AN 1} {O...., L - 1). The spacial step lengths are (h hy, h2) = (2/M 2/N,2/L), The coordinate of each sample point (I,, Vy, 24. is (I,, My. z4) (-1 hi + 1/2), -1 h (j + 1/2), -1
fg + 1/2)). The periodic boundary condition is then formulated as
Ui.j = W.4aM,44NNAt4. a, a, Z.
(15)
Finite Difference Differential Operator We use the standard central differences compute the differential operators. The first order derivative D2 is approximated by
Dzuij. Ui+1.j.k Un-1.j.k
2h
where the index + means i modulus M. The operators D,, D. are defined in a similar way. The second order derivative operator Drd and Dzy approximated by
Disu U.41,J. + 1.1k 24.
h2
D2y M4.ik M441.44K 1--15-11 M 1111k
4hghy
The other operators Dy Dzz. Dyy and Drs are defined similarly.

---

Published as a conference paper at ICLR 2023
Discrete Fourier Transformation
The discrete Fourier transformation (DFT) of Wij.k is given by
M
M-IN-I
Um. 1=0
(16)
U..sk MNL ΣEΣ
(17)
a=0 1=0
where Wmml =
E 2rrmn
2mmy
C - Wmml = e² 2rms
M
e'
2mma 20tk
e¹
and < == {m, are
the indices of the frequency coefficients. By using DFT, the differential operators converted to algebraic operators in the frequency domain.
Lemma 4. Suppose the discrete function is wi with the discrete Fourier transformation Eqn. TT and Equ. (7), by using the central difference scheme, the first order differential operator is given by
Dauijk MNL ΣΣΣ
M-IN-11-
sin
2xm M
Wannd
hg
m=0 n=0 1=0
the second order differential operators are represented by
M-IN-11- N-1L-1 M
2am (cos M
1)
D2 32 Ui.j.
MNL Σ ΣE
um.n.
Wmml
h2
M- N-IN-1 --11
2mm sin
sin
2mn
D2 Wiihk
MNL ΣΣΣ m=0
um.n.
M
N
Wmm!
liglig
n= 1=0
The other differential operators Dy, D., D pyy. D.., Dyz and Dz. also represented accordingly. The detailed proof' can be found in the supplementary material.
FFT Solver Suppose want to solve an elliptic PDE with constant coefficients on f C R3
2
2
L(u) = Σ a a.
+ -v0,, + u(z)= f(x),
(18)
p=0 9=0
1=0
with the periodic boundary condition, where q²'1,b",, are constants, the matrix (a²-9) is positive definite, namely the PDE is uniformly elliptic. By the discrete Fourier transformation F, we convert the differential equation to an algebraic equation in the frequency domain,
2
a F(FOOO)
+ Σ F(F(,,)) + =F() = F())
p=0 9=0
r=0
By applying Lemma and defining
=0 0.u 2(cos
28m M
1)
2(cos
2xm N
1)
Am.nt
+ a
h2
h
2.2 2(cos
2ml
1)
sin 2nm 2rm sin
0
L
0.2 + a..0
34
N
h2
h
(19)
sin 2an
sin 2nd
sin 2rl sin 2am
2
N
(a + a 2000
M
hygh
hgh
+ ois
2mm M
+ sin 2xm
N + 8
sin
2rd T.
+
hg
by
h
We have the algebraic equations in frequency domain,
imm...AAmm..' = fm...
With @m. ss easily obtain u..jk 's by the Inverse Discrete Fourier Transform (IDFT), which means solving the constant coefficient elliptic equation. The algorithm is described in Alg. 2 in Appendix C
The FFT for solving the constant coefficient elliptic PDE can be efficiently computed with GPUs. Moreover, the algorithm Alg. soovees the constant coefficient elliptic PDEs with periodic bound- ary condition, which can be generalized to solving the same type of PDEs with Neumann boundary condition by extending the PDE to the flat torus T using mirror reflection Eqn. 14).

---

Published as a conference paper at ICLR 2023
4 EXPERIMENTAL RESULTS
In this section, we firstly show that the our proposed FFT-OT algorithm converges linearly and runs 100x faster than the conventional convex geometry based DolverLevvy (0133) then demonstrate the method in two applications: 3D adaptive sampling and Volume Magnifier. All the algorithms are developed using generic C++ with CUDA Toolkit. All the experiments are conducted on a Windows laptop with Intel Core 17-7700HQ CPU with 16 GB memory and NVIDIA GeForce GTX 1060 Graphies Cards. More experiments can be found in Appendix D
4.1 RUNNING TIME AND CONVERGENCE ANALYSIS
To show the performance of the proposed method, we experiment on the density functions defined by the Gaussian mixture models. To be specific, the domain is a cube Q [0, 1] the 3-dimensional density function defined Ω is set to be f(x Σiz PtNV Ej), where N Σ,) represents Gaussian distribution with mean 14; and variance Σ; diag (a an oh 022) u4 E RR is uniformly sampled from 1]s, Jiz is uniformly sampled from [0, 0.5]. p3 E R is uniformly sampled from [0.2, 1] and normalized such that fg fzz)dz 1. Thus the source distribution μ is a complicated Gaussian mixture distribution restricted on Ω. Then by mirror reflection in Sec. 3.22 we obtain the
complex density function which is defined on [-1, 1]3 and satisfies the periodic boundary condition.
We directly use the FFT-OT algorithm Alg. to solve the linearized Monge-Ampère equation. With the approxi- mation error threshold æ = 1.0 x 10-6 and the resolution 256 256 256, the running time for our FFT-OT al-
Convergence FFT-OT
logillf Beillf--ily -11)))
gorithm with double precision on GPU is less than 175
seconds. The conventional convex geometry based al- gorithm for 3D optimal transportation Levy (2015) can neither handle such large data sets nor implemented on GPUs It can only compute OT map with resolution no greater than 100 x 100 x 100 on our system, which takes about 2700 seconds. When handling problem with 128 x 128 x 128 resolution, our FFT-OT consumes about 20.3 seconds, which is 130 faster than the power diagram based method Levy (2015).
-10
-12
10
15
Iterations Figure 1: Convergence Analysis.
Fig. shows the approximation error for the above Gaussian mixture density with respect to itera- tions, namely log [f Pm.l1 Our algorithm does converge linearly and the result is consistent with the prediction Eqn. 13 in Thm. Therefore, this experiment validates the theorem.
4.2 3D ADAPTIVE SAMPLING
Generating random samples matching a given density function plays an essential role in the appli- cations like Monte-Carlo integration stippling. Efficiently obtaining high quality samples is still an on-going research topic Bauer et al. (2013): Perrier et al. (2018). And optimal transportation has been successfully applied for generating high quality 2D samples de Goes et al. (2012): Nader & Guennebaud 2018). Most of the current research focuses on generating 2D samples fitting the given density function. Here we apply the proposed 3D FFT-OT method to generate high quality 3D samples according to the given complex density functions. To the best of our knowledge, it is the first work that uses OT to sample from 3D density functions.
Suppose the source probability distribution dp(z)) fz)dd is defined on Ω = |0. 1]3 with (( The target distribution dv(y) (y) = dy is the uniform distribution. We use the FFT-OT algorithm
Alg. to compute the OT map T : Ω + Ω, T = v. The domain is tessellated to 256 x 256 x 256 grid. For each Tisk. j. E [0, 1,..., 255). the image T(æ.jz) can be obtained. We use (T(x,)a) as vertices to compute the Delaunay triangulation of Ω2 Then representing the OT map 7 (2, n)) (Ω, v) as a piecewise linear map. the restriction of I on each tetrahedron is a linear map. Then the inverse OT map - (2) + (Ω,, μ) is also piecewise linear map. Namely, given grid point Ymn!, we can find a tetrahedron containing it. Suppose the vertices of the tetrahedron are (T(:;), T(xx), T(x)), then Yoand is computed
Yontt (xx) + >₃T(x)) + TxTA) + (TT))))
7

---

Published as a conference paper at ICLR 2023
(a) Density
(b) Rejection
(c) MH
(d) Slice
(c) Ours-R
(f) Ours-G
Figure 3D density function sampling. (a) The density functions in slice. The slices in each come from two different density functions. (b)-(f) The samples obtained by different sampling methods. (b) Rejection sampling. (c) Metropolis-Hasee Hastinggs (MH) algorithm Bishop (2006). (d) Slice sampling Neal (c) The sampling, results by mapping the random samples from the uniform distribution back to the desired distribution with 7- (f) The sampling results by mapping the grid centers back with T-1 The scores of the top right give the results of the Chi-square goodness-of-fit test. Smaller means better.
where the non-negative barycenter coordinates satisfy A; + X3 + A + A 1. Then the image of the inverse OT map given by
T (ymmt) AuT; + Λymz AE + ApZi-
(20)
We generate random samples } according to uniform distribution v on Q, then their images {T (yx)) are the desired random samples following the distribution μs.
In our experiment, we use the same Gaussian mixture settings of the density function as Sec. 4.1
Fig. 2 visualizes the generated samples. We randomly pick the k-th slice along the z-direction from the discretized volume, draw the source density function on this slice, and use pixel intensity to represent the density in Fig. a). (i) We uniformly generate 100k random samples {yk Ω and obtain the desired random samples by applying the inverse OT map {T [y=') (ii) We also set {yk} as the grid centers of Ω and obtain the corresponding samples of the desired distribution μ.
The samples around the k-thh slice of both sampling strategies are plotted in Fig. ([c) and Fig. Dt..
By visual comparison,
obvious that the distributions of Fig, ((e) and Fig. f are consistent
with the density function in Fig. aa... The consistency of the boundary of Fig. e) and (f) and Fig.
2)a) also verifies the obliqueness boundary condition of the Monge-Ampere equation. To further show the performance of the proposed method, we compare with the classical sampling methods, namely rejection sampling, the Metropolis- Hastings algorithm [Bishop (2006) and the slice sampling Neal (000)), shown in Fig. 2tb). Fig. ([ and Fig. dd. To quantitatively compare the sampling results, we use the Chi-square goodness-of-fit test, which firstly groups the data and then computes the 2 norm of the difference between the actual number of observations in each group and the expected number of observations. In our experiment, we set the group number to 64 x 64 x 64 and 500K samples make the comparison. The corresponding T.2 norm of each method is shown in the top-right the corresponding figure. We can see that the both sampling strategies of our method give smaller scores than the classical ones.
4.3 VOLUMETRIC MAGNIFIER
In reality, physical magnifiers can only magnify planar images. In medical image processing, it is highly desirable to magnify certain regions of the 3D MRIs or CT images. Our algorithm can address such requests with the user prescribed region of interest (ROI) and magnifying factor. Suppose the ROI is a symmetric region with the center (#, T7 Æ) E Ω and the radius oas oy, 0g in different directions. The density function f of the source measure μ is defined
f(x, y, =) = 0.5 0.5e ((x /2/ +vv (v- (==
We compute OT map T (S. (t) v). where is the uniform distribution. Similar to the method in 3D adaptive sampling, we compute the Delaunay triangulation of the images (T(z:a)), then the OT map 7' is represented piecewise linear map. The inverse optimal transportation map
8

---

Published as a conference paper at ICLR 2023
Figure 3: The volume magnifier of an aneurysm. The first column shows the original volumetric data, and the last three columns give the magnified data from the same viewpoints with different magnifying ratios. The yellow circle denotes the ROl/aneurysm. To obtain the results, =o By oz. and they are 0.83, 0.75 and 0.5 respectively.
Figure The volume magnifier of the knee. The first gives the the original volumetric data with different ROIs denoted by blue boxes from different viewpoints. and second row shows corresponding mag- nified results. In the experiments we set dg By σ. 0.75.
T-1; (2, + (S2, μ)) is also piccewise linear. For each grid point Vmn! E S2, we use Eqn. 20 to find its pre-image. Similarly, its corresponding intensity mnd is computed by linear interpolation. Then we obtain the new volumetric data {mmmt} with the magnified ROI and visualize the result with Voreen Meyer-Spradow al. (009).
Fig. demonstrates our volumetric magnifier by magnifying an aneurysm on blood vessel Hansen & Johnson (2004). We choose the aneurysm region the ROI. The first column gives the snapshot of the blood vessel, and the yellow circle denotes the location of the aneurysm. The last three columns show the magnified aneurysm with different magnifying ratio from the same viewpoints. Moreover, we show the magnified volumetric knee from different viewpoints with different ROIs denoted by the blue boxes in Fig. Our method only magnifies the ROIs and keeps other regions unchanged. Compared with the traditional method requiring tedious zoom in/out, our method only magnifies the ROI region and keeps the whole subject the field of view, which enables doctors to visualize the overall anatomy while scrutinize detailed anatomical structure at the same time.
CONCLUSION
In this paper, we propose the FFT-OT method to solve the optimal transportation problem. Accord- ing to the Brenier theory, under the quadratic distance cost, finding the solution to the OT problem is equivalent to solving the Monge-Ampère equation, which can be linearized as a sequence of variant coefficient elliptic PDEs. Later, the variant coefficient PDEs are approximated by constant coeffi- cient PDEs and solved by Fast Fourier Transfommatin. We also prove that the proposed method converges linearly. Experiments on volumetric data show that the FFT-OT can be used to sample from complex 3D density functions and magnify the volumetric data in medical images.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENT
This research was partially supported by National Key R&D Program of China A02IYFA1003003 and NSFC No. 61936002, T2225012. This work was also partially supported by NIH 3R0ILM011434-05S,, 1R IR21EB029733-01AL, NSF FAIN-2115095 and NSF CMMI-1762287.
REFERENCES
Mokhtar Z. Alaya, Maxime Berar, Gilles Gasso, and Alain Rakotomamonjy. Screening sinkhorn al- gorithm for regularized optimal transport. In Advances in Neural Information Processing Systems 32, 2019.
José I. Aliaga, Ernesto Dufrechou, Pablo Ezzatti, and Enrique S. Quintana-Orti, An efficient gpu version of the preconditioned gmres method. The Journal of Supercompuling, 75, 2019.
Jason Altschuler, Jonathan Niles-Weed, and Philippe Rigollet. Near-linear time approximation algo- rithms for optimal transport via sinkhorn iteration. In Advances in Neural Information Processing Systems 30, 2017.
Jason Altschuler, Francis Bach, Alessandro Rudi, and Jonathan Niles-Weed. Massively scalable sinkhorn distances via the nyström method. In H. Wallach, H. Larochelle, A. Beygelzimer,
F. d'Alché-Bus E. Fox, and R. Garnett (eds.). Advances in Neural Information Processing Sys- tems, volume 32. Curran Associates, Inc., 2019. URL http:tttt///ppcceeiiiiig neurips.
Dongsheng An, Yang Guo, Na Lei, Zhongxuan Luo, Shing-Tung Yau, and Xianfeng Gu. Ae-ot: A new generative model based on extended semi-discrete optimal transport. In International Conference on Learning Representations, 2020.
Dongsheng An. Na Lei, and Xianfeng Gu. Efficient optimal transport algorithm by accelerated gradient descent. In The Thirty-Sixth AAAI Conference on Artificial Intelligence (AAAI), 2022.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In ICML, pp. 214-223, 2017.
Aurenhammer, F. Hoffmann, and B. Aronov. Minkowski-type theorems and least-squares clus- tering. Algorithmica, 1998.
Martin Bauer, Sarang Joshi, and Klas Modin. Diffeomorphic density matching by optimal informa- tion transport. SIAM Journal on Imaging Sciences, 8, 2015.
J.D. Benamou, Y. Brenier, and K. Guittet. The Monge-Kantorows mass transfer and its com- putational fluid mechanics formulation. International Journal for Numerical Methods in Fluids, 2002.
Jean-David Benamou, Brittany D. Froese, and Adam M. Oberman. Numerical solution of the opti- mal transportation problem using the ampére equation. J. Comput. Phys, 2014.
Christopher M. Bishop. Pattern Recognition and Machine Learning. Springer, 2006.
Y. Brenier. Polar decomposition and increasing rearrangement of vector fields. Acad. Sci. Paris Sr. I Math., 305(19):805-808, 1987.
Y. Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Comm. Pure Appl. Math., 44(4):375-417, 1991.
Dario Cordero-Erausquin. Sur le transport de mesures périodiques monotone maps preserving pe- riodic measures. Comptes Rendus de l'Académie des Sciences Serris I Mathematics, 329: 199-202, 1999.
N. Courty, R. Flamary, D. Tuia, and A. Rakotomamonjy. Optimal transport for domain adaptation. IEEE Transactions Pattern Analysis and Machine Intelligence, 1853) 2017.
10

---

Published as a conference paper at ICLR 2023
Marco Cuturi. Sinkhor distances: Lightspeed computation of optimal transportation distances. In International Conference on Neural Information Processing Systems, 2013.
de Goes, K. Breeden, V. Ostromoukhov, and M. Desbrun. Blue noise through optimal transport. ACM Trans. Graph. (SIGGRAPH Asia), 31, 2012.
Philippe Delanoë. Classical solvability in dimension two of the second boundary-value problem associated with the Monge-Ampère operator. Annales de I'I..... Analyse non linéaire, 8(5): 443-457, 1991.
Pavel Dvurechensky, Alexander Gasnikov, and Alexey Kroshnin. Computational optimal transport: Complexity by accelerated gradient descent is better than by sinkhom's algorithm. Proceedings of the 35th International Conference on Machine Learning. PMLR, 2018.
Süli Endre. Lecture Notes on Finite Element Methods for Partial Differential Equations. University of Oxford, 2020.
David Xianfeng Gu, Feng Luo, Jian Sun, and Shing-Tune Yau. Variational principles for minkowski type problems, discrete optimal transport, and discrete monge-ampère equations. Asian Journal of Mathematics, 2016.
Charles D. Hansen and Chris R. Johnson. Visualization Handbook. Academic Press, 2004.
Jun Kitagawa, Quentin Mérigot, and Boris Thibert. Convergence of a newton algorithm for semi- discrete optimal transport. Journal of the European Mathematical Society, 2019.
Na Lei and Xianfeng Gu. Fft-ot: A fast algorithm for optimal transportation. In Proceedings of International Conference on Computer Vision (ICCV), 2021.
Bruno Levy. numerical algorithm for 12 semi-discrete optimal transport in 3d. ESAIM: M2AN, 49 16):11911 2015.
Gregorire Loeper and Francesca Rapetti. Numerical solution of the monge-ampère equation by newton's algorithm. C. Acad. Paris, pp- 319-324, 2005.
Robert J. McCann. A conerrttppiipppeeorrnneeettttgge Advances in mathematics, 128:153- 179, 1997.
Quentin Merigot. A multiscale approach to optimal transport. Computer Graphics Forum,, 2011 1.
Jennis Meyer-Spradow, Timo Ropinski, Jörg Mensmann, and Klaus H. Hinrichs. Voreen: A rapid- prototyping environment for ray-casting- caseddg--basdd volume visualizations. IEEE Computer Graphics and Applications, 2009.
Georges Nader and Gael Guennebaud. Instant transport maps on 2d grids. ACM Trans. Graph., 37 (6), 2018.
Radford M. Neal. Slice sampling. The Annals of Statistics, 2003.
Nicolas Papadakis, Gabriel Peyré, and Edouard Oudet. Optimal transport with proximal splitting. SIAM Journal on Imaging Sciences, 2014.
Hélène Perrier, David Coeurjolly, Feng Xie, Matt Pharr, Pat Hanrahan, and VictorOastrmmoukhov. Sequences with low-discrepancy blue-noise 2-d projections. Computer Graphics Forum, 2018.
Gabriel Peyré and Marco Cuturi. Computational optimal transport. Found. Trends Mach. Learn., 11 607, (667,
(5-6):355
2019.
Yousef Saad. Iterative Methods For Sparse Linear Systems. Society of Industrial and Applied Mathematics, 2003.
Filippo Santambrogio. Optimal Transport for Applied Mathematicians. Springer, 2015.
Houis-Philippo Saumier, Martial Agueh, and Boualem Khouider. An efficient numerical algorithm for the I optimal transport problem with periodic densities. IMA Journal of Applied Mathematics, 80:135-157 2013.
11

---

Published as a conference paper at ICLR 2023
Yuliy Schwartzburg, Romain Testuz, Andrea Tagliasacchi, and Mark Pauly. High-contrast com- putational caustic design. ACM Trans. Craph., 33(4), July 2014. ISSN 0730-0301. doi: 10.1145/200109:22211000. https:///ddi
11w.111111111111.222222 2601000
Justin Solomon, Fernando de Goes, Gabriel PeyrÅ, Marco Cuturi, Adrian Butscher, Andy Nguyen, Tao Du, and Leonidas Guibas. Convolutional wasserstein distances: Efficient optimal transporta- tion on geometric domains. ACM Transactions Graphies (TOG). 2015.
Kehua Su, Wei Chen, Na Lei, Junwei Zhang, Kun Qian, and Xianfeng Gu. Volume preserving mesh parameterization based on optimal mass transportation. Comput. Aided Des., 82:42-56, 2017.
Ilya Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schoelkopf. Wasserstein auto- encoders. ICLR, 2018
Cédric Villani. Topics in Optimal transportation. AMS, 2003.
Cédric Villani. Optimal transport: old and new, volume 338. Springer Science & Business Media, 2008.
12

---

Published as a conference paper at ICLR 2023
A RELATED WORK
There is a huge literature about optimal transportation. Here we will only briefly review the most re- lated works. For detailed reviews, we refer readers to Santtmmbrooio (2013): Peyré & Cuturi (2019).
The first type of algorithms is based on the Kantorovich theory. When both the input and output domains are Dirac masses, the Kantorovich problem can be treated as standard linear program- ming (LP) task. In order to tackle large data sets, Cuturi (2013) adds entropic regularizer to the original LP problem and the regularized problem can be quickly solved by the Sinkhorn al- gorithm. Recently, various algorithms have been proposed to further accelerate the computation by improving the efficiency of matrix-vector multiplications, including the Greenkhorn Altscchllrr at. (017), Screenkhorn (Alaya et al. (2019) and the NYS-SINK Altschuler et al. (2019) algorithms. Dvurechensky et al. [Duurcchensky et al. (2018) also propose the adaptive primal-dus
accelerated
gradient descent algorithm (APDAGD) to solve the discrete OT problem. An et al. An et al.) (2022) compute the approximate OT plan by smoothing the dual Kantorovich problem and solving it with the FISTA method. This kind of methods have limitations: (i) they only give transport plans and cannot produce the bijective transportation maps; and (ii) the computational complexity is too high to apply them in the scenarios with huge number of samples.
The second type of algorithms is based on the Brenier theory Brenier (987) and its intrinsic con- nection with convex geometry [Gu et al.] (2016). The semi-discrete OT algorithm proposed in Au- renhammer et al. (1998) finds the transport map between continuous distribution and discrete measure via a variational approach by dynamically constructing the power diagrams. Its efficiency can be further improved Levy) (2155; Merigot (2011) by the multi-resolution strategy. The algo- rithms proposed in Kitagawa et al. (201); Su et al. (2017) also improve the efficiency by applying the Newton's method. When both the source and target measures are continuous, some interpolation methods are necessary [Schwartzburg et al. (014). The major drawback of this type of algorithms is the high computational complexity of constructing the dynamic power diagram, which prevents them from handling high dimensional tasks. For example, for the 3D OT problems, these algorithms usually run very slow.
The third type of algorithms is based on computational fluid dynamics Benamou et al. 2002); Pa- padakis al. (2014). These methods aim at finding special camporal-spaciall flow field that trans- ports the initial source density to the target density with the minimal total kinetic energy. Then the diffeomorphism induced by the flow gives the optimal transport map under the quadratic Euclidean distance cost. However, this kind of algorithms are difficult to extend to high dimensional space.
The fourth type of algorithms directly solve the Monge-Ampère equation using numerical meth- ods. Loeper and Rapetti Loeper & Rapetti (2005) propose to solve the linearized Monge-Amper equation defined on flat torus in each iteration. Its corresponding variant coefficient elliptic PDE is converted to a positive definite linear system using the finiteddifference scheme, which can be solved by the BiCG algorithm Endre 2020). Benamou et al. Benamou et al.] (2014) propose to solve the linearized Guennebaud Nader & Guennebaud (2018) apply the similar discretization strategy and solve the Monge-Ampère equation by conjugate gradient method. Saumier et al. (Sumiier al. 2013) pro- pose to solve the linearized Monge-Ampère equation using FFT. In each iteration the elliptic PDE with spacial and temporal variant coefficients is converted to group of linear equations in the fre- quency domain, which is solved by the GMRES algorithm. Although the GMRES algorithm can be implemented on GPUs Aliaga et al. (2019) there is available open source code. The work in Saumier et al.] (2013) focuses on periodic boundary condition. but this our proposed work focuses on general second boundary condition; the work in Saumier al. (2013) concerns planar OT maps, ours emphasizes on volumetric OT maps, which has higher complexity. The work in (Saumier et al. (2013) can handle more general target measures, the proposed work currently only deals with the Lebesgue target measure. Nevertheless, the current work can be directly generalized to handle gen- eral target measures as well. Lei and Gu Lei & Gu (202I) use the fixed point method to compute the 2-dimensional OT problem based on FFT, but cannot be extended to solve the 3-dimensional problems.
Aange-Ampère more general domains using Newton's method. Nader and
In this work, we combine the idea of linearizing the Monge-Ampère equation Loeper & Rapetti (2005) and the idea of FFT Saumier et al. (2013). The key novelty of our proposed method is to use the mean linearized Monge-Amper operator Eqn. (12) replace the conventional linearized
13

---

Published as a conference paper at ICLR 2023
Monge-Ampere operator Eqn. (TO). This replacement allows the algorithm to be implemented GPUs and makes the algorithm hundreds of times faster. the following, we compute the 3- dimensional optimal transport problem by applying the proposed algorithm. Our method also runs more than 100x faster than the convex geometry based method Levy (015)).
B APPENDIX THEORY
In the section, we give the detailed proofs for several lemmas and theorems. Some of them are well known in the Mopée-Ampère PDE field and the applied mathematics field, we include them for the completeness.
B.1 EXISTENCE OF THE SOLUTION TO THE TIME DEPENDENT MONGE-AMPÈRE EQNUATION
Let Th / the n-dimensional flat torus. Below we sometimes identify with f = [0, 1 and assume all data are periodic. The existence and regularity of solutions to the Monge-Ampère equation are given by the following theorem.
Theorem 5. Suppose positive density function f Ω I. is defined on Ω = 1]", such that
Sa (())da = 1, and I CO))) then the solution Ω [0, 1] to the time-dependent Monge- Ampère equation
detD",u(x, t) = (1 (1-t)+ +tf(z,, Vzu(x, (S2) exists and is unique up to constant. Furthermore, there exist constants 0 < A < A, such that
(21)
7
u (x,tt))
Sp VEER", V(r 2x [0,1].
(22)
1=1
p-q=1
We refer readers to Cordero-Erausquin (1999) for detailed proof.
Weak Solution In practice, we compute the weak solution of the linearized Monge-Ampère Eqn. (6) using numerical methods. We first rewrite the differential operator to a divergence form, then define bi-linear form.
Since (x, f)) is the adjoint matrix of D2 t), by direct computation, we obtain
Ox (x, 1)= =0 V(x, () x
Vg 1, 1,...,
(23)
so Eqn. (6)
can be converted into the divergence form:
n Opu00
n
Σ
ap
u
pa
u a, U+
0, Σ
u a,000
p=1
4=1
p.9=1
9=1
p.4=1
we obtain
Σ 0p p=1
Σ
2 (x, ()0,((x))) = -(xx) 1
(24)
1=11
with Neumann boundary condition
Óv(2, 1) 0, V(r,) E an x [0, 1].
(25)
Ön
For any w E H (2), by differentiation of product, we obtain
Σ O p=1
Σ u
Σ
Σ
0,w0 = Σ Op
w +
w
1=1
p=1
9=1
p=1
1=1
by integrating both sides, and from the fact that satisfies the Neumann boundary condition, we deduce
/
n
Op
w
Σ
Σ
u ur
0. (26)
p=1
P.
80 p=1
1==1
14

---

Published as a conference paper at ICLR 2023
For any fixed time [0, by the divergence form, construct bilinear form ( x H (I2) and a linear form H (52) - R
n alv, w) Σ
/
up O 1(w) = 1)wd..
(27)
p-9=1
weak solution Eqn. (24 a functton v E H (12) such that
a(v, w) = I(W), Yw E H (2).
(28)
By the uniform ellipticity Eqn. the Lax-Milgram theorem Endre (2020) shows the existence of the weak solution.
B.2 DISCRETE LINEARIZED MONGE-AMPÈRE EQUATION SOLVABILITY
Galerkin Method In practice, we construct triangulation T of f2, such that the ratio between the diameter and mnsrribe-pphere radius of each simplex bounded, and variation of the diameters of all the simplexes is small. We call such kind of T quasi-uniform triangulation, and denote the largest diameter as For each vertex v; E T, we construct piecewise linear base function yPs such that pp is linear on each triangle, (vy) is dig. We define a finite dimensional subspace Vn H (22,
Vn
UA (x(x = Σ Ax E EET
Given function (S),
uA V, to denote its approximation in V.. Furthermore,
2A = Σ, Acip, we also use 4n to represent the coefficient vector ........ )x) depending on
the context. The weak solution Eqn. (28) the Monge-Ampère equation (6) is equivalent to find
a E (?)) such that a(v, w) [w) for all w E H²()). In discrete cases, we want to find
2h V/, such that
a(vr, wn) ((wn,, Vwn Vn.
(29)
Eqn.
29h is equivalent to the linear system,
a(v)1,91) 4(2)),422 :
a(p2, a((2,)42)
a(PNN)PI) a(pN, 422)
I((p1) ((p2))
A2
(30)
((
alp2, pN)
a(pp,,
AN
I(pNN)
From the weak solution to the lincarized Monge- Ampére equation [0p. we obtain the linear system
Eqn. 330
We denote the stiffness matrix (alyi, :p))) By the uniform ellipticity Eqn. (22) and
Vn
alv, v) X/ Ve
Assume Ja vdr 0, by Poincaré inequality,
|ell2 2 (1))) Ve E H (1),, / vdr 0, W
where the constant (12) depends on f. Combine the above two inequalities, we obtain
/ vdx 0. Similarly, By the uniform ellipticity Eqn. 22 and Vn C (
a v) 2 c|l Ve E H²)),,
(31)
a v) S A Vell2
For linear finite element and quasi-uniform triangulation, we have the inverse Poincaré inequality,
|/Venll C2(DA) 110m111
15

---

Published as a conference paper at ICLR 2023
where h is the diameter of each element. Combine the above two inequalities, obtain
a(vh: uh) S Cllealllaaa) Von E VA.
(32)
By combining the inequalities Eqn. 31 and Eqn. 32 we obtain
Von Vh, S. un =0, 0,
| |oalliiaca S sa(u,, vn) - Callun (2))
(33)
C3
where Cg > 1 a constant. Suppose 2/2 Di= Exp,, then
of Σ Sis J (()))) = = Et id=1
n
110 (2)
where { (E:) and the matrix
(Jn pp(p)) is positive definite. Therefore,
|El|
< 5" क5 Calle11
(34)
C4
By a(v), Uk) AE, combing incqualities Eqn.
(3 and Eqn.
34
obtain
C3C4 |EII S S C,CAlEII, YEERS
R
Σ & = 0,
(35)
where CgC4 1. This proves the following lemma, Lemma 6. By using Galerkin method using linear elements to numerically approximate the weak
solution Eqn. (28) to the linearized Monge-Ampère Eqn. (0. if the uniform ellipticity Eqn. (22 holds,
and the triangulation T quasi-uniform, then the stiffness matrix of the linear system Eqn. 30
positive definite on the space Ex & 0.
1EI < AECCLIII YEER" Σ 6 =0 CgC
(36)
1=1
where C3CAs 1.
Since the uniform ellipticity Eqn. (2 holds for any time E [0, 1). then we obtain
Corollary 7. By using Galerkin method with linear elements quasi-uniform triangulations, the
linearized Monge-Ampère equation in the continuity method Eqn. (6) always has a solution Uh. E Vn for any t E [0.
Please note that the central differential scheme can be treated Galerkin's method on special uniform triangulation. Therefore, the above estimates still hold.
B.3 CONVERGENCE RATE
Theorem 8 (main). Given domain Ω TR.", which is canonical cuboid Ω. [-1, and positive density function f : + R. with the balance condition
[ f(x)do
1
1.dr, dr,
suppose the mirror reflection extension Eqn. then Monge-Ampère equation.
TT + RRis C", (0,1), 14 of f to the flat torus f: T
det deiL)2(()) f(x),
Vu(l) Ω
can be solved using FFT-OT Algorithm Alg. 7 In particular, one can choose the step length parameter T, such that there is constant Y 1, the approximation error satisfies
| - 2 C
namely the algorithm has linear convergence rate.
16

---

Published as a conference paper at ICLR 2023
Proof. Suppose at the k + 1-th iteration, Ph-1 det(I + D2ux+1); ||x/11 O(-
f Pk+1 = -det(I I Duuu + D2vz))
f -eet( [ +DDuxx)
Σ 020 + +
= (ʃ - Pk) - Ix [vx] + o(T
where LR(Un] = Epy u/ 0 Ogk Hence by integration by parts Equ. (27)
P - pk+1 [2(0)= = ||ʃ - Pk: ( - 2 [x(vx](S - pk) ) + o(T
(p))
= [l - Pk (2) + 20x (f Pk: ok) + o(T
where ag is the bilinear form in Equ. 27 In the discrete case, all functions are in Vh. we denote
U |1 = || (()) = uz. Pun,
| = u uns
|| 1la w Aun.
by the inequality Eqn. (34) and Eqn. 35
|| I |lun = Calluall, C1
|| = 11uu1122 S C3Callual1.
C2Ca
Therefore
Il Ph, llf - Ph, 11 2T 2f At. (fh Ph,ks) + + d(-
(37)
where A& is the stiffness matrix in Eqn.(30). and A is the mean stiffnees matrix. (By the uniform ellipticity Eqn. (22), the eigen values of the adjoint matrix (xr f) is uniformly bounded away from zero in the space H : {5 E RR E, & = O}. so the eigen value of the mean adjoint matrix (1) is bounded away from zero in H. After discretization, the eigen values of A is strietly positive in H. hence A& invertible in H. In the following discussion, the term o(T will be ignored.) Remark that the following displayed equation is scalar
(Jh ph,x) A5 pb.k) tr((Sh pm, As A ph.k))
Since Aş and A are symmetric, positive definite on the space E, = 0. ||A S CyCa and ||Axll2 I C3C₁, so are their inverses. Since An and An are symmetric, positive definite on the space orthogonal to (1, 11,,, 1)"", by Eqn. (35) and ||.Az < || A 21 we have
(n 1) I1 - ph. S (fn ?T. A (fr ph,k).
C3
Plug into Eqn. 37 we have
1 (n (m
(n 1) K
Ilf 11 < 1
llf - ph. < (1-
l1fr - ph.ulla (38)
C3C2
We can choose the step-length -- such that y (0, 1). where
1)
7 = - CGg
Therefore
1/5n - ph. S 2 * - pholl3 S 1I - pnoll2
(39)
B.4 DIFFERENTIAL OPERATOR USING FFT
By using the Discrete Fourier Transformation, the differential operators can be converted to alge- braic operators in the frequency domain.
17

---

Published as a conference paper at ICLR 2023
Lemma 9. Suppose the discrete function is U. with discrete Fourier transformation
M-1N-1L-1
Wi.j.k MNL ΣEΣ
T2mms
2mms
2211
then the differential operator using central difference 0, is given by
0,D,uigh.
1141ik + U1- 2..
12
M-1 N-1
2xm 2 (cos
- 1) 2ems
MNL ΣΣ>
Um.n.
M
e
M c e
2 2012
h2
m= 1=0
where -I, and 0,
uijjk is given by,
U. 4highy
(ui+15+Ia + ui 15--1. Mi41,8-11 1,k Uy (.-11,+1.k)
1 M-IN-16-
Σ
Um.n.
20m sin
M
sin
2xm
N
2mgi
x
2ntk
MNL
h
n=0
Proof. By equations
cos(A + a) + cos(A a) - ccos(A)
[coo A sin Asin + (cos Acosa + sin Asin a:) - cooss A
2(cos - cos A
and
sin + a) sin(A a) si)
=(sin A cos + cos A Asn a) + (sin A cos cos Asin a) - 2 cos A
2(coss sin
obtain
(cos 2rm
Znmíi
2e
2ami
2mm
12
12
by direct computation, we have
1
0.D:ui.j.
(us+1.A U4-1,jj
20..j.k)
h2
M-1N-1L-1
2mm(s 31
(rm(1-1)
MNL ΣΣΣ
+ el
2e K3
c¹ N c
2016
h2
1=
M-1 N-IN-1 N-11- --11
(cos
2mm
1)
2omi 20m3 201k
C
e
MNL
h2
m==0
n=0 1=0
Similarly, by equations
cas(A A + a B + cos (A a B B) coss + 3) cas(A - a B+ 9)
cos(A + + cos B B) cos(A B+a 3) cos + B a
=2cos(A + cos(a
cos A B) cos(a 3)
=2cos(A + (coos
cos(a- 0))
coco(A + (cos
cos cos
sin sin 3 cas a cos sin sin
4cos(A A + B) sin 3
and
sin(A + B B) sin a B 3- -in0sin(AA a B 8) sin(A a B+
sinn(A B ++ + B
sin( (1
8) sin(A a+
=2sim(AA+ + B) cos(a + 3) sin(A B) cos(a 3)
=2sin(A1 + (cos(a ( cos(a - B))
=2sin(AA + (cos a cos / sin sin 3 - cos a cos B - sin sin
- 4sin + B) sin a ain 9
18

---

Published as a conference paper at ICLR 2023
deduce the following equation.
0,d, = 4hahy (Ui+I5+IK + 14-1,5-1,4 M1+1.j-1,K - Uy 11-1.++1.))
MNL ΣΣ
M- N-IN-I1--
Um.nd
sin
2am M
sin
2xn
N
e¹ e 2
2ems
6¹ 2
hghy
m=0 1=0 1=0
Similarly, we have the representations of other differential operators in the frequency domain,
0,0, uiisk
(1, + 11 2u.sskk
M M-IN-11-1
2 (cos
2nn - 1) 2rms 27ms 2x1k e"
ΣΣΣ
Um.n.
N
e
e
MNL
h
m=0 n=0 1=0
(Uijkt1 + Ui.jk-1 2(.jj.k) h:
M-IN-1L-1
(cos 2xl - 1)
ΣΣΣ
2mms 2nnj 2mtk
Um.nd
e"
MNL m=0 n= 1=0
h2
1
4hyybz (ui.+1,,41 Wi 1,8---18-11 Us.j 1. Ui 1,841)
1
Σ
M-IN-1L-
Σ
sin N
sin
L el 2 e¹ 2 e
2ath
Um.n.
MNT.
high
m= ==0
2=0
O.Dui..h =
4h.h2 (Ui+1.j4+1 + 1.j,-1 Mi+1_jk-1 Uj Mi-1.j.ke 1.j,++1)
Σ
M-IN-I1: 1-1
@mnnd hzhh
sin sin
2um
M
e"
Amgus
e. e
2emd 2mik
MNL
m=0
m=
C ALGORITHM PIPELINES
In this section. we give the algorithm pipeline of the FFT-OT in Alg. I and the details to solve the costant coefficient elliptic PDE through FFT Alg. 2
Algorithm 1: FFT-OT
Input: Domain 0 = [-1, 1] the source density function f > 0, the target density = 1, step length 7;
approximation error threshold
Output: Solution 1 |a=||2 u., the Monge-Ampère Eqn. 2 with the corresponding boundary
condition.
Initialize (x( 0:
while true do
Compute the Hessian matrix D2 (x);;
Compute the density function p.. (x) + det(I + D2 (x));
if [ Pa Pmll < a then
L Break;
Compute the adjoint matrix [H + Adj(T + D (x));
Compute the mean adjoint matrix [H using Eqn. [1]
Solve the constant coefficient elliptic PDE (12) using the FFT Solver Alg. Update the Brenier potential + (x) <- Un TU,;
19

---

Published as a conference paper at ICLR 2023
(a) Density
(b) Rejection
MH
(d) Slice
(c) Ours-rand
Ours-grid
Figure 5: 3D density function sampling. (a) The density functions in different slices of the model, namely the 40th, 56th, 72th 80th. (b)-(1) The samples obtained by different sampling methods. (b) Rejection
sampling. (c) Metropolis-Hasting (MH) algorithm Bishhp (2006). (d) Slice sampling Neal (2003). (c) The
sampling results by mapping the random samples from the uniform distribution back to the desired distribution with T- (f) The sampling results by mapping the grid centers back with T-1 The scores of the top right give the results of the Chi-square goodness-offfit test. Smaller means better. Zoom in for better visualization.
Algorithm 2: FFT Solver for the Constant Coefficient Elliptic PDE
Input: Domain Ω [-1, 1) M, N, L. {a } U", function ff with the periodic boundary condition
Output: Solution the elliptic PDE Equ. (18)
Discretize the domain
a M N grid
Sample the function f to fal.ii
Compute FFT using Eqn. [6
{ m FFT (...))): do
for (m, n. 1) E [0, MM
[0, N 1 x 10,
Compute the factor Am.n.t using Eqn. 19
Am,r. is then Um, f- 0;
else
L üm..d. + fm.n.// Am.n.t
Compute the Inverse FFT using Equ. Return {
17
(uija) <- IFFT (()m))))))))::
D APPENDIX EXPERIMENTS
In this section, as compensation of the experiments in the main paper, we give more results on the 3D adaptive sampling and volumetric magnifier.
D. MORE RESULTS ON 3D ADAPTIVE SAMPLING
In the experiments, set the density function f(x) E (n, E,), where N(u, E,) rep- resents Gaussian distribution with mean 14 and variance Σ; diag oi o oz). μ4 R3 is uniformly sampled from [0, 1|3, oij is uniformly sampled from [0, 0.5). P R is uniformly sam- pled from [0.2, 1] and normalized such that fg fz))dr 1. Thus the source distribution is complicated Gaussian mixture distribution restricted on Ω [0, 1]3. After computing the OT map
20

---

Published as a conference paper at ICLR 2023
(a) Density
(b) Rejection
MH
(d) Slice
(c) Ours-rand
Ours-grid
Figure 6: 3D density function sampling. The density functions in different slices of the model, namely the 56th, 64th, 80th 88th. (b)-(1) The samples obtained by different sampling methods. (b) Rejection
sampling. (c) Metropolis-Hasting (MH) algorithm Bishop (2006). (d) Slice sampling Neal (2003). (c) The
sampling results by mapping the random samples from the uniform distribution back to the desired distribution with T- (f) The sampling results by mapping the grid centers back with T-1 The scores of the top right give the results of the Chi-square goodness-of-fe test. Smaller means better. Zoom in for better visualization.
T from to the uniform distribution y defined on [-1, 1Js we conduct two groups of experiments: (i) we map the cell centers of the grid {y*} of [-1, 1]ss back [-1, through the inverse OT map T-1 defined by Eqn. (20): (ii) we randomly sample 100k samples {u&} from the Uniform dis- tribution defined in 1- 11 then map them back to 1-1, 1]3 through the inverse OT map T ((yk). In order to keep the consistency with the mirror reflection process in the FFT-OT algorithm, we also reflect the the generated samples back to Q. To visualize the results of kth slice, we plot samples whose coordinates satisfy the inequality.
k/128- 1/256 S = S k/128 1/256.
In Fig. and Fig. 6 we give more sampling results of different slices correspond to the two models used in Fig. in the main paper. Fig. 5 visualize the density function restricted the 40th, 56th, 72th and 80th slices for different methods the model displayed in the first row of Fig. 6 visualize the density function restricted methods of the model displayed in the second row 2 Compared with the classical methods, the both sampling strategies of our method give decent sampling results that fit the prescribed density function well. Moreover, the number of generated samples for different slices of the same 3D model fits the density functions restricted to the corresponding slices well, namely more samples are generated in the brighter regions for different slices.
the 56th, 64th, 80th and 88th slices for different
D.2 MORE RESULTS ON VOLUMETRIC MAGNIFIER
In this experiment, we magnify the volumetric MRI image of the aneurysm by different amplifica- tion factors. In Fig. 7 show the original aneurysm viewed from difference angles in the first column. The last three columns give the magnified results with different amplification factors from the viewpoints same those in the first column. We can see that the aneurysm region is successfully magnified by different factors and the rest parts the volume nearly keeps the same.
21

---

Published as a conference paper at ICLR 2023
(a) Original
(b) Magnifying ratio 1
(c) Magnifying ratio 2 (d) Magnifying ratio
Figure 7: The volume magnitier of an aneurysm. The first column shows the original volumetric data from different viewpoints, and the last three columns give magnified data from the same viewpoints the first column with different magnifying ratios. The yellow circles denote the aneurysm or the ROIs.
22


Reviews:
[{"review_id": "iHYGqBHm1qj", "reviewer": "Reviewer_wDkb", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "The authors propose a new method based on Fast Fourier to compute 3D optimal transport problems.\n\nThe authors show great speedups and provide theoretical insights using elliptic PDEs. Applications in sampling and for magnification of volumetric data.", "summary_of_the_review": "Overall, this is novel work and valuable. The speedups and theoretical connections are impressive.", "strength_and_weaknesses": "The methods is very fast but needs more memory (storage complexity increases exponentially). The method can be extended to work with any number of dimensions.\n\nThe connection between volumetric magnification and denoising or super resolution should be clarified.\n\nFigures 3 and 4 do not show the magnification ratios used.\n\nThe following statement needs clarification :\n\"Compared with the traditional method requiring tedious zoom in/out, our method only magnifies the\nROI region and keeps the whole subject in the field of view, which enables doctors to visualize the\noverall anatomy while scrutinize detailed anatomical structure at the same time.\"\n\nNot clear why this is better than magnifying the entire volume. Isn't it puzzling for the doctor looking at some regions being zoomed?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "QSCIcOM0RF", "reviewer": "Reviewer_Q68w", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "This submission proposes an algorithm for numerically estimating the optimal transport (OT) map in dimension $3$ between a source distribution $\\mu$ admitting a (Hölder) continuous density on $[0,1]^3$, and the Lebesgue measure on $[0,1]^3$. In this context, the optimal transport map can be characterized as the solution of a nonlinear degenerate elliptic partial differential equation (PDE) known as the Monge-Ampère equation. \n\nThe proposed method consists of applying the classical method of continuity to obtain a one-parameter family of PDEs $(\\det(D_x^2u(x,t))=\\rho(x,t))_{t\\in[0,1]}$ for which $u(x,0)$ is known and $u(x,1)$ can be used to obtain the OT map. Next, this family is discretized with respect to its parameter, linearized, and the resulting variable coefficient linear elliptic PDEs are made into constant coefficient PDEs by averaging the coefficients. Given the solution at the previous step, it is updated by solving the resulting constant coefficient PDE; which is amenable to computation by a fast and parallelizable spectral method based on the Fast Fourier Transform (FFT) implemented on GPUs. \n\nThe algorithm is shown to benefit from a linear rate of convergence. Applications to density sampling and volume magnification are also included.", "summary_of_the_review": "The method proposed in this submission is of interest due to its simplicity and speed. The main novelty is to average the coefficients of the linearized PDEs to enable a spectral approach. The proposed applications of this work are also of interest. As it stands, however, the submission has a number of technical mistakes and thus I cannot recommend acceptance of this work in its current state.\n\n## Post Rebuttal Comments:\n\nIn their rebuttal, the authors have fixed many of the technical issues that were brought up in the previous review. The scores have been amended to reflect these changes. Please also address the following list of issues; some of which were mentioned in the original review, but were not fixed in the revised manuscript. I also suggest that the revised manuscript be proofread to fix typos. \n\n1. Spelling of names with accents are still inconsistent. \n2. In the appendix many of the norms should be squared, but are not.\n3. After equation (37), the sentence \"The scalar\" followed by the displayed equation does not make sense. I would recommend replacing that sentence by \"Remark that the following displayed equation is a scalar\" if this emphasis is deemed to be necessary.\n4. $o(\\tau^{-1})$ is simply dropped after equation (37) without further comment, this is somewhat sloppy.\n5. In the display before (38), the upper bound is not used in the remainder of the proof to my understanding; if that is the case, the upper bound can just be dropped. The lower bound in the same display could use a bit more explanation to make it more clear to the reader. \n6. It would be useful to add a comment on why $\\bar{A}_k$ is invertible before (37).     ", "strength_and_weaknesses": "# Strengths \n- This submission builds nicely upon the previous work of Loeper-Rapetti [1] and Saumier et al. [2] by providing a fast and parallelizable spectral approach to solving the Monge-Ampère equation characterizing the OT map in this particular setting. This algorithmic improvement is enabled by the simple approach of averaging the coefficients of the variable coefficient PDEs obtained from the linearization approach discussed in [1] and [2].\n\n- The applications to density sampling and volume magnification are of interest and the proposed method appears to give good results in the provided experiments.\n\n# Weaknesses\n\n - In contrast to [2], the present submission does not cover the case where the target distribution is not the Lebesgue measure. The source distribution is also required to satisfy some conditions which restrict the general applicability of this method.\n\n-Theorem 3 states that \"one can choose the step length parameter $\\tau$ such that [...] the algorithm has a linear convergence rate.\" However, per the proof of this Theorem, the choice of $\\tau$ depends on the uniform ellipticity properties of the solution to the problem. There is, in particular, no discussion on how $\\tau$ was chosen in the experiments, nor how one should choose $\\tau$ in general. \n\n-The submission also contains a number of technical errors which we list below\n1. The entire paper appears to work under the assumption that the Brenier potential $u\\in C^2(\\Omega)$. As discussed in Theorem 2.1 in [2], this requires some additional conditions on the density of $\\mu$ which are not mentioned in the present submission.\n\n2. In equation (26), it is claimed that the integral over the boundary is equal to 0 as $\\partial \\Omega$ has zero measure. The reasoning here is erroneous, as although $\\partial \\Omega$ has zero Lebesgue measure, integrals over the boundary are taken with respect to the relevant surface measure, not the Lebesgue measure and $\\partial \\Omega$ has positive measure with respect to the surface measure.\n\n3. Although it is not explicitly mentioned in the submission, the density of $\\mu$ with respect to the Lebesgue measure must be positive on $\\Omega$ for some of the claims to hold.\n\n4. In the introduction, the OT map is stated to define a distance between any two probability distributions. This is incorrect, in particular there exists no transport map between a Dirac mass supported on one point and a normalized sum of two Dirac masses. It is rather the Kantorovich problem and corresponding OT couplings which are used to define such a distance.\n\n5. Problem 1 is stated in terms of measure preserving maps which, in turn, are only defined in this manuscript in the context of $C^1$ maps. As such the statement of Brenier's theorem provided (Theorem 2) is incorrect, as the OT map provided by the true Brenier theorem need not be $C^1$, rather it is only guaranteed to be differentiable $\\mu$-a.e. by Alexandrov's theorem under the present assumptions (cf. eg. Theorem 3.1 and Theorem 3.2 in [3]). Similarly, Problem 1 is typically stated as an optimization problem over arbitrary measurable maps satisfying the pushforward condition $T_{\\sharp} \\mu =\\nu$ i.e. $\\mu(T^{-1}(A))=\\nu(A)$ for every Borel set $A$.  \n\n6. The proof of the main theorem is difficult to follow to the point where I cannot guarantee its accuracy. In particular, the notation $a_k, A_k,\\bar A_k$ is undefined despite these objects playing a central role in the proof. Moreover, I do not understand the meaning of the display after equation (37),\n$(f_h-\\rho_{h,k})^TA_k\\bar A_k^{-1}(f_h-\\rho_{h,k})=\\text{tr} ((f_h-\\rho_{h,k})^TA_k\\bar A_k^{-1}(f_h-\\rho_{h,k}))$, it is obvious that the trace of a scalar is equal to itself.\n\n## References\n\n[1]: G. Loeper and F. Rapetti. Numerical solution of the Monge-Ampère equation by a Newton's algorithm. _C. R. Acad. Sci. Paris_, Ser. I 340(4), 2005.\n\n[2]: L.-P. Saumier, M. Agueh, and B. Khouider. An efficient numerical algorithm for the L2 optimal transport problem with periodic densities. _IMA J. Appl. Math._, 80(1), 2015.\n\n[3]: G. De Philippis and A. Figalli. The Monge–Ampère equation and its link to optimal transportation. _Bull. Am. Math. Soc_, 51(4), 2014.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "7F8ExRCm_4e", "reviewer": "Reviewer_AQZc", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper develops a fast method for solving the optimal transport problem on 3d densities. The authors convert a linearized functional equation to a set of elliptic differential equations, and solve those using an FFT solver---converting finite differences scheme to an algebraic system by passing to the frequency space. The authors perform a convergence analysis and show applications that highlight the usefulness of the method.", "summary_of_the_review": "Contributions in this field are important for a wide range of machine learning problems. I think that the result in this paper are promising and worthy of publication in some forum. However, the experimental analysis and is not in a condition to be published presently. More extensive analysis on the use of 3d FFT-OT for sampling could result in a more appropriate submission in my opinion, including comparison of runtime and more thorough empirical and theoretical analysis of the algorithm used as a sampler.\n\n--- \nI have updated more overall assessment of the contribution after more carefully reviewing the literature. The authors addressed my concerns with the g=1 assumption, which I believe that they should make more explicit in their manuscript. Other than that I think the work is acceptable, if slightly below my standards for rigorous empirical analysis. Time permitting, experiments assessing the quality of the sampling algorithm they propose on nD canonical distributions (e.g. Cauchy, Gaussian, comb) from the standpoint of information theoretic metrics would overcome my reservations here, but I don't think they're strictly necessary for publication.", "strength_and_weaknesses": "The main strength is that the method converges quickly and allows parallel calculation via gpus. It also provides a novel algorithm for nD approximate optimal transport (using the g=1, inversion trick). \n\nA primary weakness in this work is that the method is not rigorously evaluated on relevant application areas for this conference. The most relevant experiments in Section 4.2, on adaptive sampling, are very limited. Only a single density is shown, and while the numerical comparisons presented to common sampling algorithms look good I don't think it can be said that this approach is much better for sampling. In part because of obvious patterns that appear in the transported grid centers. Perhaps this could be cleaned up with a more sophisticated algorithm, but then a more thorough analysis would be called for. I think the paper could benefit from that. The 3d magnification example is also lacking in my opinion, I cannot understand the need to perform 3d magnification with this method, as opposed to restricting a new viewing window to the ROI. For medical imaging, I suspect that the resulting distortion would make whatever benefits could be obtained through magnification useless.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster