summaries: The paper introduces a novel unbalanced Gromov-Wasserstein type problem. The Gromov-Wasserstein distance is very useful in practice for comparing probability distributions that do not lie in the same metric spaces. It has recently found several successful applications in ML for computational chemistry, graphs comparisons or NLP. Following previous works on unbalanced optimal transport (i.e. soft constraints over marginals enforcement of the coupling matrix), and the rationale that disposing of unbalanced versions of transport problems can alleviate in some ways presence of outliers or noise in the distributions, the authors propose two ‘unbalanced’ variants of the Gromov-Wasserstein (GW) problem, that allow comparison of metric spaces with arbitrary positive measures up to isometries (I.e. rigid transformations).

summaries: The authors consider the Gromov Wasserstein (GW) problem for metric measure spaces having different masses (i.e., Unbalanced GW). Similar to the ideas of unbalanced optimal transport (UOT), they proposed to use a quadratic divergence to relax the marginal constraints (instead of divergence as in UOT). When divergence is KL, the authors derive a GPU-friendly algorithm for the UGW  which relies on the unbalanced Sinkhorn algorithm. Additionally, the authors also propose a variant of UGW, namely Conic Gromov-Wasserstein (CGW) to address the different masses of metric measure spaces. The authors propose that CGW has nice properties (Theorem 1). However, there is no algorithm to solve the CGW yet.

summaries: In the paper, the authors propose two versions of Gromov-Wasserstein distance when the weights of measures do not need to sum up to one. The first version, named Unbalanced Gromov-Wasserstein (UGW), is a direct application of unbalanced optimal transport (UOT) to the setting of Gromov-Wassertstein. The second version, named Conic Gromov-Wasserstein (CGW), is an extension of the conic formulation of UOT to the setting of Gromov-Wasserstein. The authors also show that CGW is a distance in the metric-measure spaces and is an lower bound of the UGW. Finally, the authors also provide some experiments with their proposed divergences.

abstractText: Comparing metric measure spaces (i.e. a metric space endowed with a probability distribution) is at the heart of many machine learning problems. The most popular distance between such metric measure spaces is the Gromov-Wasserstein (GW) distance, which is the solution of a quadratic assignment problem. The GW distance is however limited to the comparison of metric measure spaces endowed with a probability distribution. To alleviate this issue, we introduce two Unbalanced Gromov-Wasserstein formulations: a distance and a more tractable upper-bounding relaxation. They both allow the comparison of metric spaces equipped with arbitrary positive measures up to isometries. The first formulation is a positive and definite divergence based on a relaxation of the mass conservation constraint using a novel type of quadratically-homogeneous divergence. This divergence works hand in hand with the entropic regularization approach which is popular to solve large scale optimal transport problems. We show that the underlying non-convex optimization problem can be efficiently tackled using a highly parallelizable and GPU-friendly iterative scheme. The second formulation is a distance between mm-spaces up to isometries based on a conic lifting. Lastly, we provide numerical experiments on synthetic examples and domain adaptation data with a Positive-Unlabeled learning task to highlight the salient features of the unbalanced divergence and its potential applications in ML.

authors:
Alexander M Bronstein
Michael M Bronstein
Ron Kimmel

title: Generalized multidimensional scaling: a framework for isometry-invariant partial surface matching

venue: Proceedings of the National Academy of Sciences,

authors:
Hongteng Xu
Dixin Luo
Lawrence Carin

title: Scalable gromov-wasserstein learning for graph partitioning and matching

venue: In Advances in neural information processing systems,

authors:
Justin Gilmer
Samuel S Schoenholz
Patrick F Riley
Oriol Vinyals
George E Dahl

title: Neural message passing for quantum chemistry

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

authors:
Edouard Grave
Armand Joulin
Quentin Berthet

title: Unsupervised alignment of embeddings with wasserstein procrustes

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

authors:
David Alvarez-Melis
Tommi S Jaakkola

title: Gromov-wasserstein alignment of word embedding spaces

venue: arXiv preprint arXiv:1809.00013,

title: The space of spaces: curvature bounds and gradient flows on the space of metric measure spaces

venue: arXiv preprint arXiv:1208.0434,

title: Gromov–wasserstein distances and the metric approach to object matching

venue: Foundations of computational mathematics,

authors:
Matthias Liero
Alexander Mielke
Giuseppe Savaré

title: Optimal entropy-transport problems and a new hellinger–kantorovich distance between positive measures

venue: Inventiones mathematicae,

authors:
Nicolas Courty
Rémi Flamary
Devis Tuia

title: Domain adaptation with regularized optimal transport

venue: In Joint European Conference on Machine Learning and Knowledge Discovery in Databases,

authors:
Charlie Frogner
Chiyuan Zhang
Hossein Mobahi
Mauricio Araya
Tomaso A Poggio

title: Learning with a Wasserstein loss

venue: In Advances in Neural Information Processing Systems,

authors:
Lénaïc Chizat
Gabriel Peyré
Bernhard Schmitzer
François-Xavier Vialard

title: Unbalanced optimal transport: Dynamic and kantorovich formulations

venue: Journal of Functional Analysis,

authors:
Jean Feydy
Pierre Roussillon
Alain Trouvé
Pietro Gori

title: Fast and scalable optimal transport for brain tractograms

venue: In International Conference on Medical Image Computing and ComputerAssisted Intervention,

authors:
John Lee
Nicholas P Bertrand
Christopher J Rozell

title: Parallel unbalanced optimal transport regularization for large scale imaging problems

authors:
Yogesh Balaji
Rama Chellappa
Soheil Feizi

title: Robust optimal transport with applications in generative modeling and domain adaptation

venue: Advances in Neural Information Processing Systems,

authors:
Lenaic Chizat
Francis Bach

title: On the global convergence of gradient descent for over-parameterized models using optimal transport

authors:
Grant Rotskoff
Samy Jelassi
Joan Bruna
Eric Vanden-Eijnden

title: Global convergence of neuron birth-death dynamics

venue: arXiv preprint arXiv:1902.01843,

title: Sinkhorn distances: Lightspeed computation of optimal transport

venue: In Adv. in Neural Information Processing Systems,

authors:
Lenaic Chizat
Gabriel Peyré
Bernhard Schmitzer
François-Xavier Vialard

title: Scaling algorithms for unbalanced transport problems. to appear in Mathematics of Computation, 2018b

authors:
Thibault Séjourné
Jean Feydy
François-Xavier Vialard
Alain Trouvé
Gabriel Peyré

title: Sinkhorn divergences for unbalanced optimal transport

venue: arXiv preprint arXiv:1910.12958,

authors:
Charlotte Bunne
David Alvarez-Melis
Andreas Krause
Stefanie Jegelka

title: Learning generative models across incomparable spaces

venue: arXiv preprint arXiv:1905.05461,

authors:
Hongteng Xu
Dixin Luo
Ricardo Henao
Svati Shah
Lawrence Carin

title: Learning autoencoders with relational regularization

venue: arXiv preprint arXiv:2002.02913,

authors:
Ivegen Redko
Titouan Vayer
Rémi Flamary
Nicolas Courty

title: Co-optimal transport

venue: arXiv preprint arXiv:2002.03731,

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

title: Fused gromovwasserstein distance for structured objects: theoretical foundations and mathematical properties

venue: arXiv preprint arXiv:1811.02834,

authors:
Samir Chowdhury
Tom Needham

title: Gromov-wasserstein averaging in a riemannian framework

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

authors:
David Alvarez-Melis
Stefanie Jegelka
Tommi S Jaakkola

title: Towards optimal transport with global invariances

authors:
Rainer E Burkard
Eranda Cela
Panos M Pardalos
Leonidas S Pitsoulis

title: The quadratic assignment problem

venue: In Handbook of combinatorial optimization,

authors:
Gabriel Peyré
Marco Cuturi
Justin Solomon

title: Gromov-wasserstein averaging of kernel and distance matrices

venue: In International Conference on Machine Learning,

authors:
Steven Gold
Anand Rangarajan

title: A graduated assignment algorithm for graph matching

venue: IEEE Transactions on pattern analysis and machine intelligence,

authors:
Nicoló De Ponti
Andrea Mondino

title: Entropy-transport distances between unbalanced metric measure spaces

venue: arXiv preprint arXiv:2009.10636,

title: On the geometry of metric measure spaces

venue: Acta mathematica,

authors:
Laetitia Chapel
Mokhtar Z Alaya
Gilles Gasso

title: Partial gromov-wasserstein with applications on positive-unlabeled learning

venue: arXiv preprint arXiv:2002.08276,

title: The optimal partial transport problem

venue: Archive for rational mechanics and analysis,

authors:
Luca Cosmo
Emanuele Rodolà
Michael M Bronstein
Andrea Torsello
Daniel Cremers
Y Sahillioglu

title: Shrec’16: Partial matching of deformable shapes

venue: Proc. 3DOR,

authors:
Emanuele Rodola
Alex M Bronstein
Andrea Albarelli
Filippo Bergamasco
Andrea Torsello

title: A game-theoretic approach to deformable shape matching

venue: IEEE Conference on Computer Vision and Pattern Recognition,

authors:
Dmitri Burago
Iu D Burago
Yuri Burago
Sergei A Ivanov
Sergei Ivanov

title: A course in metric geometry, volume 33

venue: American Mathematical Soc.,

title: Metric properties of homogeneous and spatially inhomogeneous f -divergences

venue: IEEE Transactions on Information Theory,

title: Softmax to softassign: Neural network algorithms for combinatorial optimization

venue: Journal of Artificial Neural Networks,

authors:
J. Solomon
G. Peyré
V. Kim
S. Sra

title: Entropic metric alignment for correspondence problems

venue: ACM Transactions on Graphics (TOG),

title: Convergence of a block coordinate descent method for nondifferentiable minimization

venue: Journal of optimization theory and applications,

authors:
Haggai Maron
Yaron Lipman

title: probably) concave graph matching

venue: arXiv preprint arXiv:1807.09722,

title: Maximization of a convex quadratic function under linear constraints

venue: Mathematical programming,

authors:
Aasa Feragen
Francois Lauze
Soren Hauberg

title: Geodesic exponential kernels: When curvature and linearity conflict

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

authors:
Jason Altschuler
Francis Bach
Alessandro Rudi
Jonathan Weed

title: Massively scalable Sinkhorn distances via the nystr\" om method

venue: arXiv preprint arXiv:1812.05189,

authors:
Khiem Pham
Khang Le
Nhat Ho
Tung Pham
Hung Bui

title: On unbalanced optimal transport: An analysis of sinkhorn algorithm

venue: arXiv preprint arXiv:2002.03293,

authors:
Alexis Thibault
Lénaïc Chizat
Charles Dossal
Nicolas Papadakis

title: Overrelaxed Sinkhorn-Knopp algorithm for regularized optimal transport

venue: arXiv preprint arXiv:1711.01851,

title: Iterative procedures for nonlinear integral equations

venue: Journal of the ACM (JACM),

authors:
Damien Scieur
Alexandre d’Aspremont
Francis Bach

title: Regularized nonlinear acceleration

venue: In Advances In Neural Information Processing Systems,

authors:
Gregory Griffin
Alex Holub
Pietro Perona

title: Caltech-256 object category dataset

authors:
Kate Saenko
Brian Kulis
Mario Fritz
Trevor Darrell

title: Adapting visual category models to new domains

venue: In European conference on computer vision,

