summaries: This paper proposed a deep network for point cloud sequence super-resolution/upsampling. Building on the basis of pointNet and PU-net, the main contribution of the paper is identifying the problem of temporal incoherence in the process of upsampling a point cloud shape representation as well as a training loss to encourage temporal coherence. In the cases showed in the paper, the proposed method seems effective comparing to previous work which is not done on sequence data. My main concern about the work is that the experimental evaluation is limited.

summaries: The paper addresses the task of learning temporally stable features for point clouds with an application to upsampling point clouds. Learning point-based descriptors has been a major topic of research in the recent vision and graphics meetings, where approaches have been proposed focusing semantic labeling, geometry-oriented tasks (e.g. normal estimation), and point-based graphics. However, as the paper states, and to the best of my knowledge, no methods have been proposed to learn features in fourth dimension in a temporally stable way. Thus, the very topic of research is significantly novel and promising.

summaries: The paper tries to learn temporally stable representations for point-based data sets and focus on varying size and dynamic point sets, and demonstrate its usefulness in the context of super-resolution. To deal with a difficult target that dynamically moves and deforms over time with variable input and output size, they take a novel temporal loss function for temporally coherent point set generation and siamese network setup for temporal loss calculation. Their novel temporal loss is based on EMD to minimize differences between an estimated point cloud and a desired super-resolution point cloud. The discussion and evolution on multiple loss functions are mostly well done. Except spatial loss is considered, taking the ground truth acceleration and estimated velocity into account is beneficial to this task. Their main contribution is taking permutation invariant loss terms and a siamese training setup and generator architecture, enabling improved output variance by allowing for dynamic adjustments of the output size, and identifying a specialized form of mode collapse for temporal point networks.

abstractText: Point clouds, as a form of Lagrangian representation, allow for powerful and flexible applications in a large number of computational disciplines. We propose a novel deep-learning method to learn stable and temporally coherent feature spaces for points clouds that change over time. We identify a set of inherent problems with these approaches: without knowledge of the time dimension, the inferred solutions can exhibit strong flickering, and easy solutions to suppress this flickering can result in undesirable local minima that manifest themselves as halo structures. We propose a novel temporal loss function that takes into account higher time derivatives of the point positions, and encourages mingling, i.e., to prevent the aforementioned halos. We combine these techniques in a super-resolution method with a truncation approach to flexibly adapt the size of the generated positions. We show that our method works for large, deforming point sets from different sources to demonstrate the flexibility of our approach.

authors:
R Qi Charles
Hao Su
Mo Kaichun
Leonidas Guibas

title: Pointnet: Deep learning on point sets for 3d classification and segmentation

venue: pages 77–85,

authors:
Charles Ruizhongtai Qi
Li Yi
Hao Su
Leonidas J Guibas

title: Pointnet++: Deep hierarchical feature learning on point sets in a metric space

authors:
Pedro Hermosilla
Tobias Ritschel
Pere-Pau Vázquez
Àlvar Vinacua
Timo Ropinski

title: Monte carlo convolution for learning on non-uniformly sampled point clouds

venue: ACM Trans. Graph.,

authors:
Binh-Son Hua
Minh-Khoi Tran
Sai-Kit Yeung

title: Pointwise convolutional neural networks

venue: In Computer Vision and Pattern Recognition (CVPR),

authors:
Lequan Yu
Xianzhi Li
Chi-Wing Fu
Daniel Cohen-Or
Pheng-Ann Heng

title: Pu-net: Point cloud upsampling network

venue: In CVPR,

authors:
Matan Atzmon
Haggai Maron
Yaron Lipman

title: Point convolutional neural networks by extension operators

title: Ec-net: an edge-aware point set consolidation network

venue: In European Conference on Computer Vision,

authors:
Siamak Ravanbakhsh
Jeff G. Schneider
Barnabás Póczos

title: Deep learning with sets and point clouds

authors:
Roman Klokov
Victor S. Lempitsky

title: Escape from cells: Deep kd-networks for the recognition of 3d point cloud models

venue: In IEEE International Conference on Computer Vision, ICCV 2017,

authors:
Riccardo Roveri
A. Cengiz Öztireli
Ioana Pandele
Markus Gross

title: Pointpronets: Consolidation of point clouds with convolutional neural networks

venue: Computer Graphics Forum,

authors:
Paul Guerrero
Yanir Kleiman
Maks Ovsjanikov
Niloy J. Mitra

title: PCPNet: Learning local shape properties from raw point clouds

authors:
Kangxue Yin
Hui Huang
Daniel Cohen-Or
Hao Zhang

title: P2p-net: Bidirectional point displacement net for shape transform

authors:
Hang Su
Varun Jampani
Deqing Sun
Subhransu Maji
Evangelos Kalogerakis
Ming-Hsuan Yang
Jan Kautz

title: Splatnet: Sparse lattice networks for point cloud processing

authors:
Weiyue Wang
Ronald Yu
Qiangui Huang
Ulrich Neumann

title: Sgpn: Similarity group proposal network for 3d point cloud instance segmentation

authors:
Yifan Xu
Tianqi Fan
Mingye Xu
Long Zeng
Yu Qiao

title: Spidercnn: Deep learning on point sets with parameterized convolutional filters

venue: Lecture Notes in Computer Science,

authors:
Wenxuan Wu
Zhongang Qi
Fuxin Li

title: Pointconv: Deep convolutional networks on 3d point clouds

authors:
Jiaxin Li
Ben M. Chen
Gim Hee Lee

title: So-net: Self-organizing network for point cloud analysis

authors:
Xiaoqing Ye
Jiamao Li
Hexiao Huang
Liang Du
Xiaolin Zhang

title: 3d recurrent neural networks with context fusion for point cloud semantic segmentation

authors:
F. Groh
P. Wieschollek
H.P.A. Lensch

title: Flex-convolution (million-scale point-cloud learning beyond grid-worlds)

venue: In Computer Vision - ACCV 2018 - 14th Asian Conference on Computer Vision,

authors:
Loïc Landrieu
Martin Simonovsky

title: Large-scale point cloud semantic segmentation with superpoint graphs

authors:
Dario Rethage
Johanna Wald
Jürgen Sturm
Nassir Navab
Federico Tombari

title: Fullyconvolutional point networks for large-scale point clouds

venue: In Computer Vision – ECCV

authors:
Yue Wang
Yongbin Sun
Ziwei Liu
Sanjay E. Sarma
Michael M. Bronstein
Justin M. Solomon

title: Dynamic graph cnn for learning on point clouds

authors:
Huan Lei
Naveed Akhtar
Ajmal Mian

title: Spherical convolutional neural network for 3d point clouds

authors:
Yingxue Zhang
Michael G. Rabbat

title: A graph-cnn for 3d point cloud classification

venue: IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP),

title: Ursa: A neural network for unordered point clouds using constellations

authors:
Martin Simon
Stefan Milz
Karl Amende
Horst-Michael Groß

title: Complex-yolo: Real-time 3d object detection on point clouds

authors:
Yin Zhou
Oncel Tuzel

title: Voxelnet: End-to-end learning for point cloud based 3d object detection

authors:
Liuhao Ge
Yujun Cai
Junwu Weng
Junsong Yuan

title: Hand pointnet: 3d hand pose estimation using point sets

authors:
Nathaniel Thomas
Tess Smidt
Steven M. Kearnes
Lusann Yang
Li Li
Kai Kohlhoff
Patrick Riley

title: Tensor field networks: Rotation- and translation-equivariant neural networks for 3d point clouds

authors:
Yaoqing Yang
Chen Feng
Yiru Shen
Dong Tian

title: Foldingnet: Point cloud auto-encoder via deep grid deformation

authors:
Haowen Deng
Tolga Birdal
Slobodan Ilic

title: Ppfnet: Global context aware local features for robust 3d point matching

authors:
Yongbin Sun
Yue Wang
Ziwei Liu
Joshua E. Siegel
Sanjay E. Sarma

title: Pointgrow: Autoregressively learned point cloud generation with self-attention

authors:
Xingyu Liu
Charles Ruizhongtai Qi
Leonidas J. Guibas

title: Learning scene flow in 3d point clouds

authors:
Panos Achlioptas
Olga Diamanti
Ioannis Mitliagkas
Leonidas J. Guibas

title: Representation learning and adversarial generation of 3d point

venue: clouds. CoRR,

authors:
Haoqiang Fan
Hao Su
Leonidas J. Guibas

title: A point set generation network for 3d object reconstruction from a single image

authors:
Manuel Dahnert
Angela Dai
Leonidas Guibas
Matthias Nießner

title: Joint embedding of 3d scan and cad objects

venue: In ICCV

authors:
Yunzhu Li
Jiajun Wu
Russ Tedrake
Joshua B Tenenbaum
Antonio Torralba

title: Learning particle dynamics for manipulating rigid bodies, deformable objects, and fluids

venue: In ICLR,

authors:
Markus Ihmsen
Jens Cornelis
Barbara Solenthaler
Christopher Horvath
Matthias Teschner

title: Implicit incompressible sph

venue: IEEE Transactions on Visualization and Computer Graphics,

heading: 1 INTRODUCTION

text:
Deep learning methods have proven themselves as powerful computational tools in many disciplines, and within it a topic of strongly growing interest is deep learning for point-based data sets. These Lagrangian representations are challenging for learning methods due to their unordered nature, but are highly useful in a variety of settings from geometry processing and 3D scanning to physical simulations, and since the seminal work of Qi Charles et al. (2017), a range of powerful inference tasks can be achieved based on point sets. Despite their success, interestingly, no works so far have taken into account time. Our world, and the objects within it, naturally move and change over time, and as such it is crucial for flexible point-based inference to take the time dimension into account. In this context, we propose a method to learn temporally stable representations for point-based data sets, and demonstrate its usefulness in the context of super-resolution.
An inherent difficulty of point-based data is their lack of ordering, which makes operations such as convolutions, which are easy to perform for Eulerian data, unexpectedly difficult. Several powerful approaches for point-based convolutions have been proposed (Qi et al., 2017; Hermosilla et al., 2018; Hua et al., 2018), and we leverage similar neural network architectures in conjunction with the permutation-invariant Earth Mover’s Distance (EMD) to propose a first formulation of a loss for temporal coherence.
In addition, several works have recognized the importance of training point networks for localized patches, in order to avoid having the network to rely on a full view of the whole data-set for tasks that are inherently local, such as normal estimation (Qi Charles et al., 2017), and super-resolution (Yu et al., 2018a). This also makes it possible to flexibly process inputs of any size without being limited by memory requirements. Later on we will demonstrate the importance of such a patch-based approach with sets of changing cardinality in our setting. A general challenge here is to deal with varying input sizes, and for super-resolution tasks, also varying output sizes. Thus, in summary we target an extremely challenging learning problem: we are facing permutation-invariant inputs and targets of varying size, that dynamically move and deform over time. In order to enable deep learning approaches in this context, we make the following key contributions: Permutation invariant loss terms for temporally coherent point set generation; A Siamese training setup and generator architecture for point-based super-resolution with neural networks; Enabling improved output variance by allowing for dynamic adjustments of the output size; The identification of a specialized form of mode collapse for temporal point networks, together with a loss term to remove them. We demonstrate that these contributions together make it possible to infer stable solutions for dynamically moving point clouds with millions of points.
More formally, we show that our learning approach can be used for generating a point set with an increased resolution from a given set of input points. The generated points should provide an improved discretization of the underlying ground truth shape represented by the initial set of points. For the increase, we will target a factor of two to three per spatial dimension. Thus, the network has the task to estimate the underlying shape, and to generate suit-
able sampling positions as output. This is generally difficult due to the lack of connectivity and ordering, and in our case, positions that move over time in combination with a changing number of input points. Hence it is crucial that the network is able to establish a temporally stable latent space representation. Although we assume that we know correspondences over time, i.e., we know which point at time t moved to a new location at time t+ ∆t, the points can arbitrarily change their relative position and density over the course of their movement, leading to a substantially more difficult inference problem than for the static case.