authors:
Jeff Donahue
Yangqing Jia
Oriol Vinyals
Judy Hoffman
Ning Zhang
Eric Tzeng
Trevor Darrell

title: Decaf: A deep convolutional activation feature for generic visual recognition

venue: In International conference on machine learning,

authors:
Masahiro Kato
Takeshi Teshima
Junya Honda

title: Learning from positive and unlabeled data with a selection bias

venue: In International Conference on Learning Representations,

authors:
Yu-Guan Hsieh
Gang Niu
Masashi Sugiyama

title: Classification from positive, unlabeled and biased negative data

authors:
Marthinus C Du Plessis
Gang Niu
Masashi Sugiyama

title: Analysis of learning from positive and unlabeled data

venue: Advances in neural information processing systems,

authors:
Jean Feydy
Thibault Séjourné
François-Xavier Vialard
Shun-ichi Amari
Alain Trouvé
Gabriel Peyré

title: Interpolating between optimal transport and mmd using sinkhorn divergences

title: Optimal transport in competition with reaction: The hellinger–kantorovich distance and geodesic curves

venue: SIAM Journal on Mathematical Analysis,

title: An interpolating distance between optimal transport and fisher–rao metrics

venue: Foundations of Computational Mathematics,

title: Topics in C. Transportation

venue: Graduate studies in Math. AMS,

authors:
Rémi Flamary
Nicolas Courty

title: Pot python optimal transport library

authors:
Christian Berg
Jens Peter Reus Christensen
Paul Ressel

title: Harmonic analysis on semigroups: theory of positive definite and related functions, volume

heading: 1 Introduction

text:
Comparing data distributions on different metric spaces is a basic problem in machine learning. This class of problems is for instance at the heart of surfaces [Bronstein et al., 2006] or graph matching [Xu et al., 2019] (equipping the surface or graph with its associated geodesic distance), regression problems in quantum chemistry [Gilmer et al., 2017] (viewing the molecules as distributions of points in R3) and natural language processing [Grave et al., 2019, Alvarez-Melis and Jaakkola, 2018] (where texts in different languages are embedded as points distributions in different vector spaces).
Metric measure spaces. The mathematical way to formalize these problems is to model the data as metric measure spaces (mm-spaces). A mm-space is denoted as X = (X, d, µ) where X is a complete separable set endowed with a distance d and a positive Borel measure µ ∈M+(X). For instance, if X = (xi)i is a finite set of points, then µ = ∑ imiδxi (here δxi is the Dirac mass at xi) is simply a set of positive weights mi = µ({xi}) ≥ 0 associated to each point xi, which accounts for its mass or importance. For instance, setting some mi to 0 is equivalent to removing the point xi. We
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
refer to Sturm [2012] for a mathematical account on the theory of mm-spaces. In all the applications highlighted above, it makes sense to perform the comparisons up to isometric transformations of the data. Two mm-spaces X = (X, dX , µ) and Y = (Y, dY , ν) are considered to be equal (denoted X ∼ Y) if they are isometric, meaning that there is a bijection ψ : spt(µ)→ spt(ν) (where spt(µ) is the support of µ) such that dX(x, y) = dY (ψ(x), ψ(y)) and ψ]µ = ν. Here ψ] is the push-forward operator, so that ψ]µ = ν is equivalent to imposing ν(A) = µ(ψ−1(A)) for any set A ⊂ Y . For discrete spaces where µ = ∑ imiδxi , then one should have ν = ψ]µ = ∑ imiδψ(xi). As highlighted by Mémoli [2011], considering mm-spaces up to isometry is a powerful way to formalize and analyze a wide variety of problems such as matching, regression and classification of distributions of points belonging to different spaces. Most often, the objects of interest come with a natural distance such as an intrinsic or extrinsic distance and the uniform measure is the usual choice to make mm-spaces widely applicable. The key to unlock all these problems is the computation of a distance between mm-spaces up to isometry. So far, existing distances (reviewed below) assume that µ is a probability distribution, i.e. µ(X) = 1. This constraint is not natural and sometimes problematic for most of the practical applications to machine learning. The goal of this paper is to alleviate this restriction. We define for the first time a class of distances between unbalanced metric measure spaces, these distances being upper-bounded by divergences which can be approximated by an efficient numerical scheme.
Csiszár divergences The simplest case is when X = Y and one simply ignores the underlying metric. One can then use Csiszár divergences (or ϕ-divergences), which perform a pointwise comparison (in contrast with optimal transport distances, which perform a displacement comparison). It is defined using an entropy function ϕ : R+ → [0,+∞], which is a convex, lower semi-continuous, positive function with ϕ(1) = 0. The Csiszár ϕ-divergence reads Dϕ(µ|ν) , ∫ X ϕ ( dµ dν ) dν +
ϕ′∞ ∫ X dµ⊥, where µ = dµdν ν + µ ⊥ is called the Radon-Nikodym or the Lebesgue decomposition of µ with respect to ν and ϕ′∞ = limr→∞ ϕ(r)/r ∈ R ∪ {+∞} is called the recession constant. This divergence Dϕ is convex, positive, 1-homogeneous and weak* lower-semicontinuous, see Liero et al. [2015] for details. Particular instances of ϕ-divergences are Kullback-Leibler (KL) for ϕ(r) = r log(r)− r + 1 (note that ϕ′∞ =∞) and Total Variation (TV) for ϕ(r) = |r − 1|.
Balanced and unbalanced optimal transport. If the common embedding space X is equipped with a distance d(x, y), one can use more elaborated methods such as optimal transport (OT) distances, which are computed by solving convex optimization problems. This type of methods has proven useful for ML problems as diverse as domain adaptation [Courty et al., 2014], supervised learning over histograms [Frogner et al., 2015] and unsupervised learning of generative models [Arjovsky et al., 2017]. In this case, the extension from probability distributions to arbitrary positive measures (µ, ν) ∈ M+(X)2 is now well understood and corresponds to the theory of unbalanced OT. Following Liero et al. [2015], Chizat et al. [2018a], a family of unbalanced Wasserstein distances is defined by solving
UW(µ, ν)q , inf π∈M(X×X)
∫ λ(d(x, y))dπ(x, y) + Dϕ(π1|µ) + Dϕ(π2|µ). (1)
Here (π1, π2) are the two marginals of the joint distribution π, defined by π1(A) = π(A× Y ) for A ⊂ X . The mapping λ : R+ → R and exponent q ≥ 1 should be chosen wisely to ensure for instance that UW defines a distance (see Section 2.2.1). It is frequent to take ρDϕ instead of Dϕ (i.e. take ψ = ρϕ) to adjust the strength of the marginals’ penalization. Balanced OT is retrieved with the convex indicator ϕ = ι{1} (i.e. ϕ(1) = 0 and ϕ(x) = +∞ otherwise) or by taking the limit ρ → +∞, which enforces π1 = µ and π2 = ν. When 0 < ρ < +∞, unbalanced OT operates a trade-off between transportation and creation of mass, which is crucial to be robust to outliers in the data and to cope with mass variations in the modes of the distributions. For supervised tasks, the value of ρ should be cross-validated to obtain the best performances. Its use is gaining popularity in applications, such as medical imaging registration [Feydy et al., 2019a], videos [Lee et al., 2019], generative learning [Balaji et al., 2020] and gradient flow to train neural networks [Chizat and Bach, 2018, Rotskoff et al., 2019]. Furthermore, existing efficient algorithms for balanced OT extend to this unbalanced problem. In particular Sinkhorn’s iterations, introduced in ML for balanced OT by Cuturi [2013], extend to unbalanced OT [Chizat et al., 2018b, Séjourné et al., 2019], as detailed in Section 3.
The Gromov-Wasserstein distance and its applications. The Gromov-Wasserstein (GW) distance [Mémoli, 2011, Sturm, 2012] generalizes the notion of OT to the setting of mm-spaces up to
isometries. It replaces the linear cost ∫ λ(d)dπ of OT by a quadratic function. It reads
GW(X ,Y)q , min π∈M+(X×Y )
{∫ λ(|dX(x, x′)− dY (y, y′)|)dπ(x, y)dπ(x′, y′) :
π1 = µ π2 = ν
} . (2)
It is proved in Mémoli [2011], Sturm [2012] that GW defines with λ(t) = tq a distance up to isometries on balanced mm-spaces (i.e. the measures are probability distributions). The GW distance is applied successfully in natural language processing for unsupervised translation learning [Grave et al., 2019, Alvarez-Melis and Jaakkola, 2018], in generative learning for objects lying in spaces of different dimensions [Bunne et al., 2019] and to build VAE for graphs [Xu et al., 2020]. It has been adapted for domain adaptation over different spaces [Redko et al., 2020]. It is also a relevant distance to compute barycenters between graphs or shapes [Vayer et al., 2018, Chowdhury and Needham, 2020]. When (X ,Y) are Euclidean spaces, this distance compares distributions up to rigid isometry, and is closely related (but not equal) to metrics defined by procrustes analysis [Grave et al., 2019, Alvarez-Melis et al., 2019]. The problem (2) is non convex because the quadratic form∫ λ(|dX − dY |)dπ ⊗ π is not positive in general. It is in fact closely related to quadratic assignment problems [Burkard et al., 1998], which are used for graph matching problems, and are known to be NP-hard in general. Nevertheless, non-convex optimization methods have been shown to be successful in practice to use GW distances for ML problems. This includes for instance alternating minimization [Mémoli, 2011, Redko et al., 2020] and entropic regularization [Peyré et al., 2016, Gold and Rangarajan, 1996].
Related works and contributions. The concomitant work of De Ponti and Mondino [2020] extends the Lp transportation distance defined in Sturm et al. [2006] to unbalanced mm-spaces and studies its geometric properties. This distortion distance is not equivalent to the GW distance, and is more difficult to estimate numerically because it explicitly imposes a triangle inequality constraint in the optimization problem. The work of Chapel et al. [2020] relaxes the GW distance to the unbalanced setting by hybridizing GW with partial OT [Figalli, 2010] for unsupervised labeling. It ressembles one particular setting of our formulation, but with some important differences, detailed in Section 2. Our construction is also connected to partial matching methods, which find numerous applications in graphics and vision [Cosmo et al., 2016]. In particular, Rodola et al. [2012] introduces a mass conservation relaxation of the GW problem.
The two main contributions of this paper are the definition of two formulations relaxing the GW distance. The first one is called the Unbalanced Gromov-Wasserstein (UGW) divergence and can be computed efficiently on GPUs. The second one is called the Conic Gromov-Wasserstein distance (CGW). It is proved to be a distance between mm-spaces endowed with positive measures up to isometries, as stated in Theorem 1 which is the main theoretical result of this paper. We also prove in Theorem 1 that UGW can be used as a surrogate upper-bounding CGW. We present those concepts and their properties in Section 2. We also detail in Section 3 an efficient computational scheme for a particular setting of UGW. This method computes an approximate stationary point of a biconvex relaxation of our formulations. Even though it is a lower bound of the original problem, we provide in Theorem 3 conditions ensuring the tightness of this relaxation in many cases of interest. The algorithm leverages the strength of entropic regularization and the Sinkhorn algorithm, namely that it is GPU-friendly and defines smooth loss functions amenable to back-propagation for ML applications. Section 4 provides some numerical experiments to highlight the qualitative behavior of this algorithm and its ability to cope with outliers and mass variations in the modes of the distributions. We illustrate numerically the tightness of the relation between UGW and CGW, showing that UGW is a reasonnable proxy of a distance, at least locally. We provide an application of our divergence in the positive unlabeled learning setting, using domain adaptation data, and display results which are at par or outperform the computable competitor Chapel et al. [2020].