heading: 2 RELATED WORK

text:
Deep learning with static point sets was first targeted in PointNet (Qi Charles et al., 2017) via order-invariant networks, while PointNet++ (Qi et al., 2017) extended this concept to generate features for localized groups similar to a convolution operation for grid-based data. This concept can be hierarchically applied to the generated groups, in order to extract increasingly abstract and global features. Afterwards, the extracted features can be interpolated back to the original point cloud. The goal to define point convolutions has been explored and extended in several works. The MCNN approach (Hermosilla et al., 2018) phrased convolution in terms of a Monte Carlo integration. PointCNN (Hua et al., 2018) defined a pointwise convolution operator using nearest neighbors, while extension-restriction operators for mapping between a point cloud function and a volumetric function were used in Atzmon et al. (2018). The PU-Net (Yu et al., 2018a) proposed a network for upsampling point clouds, and proposed a similar hierarchical network structure of PointNets along the lines of PointNet++ to define convolutions. Being closer to our goals, we employ this approach for convolutional operations in our networks below. We do not employ the edge-aware variant of the PU-Net (Yu et al., 2018b) here, to keep it as simple and general as possible as we focus on temporal changes in our work.
Permutation invariance is a central topic for point data, and was likewise targeted in other works (Ravanbakhsh et al., 2016; Zaheer et al., 2017). The Deep Kd-network (Klokov and Lempitsky, 2017) defined a hierarchical convolution on point clouds via kd-trees. PointProNets (Roveri et al., 2018) employed deep learning to generate dense sets of points from sparse and noisy input points for 3D reconstruction applications. PCPNet (Guerrero et al., 2018), as another multi-scale variant of
PointNet, has demonstrated high accuracy for estimating local shape properties such as normal or curvature. P2PNet (Yin et al., 2018) used a bidirectional network and extends PointNet++ to learn a transformation between two point clouds with the same cardinality.
Recently, the area of point-based learning has seen a huge rise in interest. One focus here are 3D segmentation problems, where numerous improvements were proposed, e.g., by SPLATNet (Su et al., 2018), SGPN (Wang et al., 2018a), SpiderCNN (Xu et al., 2018), PointConv (Wu et al., 2018), SONEt(Li et al., 2018a) and 3DRNN (Ye et al., 2018). Other networks such as Flex Convolution (Groh et al., 2018), the SuperPoint Graph (Landrieu and Simonovsky, 2018), and the fully convolutional network (Rethage et al., 2018) focused on large scale segmentation. Additional areas of interest are shape classification (Wang et al., 2018b; Lei et al., 2018; Zhang and Rabbat, 2018; Skouson, 2018) and object detection (Simon et al., 2018; Zhou and Tuzel, 2018), and hand pose tracking (Ge et al., 2018). Other works have targeted rotation and translation invariant inference (Thomas et al., 2018), and point cloud autoencoders (Yang et al., 2018; Deng et al., 2018). A few works have also targeted generative models based on points, e.g., for point cloud generation (Sun et al., 2018), and with adversarial approaches (Li et al., 2018b). It is worth noting here that despite the huge interest, the works above do not take into account temporally changing data, which is the focus of our work. A notable exception is an approach for scene flow (Liu et al., 2018), in order to estimate 3D motion directly on the basis of point clouds. This work is largely orthogonal to ours, as it does not target generative point-based models.

heading: 3 METHODOLOGY

text:
We assume an input point cloud X = {x1, x2, ..., xk} of size k ∈ [1, kmax]. It consists of points xi ∈ Rd, where d includes 3 spatial coordinates and optionally additional features. Our goal is to let the network fs(X) infer a function Ỹ which approximates a desired super-resolution output point cloud Y = {y1, y2, ..., yn} of size n ∈ [1, nmax] with yi ∈ R3, i.e. fs(X) = Ỹ ≈ Y . For now we assume that the number of output points n is defined by multiplying k with a user-defined upsampling factor r, i.e. n = rk. Figure 2a) illustrates the data flow in our super-resolution network schematically. We treat the upsampling problem as a local one, i.e., we assume that the inference problem can be solved based on a spatially constrained neighborhood. This allows us to work with individual patches extracted from input point clouds. At the same time, it makes it possible to upsample adaptively, for example, by limiting the inference to relevant areas, such as complex surface structures. For the patch extraction we use a fixed spatial radius and normalize point coordinates within each patch to lie in the range of [−1, 1]. Our first building block is a measure for how well two point clouds represent the same object or scene by formulating an adequate spatial loss function. Following Achlioptas et al. (2017), we base our spatial loss LS on the Earth Mover’s Distance (EMD), which solves an assignment problem to obtain a differentiable bijective mapping φ : ỹ → y. With φ we can minimize differences in position for arbitrary orderings of the points clouds via:
LS = min φ:ỹ→y ∑ ỹi∈Ỹ ‖ỹi − φ(ỹi)‖22 (1)
3.1 TEMPORAL COHERENCE
When not taking temporal coherence explicitly into account, the highly nonlinear and ill-posed nature of the super-resolution problem can cause strong variations in the output even for very subtle changes in the input. This results in significant temporal artifacts that manifest themselves as flickering. In order to stabilize the output while at the same time keeping the network structure as small and simple as possible, we propose the following training setup. Given a sequence of high resolution point clouds Y t, with t indicating time, we can compute a velocity V t = {vt1, vt2, ..., vtk}, where vti ∈ R3. For
this we use a finite difference (yt+1i − yti), where we assume, without loss of generality, ∆t = 1, i.e. the time step is normalized to one. For training, the low resolution inputs X can now be generated from Y via down-sampling by a factor of r, which yields a subset of points with velocities. Details of our data generation process will be given below.
To train a temporally coherent network with the Y t sequences, we employ a Siamese setup shown in Figure 2b. We evaluate the network several times (3 times in practice) with the same set of weights, and moving inputs, in order to enforce the output to behave consistently. In this way we avoid recurrent architectures that would have to process the high resolution outputs multiple times. In addition, we can compute temporal derivatives from the input points, and use them to control the behavior of the generated output.
Under the assumption of slowly moving inputs, which theoretically could be ensured for training, a straightforward way to enforce temporal coherence would be to minimize the movement of the generated positions over consecutive time steps in terms of an L2 norm:
L2V = n∑ i=1 ‖ỹt+1i − ỹ t i‖22. (2)
While this reduces flickering, it does not constrain the change of velocities, i.e., the acceleration. This results in a high frequency jittering of the generated point positions. The jitter can be reduced by also including the previous state at time step t− 1 to constrain the acceleration in terms of its L2 norm:
L2A = n∑ i=1 ‖ỹt+1i − 2ỹ t i + ỹ t−1 i ‖ 2 2 (3)
However, a central problem of a direct temporal constraint via Equations (2) and (3) is that it consistently leads to a highly undesirable clustering of generated points around the center point. This is caused by the fact that the training procedure as described so far is unbalanced, as it only encourages minimizing changes. The network cannot learn to reconstruct realistic, larger motions in this way, but rather can trivially minimize the loss by contracting all outputs to a single point. For this reason, we instead use the estimated velocity of the ground truth point cloud sequence with a forward difference in time, to provide the network with a reference. By using the EMD-based mapping φ established for the spatial loss in Equation (1), we can formulate the temporal constraint in a permutation invariant manner as
LEV = n∑ i=1 ‖ ( ỹt+1i − ỹ t i)− (φ(ỹt+1i )− φ(ỹ t i) ) ‖22. (4)
Intuitively, this means the generated outputs should mimic the motion of the closest ground truth points. As detailed for the L2-based approaches above, it makes sense to also take the ground truth acceleration into account to minimize rapid changes of velocity over time. We can likewise formulate this in a permutation invariant way w.r.t. ground truth points via:
LEA = n∑ i=1 ‖ ( ỹt+1i − 2ỹ t i + ỹ t−1 i )− (φ(ỹ t+1 i )− 2φ(ỹ t i) + φ(ỹ t−1 i ) ) ‖22. (5)
We found that a combination of LEV and LEA together with the spatial loss LS from Eq. 1 provides the best results, as we will demonstrate below. First, we will introduce the additional loss terms of our algorithm.

heading: 3.2 VARIABLE POINT CLOUD SIZES

text:
Existing network architectures are typically designed for processing a fixed amount of input and output points. However, in many cases, and especially for a localized inference of super-resolution, the number of input and output points varies significantly. While we can safely assume that no patch exceeds the maximal number of inputs kmax (this can be ensured by working on a subset), it can easily happen that a certain spatial region has fewer points. Simply including more distant points could guarantee that we have a full set of samples, but this would mean the network has to be invariant to scaling, and to produce features at different spatial scales. Instead, we train our
Published as a conference paper at ICLR 2020
Sheet2
Page 1
200
250
Ref Particle Cnt Frequency
1 31 61 917 13 19 25 37 43 49 55 67 73 79 85 97 10 3 10 9 11 5 12 1 12 7
0
200
400
600
800
1000
1200
1400
1600
1800
1 31 61 917 13 19 25 37 43 49 55 67 73 79 85 97 10 3 10 9 11 5 12 1 12 7
0
200
400
600
800
1000
1200
1400
1600
1800
1 of 1 22/01/2019, 11.28
file:///Users/thuerey/LRZ Sync+Share/NeuralParticles/2019_01...
1 of 1 22/01/2019, 11.28
file:///Users/thuerey/LRZ Sync+Share/NeuralParticles/2019_01...
1 of 1 22/01/2019, 11.30
file:///Users/thuerey/LRZ Sync+Share/NeuralParticles/2019_01...
1 of 1 22/01/2019, 11.30
network for a fixed spatial size, and ensure that it can process varying numbers of inputs. For inputs with fewer than kmax points, we pad the input vector to have a fixed size. Here, we ensure that the padding values are not misinterpreted by the network as being point data. Therefore, we pad X with p ∈ {−2}d, which represents a value outside the regular patch coordinate range [−1, 1]: X̄ = {x1, x2, ..., xk, p, p, ..., p︸ ︷︷ ︸
kmax−k
}. The first convolutional layer in our network now filters out the
padded entries using the following mask: Min = {mi∈[0,k]} = {1, 1, ..., 1︸ ︷︷ ︸ k , 0, 0, .., 0︸ ︷︷ ︸ kmax−k }. The entries of p allow us to compute the mask on the fly throughout the whole network, without having to pass through k. For an input of size k, our network has the task to generate ñ = rk points. As the size of the network output is constant with rkmax, the outputs are likewise masked with Mout to truncate it to length ñ for all loss calculations, e.g., the EMD mappings. Thus, as shown in Figure 2a, ñ is used to truncate the point cloud Ȳ = {ȳ1, ȳ2, ..., ȳnmax} via a mask Mout to form the final output Ỹ = {ȳi|i ∈ [1, ñ]}. Note that in Sec. 3.1, we have for simplicity assumed that n = rk, however, in practice the number of ground truth points n varies. As such, ñ only provides an approximation of the true number of target points in the ground truth data. While the approximation is accurate for planar surfaces and volumes, it is less accurate in the presence of detailed surface structures that are smaller than the spatial frequency of the low-resolution data.
We have analyzed the effect of this approximation in Fig. 3. The histograms show that the strongly varying output counts are an important factor in practice, and Fig. 4 additionally shows the improvement in terms of target shape that results from incorporating variable output sizes. In general, ñ provides a good approximation for our data sets. However, as there is a chance to infer an improved estimate of the correct output size based on the input points, we have experimented with training a second network to predict ñ in conjunction with a differentiable output masking. While this could be an interesting feature for future applications, we have not found it to significantly improve results. As such, the evaluations and results below will use the analytic calculation, i.e., ñ = rk.

heading: 3.3 PREVENTING HALO ARTIFACTS

text:
For each input point the network generates r output points, which can be seen as individual groups g: ψ(g) = {Ỹi|i ∈ [rg + 1, (r + 1)g]}. These groups of size r in the output are strongly related to the input points they originate from. Networks that focus on maintaining temporal coherence for the dynamically changing output tend to slide into local minima where r output points are attached as a
fixed structure to the input point location. This manifests itself as visible static halo-like structures that move along with the input. Although temporal coherence is good in this case, these cluster-like structures lead to gaps and suboptimal point distributions in the output, particularly over time. These structures can be seen as a form of temporal mode collapse that can be observed in other areas of deep learning, such as GANs. To counteract this effect, we introduce an additional mingling loss term to prevent the formation of clusters by pushing the individual points of a group apart:
LM = 1
d ñr e d ñr e∑ i
|ψ(i)|∑ ỹg∈ψ(i) ‖ ∑ ψ(i) |ψ(i)| − ỹg‖2
(6)
Figure 6: Left, a result without the mingling loss from Eq. 6, right with (a single point group highlighted in orange). The former has many repeated copies of a single pattern, which the mingling loss manages to distribute as can be seen in the right picture.
Note that in contrast to previously used repulsion losses (Yu et al., 2018a), LM encourages points to globally mix rather than just locally repelling each other. While a repulsion term can lead to a deterioration of the generated outputs, our formulation preserves spatial structure and temporal coherence while leading to well distributed points, as is illustrated in Fig. 6.
In combination with the spatial and temporal terms from above, this leads to our final loss function Lfinal = LS + γLEV + µLEA + νLM , with weighting terms γ, µ, ν.

heading: 4 EVALUATION AND RESULTS

text:
We train our network in a fully supervised manner with simulated data. To illustrate the effect of our temporal loss functions, we employ it in conjunction with established network architectures from previous work (Qi Charles et al., 2017; Yu et al., 2018a). Details of the data generation and network architectures are given in the appendix. We first discuss our data generation and training setup, then illustrate the effects of the different terms of our loss function, before showing results for more complex 3D data sets. As our results focus on temporal coherence, which is best seen in motion, we refer readers to the supplemental materials at https://ge.in.tum.de/publications/ 2020-iclr-prantl/ in order to fully evaluate the resulting quality.
Ablation Study We evaluate the effectiveness of our loss formulation with a two dimensional ablation study. An exemplary patch of this study is shown in Fig. 5. In order to compare our method to previous work, we have trained a previously proposed method for point-cloud super-resolution, the PU-Net (Yu et al., 2018a) which internally uses a PointNet++ (Qi et al., 2017), with our data set, the only difference being that we use zero-padding here. This architecture will be used in the following comparisons with previous work. Fig. 5a) shows a result generated with this network. As this figure contains an average of multiple frames to indicate temporal stability, the blurred regions, esp. visible on the right side of Fig. 5a), indicate erroneous motions in the output. For this network the difficulties of temporally changing data and varying output sizes additionally lead to a suboptimal approximation
of the target points, that is also visible in terms of an increased LS loss in Table 1. While Fig. 5b) significantly reduces motions, and leads to an improved shape as well as LS loss, its motions are overly constrained. E.g., at the bottom of the shown patch, the generated points should follow the black input points, but in (b) the generated points stay in place. In addition, the lack of permutation invariance leads to an undesirable clustering of generated points in the patch center. Both problems are removed with LEV in Fig. 5c), which still contains small scale jittering motions, unfortunately. These are removed by LEA in Fig. 5d), which shows the result of our full algorithm. The success of our approach for dynamic output sizes is also shown in the LN column of Table 1, which contains an L2 error w.r.t. ground truth size of the outputs.
Temporally Coherent Features A central goal of our work is to enable the learning of features that remain stable over time. To shed light on how our approach influences the established latent space, we analyze its content for different inputs. The latent space in our case consists of a 256-dimensional vector that contains the features extracted by the first set of layers of our network. Fig. 7 contains a qualitative example for 100 randomly selected patch sequences from our test data set, where we collect input data by following the trajectory of each patch center for 50 time steps to extract coherent data sets. Fig. 7a) shows the averaged latent space content over time for these sequences. While the model trained with temporal coherence (green curve) is also visually smoother, the difference becomes clearer when considering temporal frequencies. We measure averaged frequencies of the latent space dimensions over time, as shown in Fig. 7b,c). We quantify the differences by calculating the integral of the frequency spectrum f̃ , weighted by the frequency x to emphasize high frequencies, i.e, ∫ x x · f̃(x)dx. Hence, small values are preferred. As shown in Fig. 7b), the version trained without our loss formulations contains significantly more high frequency content. This is also reflected in the weighted integrals, which are 36.56 for the method without temporal loss, and 16.98 for the method with temporal loss. To verify that our temporal model actually establishes a stable temporal latent space instead of ignoring temporal information altogether, we evaluate the temporal frequencies for the same 100 inputs as above, but with a randomized order over time. In this case, our model correctly identifies the incoherent inputs, and yields similarly high frequencies as the regular model with 28.44 and 35.24, respectively. More details in Appendix C.
In addition, we evaluated the changes of generated outputs over time w.r.t. ground truth motion. For this we mapped the generated point clouds Ỹ t = {ỹt1, ỹt2, ..., ỹtñ} for 100 frames to evenly and dense sampled ground-truth points on the original mesh Y t = {yt1, yt2, ..., ytn} (the moving man shown in Fig. 1). This gives us a dense correlation between the data and the generated point clouds. For the mapping we used an assignment based on nearest neighbors: γ : ỹ → y. Using γ we divide Ỹ t into n subsets Ŷi = {ỹj |γ(ỹj) = yi} which correlate with the corresponding ground-truth points. For each subset we can now compute the mean position 1|Ŷ i | ∑ ŷ∈Ŷi ŷ and the sample density |Ŷi| measured by the number of points assigned to a ground-truth sample position. The temporal change of these values are of particular interest. The change of the mean positions should correspond to the ground-truth changes, while the change of the density should be one. We have evaluated the error of the first and second derivative of positions, as well as the first and second derivative of density (see Fig. 8 and Table 2). As can be seen from the plots, our method leads to clear improvements for all measured quantities. The individual spikes that are visible for both versions in the position errors (c,d) most likely correspond to sudden changes of the input motions for which our networks undershoots by producing a smooth version of the motion.
3D Results Our patch-based approach currently relies on a decomposition of the input volumes into patches over time, as outlined in Appendix A. As all of the following results involve temporal data, full sequences are provided in the accompanying video. We apply our method to several complex 3D models to illustrate its performance. Fig. 9 shows the input as well as several frames generated with our method for an animation of a spider. Our method produces an even and temporally stable reconstruction of the object. In comparison, Fig. 9b) shows the output from the previous work architecture (Yu et al., 2018a). It exhibits uneven point distributions and outliers, e.g., above the legs of the spider, in addition to uneven motions.
A second example for a moving human figure is shown in Fig. 1. In both examples, our network covers the target shape much more evenly despite using much fewer points, as shown in Table 3. Thanks to the flexible output size of our network, it can adapt to sparsely covered regions by generating correspondingly fewer outputs. The previous work architecture, with its fixed output size, needs to concentrate the fixed number of output points within the target shape, leading to an unnecessarily large point count. In order to demonstrate the flexibility of our method, we also apply it to a volumetric moving point cloud obtained from a liquid simulation. Thanks to the patch-based evaluation of our network, it is agnostic to the overall size of the input volume. In this way, it can be used to generate coherent sets with millions of points. These examples also highlight our method’s capabilities for generalization. While the 3D model was only trained on data from physics simulations, as outlined above, it learns stable features that can be flexibly applied to volumetric as well as to surface-based data. The metrics in Table 1 show that for both 2D and 3D cases, our method leads to significantly improved quality, visible in lower loss values for spatial as well as temporal terms.
Another interesting field of application for our algorithm are physical simulations. Complex simulations such as fluids, often employ particle-based representations. On the one hand, the volume data is much larger than surface-based data, which additionally motivates our dynamic output. On the other hand, time stability plays a very important role for physical phenomena. Our method produces detailed outputs for liquids, as can be seen in our supplemental video.
Convergence graphs for the different versions are shown in Fig. 12 of the supplemental material. These graphs show that our method not only successfully leads to very low errors in terms of temporal coherence, but also improves spatial accuracy. The final values of LS for the 2D case are below 0.05 for our algorithm, compared to almost 0.08 for previous work. For 3D, our approach yields 0.04 on average, in contrast to ca. 0.1 for previous work.

heading: 5 CONCLUSION

text:
We have proposed a first method to infer temporally coherent features for point clouds. This is made possible by a novel loss function for temporal coherence in combination with enabling flexible truncation of the results. In addition we have shown that it is crucial to prevent static patterns as easy-to-reach local minima for the network, which we avoid with the proposed a mingling loss term. Our super-resolution results above demonstrate that our approach takes an important first step towards flexible deep learning methods for dynamic point clouds.
Looking ahead, our method could also be flexibly combined with other network architectures or could be adopted for other applications. Specifically, a combination with PSGN (Fan et al., 2016) could be used to generate point clouds from image sequences instead of single images. Other conceivable applications could employ methods like Dahnert et al. (2019) with our approach for generating animated meshes. Due to the growing popularity and ubiquity of scanning devices it will, e.g., be interesting to investigate classification tasks of 3D scans over time as future work. Apart from that, physical phenomena such as elastic bodies and fluids (Li et al., 2019) can likewise be represented in a Lagrangian manner, and pose interesting challenges and complex spatio-temporal changes.

heading: ACKNOWLEDGMENTS

text: This work is supported by grant TH 2034/1-1 of the Deutsche Forschungsgemeinschaft (DFG).

heading: A TRAINING AND EVALUATION MODALITIES