heading: 2 Unbalanced Gromov-Wasserstein formulations

text: We present in this section our two new formulations and their properties. The first one, called UGW, is exploited in Sections 3 and 4 to derive an efficient algorithm used in numerical experiments. The second one, called CGW, defines a distance between mm-spaces up to isometries. Those results build upon the work of Liero et al. [2015], and a summary of the construction of UOT is detailed in Appendix A. In all what follows, we consider complete separable mm-spaces endowed with a metric and a positive measure.

heading: 2.1 The unbalanced Gromov-Wasserstein divergence

text:
This new formulation makes use of quadratic ϕ-divergences, defined as D⊗ϕ (ρ|ν) , Dϕ(ρ⊗ ρ|ν⊗ ν), where ρ ⊗ ρ ∈ M+(X2) is the tensor product measure defined by d(ρ ⊗ ρ)(x, y) = dρ(x)dρ(y). Note that D⊗ϕ is not a convex function in general.
Definition 1 (Unbalanced GW). The Unbalanced Gromov-Wasserstein divergence is defined as UGW(X ,Y) = infπ∈M+(X×Y ) L(π) , G(π) + D⊗ϕ (π1|µ) + D⊗ϕ (π2|ν).
This definition can be understood as an hybridation between (1) and (2) but with a twist: one needs to use the quadratic divergence D⊗ϕ in place of Dϕ. To the best of our knowledge, it is the first time such quadratic divergences are being used and studied. In the TV case, this is the most important distinction between UGW and partial-GW [Chapel et al., 2020]. Note also that the balanced GW distance (2) is recovered as a particular case when using ϕ = ι{1} or by letting ρ → +∞ for an entropy ψ = ρϕ.
Using quadratic divergences results in UGW being 2-homogeneous: for θ ≥ 0, writing (Xθ,Yθ) equiped with (θµ, θν), one has θ−2UGW(Xθ,Yθ) = UGW(X ,Y). When using non tensorized ϕ-divergences, the resulting unbalanced Gromov-Wassertein functional between Xθ and Yθ have very different and inconsistent behaviors when θ → 0 and θ → +∞. Indeed, once normalized by θ−2 and θ−1, one obtains respectively balanced GW and a Hellinger-type distance. Using tensorized divergences ensures that the behavior does not depends on θ. It is also fundamental to connect UGW with our distance CGW, see Theorem 1 and Appendix B.
We first prove the existence of optimal plans π minimizing L, which holds for the three key settings of Section 2.2.1, namely for KL, TV, and for compact metric spaces (such as finite pointclouds and graphs). All proofs are deferred in Appendix B.
Proposition 1 (Existence of minimizers). We assume that (X,Y ) are compact and that either (i) ϕ superlinear, i.e ϕ′∞ = ∞, or (ii) λ has compact sublevel sets in R+ and 2ϕ′∞ + inf λ > 0. Then there exists π ∈M+(X × Y ) such that UGW(X ,Y) = L(π).
The following proposition ensures that the functional UGW can be used to compare mm-spaces.
Proposition 2 (Definiteness of UGW). Assume that ϕ−1({0}) = {1} and λ−1({0}) = {0}. Then UGW(X ,Y) ≥ 0 and is 0 if and only if X ∼ Y .
We end this section with a reformulation of UGW which is important to make the connection with the second formulation CGW of the following section. It splits UGW into two parts: the term ϕ(0)(|(µ⊗µ)⊥|+ |(ν⊗ν)⊥|) accounts for the pure creation/destruction of mass and a new transport cost Lc accounts for the remaining part (partial/pure transport and partial creation/destruction of mass).
Lemma 1. Defining Lc(a, b) , c + aϕ(1/a) + bϕ(1/b), and writing (f , dµdπ1 , g , dν dπ2 ) the Lebesgue densities of (µ, ν) w.r.t. (π1, π2) such that µ = fπ1 + µ⊥ and ν = gπ2 + ν⊥, one has
L(π) = ∫ X2×Y 2 Lλ(|dX−dY |)(f ⊗ f, g ⊗ g)dπdπ + ϕ(0)(|(µ⊗ µ) ⊥|+ |(ν ⊗ ν)⊥|). (3)
Proof. Write f = dµdπ1 and g = dν dπ2 . The Lebesgue decompositions read µ ⊗ µ = (f ⊗ f)π1 ⊗ π1 + (µ ⊗ µ)⊥ and ν ⊗ ν = (g ⊗ g)π2 ⊗ π2 + (ν ⊗ ν)⊥, thanks to the tensorized structure of the decomposed plans. To prove Equation (3), we need to define the reverse entropy Liero et al. [2015] such that Dϕ(α|µ) = Dψ(µ|α), where ψ(x) , xϕ( 1x ) is also an entropy function satisfying ψ′∞ = ϕ(0). One then has
L(π) = ∫ X2×Y 2 λ(Γ)dπdπ + D⊗ϕ (π1|µ) + D⊗ϕ (π2|ν)
= ∫ X2×Y 2 λ(Γ)dπdπ + D⊗ψ (µ|π1) + D ⊗ ψ (ν|π2)
L(π) = ∫ X2×Y 2 λ(Γ)dπdπ + ∫ X2 ψ(f ⊗ f)dπ1dπ1 + ∫ Y 2 ψ(g ⊗ g)dπ2dπ2
+ ϕ(0)(|(µ⊗ µ)⊥|+ |(ν ⊗ ν)⊥|)
= ∫ X2×Y 2 Lλ(Γ)(f ⊗ f, g ⊗ g)dπdπ + ϕ(0)(|(µ⊗ µ)⊥|+ |(ν ⊗ ν)⊥|).
Using the definition of ψ in Lc ends the proof.

heading: 2.2 The conic Gromov-Wasserstein distance

text: We introduce a second “conic” formulation of unbalanced GW, which is connected to UGW, and whose construction is inspired by the conic formulation of UOT (see Appendix A for an overview).

heading: 2.2.1 Background on cone sets and distances

text:
The conic formulation lifts a point x ∈ X to a couple (x, r) ∈ X × R+ where r encodes some (power of a) mass. Then we seek optimal transport plans defined over C[X] , X × R+/(X × {0}), where coordinates (x, r = 0) with no mass are merged into a single point 0X called the apex of the cone. In the sequel, points of X × R+ are noted (x, r), while [x, r] are quotiented points of C[X]. While transport plans depend on variables ([x, r], [y, s]) and ([x′, r′], [y′, s′]) in C[X]×C[Y ], the transportation cost involved in our conic formulation only makes use of the 2-D cone C[R+] over R+ endowed with the distance |u−v| (note that any other distance on R could be used as well). More specifically, we consider coordinates of the form ([u, a], [v, b]) = ([dX(x, x′), rr′], [dY (y, y′), ss′]) ∈ C[R+] × C[R+]. Thus we now describe conic discrepancies D on C[R+], which are defined for (p, q) ≥ 1 as D([u, a], [v, b])q , Hλ(|u−v|)(ap, bp), where Hc(ap, bp) , infθ≥0 θLc(a p θ , bp
θ ) is the perspective transform of Lc introduced in Lemma 1. The intuition underpinning the definition of this cost is that the perspective transform accounts for the possibility to rescale a transport plan π by a scalar θ but the scaling is performed pointwise instead of globally. In general D is not a distance, but it is always definite as stated by this result proved in Appendix A. Proposition 3. Assume λ−1({0}) = {0}, ϕ−1({0}) = {1} and ϕ is coercive. Then D is definite on C[R+], i.e. D([u, a], [v, b]) = 0 if and only if (a = b = 0) or (a = b and u = v).
Of particular interest are those ϕ where D is a distance, which necessitates a careful choice of λ, p and q. We now detail three examples where this is the case. Gaussian Hellinger distance (GH). When Dϕ = KL, λ(t) = t2 and q = p = 2, then one has D([u, a], [v, b])2 = a2 + b2 − 2abe−|u−v|/2. This cone distance [Burago et al., 2001] is further generalized by De Ponti [2019] who shows that D is a distance for power entropies ϕ(s) = s
p−p(s−1)−1 p(p−1)
if p ≥ 1 (the case p = 1 corresponding to Dϕ = KL). Hellinger-Kantorovich (HK) / Wasserstein-Fisher-Rao distance (WFR). When Dϕ = KL, λ(t) = − log cos2(t∧ π2 ) and q = p = 2, then one hasD([u, a], [v, b])
2 = a2 +b2−2ab cos(π2 ∧|u− v|). This construction, which might seem peculiar, corresponds to the one used to make unbalanced OT a geodesic distance, as detailed in [Liero et al., 2015, Chizat et al., 2018a]. Partial optimal transport distance (PT). When Dϕ = TV, λ(t) = tq, q ≥ 1 and p = 1, then D([u, a], [v, b])q = a+ b− (a ∧ b)(2− |u− v|q)+ defines a cone distance [Chizat et al., 2018a].

heading: 2.2.2 Definitions and properties

text:
The conic formulation consists in solving a GW problem on the cone, with the addition of two linear constraints. Informally speaking, Lc from Lemma 1 becomes D, the term (|(µ⊗ µ)⊥|+ |(ν ⊗ ν)⊥|) is taken into account by the constraints (5) below, and the variables (f, g) are replaced by (rp, sp). It reads CGW(X ,Y) , infα∈Up(µ,ν)H(α) where
H(α) , ∫ D([dX(x, x′), rr′], [dY (y, y′), ss′])q dα([x, r], [y, s])dα([x′, r′], [y′, s′]), (4)
and Up(µ, ν) is defined as the set Up(µ, ν) , { α ∈M+(C[X]× C[Y ]), ∫ R+ rpdα1(·, r) = µ, ∫ R+ spdα2(·, s) = ν } . (5)
It is similar to the conic formulation of UW, see Appendix A. Note that similarly to the GW formulation (2) – and in sharp contrast with the conic formulation of UW – here the transport plans are defined on the cone C[X]× C[Y ] but the cost D is a distance on C[R+]. We present now the main contributions of this paper, proved in Appendix C. We state that CGW defines a distance under conditions that hold for the settings of Section 2.2.1, and that it is upperbounded by UGW. The divergence UGW can be approximated with efficient numerical schemes as detailed in Section 3.
Theorem 1. (i) The divergence CGW is symmetric, positive and definite up to isometries. (ii) If D is a distance on C[R+], then CGW1/q is a distance on the set of mm-spaces up to isometries. (iii) For any (Dϕ, λ, p, q) with associated cost D on the cone, one has UGW ≥ CGW.

heading: 3 Algorithms

text: We focus in this section on the numerical computation of the upper bound UGW using a bi-convex relaxation and derive an alternate minimization scheme coupled with entropic regularization. We also propose to approximate CGW by doing a similar alternate minimization, as detailed in Appendix E. We provide guarantees of tightness on the bi-convex relaxation for CGW (see Theorem 3). The computation of the distance CGW is heavy in practice because it requires an optimization over a lifted conic space, which needs to be discretized. Thus it does not scale to large problem for CGW, but allows to explore numerically how tight is the upper bound UGW ≥ CGW, see Section 4. The algorithm for UGW is presented on arbitrary measures, the special case of discrete measures being a particular case. The discretized formulas and algorithms are detailed in Appendix D, see also Chizat et al. [2018b], Peyré et al. [2016]. All implementations are available at https: //github.com/thibsej/unbalanced_gromov_wasserstein, and installable in Python with the command pip install unbalancedgw.