text:
Data Generation We employ a physical simulation to generate our input and output pairs for training. This has the advantage that it leads to a large variety of complex motions, and gives full control of the generation process. More specifically, we employ the IISPH (Ihmsen et al., 2014) algorithm, a form of Lagrangian fluid simulator that efficiently generates incompressible liquid volumes. These simulations also have the advantage that they inherently control the density of the point sampling thanks to their volume conserving properties. In order to generate input pairs for training, we randomly sample regions near the surface and extract points with a given radius around a central point. This represents the high-resolution target. To compute the low-resolution input, we downsample the points with a Poisson-disk sampling to compute a point set with the desired larger spacing. In order to prevent aliasing from features below the coarse resolution, we perform a pass of surface fairing and smoothing before downsampling. Due to the large number of patches that can be extracted from these simulations, we did not find it necessary to additionally augment the generated data sets. Examples of the low- and high-resolution pairs are shown in the supplemental material.
Below we will demonstrate that models trained with this data can be flexibly applied to moving surface data as well as new liquid configurations. The surface data is generated from animated triangle meshes that were resampled with bicubic interpolation in order to match a chosen average per-point area. This pattern was generated once and then propagated over time with the animation. When applying our method to new liquid simulations, we do not perform any downsampling, but rather use all points of a low-resolution simulation directly, as a volumetric re-sampling over time is typically error prone, and gives incoherent low resolution inputs.
Given a moving point cloud, we decompose it into temporally coherent patches in the following manner: We start by sampling points via a Poisson-disk sampling in a narrow band around the surface, e.g., based on a signed distance function computed from the input cloud. These points will persist as patch centers over time, unless they move too close to others, or too far away from the surface, which triggers their deletion. In addition, we perform several iterations for every new frame to sample new patches for points in the input cloud that are outside all existing patches. Note that this resampling of patches over time happens instantaneously in our implementation. While a temporal fading could easily be added, we have opted for employing transitions without fading, in order to show as much of the patch content as possible.
Network Architecture and Training Our architecture heavily relies on a form of hierarchical point-based convolutions. I.e., the network extracts features for a subset of the points and their nearest neighbors. For the point convolution, we first select a given number of group centers that are evenly distributed points from a given input cloud. For each group center, we then search for a certain number of points within a chosen radius (a fraction of the [-1,1] range). This motivates our choice for a coordinate far outside the regular range for the padded points from Sec. 3.2. They are too far away from all groups by construction, so they are filtered out without any additional overhead. In this way, both feature extraction and grouping operations work flexibly with the varying input sizes. Each group is then processed by a PointNet-like sub-structure (Qi Charles et al., 2017), yielding one feature vector per group.
The result is a set of feature vectors and the associated group position, which can be interpreted as a new point cloud to repeatedly apply a point convolution. In this way, the network extracts increasingly abstract and global features. The last set of features is then interpolated back to the original points of the input. Afterwards a sub-pixel convolution layer is used to scale up the point cloud extended with features and finally the final position vectors are generated with the help of two additional shared, fully-connected layers. While we keep the core network architecture unmodified to allow for comparisons with previous work, an important distinction of our approach is the input and output masking, as described in Sec. 3.2.
Our point data was generated with a mean point spacing, i.e., Poisson disk radius, of 0.5 units. For the 2D tests, an upscaling factor of r = 9 was used. For this purpose, patches with a diameter of 5 were extracted from the low-resolution data and patches with a diameter of 15 from the high-resolution
data. We used the thresholds kmax = 100 and nmax = 900. For the loss, we used γ = 10, µ = 10, and ν = 0.001. The network was trained with 5 epochs for a data set with 185k pairs, and a batch size of 16, the learning rate was 0.001 with a decay of 0.003. For the 3D results below, the scaling factor r was set to 8. The diameter of the patches was 6 for the low-resolution data and 12 for the high-resolution data, with kmax = 1280 and nmax = 10240. The loss parameters were γ = µ = 5, with ν = 0.001. Learning rate and decay were the same for training, but instead we used 10 epochs with 54k patches in 3D, and a batch size of 4.

heading: B NETWORK ARCHITECTURE DETAILS

text:
The input feature vector is processed in the first part of our network, which consists of four point convolutions. We use (ng, rg, [l1, ..., ld]) to represent a level with ng groups of radius rg and [l1, ..., ld] the d fully-connected layers with the width li(i = 1, ..., d). The parameters we use are (kmax, 0.25, [32, 32, 64]), (kmax/2, 0.5, [64, 64, 128]), (kmax/4, 0.6, [128, 128, 256]) and (kmax/8, 0.7, [256, 256, 512]). We then use interpolation layers to distribute the features of each convolution level among the input points. In this step, we reduce the output of each convolution layer with one shared, fully-connected layer per level, to a size of 64 and then distribute the features to all points of the input point cloud depending on their position. This extends the points of our original point cloud with 256 features. Fig. 11 shows a visual overview of the data flow in our network.
Afterwards, we process the data in r separate branches consisting of two shared, fully interconnected layers with 256 and 128 nodes. The output is then processed with two shared fully-connected layers of 64 and 3 nodes. Finally, we add our resulting data to the input positions that have been repeated r times. This provides an additional skip connection which leads to slightly more stable results. All convolution layers and fully interconnected layers use a tanh() activation function.
For the input feature vector, we make use of additional data fields in conjunction with the point positions. Our network also accepts additional features such as velocity, density and pressure of the SPH simulations used for data generation. For inputs from other sources, those values could be easily computed with suitable SPH interpolation kernels. In practice, we use position, velocity and pressure fields. Whereas the first two are important (as mentioned in Sec. 3.1), the pressure fields turned out to have negligible influence.

heading: C FREQUENCY EVALUATION OF LATENT SPACE

text: In this section we give details for the frequency evaluation of Sec. 4. In order to measure the stability of the latent space against temporal changes, we evaluated the latent space of our network with and without temporal loss, once for 100 ordered patch sequences and once for 100 un-ordered ones. The central latent space of our network consists of the features generated by the point-convolution layers in the first part of the network and is 256 dimensional (see Fig. 11). To obtain information about its general behavior, we average the latent space components over all 100 patch sequences, subtract the mean, and normalize the resulting vector w.r.t. maximum value for each data set. The result is a time sequence of scalar values representing the mean deviations of the latent space. The Fourier transform of these vectors f̃ , are shown in Fig. 7, and were used to compute the weighted frequency content∫ x x · f̃(x)dx. Here, large values indicate strong temporal changes of the latent space dimensions. The resulting values are given in the main document, and highlight the stability of the latent space learned by our method.

heading: D TRAINING DATA AND GRAPHS

text:
Two examples with ground truth points and down-sampled input versions are shown in Fig. 10.
Additionally, Fig. 12 shows loss graphs for the different versions shown in the main text: 2D previous work, our full algorithm in 2D, as well as both cases for 3D. The mingling loss LM is only shown as reference for the previous work versions, but indicates the strong halo-like patterns forming for the architectures based on previous work.
pointNet
MLP
interpolate
kmax points (pos, vel, dens, press)
interpolate interpolate interpolate
64 64 64 64 kmax times
input 256 features 2 5 6 1 2 8×
128
2 5 6 1 2 8×
128
2 5 6 1 2 8×
128
2 5 6 1 2 8×
128
r branches
n max times
PointNet PointNet PointNet
MLP MLP MLP
kmax 2 )(
kmax 4 )(
kmax 8 )(
kmax 16 )(
kmax 4 )( kmax 2 )(
kmax 8 )( , 64 )(
, 64 , 128 , 256
, 512
, 64, 64, 64
6 4 3× output
input position
+
kmax 16
} kmax
128
Page 1
Page 2
Page 1
2 26 50 74 98 12 2 14 6 17 0 19 4 21 8 24 2 26 6 29 0 31 4 33 8 36 2 38 6 41 0 43 4 45 8 48 2 50 6 53 0 55 4 57 8
0
0.05
0.1
0.15
0.2
Page 1
Page 2
2 24 46 68 90 11 2 13 4 15 6 17 8 20 0 22 2 24 4 26 6 28 8 31 0 33 2 35 4 37 6 39 8 42 0 44 2 46 4 48 6 50 8 53 0 55 2 57 4
0
1
2
3
4
5
Page 2
Page 1
2 26 50 74 98 12 2 14 6 17 0 19 4 21 8 24 2 26 6 29 0 31 4 33 8 36 2 38 6 41 0 43 4 45 8 48 2 50 6 53 0 55 4 57 8
0
5
10
15
20
25
Page 1
Sheet2
0.02
0.04
0.06
0.08
0.1
0.12
0.14
Full Loss
Figure 1 : Convergence plots for the training runs of our different 2D and 3D versions. The combined loss only illustrates convergence behavior for each method separately, as weights and terms differ across the four variants. LM for previous work is not minimized, and only given for reference.
16

title: TRANQUIL CLOUDS: NEURAL NETWORKS FOR LEARNING TEMPORALLY COHERENT FEATURES IN POINT CLOUDS


summaries: The authors consider the alignment problem for multiple datasets with side information via entropic optimal transport (Sinkhorn). The authors formulate it as a transport cost learning in optimal transport framework with constraints giving by side information. Empirically, the authors illustrated the effectiveness of the proposed approach over state-of-the-art on several applications (e.g. single-cell RNA-seq, marriage-matching, MNIST with its perturbed one.

summaries: This paper proposes a new way to learn the optimal transport (OT) cost between two datasets that can utilize subset correspondence information. The main idea is to use a neural network to define a transport cost between two examples instead of using the popular Euclidean distance and have a L2 regularizer that encourages the resulting transport plan to be small if two examples belong to the same subset. The results on synthetic datasets (2D, MNIST image) and real-world datasets (RNA sequence) show that the learned transport cost outperforms the Euclidean distance and Procrustes-based OT.

summaries: The paper presents a gradient-based method for learning the cost function for optimal transport, applied to dataset alignment. The algorithm is based on the Sinkhorn-Knopp iterative procedure for approximating the optimal transport plan. The current paper proves that the Sinkhorn transport plan with a parameterized cost function is infinitely differentiable. The cost function can thus be optimized by unrolling the Sinkhorn iteration for a fixed number of steps and then backpropagating through the iterative updates. The cost function is optimized via a loss function derived from side information, specifically subsets of the two datasets to be aligned that contain elements that should be matched in the optimal transport plan. Experiments on a variety of synthetic and real datasets show that the proposed method is often better than, and almost always competitive with, other modern approaches.

abstractText: Learning to align multiple datasets is an important problem with many applications, and it is especially useful when we need to integrate multiple experiments or correct for confounding. Optimal transport (OT) is a principled approach to align datasets, but a key challenge in applying OT is that we need to specify a transport cost function that accurately captures how the two datasets are related. Reliable cost functions are typically not available and practitioners often resort to using hand-crafted or Euclidean cost even if it may not be appropriate. In this work, we investigate how to learn the cost function using a small amount of side information which is often available. The side information we consider captures subset correspondence—i.e. certain subsets of points in the two data sets are known to be related. For example, we may have some images labeled as cars in both datasets; or we may have a common annotated cell type in single-cell data from two batches. We develop an end-to-end optimizer (OT-SI) that differentiates through the Sinkhorn algorithm and effectively learns the suitable cost function from side information. On systematic experiments in images, marriage-matching and single-cell RNA-seq, our method substantially outperform state-of-the-art benchmarks.

authors:
Jason Altschuler
Jonathan Weed
Philippe Rigollet

title: Near-linear time approximation algorithms for optimal transport via sinkhorn iteration

venue: In Advances in Neural Information Processing Systems,

authors:
David Alvarez-Melis
Tommi Jaakkola
Stefanie Jegelka

title: Structured optimal transport

venue: In International Conference on Artificial Intelligence and Statistics,

authors:
David Alvarez-Melis
Stefanie Jegelka
Tommi S Jaakkola

title: Towards optimal transport with global invariances

venue: In The 22nd International Conference on Artificial Intelligence and Statistics,

authors:
Matthew Amodio
Smita Krishnaswamy. Magan

title: Aligning biological manifolds

venue: In International Conference on Machine Learning,

authors:
Mikhail Bilenko
Sugato Basu
Raymond J Mooney

title: Integrating constraints and metric learning in semi-supervised clustering

venue: In Proceedings of the twenty-first international conference on Machine learning,

authors:
Andrew Butler
Paul Hoffman
Peter Smibert
Efthymia Papalexi
Rahul Satija

title: Integrating single-cell transcriptomic data across different conditions, technologies, and species

venue: Nature biotechnology,

authors:
Chao Chen
Kay Grennan
Judith Badner
Dandan Zhang
Elliot Gershon
Li Jin
Chunyu Liu

title: Removing batch effects in analysis of expression microarray data: an evaluation of six batch adjustment methods

venue: PloS one,

authors:
Yunjey Choi
Minje Choi
Munyoung Kim
Jung-Woo Ha
Sunghun Kim
Jaegul Choo

title: Stargan: Unified generative adversarial networks for multi-domain image-to-image translation

venue: In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition,

authors:
Roberto Cominetti
Jaime San Martín

title: Asymptotic analysis of the exponential penalty trajectory in linear programming

venue: Mathematical Programming,

authors:
Nicolas Courty
Rémi Flamary
Mélanie Ducoffe

title: Learning wasserstein embeddings

venue: arXiv preprint arXiv:1710.07457,

authors:
Nicolas Courty
Rémi Flamary
Devis Tuia
Alain Rakotomamonjy

title: Optimal transport for domain adaptation

venue: IEEE transactions on pattern analysis and machine intelligence,

authors:
Paolo Cremonesi
Yehuda Koren
Roberto Turrin

title: Performance of recommender algorithms on top-n recommendation tasks

venue: In Proceedings of the fourth ACM conference on Recommender systems,

title: Sinkhorn distances: Lightspeed computation of optimal transport

venue: In Advances in neural information processing systems,

authors:
Marco Cuturi
David Avis

title: Ground metric learning

venue: The Journal of Machine Learning Research,

authors:
Joel Franklin
Jens Lorenz

title: On the scaling of multidimensional matrices

venue: Linear Algebra and its applications,

authors:
Alfred Galichon
Bernard Salanié

title: Matching with trade-offs: Revealed preferences over competing characteristics

authors:
Aude Genevay
Gabriel Peyré
Marco Cuturi

title: Learning generative models with sinkhorn divergences

venue: arXiv preprint arXiv:1706.00292,

authors:
Pascal Germain
Amaury Habrard
François Laviolette
Emilie Morvant

title: A pac-bayesian approach for domain adaptation with specialization to linear classifiers

venue: In International conference on machine learning,

authors:
Edouard Grave
Armand Joulin
Quentin Berthet

title: Unsupervised alignment of embeddings with wasserstein procrustes

authors:
Laleh Haghverdi
Aaron TL Lun
Michael D Morgan
John C Marioni

title: Batch effects in single-cell rna-sequencing data are corrected by matching mutual nearest neighbors

authors:
Jihun Ham
Daniel D Lee
Lawrence K Saul

title: Semisupervised alignment of manifolds

venue: In AISTATS, pp

authors:
Yehuda Koren
Robert Bell
Chris Volinsky

title: Matrix factorization techniques for recommender systems

authors:
Ruilin Li
Xiaojing Ye
Haomin Zhou
Hongyuan Zha

title: Learning to match via inverse optimal transport

venue: Journal of Machine Learning Research,

authors:
Mingsheng Long
Han Zhu
Jianmin Wang
Michael I Jordan

title: Deep transfer learning with joint adaptation networks

venue: In Proceedings of the 34th International Conference on Machine Learning-Volume

authors:
Giulia Luise
Alessandro Rudi
Massimiliano Pontil
Carlo Ciliberto

title: Differential properties of sinkhorn approximation for learning with wasserstein distance

authors:
Andriy Mnih
Ruslan R Salakhutdinov

title: Probabilistic matrix factorization

authors:
Adam Paszke
Sam Gross
Soumith Chintala
Gregory Chanan
Edward Yang
Zachary DeVito
Zeming Lin
Alban Desmaison
Luca Antiga
Adam Lerer

title: Automatic differentiation in pytorch

authors:
Anand Rangarajan
Haili Chui
Fred L Bookstein

title: The softassign procrustes matching algorithm

venue: In Biennial International Conference on Information Processing in Medical Imaging,

title: Factorization machines with libfm

venue: ACM Transactions on Intelligent Systems and Technology (TIST),

title: Optimal Transport for Applied Mathematicians: Calculus of Variations, PDEs, and Modeling, volume 87

authors:
Nicholas Schaum
Jim Karkanias
Norma F Neff
Andrew P May
Stephen R Quake
Tony WyssCoray
Spyros Darmanis
Joshua Batson
Olga Botvinnik
Michelle B Chen

title: Single-cell transcriptomics of 20 mouse organs creates a tabula muris: The tabula muris

venue: consortium. Nature,

authors:
Vivien Seguy
Bharath Bhushan Damodaran
Rémi Flamary
Nicolas Courty
Antoine Rolet
Mathieu Blondel

title: Large-scale optimal transport and mapping estimation

venue: arXiv preprint arXiv:1711.02283,

authors:
Richard Sinkhorn
Paul Knopp

title: Concerning nonnegative matrices and doubly stochastic matrices

venue: Pacific Journal of Mathematics,

authors:
Matthew Staib
Sebastian Claici
Justin M Solomon
Stefanie Jegelka

title: Parallel streaming wasserstein barycenters

authors:
Marlon Stoeckius
Christoph Hafemeister
William Stephenson
Brian Houck-Loomis
Pratip K Chattopadhyay
Harold Swerdlow
Rahul Satija
Peter Smibert

title: Simultaneous epitope and transcriptome measurement in single cells

venue: Nature methods,

authors:
Vayer Titouan
Nicolas Courty
Romain Tavenard
Chapel Laetitia
Rémi Flamary

title: Optimal transport for structured data with application on graphs

title: Optimal transport: old and new, volume 338

venue: Springer Science & Business Media,

authors:
Chang Wang
Sridhar Mahadevan

title: Manifold alignment using procrustes analysis

venue: In Proceedings of the 25th international conference on Machine learning - ICML

title: Manifold alignment without correspondence

venue: In IJCAI,

authors:
Eric P Xing
Michael I Jordan
Stuart J Russell
Andrew Y Ng

title: Distance metric learning with application to clustering with side-information

authors:
Amir R Zamir
Alexander Sax
William Shen
Leonidas J Guibas
Jitendra Malik
Silvio Savarese

title: Taskonomy: Disentangling task transfer learning

authors:
Jun-Yan Zhu
Taesung Park
Phillip Isola
Alexei A Efros

title: Unpaired image-to-image translation using cycle-consistent adversarial networks

venue: In Proceedings of the IEEE international conference on computer vision,

heading: 1 INTRODUCTION

text:
In many applications, we have multiple related datasets from different sources or domains, and learning efficient computational mappings between these datasets is an important problem (Long et al., 2017; Zamir et al., 2018). For example, we might have single-cell RNA-Seq datasets generated for the same tissue type from two different labs. Since data come from the same type of tissue, we would like to map cells between the two datasets to merge them, so that we could analyze them jointly. However, there are often complex nonlinear batch artifacts generated by the different labs. Moreover the cells are not paired—for each cell measured in the first lab, there is not an identical clone in the second lab. How to integrate or align these two datasets is therefore a challenging problem.
Optimal transport (OT) is a principled analytical framework to align heterogeneous datasets (Santambrogio, 2015). It has been increasingly applied to problems in domain adaptation and transfer learning (Seguy et al., 2017; Genevay et al., 2017; Courty et al., 2017b; Li et al., 2019). Optimal transport is an approach for taking two datasets, and computing a mapping between them in the form of a "transport plan" γ. The mapping is optimal in the sense that among all reasonable mappings (precisely defined in Section 2), it minimizes the cost of aligning the two datasets. The transport cost is given by the user and encodes expert knowledge about how datasets relate to each other. For example, if the expert believes that one data Y is essentially data X with added Gaussian noise, then Euclidean cost could be natural. If the cost is correctly specified, then there are powerful methods for finding the global optimal transport (Villani, 2008). A major challenge in practice, e.g. for single-cell RNA-seq, is that we and experts do not know what cost is appropriate. Users often resort to using Euclidean or other hand-crafted cost functions, which could give misleading mappings.
Our contributions. We propose a novel approach to automatically learn good transport costs by leveraging side information we may have about the data. In particular, aligning datasets across different conditions is a very important problem in biology, especially in single cell analysis (Butler et al., 2018), where we often have annotations that certain cluster of cells in a dataset corresponds to one cell type (e.g. B cells and T cells) based on known marker genes (Schaum et al., 2018). Then we can deduce that T cells from dataset X should be at least mapped to T cells from dataset Y . We only need T cells to be crudely annotated in both datasets, which is reasonable; we don’t need to know that a particular T cell should be mapped to another specific T cell. This gives rises to the side information that we can leverage in our algorithm — a certain subset of points in dataset X should be mapped to another subset of points in dataset Y .
We present the first algorithm, OT-SI, to leverage subset correspondence as a general form of side information. In contrast, previous works mainly focus on pair matching problems (Li et al., 2019; Galichon & Salanié, 2010) — the extreme case of subset correspondence when the subset sizes are 1. In practice, exact one-to-one matching labels are often expensive to obtain or even intractable. OT-SI is an end-to-end framework that learns the transport cost. The intuition is to optimize over a parametrized family of transport costs, to identify the cost under which the annotated subsets are naturally mapped to each other via optimal transport. OT-SI efficiently leverages even a small amount of side information and it generalizes well to new, unannotated data. The learned transport cost is also interpretable. We demonstrate in extensive experiments across image, single-cell, marriage and synthetic datasets that our method OT-SI substantially outperform state-of-the-art methods for mapping datasets.
Related Work Optimal transport been well studied in the mathematics, statistics and optimization literature (Villani, 2008; Courty et al., 2017a; Li et al., 2019; Courty et al., 2017b). OT can be used to define a distance metric between distributions (e.g. Wasserstein distance) or to produce an explict mapping across two datasets. The latter is the focus of our paper. In machine learning, there has been significant work on developing fast algorithms for efficient computation of the optimal transport plan (Cuturi, 2013; Altschuler et al., 2017; Staib et al., 2017), and analyzing the properties of the transport plan under various structures and constraints on the optimization problem (Alvarez-Melis et al., 2018; Titouan et al., 2019). The previous work on learning the transport cost is done on a very different setting from ours – learning feature histogram distances between many pairs of datapoints (Cuturi & Avis, 2014). Some classical clustering (Xing et al., 2003; Bilenko et al., 2004) and alignment methods (Ham et al., 2005; Wang & Mahadevan, 2008; 2009) have realized benefits by including side information, but these semi-supervised approaches differ from our explicit parametrization and optimization of the transport cost function.
Separately, there have been recent efforts to directly map between datasets, without learning a transport cost. The standard alignment methods can be divided into two categories: GANs-based (Zhu et al., 2017; Choi et al., 2018) and OT-based (Grave et al., 2019; Alvarez-Melis et al., 2019). GAN-based approaches have been used to align single-cell RNA-seq data when pairs of cells are known to be related (Amodio & Krishnaswamy, 2018). However the exact pairing of individual cells is always not readily available or even intractable. To address this issue, our method OT-SI allows for more general correspondence between subsets, i.e., clusters, cell types and also individual cells. In the meantime, the OT-based methods always rely on Procrustes analysis (Rangarajan et al., 1997) — a linear transformation between the datasets is assumed, which lacks the flexibility to handle nonlinear artifacts and the side information cannot be utilized. In contrast, a major benefit of our approach is its graceful adaptation to partial subset correspondence information, where we frame the problem as semi-supervised.

heading: 2 LEARNING COST METRICS

text: A good choice of the cost function for optimal transport is the key to a successful mapping between two datasets. In this section, we present the algorithm OT-SI, which parametrizes the cost function with weight θ and adaptively learns θ using side information about the training data. The side information we consider is subset correspondence — a common situation when some subsets of training points are known to be related; pair matching is included as an extreme case. The learned cost function is further evaluated on the unseen test data to prove generalizability.

heading: 2.1 OPTIMAL TRANSPORT

text:
Consider learning a mapping between two datasets X = {x(1), ..., x(nX)} and Y = {y(1), ..., y(nY )}. Here we use nX and nY to denote the number of datapoints; each sample x(i) or y(j) could be a vector as well. We briefly recall the optimal transport framework in this setting. Given probability vectors µX and µY , the transport polytope is defined as
U(µX , µY ) := {γ ∈ RnX×nY+ |γ1nX = µX , γT1nY = µY },
where 1nX (1nY ) is the nX (nY ) dimensional vector of ones. Here the probability vector µX (µY ) is in the simplex ∑ n := {p ∈ Rn+ : pT1n = 1} for n = nX (nY ). For two random variables with distribution µX and µY , the transport polytope U(µX , µY ) represents the set of all possible joint probabilities of the two variables. In this paper, we consider µX and µY to represent the empirical distributions of the samples X and Y, respectively, and set µ = (1/n)1n.
Given a nX×nY cost matrix C, the classical optimal transport plan between µX and µY is defined as γ∗ = arg min γ∈U(µX ,µY )〈γ,C〉, where 〈·, ·〉 denotes the Frobenius inner product. γ
∗ is also called a coupling. Despite its intuitive formulation, the computation of this linear program quickly becomes prohibitive especially in the common situation when nX and nY , the sizes of the datasets, exceed a few hundred. For computational efficiency, Sinkhorn-Knopp iteration is widely used to compute the optimal transport (Cuturi, 2013). Sinkhorn-Knopp is a fast iterative algorithm for approximately solving the optimization problem with entropy regularization Santambrogio (2015):
γλ = arg min γ∈U(µX ,µY ) 〈γ,C〉 − 1 λ h(γ). (1)
where λ > 0 is a regularization parameter and h(γ) = − ∑nX i=1 ∑nY j=1 γij log γij denotes the entropy. The regularized solution γλ converges to the classical one γ∗ when the regularization diminishes, i.e., λ→∞, with exponential convergence rate (Cominetti & San Martín, 1994). The transport γλ treats X and Y symmetrically.
The cost matrix is retrieved from the cost function Cij = c(x(i), y(j)). A good choice of the cost function is the key to influencing the learned mapping γλ. However, reliable cost functions are typically not available and Euclidean cost is mostly used. In this paper, the representation of the cost function is adaptively learned using side information about the data.

heading: 2.2 SIDE INFORMATION

text:
Subset correspondence describes a common situation when certain subsets of points are known to be related. For example, images with the same objects should always be mapped together in domain adaptation tasks (Courty et al., 2017b), while cells in a single-cell dataset need to be aligned to those with the same cell type, where the cell type annotation is available.
Given m corresponding subsets, we write S(k)X and S (k) Y , k = 1, ...,m, to denote the sets of data indices in the corresponding subsets, i.e., {x(i) | i ∈ S(k)X } and {y(i) | i ∈ S (k) Y }. Note that S (k) X and S (k) Y could have different probability mass. If |S(k)X | nX ≤ |S (k) Y | nY , we take this side information to be that S (k) X should be mapped into S (k) Y . In other words, all the other entries of the transport matrix γ that maps S(k)X to outside of S (k) Y should be 0. Everything is swapped if |S(k)Y | nY ≤ |S (k) X | nX
. Mathematically, this side information corresponds to the constraint that γij = 0 for (i, j) ∈ S0, where
S0 = m⋃ k=1 {(i, j)|i ∈ S(k)X , j /∈ S (k) Y , |S(k)X | nX ≤ |S(k)Y | nY }∪{(i, j)|i /∈ S(k)X , j ∈ S (k) Y , |S(k)Y | nY ≤ |S(k)X | nX }
(2)
Note that pair matching is an extreme case of subset correspondence with subset size 1 — that is, the exact pairwise relation is known. The pair matching problem has been addressed in literature (Li et al., 2019; Galichon & Salanié, 2010). However, in practice, exact one-to-one matching labels are often expensive to obtain or even intractable. In this paper, we show that subset correspondence, as a small amount of side information, can significantly aid cost function learning.
We investigate how to learn the cost function using subset correspondence information between training datasets Xtrain and Ytrain. The learned cost function is evaluated via the mapping quality on the test datasets Xtest and Ytest. Note that the training and test sets are not necessarily under the same distribution. We demonstrate the power of OT-SI in generalizing to new subsets that were not seen in the training process.

heading: 2.3 THE OT-SI ALGORITHM

text:
Our ultimate goal is to learn a cost function c(·), such that the computed optimal transport γλtrain satisfies the side information given in Eq. (2) as faithfully as possible.
Cost function parametrization. When the cost function c(·) is Euclidean, the entry of the cost matrix is computed asCij = c(x(i), y(j)) = ∑d k=1(x (i) k −y (j) k )
2, where d is the data dimension. To learn the cost function systematically, we parametrize it as Cij(θ) = c(x(i), y(j); θ) with weight θ. Here the function form c(·) can be chosen by users. To illustrate the improvement over the commonly-used Euclidean cost, we parametrize c(x(i), y(j); θ) as a polynomial in (x(i)1 , ..., x (i) dx , y (j) 1 , ..., y (j) dy
) with coefficients θ and degree 2 for low-dimensional data. The Euclidean cost is equivalent to a specific choice of θ0 which is set as the initialization; see Appendix for more discussions. For high-dimensional data, the memory required to store the second order polynomials becomes too large and we use a fully connected neural network to parametrize c(x(i), y(j); θ) with input (x(i)1 , ..., x (i) dx , y (j) 1 , ..., y (j) dy
) and weights θ. Throughout this paper, the polynomial parametrization is used if not specified.
The optimal transport solution is characterized by θ as
γλ(θ) = arg min γ∈U(µX ,µY ) 〈γ,C(θ)〉 − 1 λ h(γ). (3)
Then the problem can be formulated as optimizing θ to make the transport γλ(θ) approximately satisfy the conditions defined in Eq. (2), penalizing deviation of the solution from these constraints with the loss
Lλ(θ) = ∑
(i,j)∈S0,train
||γλij,train(θ)||22. (4)
Theorem 1. For any λ > 0, the optimal transport plan γλ(θ): Rdθ → U(µX , µY ) is C∞ in the interior of its domain.
The infinite differentiability of the Sinkhorn distance is previously-known (Luise et al., 2018); Thm. 1 proves that the Sinkhorn transport plan also has this desirable property. Because Thm. 1 guarantees that γλ(θ) is infinitely differentiable, we are able to optimize Lλ(θ) in Eq. (4) by gradient descent. In practice, we iterate Sinkhorn’s update a sufficient number of times to converge to γλ. Each iteration is a matrix operation involving the cost matrix C(θ), and when the number of iterations is fixed, we can propagate the gradient∇θ through all of the iterations using the chain rule. Updating θ by one forward and backward pass has complexity of O(n2) up to logarithmic terms. Hence OT-SI has the same complexity as the Procrustes-based OT methods which alternatively optimize over coupling matrix and linear transformation matrix (Grave et al., 2019; Alvarez-Melis et al., 2019). To further boost the performance, we propose to use a mimic learning method for initialization, which does not need to propagate the gradient. The pseudocode for OT-SI is in Algorithm 1 and details are in the Appendix. The proof of Thm. 1, mimic learning algorithm, details, and discussions about convergence are also in the Appendix.

heading: 2.4 EXPERIMENTS SETUP

text:
The OT-SI algorithm is carried out in Pytorch (Paszke et al., 2017) and trained with GPU. The model is fitted on training set, and evaluated on test set. We use validation set for hyperparameter selection and early stopping. We evaluate OT-SI with different types of data and correspondence information.
Comparison methods. We use optimal transport with Euclidean cost function as a baseline for comparison, referred as "OT-baseline". We also compare our result with state-of-the-art GAN-based data alignment methods, MAGAN (Amodio & Krishnaswamy, 2018) and CycleGAN (Zhu et al.,
Algorithm 1 OT-SI
Require: training datasets Xtrain and Ytrain, corresponding subsets index S (k) Xtrain and S(k)Ytrain (k = 1, ...,mtrain), step size α, total training steps T , weights θ0 from initialization procedure, Sinkhorn regularization parameter λ and number of Sinkhorn iterations nSinkhorn.
1: nX = length(Xtrain), nY = length(Ytrain), µX = (1/nX)1nX , µY = (1/nY )1nY 2: Initialize θ = θ0 3: for t = 1 to T do 4: Compute cost matrix C(θ) with entries Cij(θ) = c(x(i), y(j); θ). 5: Solve γλtrain(θ) = arg min γ∈U(µX ,µY )〈γ,C(θ)〉 − 1 λh(γ) with Sinkhorn’s update with
nSinkhorn iterations. 6: Derive ∇θγλtrain(θ) by backpropagating the gradient through all Sinkhorn-Knopp iterations. 7: Update weights θ := θ − α ∑ (i,j)∈S0,train γ λ ij,train(θ)∇θγλij,train(θ). 8: end for
2017), as well as the OT-based methods, RIOT (Li et al., 2019) which is developed for specific pair matching applications and the Procrustes-based OT (Grave et al., 2019), referred as "OT-Procrustes". For MAGAN and CycleGAN, the matching point for a source sample is set as its nearest neighbor in the target after mapping. Among the five comparison methods, OT-baseline, OT-Procrustes and CycleGAN do not use any side information; MAGAN makes use of matching pairs; RIOT requires the one-to-one matching labels for all the datapoints. Because MAGAN and RIOT requires pairwise correspondence, they are not applied in some experiments and these are marked as N/A. We use the same settings and hyperparameters for the comparison methods as in their original implementations.
Evaluation metrics. When the subset correspondence is known on the test set (not shown to the algorithm), we evaluate a transport plan γ by how much it satisfies the correspondence. Mathematically, we define subset matching accuracy:
Accuracy =
∑m k=1 ∑ i∈S(k)Xtest ∑ j∈S(k)Ytest
γij,test∑m k=1 min{|S (k) Xtest |/nXtest , |S (k) Ytest |/nYtest} .
From the definition, 0 6 Accuracy 6 1 gives the probability of mapping to the correct corresponding subsets. When all the test datapoints are mapped into the correct subsets, the accuracy is 1; when all the data are matched to the wrong subsets, accuracy is 0. As an extreme example, pair matching is equivalent to subset correspondence with subset sizes 1, referred to as pair matching accuracy. In the next few sections, we thoroughly evaluate OT-SI and several state-of-the-art methods in extensive and diverse experiments—aligning single-cell RNA-seq data to correct for batch effects, aligning single-cell gene expression and protein abundances, a marriage data, an image dataset, and the synthetic twin-moon data for illustration.

heading: 3 BENCHMARK ON SYNTHETIC DATASETS

text:
We first experiment with the benchmark toy example for domain adaptation — two moon datasets—to illustrate the challenges of data alignment, before we move onto complex real-world data (Germain et al., 2013; Courty et al., 2017b). The dataset is simulated with two domains, source and target. As shown in Fig. 1a, each domain contains two standard entangled moons. The two moons are associated with two different classes, denoted by circle and crossing respectively. The target (colored in orange), is built by adding noise to the source (colored in orange) and rotating by 60 degrees. In the experiment, we generate the training, test and validation datasets with 100, 100, and 50 samples of each moon. We set the parameter λ = 103 and the number of Sinkhorn-Knopp iterations N = 200. The algorithm is run for 100 epochs with step size 1. There are two types of side information available for OT tasks: (i) subset correspondence — datapoints are known to be mapped into the corresponding moon class; (ii) pair matching — known matched datapoints after rotation. The result is averaged over 10 (50) independent runs when the side information is subset correspondence (pair matching).
Baseline performance. The optimal transport plan under Euclidean cost function is depicted in Fig. 1b. Datapoints learned to be matched are connected by solid curves. The red curves indicate wrong
transports which map the data into the wrong subset, i.e., the other moon. As shown by Fig. 1b, most wrong transports are between the points at the edges of the moons. In the euclidean space, the edge of one moon becomes “closer" to the other moon after rotation, which leads to a small cost between datapoints in different moon classes. A new cost function which captures the rotation property is expected. Quantitatively, only 72% of the data are mapped into the correct moon it belongs to and only 2% of the data are matched to their corresponding pairs, as given in Table 1.
Subset correspondence. We first evaluate our method OT-SI when the side information is only that the data on the corresponding moons are known to be related—i.e. S(1)X and S (2) X are the two moons in dataset one, and S(1)Y and S (2) Y are the two moons in dataset two. With the learned cost function, almost all the datapoints are mapped into the corresponding moon, i.e., the subset matching accuracy on the test data achieves 100% for both methods, as shown in Table 1. The results are averaged for 10 independent runs. Interestingly, although there is no pair matching information provided during training, the learned cost function significantly improves the matching performance. As shown in Table 1, 87% datapoints are transported into their exact matching points after rotation. The learned mapping of OT-SI is depicted in Fig. 1c. The rotation property of the datasets is correctly captured. In contrast, OT-Procrustes sometimes learns as good as OT-SI, similar to Fig. 1c, but sometimes mistakenly learns rotation as flipping, indicated by Fig. 1d. This results in overall worse accuracy and large variance for OT-Procrustes.
Pair matching. OT-SI demonstrates substantial improvement when the pair matching information is provided—only 1 and 10 pairs are known out of the total 100 training pairs. The matching pairs are randomly selected from the training data. OT-SI significantly outperforms all four bench methods particularly when the number of known pairs is very limited. Even when only 1 matching pair is provided, the learned cost function greatly improves the OT performance, as given in Table 1. The improvement here is largely attributable to the unlabeled data, i.e., the datapoints without any pair
matching information. For comparison, we carry out another experiment with only 3 unlabeled data, with all other settings unchanged. The algorithms are not able to learn the right cost function anymore — the test pair matching accuracy is only 0.5% after learning, even worse than the Euclidean baseline. With 199 unlabeled data, the accuracy achieves 53%. In contrast, the competing methods MAGAN and RIOT learn barely any patterns, because too few labeled datapoints are available and the unlabeled ones are wasted.

heading: 4 BIOLOGICAL MANIFOLD ALIGNMENT

text: In this section, we implement our method OT-SI to learn a cost function that aligns biological manifolds with partial supervision — annotations of some cell types or clusters, which is the common situation in biological studies. The pair matching methods, MAGAN and RIOT, are not compared here because the the cell-to-cell matching information is not available. Similar to Sec. 3, the CycleGAN does not learn correctly for these data types and is not presented here. We show that OT-SI has substantial improvement for aligning datasets with different data types and aligning data from different batches.

heading: 4.1 ALIGNMENT OF PROTEIN AND RNA SEQUENCING DATA

text:
How to align datasets with different data types has been a major topic in many fields. For example, in single-cell studies, RNA and protein sequencing can both be done at cellular resolution. How to map between those two types of data, i.e., map cells with certain mRNA level to cells with certain protein level, becomes critical for downstream studies such as RNA and protein co-expression analysis.
We demonstrate the power of OT-SI in learning cost function that aligns two different data types — RNA and protein expression in the CITE-seq cord blood mononuclear cells (CBMCs) experiments (Stoeckius et al., 2017). The dataset is subset to 8,005 human cells with the expression of 5,001 highly variable genes and 13 proteins. Fifteen clusters are identified using the Louvain modularity clustering. The CITE-seq technology has enabled the simultaneous measurement of RNA and protein expression at single-cell level, hence the ground truth about the cell pairing is available. To emulate the common situation, we only use the information of cluster correspondence in the training and report the performance of both subset (cluster) matching and pair matching for test. We randomly sampled 500 and 500 cells for validation and test purpose. When OT-SI is learned in the original data space with the expression of 5001 mRNA and 13 proteins for each cell, we use a fully connected neural network to parametrize the cost function, with two hidden layers of 100 and 5 neurons.
We align the RNA and protein expression datasets in two scenarios: i) the embedding space where we use the first 10 PCs for both datasets; ii) the original expression space. Table 2 shows that OT-SI substantially outperforms other methods. OT-SI is able to learn good cost function when the dimensions of the two datasets are highly unbalanced as 5001:13. The learned cost metrics in the original expression space can be used for future biological analysis on the effect and relation between RNA and protein expressions. Although the single-cell sequencing data are noisy and the algorithm is not informed of any matching pair during training, the accuracy of test pair matching is improved.