heading: 3.1 Bi-convex relaxation and tightness

text:
In order to derive a simple numerical approximation scheme, following Mémoli [2011], we introduce a lower bound obtained by introducing two transportation plans. To further accelerate the method and enable GPU-friendly iterations, similarly to Gold et al. [1996], Solomon et al. [2016], we consider an entropic regularization. It reads, for any ε ≥ 0,
UGWε(X ,Y) , inf π L(π) + εKL⊗(π|µ⊗ ν) ≥ inf π,γ F(π, γ) + εKL(π ⊗ γ|(µ⊗ ν)⊗2), (6) and F(π, γ) , ∫ X2×Y 2 λ(|dX − dY |)dπ ⊗ γ + Dϕ(π1 ⊗ γ1|µ⊗ µ) + Dϕ(π2 ⊗ γ2|ν ⊗ ν) ,
where (γ1, γ2) denote the marginals of the plan γ. In the sequel we write Fε = F + εKL⊗. Note that in contrast to the entropic regularization of GW Peyré et al. [2016], here we use a tensorized entropy to maintain the overall homogeneity of the energy. A simple method to approximate this lower bound is to perform an alternate minimization on π and γ, which is known to converge for smooth ϕ to a stationary point since the coupling term in the functional is smooth [Tseng, 2001]. Note that if π ⊗ γ is optimal then so is (sπ)⊗ ( 1sγ) with s ≥ 0. Thus without loss of generality we can optimize under the constraint m(π) = m(γ) by setting s = √ m(γ)/m(π).
We now discuss the tightness of the bi-convex relaxation by generalizing a result of Konno. We first present a result which applies to general quadratic assignment problems, then state its application to our setting.
Theorem 2 (Tight relaxation). Let B a Banach space, let f : B 7→ R ∪ {+∞} be a function and let L : C ⊂ B 7→ R the function defined on the convex set C ⊂ B by L(π) = 12 〈π, k(π)〉+2f(π) where k is a symmetric bilinear map which is negative (not necessarily definite) on ∆C , Span({π − γ ; (π, γ) ∈ C}), that is, for any z ∈ ∆C, 〈z, kz〉 ≤ 0. Assume that there exists π0 ∈ C such that L(π0) < +∞, and define F(π, γ) , 12 〈π, k(γ)〉 + f(π) + f(γ). Then, for any (π∗, γ∗) ∈ arg minF(π, γ), we have F(π∗, π∗) = F(γ∗, γ∗) = F(π∗, γ∗). Moreover, if one assumes either that k is a definite kernel or f is strictly convex, one gets π∗ = γ∗.
The above Theorem 2 is proved in Appendix D. As an application, we now state our tightness result for GWε and CGW. In those settings the optimizers of the bi-convex relaxation are also optimal for the original problem. Theorem 3. For GWε with ε ≥ 0 or for CGW, assume that λ(t) = t2 and that (dX , dY ) are both conditionnally negative (or conditionally positive) kernels. Then the bi-convex relaxation of both problems is tight.
Proof. The proof for CGW is detailed in Appendix E, we prove the tightness for GWε. When λ(t) = t2 the kernel k = λ(|dX − dY |) is conditionally negative on the set {(π, γ), π1 = γ1 and π2 = γ2}, i.e. we have 〈(π−γ), k(π−γ)〉 ≤ 0 (see Maron and Lipman [2018]). For GWε one has π1 = γ1 = µ and π2 = γ2 = ν thanks to the constraints on marginals. Thus the kernel is negative semi-definite and the proof of Theorem 2 applies, hence the tightness of the relaxation.
Konno’s result Konno [1976] applies for unregularized (ε = 0), Balanced-GW. The novelty of Theorem 3 is its extension to both GWε and CGW. So far it is an open question whether the relaxation is tight or not for UGWε, because the above proof no longer holds. Note that in all our numerical simulations, our solvers always found solutions of UGWε such that π = γ when |dX − dY |2 is conditonally negative. The property that the kernel |dX − dY |2 is negative does not hold in general (e.g. for graph geodesic distances) and the tightness of the relaxation remains open in this setting. We know from [Maron and Lipman, 2018, Theorem 1] that it is conditionally negative semi-definite when both (dX , dY ) are conditionally negative kernels. Examples of distances which are negative kernels are tree metrics in the case of graphs, as well as Euclidean, spherical and hyperbolic distances over their respective manifolds Feragen et al. [2015]. In practice, when ε is small, we observed in the indefinite setting that the relaxation outputs more frequently spurious minima than in the negative semi-definite setting.

heading: 3.2 Alternate Sinkhorn minimization

text:
Minimizing the lower bound (6) with respect to either π or γ is non-trivial for an arbitrary ϕ. We restrict our attention to the Kullback-Leibler case Dϕ = ρKL with ρ > 0, which can be addressed by solving a regularized and convex unbalanced problem as studied in Chizat et al. [2018b], Séjourné et al. [2019]. It is explained in the following proposition. Proposition 4. For a fixed γ, the optimal π ∈ arg min
π F(π, γ) + εKL(π ⊗ γ|(µ⊗ ν)⊗2) solves
min π
∫ cεγ(x, y)dπ(x, y) + ρm(γ)KL(π1|µ) + ρm(γ)KL(π2|ν) + εm(γ)KL(π|µ⊗ ν),
where m(γ) , γ(X × Y ) is the mass of γ, and where we define the cost associated to γ as cεγ(x, y), ∫ λ(|dX(x, ·)− dY (y, ·)|)dγ+ρ ∫ log(
dγ1 dµ )dγ1+ρ
∫ log(
dγ2 dν )dγ2+ε
∫ log( dγ
dµdν )dγ.
Algorithm 1 – UGW(X , Y , ρ, ε) Input: mm-spaces (X ,Y), relax. ρ, regul. ε Output: π, γ solving (6)
Init. π = γ = µ⊗ ν/ √ m(µ)m(ν), g = 0. while (π, γ) has not converged do Update π ← γ, then c← cεπ , ρ̃← m(π)ρ, ε̃← m(π)ε while (f, g) has not converged do f ← − ε̃ρ̃
ε̃+ρ̃ log
∫ e(g(y)−c(·,y))/ε̃dν(y)
g ← − ε̃ρ̃ ε̃+ρ̃
log ∫ e(f(x)−c(x,·))/ε̃dµ(x)
end while Upd. γ(x, y)←e f(x)+g(y)−c(x,y) ε̃ µ(x)ν(y)
Rescale γ ← √ m(π)/m(γ)γ
end while Return (π, γ).
Computing the cost cεγ for spaces X and Y of n points has in general a costO(n4) in time and memory. However, as explained for instance in Peyré et al. [2016], for the special case λ(t) = t2, this cost is reduced to O(n3) in time and O(n2) in memory. This is the setting we consider in the numerical simulations. This makes the method applicable for scales of the order of 104 points. For larger datasets one should use approximation schemes such as hierarchical approaches [Xu et al., 2019] or Nyström compression of the kernel [Altschuler et al., 2018].
The resulting alternate minimization method is detailed in Algorithm 1, see Appendix D for a discretized version. It uses the unbalanced Sinkhorn algorithm of Chizat et al. [2018b], Séjourné et al. [2019] as subiterations and takes π = µ ⊗ ν/ √ m(µ)m(ν) to initialize the updates. This Sinkhorn algorithm operates over a pair of continuous functions (so-called
Kantorovitch potentials) f(x) and g(y). For discrete spaces X and Y of size n, these functions are stored in vectors of size n, and that integral involved in the updates becomes a sum. Each iteration of Sinkhorn thus has a cost n2, and all the involved operation can be efficiently mapped to parallelizable GPU routines as detailed in Chizat et al. [2018b], Séjourné et al. [2019]. Another advantage of using an unbalanced Sinkhorn algorithm is its complexity O(n2/ε) to compute an ε-approximation, as stated in Pham et al. [2020], which should be compared to O(n2/ε2) operations for balanced Sinkhorn.
Note also that balanced GW is recovered as a special case when setting ρ → +∞, so that ρ̃/(ε̃ + ρ̃) → 1 should be used in the iterations. In order to speed up Sinkhorn inner-loops, especially for small values of ε, one can use linear extrapolation [Thibault et al., 2017] or non-linear Anderson acceleration [Anderson, 1965, Scieur et al., 2016].
There is an extra scaling step after computing γ involving the mass m(π). It corresponds to the scaling s of π⊗ γ such that m(π) = m(γ), and we observe that this scaling is key not only to impose this mass equality but also to stabilize the algorithm. Otherwise we observed that m(γ) < 1 < m(π) and underflows whenever m(γ)→ 0 and m(π)→∞.

heading: 4 Numerical experiments