heading: 4.2 BATCH ALIGNMENT

text:
Another fascinating biological application of optimal transport is to align data from different batches. In biological studies, the samples processed or measured in different batches usually result in nonbiological variations, known as batch effect (Chen et al., 2011; Haghverdi et al., 2018). Here we use OT to align two batches of data1 — data collected with fluorescence activated cell sorting (FACS) and droplet methods. In this case, the cell types are used as the subset correspondence information. For illustration purposes, we subsample the top 10 celltypes with 400 samples of each. There are 1,682 genes after filtering and the first 10 principal components are used for analysis. The dataset is split into training, validation and test sets with ratio 50%, 20% and 30%. The experiment setting is the same as in Sec. 3.
We demonstrate the power of our learned metric in generalizing to entirely new cell types that were not used to train the cost. This is a hard task (a zero-shot learning task), and is more realistic because in most settings we only have partial annotations for cell types and we would like the mapping to generalize to all of the data. To do this, we choose two cell types — T cell and immature T cell — as held-out and train on the rest. Among all the ground truth expert annotations, T cell and immature T cell are most difficult to be aligned. With OT-baseline, only 70% cells are mapped into the correct cell types. Substantial improvement is achieved, as shown in Table 3. From a small number of annotated cells types, OT-SI is able to learn a transport cost that captures the batch artifacts between FACS and droplet which generalized to mapping these two new cell types. We anticipate future uses of our formulation to further investigate the cost function, particularly in biological discovery applications like isolating genes that mark single-cell heterogeneity.

heading: 5 EXPERIMENTS ON IMAGES AND TABULAR DATA ALIGNMENT

heading: 5.1 IMAGE ALIGNMENT FOR NEW DIGITS

text:
To illustrate the use of OT-SI in image alignment, we use it to learn cost metrics in aligning images with partial annotation on the MNIST dataset which contains 28 × 28 images of handwritten digits. We subsample 200 images from each digit class and split them into training, validation, and test sets with ratio 50%, 20% and 30%. For illustration purposes, we use the first ten principal components for the alignment analysis. We generate four different types of perturbations to the original images, as plotted in Table 4. Then we align the original images with the perturbed ones using OT, respectively.
1https://github.com/czbiohub/tabula-muris-vignettes/blob/master/data
The experiment setting is the same as in Sec. 3. Similar to Sec. 4, we test how well our algorithm generalizes to new classes that were not used in learning the cost function. In the experiment, we hold out digits 3 and 5, and demonstrate the metric learned on the other eight digits can help the alignment of digits 3 and 5. We achieve consistent improvement over the baseline on all four distribution types, as indicated by Table 4.

heading: 5.2 MARRIAGE DATASET FOR PAIR MATCHING

text:
While OT-SI is designed for the more general form of side information — subset correspondence, it can also be used for pair matching purpose. We finally benchmark it from the comparison with other state-of-the-art pair matching methods, including RIOT, factorization machine model (FM) (Rendle, 2012), probabilistic matrix factorization model (PMF) (Mnih & Salakhutdinov, 2008), item-based collaborative filtering model (itemKNN) (Cremonesi et al., 2010), classical SVD model (Koren et al., 2009) and baseline random predictor model. These methods were also used as comparisons in Li et al. (2019). We follow the same experimental protocol as in Li et al. (2019) for the Dutch Household Survey (DHS) dataset. The exact matching matrix between 50 datapoints with 11 features are known. We note that the coupling matrix is treated as continuous and OT-Procrustes is not applicable.
The performance is evaluated by the root mean square error (RMSE) and the mean absolute error (MAE) for the predicted matching matrix, as given in Table 5. When used for pair matching purposes, OT-SI report comparable performance to state-of-the-art matching algorithms. Note that this marriage dataset was a primary motivating dataset used to design RIOT (Li et al. (2019)), and therefore we expect RIOT to perform very well for this task.

heading: 6 DISCUSSION