text:
This section presents simulations on synthetic examples to highlight the qualitative behavior of UGW and the tightness of the bound UGW ≥ CGW. Other illustrations on UGW are available in Appendix E. We end the section with a learning application of UGW in a positive-unlabeled setting, using domain adaptation data so as to compare with PGW Chapel et al. [2020]. In the synthetic experiments, µ and ν are probability distributions, which allows us to compare GW with UGW.
Robustness to imbalanced classes. In this first example, we take X = R3, Y = R2 and consider E2, E3, C and S to be uniform distributions on a 2D and 3D ellipse, a square and a sphere. We consider mm-spaces of different dimensions to emphasize the ability of (U)GW to compare different spaces. Figure 1 contrasts the transportation plan obtained by GW and UGW for a fixed µ = 0.5E3 + 0.5S and ν obtained using two different mixtures of E2 and C. The black segments show the largest entries of the transportation matrix π, for a sub-sampled set of points (to ease visibility), thus effectively displaying the matching induced by the plan. Furthermore, the width of the dots are scaled according to the mass of the marginals π1 ≈ µ and π2 ≈ ν, i.e. the smaller the point, the smaller is the amount of transported mass. This figure shows that the exact conservation of mass imposed by GW leads to a poor geometrical matching of the shapes which have different global mass. As this should be expected, UGW recovers coherent matchings. We suspect the alternate minimization algorithm is able to find the global minimum in these cases.
Tightness of the bound CGW≤UGW We propose to approximate CGW by doing a similar alternate minimization as for UGW, as detailed in Appendix E. This numerical scheme does not scale to large problems, but allows us to explore numerically how tight is the upper bound UGW ≥ CGW. Figure 2 highlights the fact that in Euclidean space X = Y = Rd, this bound seems to be tight when the two measures are sufficiently close. We consider discrete measures µ = 1n ∑ i δxi in
X = Y = Rd and νt = 1n ∑ i δyi where yi = xi + t∆i where ∆i are random perturbations and
denote (X ,Yt) the two mm-spaces associated to the Euclidean distance. As t → 0, µ and νt get closer, we observe numerically that UGW ≈ CGW. Figure 3 considers random points (xi)i and (yi)i and displays the histograms of the ratio CGW/UGW for n = 3. This shows that while the bound CGW ≤ UGW seems not tight, the ratio appears to be bounded even for points not being close. This numerical experiment suggests that UGW and CGW are locally equivalent and that UGW is in practice an acceptable proxy of the distance CGW. We leave for future works a tighter analysis of the gap between UGW and CGW.
their similar features, while negative samples would be ignored due to dissimilar features that induce a laziness to transport mass and match them.
We consider PU learning over the Caltech office dataset used for domain adaptation tasks (with domains Caltech (C) Griffin et al. [2007], Amazon (A), Webcam (W) and DSLR (D) Saenko et al. [2010]). The Caltech datasets are represented with two embeddings based on Surf and Decaf features Saenko et al. [2010], Donahue et al. [2014]. On the latter datasets, we perform PU learning over similar features (e.g. surf-C→ surf-* or decaf-C→ decaf-*) and from one feature format to the other (e.g. surf-C→ decaf-* or surf-C→ decaf-*). Those features are projected via PCA to subspaces of dimension 10 for surf features and 40 for decaf features. In the last task, one cannot use standard PU-method, and to the best of our knowledge, Unbalanced-GW methods are the only approaches for PU learning across different domains/features.
The procedure is the following. We solve the PU learning problem by computing the optimal plan π for UGW(X ,Y). We compute its first marginal π2 on Y , and predict the labels of some y ∈ Y as `(y) , sign(π2(y) − q) where q is the quantile of π2 corresponding to the proportion r of positives samples in Y . Following Chapel et al. [2020] which is adapted from Kato et al. [2018], Hsieh et al. [2019], this proportion r is assumed to be known. We report the accuracy of the prediction over the same 20 folds of the datasets, and use 20 other folds to validate the parameters of UGW. We consider 100 random samples for each fold of (X,Y ), a ratio of positive samples r = 0.1 for domains (C,A,W,D), and a ratio r = 0.2 for domains (C,A,W).
Since the GW objective is non-convex, the initialization of the minimization algorithms is key to obtain good perfor-
mances. In Chapel et al. [2020] and our experiments, for tasks where Y and Y belong to the same Euclidean space, (e.g. surf-C→ surf-*) we initialize π with the Partial-Wasserstein (PW) solution with a squared Euclidean cost. For cross-domain prediction (e.g. surf-C→ decaf-*), following Chapel et al. [2020], PGW is initialized with a list of plans built using a coarsened representation of the data with k-NN. While Chapel et al. [2020] makes use in an oracle manner of the plan providing the best accuracy, we modified their protocol and keep the plan which has the lowest PGW cost, which seems fairer, hence the difference in performance with Chapel et al. [2020]. To initialize UGW when X and Y do not belong to the same Euclidean space, we use a UOT solution of a matching between distance
histograms called FLB Mémoli [2011]. We define FLB in our UGW setting as FLB(X ,Y) , min ∫ X×Y |µ̄ ? dX − ν̄ ? dY |2dπ+ ρKL(π1|µ) + ρKL(π2|ν) + εKL(π|µ⊗ ν), (7)
where µ ? dX(x) , ∫ dX(x, x
′)dµ(x′) is the eccentricity, i.e. a histogram of aggregated distances, and µ̄ = µ/m(µ). Contrary to GW Mémoli [2011], there is a priori no link between FLB and UGW.
In the experiments we slightly generalize UGW and use two different marginal penalties ρ1KL
⊗(π1|µ) + ρ2KL⊗(π2|ν) with two parameters (ρ1, ρ2) to take into account shifts between domains/features. Note that PGW has a single parameter (which plays a role similar to (ρ1, ρ2)) which controls the cost of mass creation/destruction. We set ε = 2−9, which avoids introducing an extra parameter in the method. The value (ρ1, ρ2) ∈ {2−k, k ∈ J5, 10K}2 are cross validated for each task on the validation folds, and we report the average accuracy on the testing folds. We discuss in Appendix E the impact of reducing the number of parameters on the performance. Comparison with other methods – PU and PUSB Kato et al. [2018], Du Plessis et al. [2014] – are provided in Chapel et al. [2020] and we focus here on the comparison with PGW only.
The results are reported in Table 1. We display the performance of PGW, UGW and the initialization used for UGW to guarantee that using UGW does improve the performance. We observe that when the source and target dataset is the same (C→C tasks), the PW initialization performs better and PGW/UGW degrade the performance, so that in this setting Optimal Transport should be preferred over GW, which is to be expected. However when the domains are different, applying UGW improves the performance over the initialization (which is FLB) in almost all tasks. Note that in that case the methods PU, PUSB or PW cannot be used. Overall, this shows that GW methods are able to solve to some extent the PU learning problem across different spaces, and that using a “softer” KL penalties in UGW is at least at par with Partial GW, and performs better in some settings.

heading: 5 Conclusion and perspectives

text: This paper defines two Unbalanced Gromov-Wasserstein formulations: CGW and UGW. We prove that they are both positive and definite. We provide a scalable, GPU-friendly algorithm to compute UGW illustrate its applicability in learning tasks, and show that CGW is a distance between mmspaces up to isometry. These divergences and distances allow for the first time to blend in a seamless way the transportation geometry of GW with creation and destruction of mass. This hybridization is the key to unlock both theoretical and practical issues. This work opens new questions for future works, for instance removing the bias introduced by the use of entropic regularization, which is important for applications to ML. Note that such a debiasing was successfully applied for BalancedGW in Bunne et al. [2019] and is shown to lead to a valid divergence for balanced OT in Feydy et al. [2019b] and UW in Séjourné et al. [2019]. The design of efficient numerical solvers for CGW is also an interesting avenue for future works, as well as the study of its induced topology.

heading: Acknowledgements

text:
The works of Thibault Séjourné and Gabriel Peyré is supported by the ERC grant NORIA. The work of G. Peyré was supported in part by the French government under management of Agence Nationale de la Recherche as part of the "Investissements d’avenir" program, reference ANR19-P3IA-0001 (PRAIRIE 3IA Institute).
The authors thank Rémi Flamary for his remarks and advices, as well as Laetitia Chapel for her help to reproduce her experiments.

title: The Unbalanced Gromov Wasserstein Distance: Conic Formulation and Relaxation


summaries: The submission deals with eliminating neurons in a network where either a) all the input connections xor b) all the output connections have been pruned. When this is the case, the unpruned a) output or b) input connections are unused and can also be pruned: and the freed parameter budget used for other more useful connections. This is shown to improve the accuracy of pruned networks at a given sparsity ratio, especially for very high levels of sparsity.

summaries: The paper proposes to remove dead neurons and their connected parameters through a very simple check while reviving pruned (salient) parameters up to the prespecified sparsity level, such that the sparse network obtained could perform better. The main (and perhaps the single major) contribution of this work is in its demonstration that such a simple method is indeed effective for different pruning methods on various network architectures and datasets. The proposed method (AAP) can perhaps be considered as a generic post-processing step that could be equipped to any pruning method leaving dead neurons.

summaries: This work proposes a novel pruning method, called all-alive pruning (AAP), which is a general technique to remove dead connections from pruned neural networks. AAP is broadly applicable to various saliency-based pruning methods and model architectures. AAP equipped with existing pruning methods consistently improves the accuracy of original methods on three benchmark datasets.

abstractText: Network pruning has been widely adopted for reducing computational cost and memory consumption in low-resource devices. Recent studies show that saliencybased pruning can achieve high compression ratios (e.g., 80–90% of the parameters in original networks are removed) without sacrificing much accuracy loss. Nevertheless, finding the well-trainable networks with sparse parameters (e.g., < 10% of the parameters remaining) is still challenging to network pruning, commonly believed to lack model capacity. In this work, we revisit the procedure of existing pruning methods and observe that dead connections, which do not contribute to model capacity, appear regardless of pruning methods. To this end, we propose a novel pruning method, called all-alive pruning (AAP), producing the pruned networks with only trainable weights. Notably, AAP is broadly applicable to various saliency-based pruning methods and model architectures. We demonstrate that AAP equipped with existing pruning methods (i.e., iterative pruning, one-shot pruning, and dynamic pruning) consistently improves the accuracy of original methods at 128×–4096× compression ratios on three benchmark datasets.

authors:
Sanjeev Arora
Rong Ge
Behnam Neyshabur
Yi Zhang

title: Stronger generalization bounds for deep nets via a compression approach

venue: In International Conference on Machine Learning,

authors:
Davis W. Blalock
Jose Javier Gonzalez Ortiz
Jonathan Frankle
John V. Guttag

title: What is the state of neural network pruning

authors:
Sutskever
Dario Amodei

title: Language models are few-shot learners

title: Carreira-Perpiñán and Yerlan Idelbayev. "learning-compression" algorithms for neural net pruning

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

title: A back-propagation algorithm with optimal use of hidden units

venue: In Advances in Neural Information Processing Systems, pp

authors:
Utku Evci
Trevor Gale
Jacob Menick
Pablo Samuel Castro
Erich Elsen

title: Rigging the lottery: Making all tickets winners

authors:
Jonathan Frankle
Michael Carbin

title: The lottery ticket hypothesis: Finding sparse, trainable neural networks

authors:
Trevor Gale
Erich Elsen
Sara Hooker

title: The state of sparsity in deep neural networks

authors:
Yiwen Guo
Anbang Yao
Yurong Chen

title: Dynamic network surgery for efficient dnns

authors:
Song Han
Jeff Pool
John Tran
William J. Dally

title: Learning both weights and connections for efficient neural network

venue: In Neural Information Processing Systems

authors:
Babak Hassibi
David G. Stork
Gregory J. Wolff

title: Optimal brain surgeon and general network pruning

venue: In International Conference on Neural Networks,

authors:
Kaiming He
Xiangyu Zhang
Shaoqing Ren
Jian Sun

title: Deep residual learning for image recognition

title: Structural learning with forgetting

venue: Neural Networks,

title: A simple procedure for pruning back-propagation trained neural networks

venue: IEEE Trans. Neural Networks,

authors:
Yann LeCun
John S. Denker
Sara A. Solla

title: Optimal brain damage

authors:
Yann LeCun
Léon Bottou
Yoshua Bengio
Patrick Haffner

title: Gradient-based learning applied to document recognition

venue: Proceedings of the IEEE,

authors:
Namhoon Lee
Thalaiyasingam Ajanthan
Philip H.S. Torr

title: SNIP: single-shot network pruning based on connection sensitivity

authors:
Hao Li
Asim Kadav
Igor Durdanovic
Hanan Samet
Hans Peter Graf

title: Pruning filters for efficient convnets

authors:
Ning Liu
Xiaolong Ma
Zhiyuan Xu
Yanzhi Wang
Jian Tang
Jieping Ye

title: Autocompress: An automatic DNN structured pruning framework for ultra-high compression rates

venue: In AAAI Conference on Artificial Intelligence,

authors:
Zhuang Liu
Mingjie Sun
Tinghui Zhou
Gao Huang
Trevor Darrell

title: Rethinking the value of network pruning

authors:
Christos Louizos
Max Welling
Diederik P. Kingma

title: Learning sparse neural networks through l_0 regularization

authors:
Decebal Constantin Mocanu
Elena Mocanu
Peter Stone
Phuong H. Nguyen
Madeleine Gibescu
Antonio Liotta

title: Scalable training of artificial neural networks with adaptive sparse connectivity inspired by network science

venue: Nature Communications,

authors:
Dmitry Molchanov
Arsenii Ashukha
Dmitry P. Vetrov

title: Variational dropout sparsifies deep neural networks

authors:
Pavlo Molchanov
Stephen Tyree
Tero Karras
Timo Aila
Jan Kautz

title: Pruning convolutional neural networks for resource efficient inference

authors:
Hesham Mostafa
Xin Wang

title: Parameter efficient training of deep convolutional neural networks by dynamic sparse reparameterization

authors:
Michael Mozer
Paul Smolensky

title: Skeletonization: A technique for trimming the fat from a network via relevance assessment

authors:
Sharan Narang
Greg Diamos
Shubho Sengupta
Erich Elsen

title: Exploring sparsity in recurrent neural networks

authors:
Sejun Park
Jaeho Lee
Sangwoo Mo
Jinwoo Shin

title: Lookahead: A far-sighted alternative of magnitude-based pruning

title: Pruning algorithms-a survey

authors:
Alex Renda
Jonathan Frankle
Michael Carbin

title: Comparing rewinding and fine-tuning in neural network pruning

authors:
Mark Sandler
Andrew G. Howard
Menglong Zhu
Andrey Zhmoginov
Liang-Chieh Chen

title: Mobilenetv2: Inverted residuals and linear bottlenecks

authors:
Mingxing Tan
Quoc V. Le

title: Efficientnet: Rethinking model scaling for convolutional neural networks

authors:
Hugo Touvron
Andrea Vedaldi
Matthijs Douze
Hervé Jégou

title: Fixing the train-test resolution discrepancy: Fixefficientnet

authors:
Chaoqi Wang
Guodong Zhang
Roger B. Grosse

title: Picking winning tickets before training by preserving gradient flow

authors:
Andreas S. Weigend
David E. Rumelhart
Bernardo A. Huberman

title: Generalization by weightelimination with application to forecasting

authors:
Michael Zhu
Suyog Gupta

title: To prune, or not to prune: Exploring the efficacy of pruning for model compression

venue: In International Conference on Learning Representations (Workshop),

authors:
pruning
SNIP (Lee

heading: 1 INTRODUCTION

text:
The state-of-the-art neural networks have shown remarkable performance gains on various downstream tasks such as computer vision, natural language processing, and speech recognition. Because neural networks are typically overparameterized, they require high computational cost and memory consumption. Such a nature inherently hinders the deployment of models with excessive parameters on low-end devices such as mobile, embedded, and on-device systems.
Network pruning (Reed, 1993) is the prevalent technique to compress high-capacity models by removing unnecessary units such as weights/filters while maintaining the performance with minimal accuracy loss. Existing pruning methods can be divided into two categories. 1) The first group enforces the sparsity as model regularization (Chauvin, 1988; Weigend et al., 1990; Ishikawa, 1996; Molchanov et al., 2017a; Carreira-Perpiñán & Idelbayev, 2018; Louizos et al., 2018). It is theoretically well-investigated and does not require network retraining. 2) Another group develops saliency criteria to prune less important units (Mozer & Smolensky, 1988; LeCun et al., 1989; Karnin, 1990; Hassibi et al., 1993; Han et al., 2015; Guo et al., 2016; Lee et al., 2019; Park et al., 2020; Evci et al., 2019).
Because of its simple operation and outstanding pruning performance, magnitude pruning (MP) (Han et al., 2015; Narang et al., 2017; Zhu & Gupta, 2018) is the most popular saliency-based pruning method. Recently, the effectiveness of MP is highlighted by the success of the lottery ticket hypothesis (Frankle & Carbin, 2019) and learning rate rewinding (Renda et al., 2020), which achieves less than 1% accuracy loss even after pruning 90% of the parameters. However, all pruning methods show that the trade-off between sparsity and accuracy is significantly degraded, especially at the extremely high sparsity (Gale et al., 2019; Liu et al., 2019; Blalock et al., 2020). Considering the explosive increase in the size of state-of-the-art models for downstream tasks (e.g., GPT-3 (Brown et al., 2020) for machine translation and FixEfficientNet-L2 (Touvron et al., 2020) for image classification), the effective pruning methods at extreme compression ratios must be accomplished, and it is particularly crucial for adopting high-capacity models to low-resource devices (e.g. One Laptop per Child).
To break the performance bottleneck at high compression ratios of the state-of-the-art pruning methods, we investigate the procedure of existing pruning methods. Surprisingly, it is revealed that all existing studies overlook the existence of dead neurons after pruning – dead neurons are the nodes
20
40
60
80
100
32 × 64 × 12 8× 25 6× 51 2× 10 24 ×
0
20
40 60 T es t A cc u ra cy ( % )
% o
f d ea
d c
o n n ec
ti o n s
Compression ratio
LeNet-300-100 MNIST
20
40
60
80
100
12 8× 25 6× 51 2× 10 24 × 20 48 × 40 96 ×
0
20
40
60
T es
t A
cc u ra
cy (
% )
% o
f d ea
d c
o n n ec
ti o n s
Compression ratio
ResNet-18 CIFAR-10
0
15
30
45
60
32 × 64 × 12 8× 25 6× 51 2× 10 24 ×
0
20
40
60
T es
t A
cc u ra
cy (
% )
% o
f d ea
d c
o n n ec
ti o n s
Compression ratio
ResNet-50 Tiny-ImageNet
IMP Accuracy % of dead connections
Figure 1: Model accuracy and percentage of dead connections under varying compression ratios: LeNet-300-100 on MNIST, ResNet-18 on CIFAR-10, and ResNet-50 on Tiny-ImageNet. We used IMP with learning rate rewinding (Renda et al., 2020) as the base pruning method. The percentage of dead connections is calculated by the number of dead connections divided by the total number of connections in a given network.
with no input connections or output connections. That is, dead neurons make all their connected weights useless, called dead connections. Although Han et al. (2015) have already raised the potential issue of dead neurons, they anticipated that the dead neurons could be negligible after multiple iterative pruning. Unfortunately, we find that this problem is not as simple as their expectation.
Figure 1 depicts prediction accuracy (green line) and the percentage of dead connections (red line) over various compression ratios. For this empirical study, iterative magnitude pruning (IMP) with learning rate rewinding (Renda et al., 2020) is employed as the baseline pruning method. Then, we analyze the correlation between pruning accuracy and the occurrence of dead neurons on three benchmark datasets. We discover two meaningful findings: 1) The dead neurons cannot be successfully removed by iterative pruning, especially at high compression ratios. 2) The severe performance degradation at high compression ratios is strongly correlated with the number of dead connections, i.e., they are inversely proportional to each other.
Based on these valuable findings, we propose all-alive pruning (AAP), which improves the pruning performance by effectively eliminating dead connections. Specifically, we search dead neurons at the pruning stage by inspecting their gradient flows – when zero gradients are passing through the node, we regard them as dead neurons. Once identified, dead neurons and corresponding dead connections (any weights linked from or to dead neurons) are removed together during the pruning procedure. Note that detecting gradient-based dead neurons can be applied for complex model architectures with skip connections. If we remove more weights than a sparsity threshold, we can revive some weights with the highest saliency scores. As a result, AAP constitutes all-alive subnetwork, i.e., all connections in the subnetwork are kept as trainable weights.
To summarize, the key advantage of AAP is two-fold: 1) AAP is versatile – it is broadly applicable to various saliency-based pruning methods and model architectures by minimizing the loss of prediction accuracy. 2) AAP consistently improves the accuracy of the original pruning methods at high compression ratios and breaks the state-of-the-art performance on three benchmark datasets (i.e., MNIST, CIFAR-10, and Tiny-ImageNet).

heading: 2 RELATED WORK

text:
Network pruning utilizes the highly overparameterized nature of modern neural networks to compress the model. Also, it provides meaningful insight into what leads to the success of neural networks; small subnetworks in the original network can achieve comparable performance and improve generalization effects (Arora et al., 2018).
Recent advances. In general, network pruning can be categorized into two groups. The first approach employs the loss function with regularization terms to enforce the sparsity (Chauvin, 1988; Weigend et al., 1990; Ishikawa, 1996). Also, Molchanov et al. (2017a) proposed variational dropout to produce
highly sparse networks and Louizos et al. (2018) reparameterized network weights as the product of weight and a stochastic gate variable. Carreira-Perpiñán & Idelbayev (2018) attempted to minimize the loss using the stochastic version of projected gradient descent.
The second approach utilizes the sensitivity of the loss for weights (LeCun et al., 1989; Karnin, 1990; Hassibi et al., 1993). For saliency-based criteria, Han et al. (2015) proposed iterative pruning based on the magnitude of weights. Despite its simplicity, magnitude pruning shows significant performance gain at high compression ratios (Guo et al., 2016; Narang et al., 2017; Zhu & Gupta, 2018). Recently, Park et al. (2020) proposed lookahead pruning (LAP) that considers the neighbors of connections. While iterative pruning requires pre-trained networks, SNIP (Lee et al., 2019) and GraSP (Wang et al., 2020) proposed to prune the network at initialization. Also, dynamic pruning methods (e.g., SET (Mocanu et al., 2018), DSR (Mostafa & Wang, 2019), and RigL (Evci et al., 2019)) adjust the sparsity level with prune-and-grow cycles. Comprehensive analysis and experimental results for network pruning can be found (Gale et al., 2019; Liu et al., 2019; Blalock et al., 2020).
Detecting dead connections in network pruning. Pruning neurons/weights may result in unwanted side-effects, producing useless connections in the pruned networks. In structured pruning, Li et al. (2017) and Liu et al. (2020) attempted to eliminate dead neurons by merely considering the connections between adjacent filters. However, it is not applicable for complex networks with shortcut connections. In unstructured pruning, Han et al. (2015) discussed the existence of dead weights but neglected them. Evci et al. (2019) proposed dynamic pruning, which iteratively updates the weights to identify more useful weights under given computing constraints. Although dynamic pruning eliminates the connections with relatively low importance, it still cannot correctly handle dead connections in each pruning stage. In contrast, this paper devises a new pruning method that removes dead connections and keeps all trainable connections in general networks, including shortcut connections and batch normalization.

heading: 3 PROPOSED METHOD: ALL-ALIVE PRUNING (AAP)

text:
Given a dataset D = {x(i), y(i)}ni=1 and a sparsity threshold τ (i.e., the number of non-zero weights), unstructured pruning can be formulated as the constrained optimization problem.
min c,w L(c w;D) = min c,w
1
n n∑ i=1 `(c w;D), s.t. c ∈ {0, 1}m and ||c||0 ≤ τ, (1)
where m is the number of parameters, and n is the number of samples in D. Also, `(·) is the loss function, is the Hadamard product, and || · ||0 is the L0 norm. c ∈ {0, 1}m is the indicator variable to determine whether or not w is pruned. In other words, cj indicates whether wj is active (cj = 1) or inactive (cj = 0).
In this work, we aim at improving saliency-based pruning, where the importance of weights in the network is measured by a pre-defined saliency criterion. Popular criteria include the magnitude of weights (Han et al., 2015), the sensitivity of weights (Lee et al., 2019), and the neighboring connections of weights (Park et al., 2020). Given a saliency criterion, we calculate the importance of weight wj as a saliency score sj ; the higher the saliency score, the more important the weight is. Then, each element of c is expressed as cj = 1 if sj ∈ S and cj = 0 otherwise, where j ∈ {1, · · · ,m} and S is a set of τ -largest scores associated with w. That is, τ connections are kept in the original network, and the rest of the connections are removed from the network.

heading: 3.1 LIMITATIONS OF EXISTING SALIENCY-BASED PRUNING

text:
Depending on training efficiency, prior studies in saliency-based pruning utilize either one-shot pruning or iterative pruning. One-shot pruning solely erases selected connections at once and retrains the network. Meanwhile, iterative pruning alternates pruning and retraining cycles by gradually increasing the sparsity levels at multiple steps. Based on this process, iterative pruning can achieve higher accuracy than one-shot pruning for highly compressed networks.
However, pruning weights can lead to unexpected changes in network architecture, and it affects the importance of weights. Although the pruned network can constitute the most useful subnetwork in
terms of selected criteria, the subnetwork may include useless connections, called dead connections, as a result of pruning.
To describe the emergence of dead connections, we consider a simple fully-connected network. Here, w
(l) ij is a weight connecting two nodes a (l−1) i and a (l) j between the (l− 1)-th layer and the (l)-th layer.
Based on a backpropagation rule, the gradient of w(l)ij is rewritten as follows:
∂L(w;D) ∂w
(l) ij
= ∂a
(l) j
∂w (l) ij
∂L(w;D) ∂a
(l) j
= a (l−1) i δ (l) j and a (l) j = ψ ( d∑ i=1 w (l) ij a (l−1) i + b (l) j ) , (2)
where δ(l)j = ∂L(w;D)/∂a (l) j , and ψ(·) is the activation function. Also, a (l) j and b (l) j are an activation and a bias of the j-th node at the (l)-th layer.
After pruning, the updated network can possibly have a dead neuron which appears when the neuron has zero input connections or zero output connections, i.e., either a(l−1)i or δ (l) j is zero:
1) No input connections: When a(l−1)i = 0, a (l−1) i is apparently a dead neuron. According to Eq. (2), the gradient of all output weights w(l)i∗ are also zeros, and thus can be marked as dead connections. Note that, when a(l−1)i = b l−1 i , a (l−1) i is also regarded as a dead neuron since we can reestablish the network by transferring the bias b(l−1)i of a (l−1) i to the bias b (l) j of a (l) j .
2) No output connections: When δ(l)j = 0, all input weights w (l) ∗j connected to a (l) j and the biases b l j can be removed.
When the dead neurons are detected, we can also eliminate all useless connections, including the biases and the weights for batch normalization, associated with dead neurons. Note that the dead neurons in the fully-connected layer can be extended as filters in the convolutional networks.

heading: 3.2 BUILDING ALL-ALIVE SUBNETWORK

text:
The key idea of our pruning method, namely all-alive pruning (AAP), is to build all-alive subnetwork during the pruning procedure. Toward this goal, we first sort all weights by a saliency-based criterion. We highlight that our pruning method can be applied to various saliency-based criteria. Once the weights with τ -highest scores are determined, we check whether dead connections appear as a side effect of removing the connections. That is, additional dead connections lead to overly pruned networks. To meet the sparsity threshold τ , we can revive some connections with the highest saliency scores among previously pruned connections and update c.
Specifically, the overall process of AAP consists of the following three phases. (In Appendix, Algorithm 2 is the pseudo-code of detecting dead connections in the pruning process.)
1) Eliminating dead connections: Choose a set of high-impacted τ weights according to a saliencybased criterion, and eliminate the rest weights as well as dead connections that are associated with pruned weights.
2) Reviving new connections: If the more number of weights are eliminated than a sparsity threshold, we restore some connections with the highest saliency scores, where the restored connections have their original weight values.
3) Repeat steps 1–2 until the pruned subnetwork is converged to the all-alive subnetwork.
Figure 2 depicts AAP equipped with magnitude pruning (MP). When we prune four weights from the original network, we eliminate the least important connections (dashed lines). While existing pruning methods include dead connections, we eliminate all the incoming (or outgoing) weights of a dead neuron as dead connections (red color) and then revive alternative connections (green color). Finally, all remaining weights are active, thereby the pruned network being all-alive. Please refer to Algorithm 2 for the detailed procedure of AAP in the Appendix.
We observe that AAP has two key advantages. Firstly, all weights in the subnetwork are kept alive without dead connections. That is, we utilize all connections in the subnetwork. Owing to this
compact pruning, we can further compress networks; this is equivalent to improve the prediction accuracy under fixed compression ratios, particularly effective for pruning the networks at extremely high compression ratios. Secondly, since AAP eliminates all the weights associated with dead neurons, AAP results in the memory layout similar to those of structured pruning methods. As a positive byproduct of AAP, it is more compatible with modern memory architectures.

heading: 4 EXPERIMENTAL RESULTS

text:
We evaluate AAP with a variety of network architectures on MNIST, CIFAR-10, and Tiny-ImageNet datasets. Concretely, we adopt LeNet-300-100 (LeCun et al., 1998), ResNet (He et al., 2016), MobileNetV2 (Sandler et al., 2018), and EfficientNet (Tan & Le, 2019). We employ iterative magnitude pruning (IMP) with learning rate rewinding (Renda et al., 2020), a state-of-the-art pruning method, as our base scheme of AAP. Also, we incorporate AAP into other saliency-based pruning methods such as SNIP (Lee et al., 2019), lookahead pruning (LAP) (Park et al., 2020), and dynamic pruning inspired by RigL (Evci et al., 2019).
Implementation details. We use the same environmental setting for all experiments as proposed by Blalock et al. (2020), including initialization values, hyper-parameters, and data augmentation settings. We use Adam optimizer, the batch size of 60, and the constant learning rate of 0.0003 with 50 training epochs for MNIST dataset. For CIFAR-10 dataset, we use the batch size of 128, and the learning rate is initially set to 0.1 and is decayed by 0.1 at every 30 epochs. For Tiny-ImageNet dataset, we use the batch size of 1024, and the learning rate linearly warmed up to its maximum value 0.4 at epoch 5 and decayed by 0.1 at epochs 30, 60, and 80. We use SGD with a momentum of 0.9 and a weight decay rate of 0.0001 for both datasets and train them for 120 epochs and 90 epochs, respectively. For all experiments, we use 10% of the training set as the validation set.
For IMP, we rewind the learning rate schedule to the beginning of the training and retrain from the final values of the weights. For iterative pruning, we prune 50% of the remaining parameters in each step. Exceptionally, since we use the constant learning rate for MNIST dataset, we use weight rewinding proposed by Frankle & Carbin (2019). For MNIST dataset, we re-initialize the pruned model’s weights to the value of the original model’s initial weights.
Compression ratio. To measure the degree of network pruning, we adopt a compression ratio, which is calculated by original model size / compressed size. For instance, when only 1% of the parameters remaining in the compressed network, the compression ratio is 100×. Note that we consider the number of all parameters, including the weights for batch normalization and biases.

heading: 4.1 EFFECT OF ALL-ALIVE PRUNING

text:
We employ IMP with learning rate rewinding (Renda et al., 2020) and weight rewinding (Frankle & Carbin, 2019) as the baseline pruning method due to their state-of-the-art performances. At each pruning iteration, we apply AAP after pruning the connections with the lowest magnitude scores. To
validate AAP can alleviate performance degradation at high compression ratios, we start to apply AAP to the point (gray colors in Figure 3), where the performance of pruned networks begins to be worse than the original model. On the basis of IMP, our pruning method, namely IMP-AAP, is compared with IMP in terms of the prediction accuracy under the same compression ratio to evaluate the effectiveness of AAP.
LeNet-300-100 on MNIST. LeNet-300-100 consists of the two-layer fully connected (FC) networks with 300 and 100 nodes. We compute the test accuracies of IMP and IMP-AAP at various compression ratios, starting from 16× to 1024×. Figure 3(a) clearly showcases the performance gains by adopting AAP, particularly noticeable as the compression ratio exceeds 64×. Specifically, IMP-AAP achieves gains of 10.92% at 512× and 32.25% at 1024×. From this result, we find two positive effects of AAP: 1) The impressive performance gains at the high compression ratios show the effectiveness of AAP. 2) The success of LeNet-300-100 represents that AAP is effective in fully connected layers.
ResNet-32 and ResNet-18 on CIFAR-10. We evaluate AAP on ResNet-32, and ResNet-18 for CIFAR-10 dataset. Figure 3(b) compares the test accuracies of IMP and IMP-AAP over the compression ratios from 8× to 512× on ResNet-32. Despite the rapid drop of IMP in performance after 128×, IMP-AAP slows down the degradation of IMP and always achieves higher accuracies than IMP under the same compression ratios. Figure 3(c) depicts the experimental results from ResNet-18. Interestingly, we observe the significant benefits by introducing AAP – IMP-AAP improves IMP by 6% at an extremely high compression ratio 2048×. The performance gaps are consistently increasing at higher compression ratios (e.g., greater than 512×). From both ResNet-32 and ResNet-18
experiments, it is consistently observed that the pruned network by IMP-AAP often discards the entire ResBlock except for the skip connection. We conjecture that this is possible because the skip connection allows using the outputs from the previous layers while achieving a high compression ratio. Overall, we confirm that AAP is also effective in compressing convolution layers and the complex structure even with skip connections.
ResNet-50, MobileNetV2, and EfficientNet-B0 on Tiny-ImageNet. For Tiny-ImageNet dataset, we employ ResNet-50, MobileNetV2, and EfficientNet-B0 as backbone networks. For ResNet-50, we observe that IMP-AAP improves the trade-off between sparsity and accuracy. When the compression ratio is 512×, IMP experiences a severe performance drop, while IMP-AAP improves IMP by 7%. In Figures 3(e) and 3(f), we evaluate AAP on MobileNetV2 and EfficientNet-B0 to validate whether or not AAP is still effective in the compact backbone models (i.e. having fewer parameters). Due to the highly efficient nature of MobileNetV2 and EfficientNet-B0, both IMP and IMP-AAP suffer from performance degradation when pruning at high compression ratios (e.g., greater than 16×). However, IMP-AAP still consistently achieves a meaningful improvement on top of IMP. Concretely, IMP-AAP consistently enjoys the performance improvements about 1-2% at the compression ratios greater than 8× for MobileNetV2. For EfficientNet-B0, the gain by IMP-AAP is also noticeable – IMP-AAP shows 2% improvement on test accuracy at 64× and 6% at 128×.

heading: 4.2 APPLICABILITY ON VARIOUS SALIENCY-BASED PRUNING

text:
To evaluate the versatility of AAP, we apply AAP on top of existing pruning methods and see whether or not AAP can enjoy the performance benefits of others. We re-implement other saliencybased pruning, referring to the PyTorch implementations, i.e., SNIP1 and LAP2. We apply the same criterion as the weights to the biases for SNIP and use global pruning for LAP. We also validate AAP in dynamic pruning setting, where the prune-and-grow strategy was introduced by RigL (Evci et al., 2019). Specifically, we prune k connections with the lowest magnitude weights then grow k connections with the highest magnitude gradients for every 1500 steps until 3/4 of the full training epochs. We decay k using a cosine update schedule, where the initial k is 30% of all connections in the network. Then, we apply AAP after each prune-and-grow cycle. Originally, RigL introduced the sparsity strategy for practical training, which assigns the target sparsity level at each layer for preserving low FLOPS. In this experiment, we prune and grow the weights globally, including biases
1https://github.com/mil-ad/snip 2https://github.com/alinlab/lookahead_pruning
and weights for batch normalization unlike RigL. We employ LeNet-300-100 for MNIST dataset and ResNet-18 for CIFAR-10 dataset as backbone networks.
For LeNet-300-100, we conduct comparisons on the basis of one-shot magnitude pruning (One-shot MP), SNIP, and lookahead pruning (LAP). For ResNet-18, the experiment of LAP is not conducted since our method considers all the parameters in the network, including those in skip connections and batch normalization; LAP does not fully consider skip connections and parameters associated with batch normalization. Instead, for ResNet-18, we choose dynamic pruning as a base pruning scheme and apply AAP at each prune-and-grow cycle. For one-shot MP, we apply weight rewinding on LeNet-300-100 and learning rate rewinding on ResNet-18.
Figure 4 depicts the efficiency of AAP over various saliency-based pruning. As expected, since one-shot pruning involves a single pruning stage, any one-shot pruning methods and those with AAP are less effective than IMP. Nevertheless, applying AAP clearly improves various saliency-based, one-shot pruning methods. For one-shot MP on LeNet-300-100 and ResNet-18, we achieve 10% of the accuracy improvement at 128× (from 82.47% to 93.11%) and 6% of the accuracy improvement at 512× (from 76.6% to 82.3%), respectively. Like one-shot MP, AAP incorporates with SNIP well for both LeNet-300-100 and ResNet-18. At 256×, AAP outperforms the original SNIP by 33% on LeNet-300-100 (from 55.48% to 88.47%) and 3% on ResNet-18 (from 75.85% to 82.03%). On LeNet-300-100, LAP seems to preserve the good performance even in the high compression ratio; nonetheless, AAP can improve the performance by 12% (81.73% to 93.58%) at 256×. Unlike SNIP or LAP, prune-and-grow policy in dynamic pruning leads to architectural changes in the network model. AAP is also successful in improving dynamic pruning; AAP achieves 15% improvements at 256× for dynamic pruning on ResNet-18. From these results, we confirm the efficacy of AAP independently of saliency-based pruning even with dynamic pruning.
In addition to the accuracy gain, Figure 5 compares the weight distribution of the pruned network between one-shot MP and the same one with AAP. While all active neurons from one-shot MP are spread arbitrarily, those from AAP tend to form clusters for more columns and rows. This is notable because AAP does not require explicit or implicit constraints to cluster the weights. We observe that AAP produces a small-dense subnetwork similar to structured pruning (Li et al., 2017; Molchanov et al., 2017b; Liu et al., 2020), which is more favorable for modern memory and computing architectures even with unstructured pruning.

heading: 5 CONCLUSION

text: While existing studies mainly focus on developing new saliency-based criteria or optimizing magnitude pruning methods, existing studies do not properly address the dead connections during the pruning process. In this paper, we propose a simple-yet-effective and versatile unstructured pruning method, namely all-alive pruning (AAP), to eliminate dead connections and make all weights in the subnetwork trainable. In experimental results, AAP consistently improves various saliency-based pruning methods with different model architectures at 128×–4096× compression ratios, achieving state-of-the-art performance on several benchmark datasets.

heading: A DETAILED ALGORITHMS OF THE PROPOSED METHOD

text:
Finding dead connections. Given a binary mask vector c ∈ {0, 1}m and the corresponding parameter w ∈ Rm, we find a dead connection vector d ∈ {0, 1}m, where d(l)ij is 1 if w (l) ij is a dead connection; d(l)ij is 0 otherwise. When the weight matrix is higher than two-dimensional, e.g.filters in CNN, the first two dimensions are regarded as input and output dimensions. To simplify the notation, gi∗ = 0 is same as gi:∗:∗:∗ = 0.
Algorithm 1 is the pseudo-code of finding dead connections. First of all, we preprocess a network to check the gradient of weights. In this preprocessing stage, the activation function is regarded as an identity function. Besides, we set all weights to be positive and ignore all bias terms to avoid false positive neurons (i.e., detecting the false dead neurons, although it is not dead neurons). After that, we update the weights by considering the given mask vector and compute the gradient g(·) of w by taking the all-one vector as input. After the dead connection vector is initialized with a zero vector, we update the gradient for dead connections.
Pseudo-code of all-alive pruning. Algorithm 2 is the pseudo-code of the all-alive pruning algorithm. Given a network, we compute the saliency score using an arbitrary saliency criterion, e.g.magnitude pruning, SNIP (Lee et al., 2019), and lookahead pruning (LAP) (Park et al., 2020). Then we update c using the saliency score, i.e.the weights with the lowest scores are set to 1. For the updated sub-network, we find the dead connections and recalculate the saliency scores accordingly. That is, if the weight is identified as a dead connection, its saliency score is set to 0. We repeat eliminating dead connections and making all-alive subnetwork by given saliency score until all connections are kept alive. Then we prune the network with the finally obtained mask vector.
Algorithm 1 Finding the dead connections Input: Binary mask c ∈ {0, 1}m, Binary dead connections d ∈ {0, 1}m, model parameters w ∈ Rm Output: Updated binary dead connections d ∈ {0, 1}m corresponding to w ∈ Rm
1: Preprocess: Make all weights positive, remove all biases. 2: w̃← w c. . Update the pruned network with c. 3: Compute the gradient g(w̃;1). . Update the gradient g for all-one input data. 4: for l = 1 to N do 5: d(l)ij ← 1 if g (l) i∗ = 0 or g (l) ∗j = 0. 6: end for
Algorithm 2 All-alive pruning Input: Model parameters w ∈ Rm, sparsity threshold τ Output: All-alive subnetwork w
1: Compute saliency s for w. . Use an arbitrary saliency criterion. 2: Initialize binary dead connections d ∈ {0, 1}m as zero. 3: repeat 4: S ← top-τ connections by s. 5: Initialize c as zero and update c for S. . Set top-τ connections in S as cj = 1. 6: Update dead connections d for w and c. . Refer to Algorithm 1. 7: for l = 1 to N do 8: s(l)ij ← 0 if d (l) ij = 1. . Update the score of dead connections as s (l) ij = 0.
9: end for 10: until all connections in c is alive. 11: w← w c

heading: B MODEL ARCHITECTURE DETAILS AND HYPERPARAMETER SETTINGS

text:
We modify the input and output layers of model architectures to handle a smaller size of input data. For CIFAR-10, we use a 3 × 3 filter with a stride of 1 at the first convolutional layer of ResNet. Because the image resolution of the Tiny-ImageNet dataset is larger than that of the CIFAR-10
dataset, we double the stride of the first convolutional layer. For MobileNetV2 and EfficientNet-B0, we use a stride of 1 instead of 2 at the first convolutional layer. Also, we adjust the number of output nodes for all networks according to the number of labels for each dataset.
We re-implement SNIP and LAP, referring to the PyTorch implementations. We apply the same criterion as the weights to the biases for SNIP. For LAP, we use global pruning (i.e.prune the entire network at once) instead of layer-wise pruning for fair comparisons with other pruning methods; the original LAP chooses layer-wise pruning. Whereas SNIP and LAP do not fully consider the biases and the weights for batch normalization, our AAP equally treats and removes all parameters in the network.

heading: C DETAILED EXPERIMENTAL RESULTS

text: We report numerical results depicted in Figure 3. Note that all the results are computed by averaging three trials with different random seeds. In all tables, subscripts denote standard deviations, and bracketed numbers indicate the accuracy gains obtained by the proposed model over the original IMP.

heading: D SPARSITY PATTERNS IN LENET-300-100

text:
We visualize the change of the pruned weights in LeNet-300-100 over compression ratios (from 64× to 256×). In all figures, each dot denotes the unpruned weight on the 300× 100 matrix. To observe the differences between pruning methods, we perform one-shot MP in the same network. As we expected, all-alive pruning behaves similar to structured pruning by removing dead neurons. Note
that the structured pruning makes the pruned sub-network to be small-dense. Similarly, AAP also produces the sub-network to have a small-dense structure because the dots are clustered along with either columns or rows – if the black bars along with the vertical or horizontal axis are scarce, the sub-network becomes small-dense. Interestingly, this tendency becomes prominent with an extremely high compression ratio (e.g.256×) as shown in Figure 7. In other words, AAP achieves 89.65% accuracy even with weights connected with 36 out of 300 neurons and 29 out of 100 neurons.

heading: E THEORETICAL STUDY OF ALL-ALIVE PRUNING

text:
E.1 CONVERGENCE OF ALL-ALIVE PRUNING
In this section, we study the convergence of the proposed all-alive pruning. In all iterations of AAP (Algorithm 2), we mask all dead connections found in the previous iteration to avoid reviving the dead connections. That is, the sub-network derived from one iteration of AAP should have a smaller number of connections than the previous iteration, i.e., the number of connections in the sub-network excluding dead connections for t iteration is followed as:
0 < n(θt) = n(θt−1)− n(Dτt ) < n(θt−1), (3)
where θt and Dτt are the parameters in the sub-network and the found dead connections respectively at t iteration of AAP. Since the number of connections considered for each iteration decreases, AAP can always converge. It is observed that the 2-3 iterations are repeated in most cases.
E.2 ELIMINATING DEAD CONNECTIONS FROM LOWER COMPRESSION RATIOS
Without knowing the dependency of all connections, finding the optimal sub-network with the highest saliency scores without any dead connections is hard. All-alive pruning is proposed to obtain a more efficient sub-network with the existing saliency criteria; nevertheless, it also cannot guarantee the sub-network with the highest scores. In this section, we deal with the more efficient sub-network by applying AAP incrementally from lower compression ratios (higher sparsity thresholds). With the given sparsity threshold τ , the total saliency scores of the sub-network obtained by AAP is followed as:
Sτ = ∑ TopK(|θ −Dτ |, τ), (4)
where TopK(v, k) indicates the top-k elements of set v. Note that the saliency scores of the dead connections could be considered as zero as they cannot properly join the regular training. If we eliminate the dead connections for τ ′ > τ sparsity threshold first, then pruning up to τ again, the dead connections found for the final sub-network is changed as:
Dτ ′→τ = Dτ+n(D τ′ ) ⊆ Dτ , (5)
since the dead connections found by the previous iteration decrease the removal of the connections of the sub-network needed for satisfying the sparsity threshold, in the same manner with the motivation of AAP. It is obvious that the dead connections Dτ+n(D τ′ ) caused by fewer pruned connections than
τ is a subset of Dτ . The total saliency scores of the sub-network Sτ ′→τ with the corresponding dead connections Dτ ′→τ are changed to:
Sτ ′→τ = ∑ TopK(|θ −Dτ ′→τ |, τ)
= ∑ TopK(|θ −Dτ+n(D τ′ )|, τ)
≥ Sτ = ∑ TopK(|θ −Dτ |, τ) ∵ Dτ+n(D τ′ ) ⊆ Dτ .
(6)
According to Eq. (6), we can see that the dead connections are eliminated in smaller units, the higher the total saliency scores can be obtained; eliminating dead connections whenever prune each connection can obtain the sub-network with much higher saliency scores. However, this does not mean that we can get an optimal network with the highest saliency scores, e.g., the connection with the lowest saliency score can still be needed to alive, and the existing saliency-based pruning methods do not allow to make it alive.
To validate the effectiveness of saving the saliency scores by applying AAP incrementally from the lower compression ratios (higher sparsity thresholds), we conducted the experiments in one-shot magnitude pruning. We satisfied the sparsity threshold τ over several pruning steps without any training process, but only with eliminating newly appearing dead connections. Figure 8 depicts the experiments of the AAP with smaller units. For example, AAP(1%) denotes that we repeat the process of 1) prune n(θ)−τ100 connections, then 2) eliminate the corresponding dead connections 100 times for satisfying the sparsity threshold τ . Interestingly, we can observe that the sub-network with the higher total saliency scores can achieve better accuracy at high compression ratios. Note that the existing pruning always outputs the same sub-network even with incremental pruning since they eliminate the connections with the lowest saliency scores without considering the effect of pruning on other connections.