text:
In this paper, we study the problem of learning the transport cost using side information in the form of a small number of corresponding subsets between the two datasets. This is a new problem formulation, to the best of our knowledge. Previous works rely on more restricted information such as that specific pairs of points should be aligned. In settings such as genomics and images, it is often difficult to say that a single point in dataset one should be mapped onto a particular point in dataset two. It is more common to have partial annotation of subsets of points—e.g. T cells are annotated in two single-cell RNA-seq datasets—which motivates our generalization.
We propose a flexible and principled method to learn the transport cost with side information. Experiments demonstrate that they work significantly better than state-of-the-art methods when the side-information is very limited, which is often the case. We compare against state-of-the-art methods for the special case when the side information consists of matching pairs, since we are not aware of other published OT methods that deal with the more general subset correspondence. One interesting reason for the improved performance is that by learning the transport cost directly, our algorithms are more efficiently using all of the unannotated datapoints that are not in any pairs or subsets. These unannotated data act as regularization (similar to in semi-supervised 1learning), which enables the model to avoid overfitting to the limited side information. An interesting direction of future work is to interpret the learned cost function for insights on how the datasets differ.

heading: A COST FUNCTION

text:
In Sec. 2.3, the cost function is parametrized as Cij(θ) = c(x(i), y(j); θ) with weight θ. For lowdimensional data, we choose c(x(i), y(j); θ) as a polynomial in (x(i)1 , ..., x (i) dx , y (j) 1 , ..., y (j) dy
) with coefficients θ and degree 2:
Cij(θ) = ∑
α,β∈Zd, α,β>0 161Td α+1 T d β62
θα,β dx∏ k=1 dy∏ l=1 (x (i) k ) αk(y (j) l ) βl . (5)
Here we do not require the dimension of the two datasets dx and dy to be the same. When the two datasets X and Y have different features, the learned weights indicate the coupling between different features in the data mapping. In general, the function form c(·) can be chosen by users. We have also investigated parametrizing C(θ) as a small fully connected neural network and achieved very similar performance.
The Euclidean cost is equivalent to a specific choice of θ, when the two datasets have the same feature space dx = dy = d. The entry of the cost matrix is computed as Cij = c(x(i), y(j); θ) =∑d k=1((x (i) k ) 2 + (y (j) k ) 2 − 2x(i)k y (j) k ).

heading: B OT-SI ALGORITHM

text:
The Lagrangian dual of Eq. (3) is
max u,v
µTXu+ µ T Y v −
1
λ nX∑ i=1 nY∑ j=1 e−λ(Cij(θ)−ui−vj) (6)
By Sinkhorn’s scaling theorem (Sinkhorn & Knopp, 1967), the optimal transport plan γλ(θ) is computed as
γλ(θ) = diag(eλu ∗(θ))e−λC(θ)diag(eλv ∗(θ)), (7) where u∗(θ) and v∗(θ) are the solutions to the dual problem in Eq. (6).
B.1 PROOF OF THEOREM 1
The analysis for Theorem 1 follows the strategy of the proofs of Theorem 2 in (Luise et al., 2018).
Theorem 1. For any λ > 0, the optimal transport plan γλ(θ): Rdθ → U(µX , µY ) is C∞ in the interior of its domain.
Proof. Based on Eq. (7), the optimal transport plan γλ(θ) is a smooth function when u∗(θ), v∗(θ) and C(θ) are smooth. In the meantime, the cost C(θ) is a linear function of θ, as indicated by Eq. (5). Thus to prove the smoothness of γλ(θ), we only need to demonstrate u∗(θ) and v∗(θ) are both smooth in θ.
Here we define
σ(θ;u, v) = −µTXu− µTY v + 1
λ nX∑ i=1 nY∑ j=1 e−λ(Cij(θ)−ui−vj)
The dual problem in Eq. (6) becomes minu,v σ(θ;u, v). From the definition, σ(θ;u, v) is smooth and strictly convex in (u, v). Note that C(θ) is linear in θ. Then for any fixed θ in the interior of Rdθ , there exits (u∗(θ), v∗(θ)) such that σ(θ;u∗(θ), v∗(θ)) = minu,v σ(θ;u, v). The function ∇(u,v)σ(θ;u, v) ∈ C∞ due to the smoothness of σ(θ;u, v). Now we fix (u0, v0, x0) such that ∇(u,v)σ(θ0;u0, v0) = 0. The strict convexity of σ(θ;u, v) ensures that ∇2(u,v)σ(θ0;u0, v0) is invertible.
From implicit function theorem, we can always find a function f and a subset Uθ0 ∈ Rdθ such that i) f(θ0) = (u0, v0); ii) ∇(u,v)σ(θ; f(θ)) = 0 for any θ ∈ Uθ0 ; iii) f ∈ C∞(Uθ0). That is, f(θ) is a
stationary point of the function σ for any θ in Uθ0 . Together with the strict convexity of σ, we derive f(θ) = (u∗(θ), v∗(θ)). Recalling f ∈ C∞(Uθ0), we prove that (u∗(θ), v∗(θ)) is C∞ in the interior of its domain.
B.2 CONVERGENCE PROPERTIES
The gradient of γλ(θ) in Eq. (7) is computed as ∇θγλij(θ) = λγλij(θ) ( ∇θu∗i (θ) +∇θv∗j (θ)−∇θCij(θ) ) The convergence of scaling factors eλu(θ) and eλv(θ) is linear (i.e. exponential in nSinkhorn), with the bounded rate given by (Franklin & Lorenz, 1989); our experiments use a linear C(θ), and we do not find Sinkhorn to bottleneck convergence. Besides this point, convergence is determined by the loss landscape of C. Analyzing convergence to the global optimum, and the role of primal regularization of γ, is a relevant open question.
B.3 COMPUTING GRADIENT
In practice, we iterate Sinkhorn’s update N times to converge to γλ. More specifically, the dual solutions u∗(θ) and v∗(θ) for Eq. (6) are computed by iterating for N times
u← µX/e−λC(θ)v
v ← µY /(e−λC(θ))′u For the optimal plan γλ(θ) given in Eq. (7), its gradient∇θγλ(θ) is obtained by propagating through all the iterations using the chain rule and implemented by Pytorch. Thus updating θ by one forward and backward pass has complexity of O(n2) up to logarithmic terms. OT-SI has the same complexity as the Procrustes-based OT methods which alternatively optimize over coupling matrix and linear transformation matrix. Taking the image alignment for watering MNIST dataset as an example, the running time for OT-SI and OT-Procrustes is 229.5s and 196.7s, respectively.

heading: C MIMIC LEARNING FOR INITIALIZATION

text:
In this section, we derive a mimic learning as an fast initialization method to boost the performance and accelerate the learning. While Algorithm 1 requires to differentiate through Sinkhorn updates, the mimic learning approach does not need to propagate the gradient through all the iterations and is applicable for any kind of OT algorithm. Here we take the classical optimal transport plan γ∗ = arg min γ∈U(µX ,µY )〈γ,C〉 as an example. As discussed in Sec. 2.3, our ultimate goal is to learn a cost function, such that the optimal transport γ∗ satisfies the side information defined in Eq. (2) as faithfully as possible. From another perspective, we force an additional set of constraints on the transport plan γ to fulfill the condition in Eq. (2):
U c = m⋂ k=1 {γ|γij = 0, i ∈ S(k)X , j /∈ S (k) Y , |S(k)X | nX ≤ |S(k)Y | nY }
∩ {γ|γij = 0, i /∈ S(k)X , j ∈ S (k) Y , |S(k)Y | nY ≤ |S(k)X | nX }
(8)
To quantify how much the learned γ∗ in Eq. (1) follows the side information, we compare it with
γ̂ = arg min γ∈U(µX ,µY )∩Uc
〈γ,C〉. (9)
Here γ̂ is interpreted as the optimal transport plan when the side information is completely satisfied, and 〈γ̂, C〉 is the smallest transport distance under the constraint. With the cost function parametrized as Cij(θ) = c(x(i), y(j); θ), The optimal transport solution in Eq. (9) is characterized by θ as γ̂(θ) = arg min γ∈U(µ,ν)∩Uc〈γ,C(θ)〉.
Algorithm 2 Mimic Learning for Initialization
Require: training datasets Xtrain and Ytrain, corresponding subsets index S (k) Xtrain and S(k)Ytrain (k = 1, ...,mtrain), step size α, total steps Tinit, optimal transport solver OTSolver.
1: nX = length(Xtrain), nY = length(Ytrain), µX = (1/nX)1nX , µY = (1/nY )1nY 2: Define transport polytopes U∗ = U(µX , µY ) and constraints U c from Eq. (8). 3: Initialize θ such that the cost function is equivalent to the Euclidean cost. 4: for t = 1 to Tinit do 5: Compute cost matrix C as Cij = c(x(i), y(j); θ) 6: γ∗ = OTSolver(C,U∗), γ̂ = OTSolver(C,U(µX , µY ) ∩ U c) 7: Linit(θ) = 〈γ̂,C(θ)〉−〈γ∗,C(θ)〉 〈γ̄,C(θ)〉−〈γ∗,C(θ)〉 8: Update weights θ := θ − α∇θLinit(θ) 9: end for
Then we also expect a good cost function to make the distance under constraint 〈γ̂, C〉 to be as small as the lowest one 〈γ∗, C〉 as possible — optimize θ to minimize the loss
Linit(θ) = 〈γ̂(θ), C(θ)〉 − 〈γ∗(θ), C(θ)〉. (10)
We refer to this method as mimic learning, because its objective is to make the γ̂ mimic the cost performance of γ∗.
Note that γ∗(θ) is the optimal solution for any transport matrix in U(µX , µY ). That is, the optimal distance 〈γ∗(θ), C(θ)〉 ≤ 〈γ,C(θ)〉 for any γ ∈ U(µX , µY ). The equality holds true only when γ = γ∗(θ) for the convex transport problem. In the meantime, we have γ̂(θ) ∈ U(µX , µY )∩U c. Thus the loss Linit(θ) is always larger or equal to 0. When zero loss is achieved, we have γ∗(θ) = γ̂(θ), coinciding with the optimal solution for L(θ) in Eq. (3).
Equation (10) describes the absolute difference between the two transport distances, but a relative difference is more desirable in practice to adjust for the scale of the objective function around 〈γ∗, C〉. For example, scaling the cost matrix C(θ) by a constant does not change the solutions γ∗ and γ̂, but does scale the loss defined in Eq. (10) by the same constant. We modify the loss to be invariant to such scaling:
Linit(θ) = 〈γ̂(θ), C(θ)〉 − 〈γ∗(θ), C(θ)〉 〈γ̄, C(θ)〉 − 〈γ∗(θ), C(θ)〉 , (11)
Here γ̄ is a uniform nX × nY matrix used to stand for the averaged performance of random transport plans. Eq. (11) captures how close the distance under constraint 〈γ∗, C〉 is to the best one, compared to other random transports.
The mimic learning is approximately solved by alternating minimization. As described in Algorithm 2, we iterate over two steps: (i) compute the value of γ∗ and γ̂ while fixing θ; (ii) take one gradient step with respect to Linit(θ) with fixed γ∗ and γ̂. The computation for optimal transport plans and the optimization of θ are carried out in alternating fashion.
The OT solver is used only to estimate the value of γ∗ and γ̂ in the first step, requiring no gradient propagation. Given such estimates of transport mappings γ∗ and γ̂, the second step can be interpreted as learning a cost function which equates their transport costs, i.e., makes the behavior of γ̂ mimic that of γ∗. In the experiments, we set α = 1 and Tinit = 10 for initialization purpose.