# Robust Network Pruning With Sparse
Entropic Wasserstein Regression

Lei You
Department of Engineering Technology
Technical University of Denmark
Ballerup, DK-2750, Denmark
<EMAIL_ADDRESS>
&Hei Victor Cheng
Electrical and Computer Engineering Department
Aarhus Universrity
Aarhus, DK-8200, Denmark
<EMAIL_ADDRESS>Corresponding Author.

###### Abstract

This study unveils a cutting-edge technique for neural network pruning that
judiciously addresses noisy gradients during the computation of the empirical
Fisher Information Matrix (FIM). We introduce an entropic Wasserstein
regression (EWR) formulation, capitalizing on the geometric attributes of the
optimal transport (OT) problem. This is analytically showcased to excel in
noise mitigation by adopting neighborhood interpolation across data points.
The unique strength of the Wasserstein distance is its intrinsic ability to
strike a balance between noise reduction and covariance information
preservation. Extensive experiments performed on various networks show
comparable performance of the proposed method with state-of-the-art (SoTA)
network pruning algorithms. Our proposed method outperforms the SoTA when the
network size or the target sparsity is large, the gain is even larger with the
existence of noisy gradients, possibly from noisy data, analog memory, or
adversarial attacks. Notably, our proposed method achieves a gain of 6%
improvement in accuracy and 8% improvement in testing loss for MobileNetV1
with less than one-fourth of the network parameters remaining.

##  1 Introduction

The advent of deep learning has revolutionized various domains of artificial
intelligence, with neural networks showing remarkable performance across an
array of applications. Nonetheless, the increase in model complexity has led
to escalating computational demands and substantial memory requirements. This
poses significant challenges for deploying these models in resource-
constrained environments such as mobile or internet of things (IoT) devices.
Therefore, the concept of neural network pruning emerges as a critical
solution. It aims to optimize the network by removing less important
parameters, which reduces computational overhead while maintaining the
performance of the original model.

In the realm of state-of-the-art deep learning, the models often exhibit
substantial size and complexity, with up to trillions of parameters, as
exemplified by models such as GPT-4. The immense computational demand, energy
inefficiency, and the challenges with model interpretability associated with
these models highlight the need for innovative and efficient optimization
techniques. These techniques should ideally minimize the model size while
improving their robustness and interpretability. Taking into account the
limitations of previous work, especially those arising from the influence of
noisy gradients, the paper proposes a promising pathway for robust pruning.

Below, we inspect the network pruning problem from an optimization
perspective, with a concise introduction of the most relevant existing works.
A sketch of our approach is given following a brief discussion of the
motivations.

Related Work on Pruning as Optimization. Denote by
$\mathbf{\bar{w}}\in\mathbb{R}^{p}$ a trained model and
$\mathcal{L}(\mathbf{w})$ the loss function given arbitrary model
$\mathbf{w}$. The loss function can be locally approximated around
$\mathbf{\bar{w}}$ with Taylor Expansion as shown in equation 1.

$\mathcal{L}(\mathbf{w})=\mathcal{L}(\mathbf{\bar{w}})+\nabla\mathcal{L}(\mathbf{\bar{w}})^{\top}(\mathbf{w}-\mathbf{\bar{w}})+\frac{1}{2}(\mathbf{w}-\mathbf{\bar{w}})^{\top}\nabla^{2}\mathcal{L}(\mathbf{\bar{w}})(\mathbf{w}-\mathbf{\bar{w}})+O(\left\lVert\mathbf{w}-\mathbf{\bar{w}}\right\rVert^{3})$
(1)

Consider a neural network with a loss function
$\mathcal{L}(\mathbf{w})=\frac{1}{N}\sum_{i=1}^{N}{\bm{\ell}}_{i}(\mathbf{w})$,
where $\bm{\ell}_{i}(\mathbf{w})\in\mathbb{R}^{p}$ is the loss incurred at
data point $i$ ($i=1,\ldots,N$). The goal of network pruning is to find a set
of $\mathbf{w}$ such that there are $k$ ($k<p$) elements of $\mathbf{w}$ being
zero while keeping the newly obtained model $\mathbf{w}$’s performance as good
as possible to the original one $\mathbf{\bar{w}}$. Mathematically, we want to
find some $\mathbf{w}\in\mathbb{R}^{p}$ that satisfies both
$\mathcal{L}(\mathbf{w})\approx\mathcal{L}(\mathbf{\bar{w}})$ and
$\left\lVert\mathbf{w}\right\rVert_{0}\leq k$, with $k<p$.

This line of research can be dated back to LeCun et al. (1989), where the
approximation in equation 1 is adopted. Under the assumption that gradient
$\nabla\mathcal{L}(\mathbf{\bar{w}})\approx 0$ when the network is trained,
the network weights are pruned one-by-one in decreasing order based on the
value of
$(\mathbf{w}-\mathbf{\bar{w}})^{\top}\mathbf{H}(\mathbf{w}-\mathbf{\bar{w}})$.
In their approach, the $\mathbf{H}$ is approximated as a diagonal matrix, this
is later extended in Hassibi & Stork (1992) to include the whole Hessian
matrix and authors also proposed using the Fisher information matrix (FIM) as
an approximation to the Hessian. Later Singh & Alistarh (2020) proposed to
reduce the computation complexity by using block diagonal Hessian, and the FIM
is approximated using a small subset of the training data.

These approaches all use equation 1 to prune the network in a one-by-one
manner, namely the weight with the least importance are set to zero according
to the different approximations of equation 1. In this way, the potential
interactions of pruning multiple weights are ignored. To explore this, the
network pruning problem is formulated as a mixed integer quadratic programming
(MIQP) in Yu et al. (2022). Namely, an objective function

$f(\mathbf{w})=(\mathbf{w}-\mathbf{\bar{w}})^{\top}\mathbf{H}(\mathbf{w}-\mathbf{\bar{w}})+\lambda\left\lVert\mathbf{w}-\mathbf{\bar{w}}\right\rVert^{2}\text{
($\lambda\geq 0$)}$ (2)

is minimized, and Hessian is approximated as
$\mathbf{H}\approx\nabla^{2}\mathcal{L}(\mathbf{\bar{w}})$, subject to the
sparsity constraint $\left\lVert\mathbf{w}\right\rVert_{0}\leq k$, where
$\lambda$ is a regularization parameter. While this approach shows significant
improvements over the previous work, it suffers from scalability issues as it
requires the full Hessian matrix.

Sparse Linear Regression (LR) Formulation. To reduce the computational
complexity, the Hessian matrix can be approximated by the empirical FIM, using
$n$ samples as in Chen et al. (2022); Benbaki et al. (2023). Denote
$\mathbf{G}=[\nabla{\bm{\ell}}_{1},\ldots,{{\nabla\bm{\ell}}}_{n}]^{\top}\in\mathbb{R}^{n\times
p}$, where $\nabla\bm{\ell}_{i}=\nabla\bm{\ell}_{i}(\mathbf{\bar{w}})$. For
simplicity, $\nabla\bm{\ell}_{i}$ is used to represent the derivative of the
data point $i$’s loss at $\mathbf{\bar{w}}$ consistently in this paper unless
specified otherwise. The Hessian is approximated through the expression
$\mathbf{H}\approx(1/n)\sum_{i=1}^{n}{\nabla\bm{\ell}}_{i}{\nabla\bm{\ell}}_{i}^{\top}=(1/n)\mathbf{G}^{\top}\mathbf{G}$,
which is the so-called FIM. Denote
$x_{i}={\nabla\bm{\ell}}_{i}^{\top}\mathbf{w}$ and
$y_{i}={\nabla\bm{\ell}}_{i}^{\top}\mathbf{\bar{w}}$, equation 2 is formulated
to the sparse LR problem shown in equation 3 below.

$\min_{\mathbf{w}}\bar{Q}(\mathbf{w})=\sum_{i=1}^{n}\left\lVert
x_{i}(\mathbf{w})-y_{i}\right\rVert^{2}+n\lambda\left\lVert\mathbf{w}-\mathbf{\bar{w}}\right\rVert^{2},\text{
s.t. }\|\mathbf{w}\|_{0}\leq k$ (3)

This formulation has an advantage in computational cost, as empirical FIM
needs not to be computed explicitly. It is shown that the formulation scales
to large neural network pruning Chen et al. (2022).

Motivation of Combating Against Noise. In practice, it is not always easy to
obtain the correct gradients for pruning large neural networks. There can be
noise contained in the data samples, and the gradients can also be corrupted
due to various reasons, e.g. distributed or federated learning Turan et al.
(2022), or adversarial attacks like data poisoning Steinhardt et al. (2017).

As pointed out by Mahsereci et al. (2017); Siems et al. (2021), conditional on
the underlying true gradient $\nabla\mathcal{L}(\mathbf{\bar{w}})=0$, there
are mini-batch gradients which are not informative anymore as it can be fully
explained by sample noise and the vanishing gradients. These gradients would
not contribute to the covariance information of the empirical FIM but serve as
outliers in Hessian approximation.

In the scenarios of federated learning (FL), gradients computed by different
clients are skewed and consequently, local models move away from globally
optimal models Huang et al. (2022), imposing challenges for constructing
informative FIM. Besides, noise can be added to the gradient for privacy
concerns in communications Li et al. (2020). Additionally, the client usually
contains inevitable noisy samples and labels, making models suffer from a
significant performance drop Tuor et al. (2021). Finally, over-the-air
communications itself suffer from unavoidable noises Ang et al. (2020); Yang
et al. (2020). These lead to concerns for network pruning with noisy
gradients.

The analog memory recently gained attention for deep learning model deployment
Garg et al. (2022). When neural network parameters and data get stored in
these analog devices, they are susceptible to device-related noise, affecting
the performance of network compression Isik et al. (2023).

Approach Sketch. We revisit the MIQP network pruning optimization from a
perspective of entropic Wasserstein regression (EWR), which leverages
Wasserstein distance to model the dissimilarity between two distributions. In
our context, it measures the dissimilarity of distributions relevant to model
parameters and gradient magnitudes before and after pruning. Namely,
${\nabla\bm{\ell}}$ is a $p$ dimensional distribution, capturing geometric
properties of the loss at $\mathbf{\bar{w}}$ before pruning. Both
$\mathbf{\bar{w}}$ and $\mathbf{w}$ perform projections for
${\nabla\bm{\ell}}$ to a 1-D distribution respectively as
${\nabla\bm{\ell}}^{\top}\mathbf{\bar{w}}$ and
${\nabla\bm{\ell}}^{\top}\mathbf{w}$. Computing the distance between
${\nabla\bm{\ell}}^{\top}\mathbf{\bar{w}}$ and
${\nabla\bm{\ell}}^{\top}\mathbf{w}$ falls into the framework of sliced
probability divergence Nadjahi et al. (2020). Under this framework, pruning
optimization essentially fine-tunes $\mathbf{w}$ and selectively reserves its
elements such that the divergence is minimized subject to the sparsity
constraint.

Our approach’s effectiveness in combatting noisy gradients is established both
analytically and numerically. We demonstrate that pruning through the
Wasserstein regression implicitly enacts gradient averaging using Neighborhood
Interpolation. This entails a nuanced balance between capturing gradient
covariance and diminishing gradient noise. Notably, the sparse LR formulation
is merely a specific instance of ours. Yet, our proposed algorithm doesn’t
demand a markedly higher computational expense. This modest additional effort
bestows upon us enhanced robustness.

##  2 Problem Setup and Formulation

We first introduce the optimal transport (OT) problem in Kantorovich
formulation with entropic regularization, which measures the divergence
between two distributions, defined in equation 4 below. The Wasserstein
regression formulation as a generalization of the LR one is then proposed.

The Kantorovich Problem. Denote $\mathcal{P}_{2}$ the set of probability
measures with finite second moments. Let $\mu$, $\nu\in\mathcal{P}_{2}$ and
let $\Pi(\mu,\nu)$ denote the set of probability measures in $\mathcal{P}_{2}$
with marginal distributions equal to $\mu$ and $\nu$. The 2-Wasserstein
distance is defined as

$W_{2}^{2}(\mu,\nu)=\inf_{\pi\in\Pi(\mu,\nu)}\int_{\mathbb{R}^{d\times
d}}\left\lVert
x-y\right\rVert^{2}\operatorname{d}\\!{\pi}(x,y)+\varepsilon\int_{\mathbb{R}^{d}}\log\left(\frac{\operatorname{d}\\!{\pi}}{\operatorname{d}\\!{\mu}\operatorname{d}\\!{\nu}}\right).$
(4)

This is also referred to as the entropic OT problem, where the first term is
the transportation cost between the two measures and the second term is the
entropic regularization with multiplier $\varepsilon$.

Sparse EWR Formulation. The pruning problem formulation is defined in equation
5 below.

$\displaystyle\min_{\mathbf{w}}$ $\displaystyle\quad
Q(\mathbf{w})=W_{2}^{2}(x(\mathbf{w}),y)+\lambda\left\lVert\mathbf{w}-\mathbf{\bar{w}}\right\rVert^{2}$
(5a) s.t. $\displaystyle\quad\left\lVert\mathbf{w}\right\rVert_{0}\leq k$ (5b)

The term $W_{2}^{2}(x(\mathbf{w}),y)$ is a Wasserstein distance between the
two one-dimensional distributions $x$ and $y$ (or a sliced Wasserstein
distance for $\nabla\ell$ with two one-dimensional projections). The
optimization is to alter $\mathbf{w}$ such that the distance between the two
distributions is minimized.

Let $x$ and $y$ follow empirical distributions $\\{x_{i}\\}_{i=1}^{n}$ and
$\\{y_{i}\\}_{i=1}^{n}$. Denote by $a_{i}$ and $b_{i}$ the mass of data points
$x_{i}$ and $y_{i}$, respectively. We use $\mathbf{\Pi}$ to refer to a matrix
representing the transportation probability between $x$ and $y$, and $\Pi$ the
set of all such matrices, i.e.
$\Pi=\\{\mathbf{\Pi}|\sum_{i=1}^{n}\pi_{i}=1~{}\forall j\text{ and
}\sum_{j=1}^{n}\pi_{j}=1~{}\forall i\\}$. Then the formulation equation 5
reads as follows.

$\displaystyle\min_{\mathbf{w}}$
$\displaystyle~{}Q(\mathbf{w})=\inf_{\mathbf{\Pi}\in\Pi}\left\\{\sum_{i=1}^{n}\sum_{j=1}^{n}\left\lVert
x_{i}(\mathbf{w})-y_{i}\right\rVert^{2}\pi_{ij}+\varepsilon\sum_{i=1}^{n}\sum_{j=1}^{n}\log\left(\frac{\pi_{ij}}{a_{i}b_{j}}\right)\pi_{ij}\right\\}+\lambda\left\lVert\mathbf{w}-\mathbf{\bar{w}}\right\rVert^{2}$
(6a) s.t. $\displaystyle~{}\left\lVert\mathbf{w}\right\rVert_{0}\leq k$ (6b)

LR as a Special Case. Let $\varepsilon=0$. Once we set $\mathbf{\Pi}$ to be
the diagonal matrix with constant value $1/n$, i.e.
$\operatorname{diag}(1/n)$, the mass transportation happens only between data
point pairs $(x_{i},y_{i})$ for $i=1,\ldots,n$. Therefore we have

$Q_{\mathbf{\Pi}=\operatorname{diag}(1/n)}(\mathbf{w})=\frac{1}{n}\sum_{i=1}^{n}\left\lVert
x_{i}(\mathbf{w})-y_{i}\right\rVert^{2}+\lambda\left\lVert\mathbf{w}-\mathbf{\bar{w}}\right\rVert^{2}=\frac{1}{n}\bar{Q}(\mathbf{w}),$
(7)

i.e. the formulation equation 6 in this case degrades to the LR formulation in
equation 3.

##  3 Theoretical Aspects

This section reveals good theoretical properties of the Sparse EWR formulation
for network pruning. We start with Proposition 1 below (proof in Appendix A.1)
that states a geometry property of OT with squared Euclidean distance cost.
Second, we demonstrate the Neighborhood Interpolation mechanism that happens
implicitly in solving the EWR. Third, we show that such a mechanism strikes a
balance in capturing gradient covariance and reducing gradient noise, with a
brief discussion on the advantage of using entropic regularization in terms of
sample complexity.

######  Proposition 1 (Convex Hull Distance Equality).

Consider a set $S$ and its convex hull $\textit{Conv}(S)$ in a Euclidean
space, and an arbitrary point ${x}$ in the space. For any probability measure
$\hat{\nu}$ on $S$, we can find a point ${y}^{\prime}$ in $\textit{Conv}(S)$
as ${y}^{\prime}=\int{y}\,\mathrm{d}\nu({y})$ such that
$\|{x}-{y}^{\prime}\|^{2}=\int\|{x}-{y}\|^{2}\,\mathrm{d}\hat{\nu}({y})$,
where $\nu$ is a measure on $\textit{Conv}(S)$.

Neighborhood Interpolation. In formulation equation 5, let $W_{x}$ be the
first term of $W_{2}^{2}$ in equation 5 for an arbitrary given $x$, i.e.,
$W_{x}=\int\left\lVert
x(\mathbf{w})-y\right\rVert^{2}\operatorname{d}\\!{\pi}(\cdot|x)(y)$, where
$\pi(\cdot|x)(y)$ is a conditional measure given $x$. Now, divide the
Euclidean space $\mathbb{R}^{d}$ by sub-spaces $S^{y}_{1}$, $S^{y}_{2}$,…,
$S^{y}_{n}$ for $y$. For any conditional measure $\pi(\cdot|x)(y)$ defined on
any $S^{y}_{i}$ ($i=1,\ldots,n$), there exists a measure $\nu(y)$ defined on
$\textit{Conv}(S_{i}^{y})$ such that the weighted distance from $S^{y}_{i}$ to
$x$ equals the distance from $x$ to a point $y^{\prime}$ in
$\textit{Conv}(S^{y}_{i})$. Hence

$\displaystyle W_{x}$ $\displaystyle=\int_{S^{y}_{1}\cup
S^{y}_{2}\cup\cdots\cup
S^{y}_{n}}\left\|{x}-y\right\|^{2}\operatorname{d}\\!{\pi}(\cdot|x)(y)$
$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}\left\|{x}-y_{i}^{\prime}\right\|^{2}\text{
s.t.
}y_{i}^{\prime}=\int_{\textit{Conv}(S^{y}_{i})}{y}\,\operatorname{d}\\!{\nu(y)},\nu\in\mathbb{V}_{i}(x),~{}i=1,\ldots
n.$

where $\mathbb{V}_{i}(x)$ is the set of measures $\nu$ that make the equality
holds and $\mathbb{V}_{i}(x)\neq\phi$ by Proposition 1.

Similarly, we could define $W_{y}$ for any given $y$ and a set $S^{x}_{m}$,

$\displaystyle W_{y}$ $\displaystyle=\int_{S^{x}_{1}\cup
S^{x}_{2}\cup\cdots\cup
S^{x}_{m}}\left\|{x}-y\right\|^{2}\operatorname{d}\\!{\pi}(\cdot|y)(x)$
$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}\left\|{x}^{\prime}_{i}-y\right\|^{2}\text{
s.t.
}x^{\prime}_{i}=\int_{\textit{Conv}(S^{x}_{i})}{x}\,\operatorname{d}\\!{\mu}{(x)},~{}\mu\in\mathbb{U}_{i}(y),~{}i=1,\ldots
n.$

where $\mu$ is a measure defined on $\textit{Conv}(S^{x}_{i})$ and
$\mathbb{U}_{i}(y)\neq\phi$.

${y_{1}}$${y_{2}}$${y_{3}}$${y_{4}}$${y_{5}}$${x}$${y}^{\prime}$

We demonstrate the concept of “Neighborhood Interpolation” through an
empirical distribution example. Define  $S$ as a subset of  $y$ such that for
every element  $y_{i}\in S$ ,  $\pi_{x,i}>0$ . Without loss of generality, we
can denote  $S=\\{y_{1},y_{2},y_{3},y_{4},y_{5}\\}$ . The area shaded in gray,
denoted as  $\textit{Conv}(S)$ , represents the convex hull of  $S$ .  $W_{x}$
computes a weighted summation of distances between  $x$ and the points
$y_{1},\ldots,y_{5}$ . The weights  $\pi_{x,1},\ldots,\pi_{x,5}$ are decided
by OT. A significant  $\pi_{x,i}$ typically implies that  $y_{i}$ is in
proximity to  $x$ , indicating a neighborhood relation. By Proposition 1, this
weighted distance is analogous to the distance between  $x$ and  $y^{\prime}$
, where  $y^{\prime}$ is derived from  $\textit{Conv}(x)$ .

Revisit the EWR formulation. The integral of either $W_{x}$ or $W_{y}$
respectively on $x$ or $y$ gives the first term of $W_{2}^{2}$. One can then
reformulate equation 5 as equation 8 below.

$\min_{\mathbf{w}:\|\mathbf{w}\|_{0}\leq
k}Q(\mathbf{w})=\frac{1}{2}\inf_{\pi}\left\\{\int
W_{x}(\mathbf{w})\operatorname{d}\\!{\mu}(x)+\int
W_{y}(\mathbf{w})\operatorname{d}\\!{\nu}(y)\right\\}+\lambda\left\lVert\mathbf{w}-\mathbf{\bar{w}}\right\rVert^{2}.$
(8)

Interpretation: The objective function calculates the Euclidean distance
between a point $x$ and $n$ distinct points. These $n$ points originate from
$n$ convex hulls, each shaped by different $n$ subspaces within $y$.
Similarly, the function measures the distance between each point $y$ and $m$
unique points derived from $m$ convex hulls, each formed by distinct $m$
subspaces within $x$.

We claim that the EWR formulation is more resilient to noisy gradient than its
counterpart, the LR formulation given by equation 3. To understand this claim
better, let’s reimagine the problem using empirical distributions, as
indicated by equation 6. In this context, we use $x_{i}$ and $y_{i}$ as
substitutes for $S^{x}_{i}$ and $S^{y}_{i}$. Moreover, the integration in both
$W_{x}$ and $W_{y}$ is replaced with summations, offering a more insightful
version of our initial EWR formulation, shown as equation 9.

$\min_{\mathbf{w}:\left\lVert\mathbf{w}\right\rVert_{0}\leq
k}Q(\mathbf{w})=\inf_{\mathbf{\Pi}}\left\\{Q_{\mathbf{\Pi}}(\mathbf{w})+\varepsilon\sum_{i=1}^{n}\sum_{j=1}^{n}\log\left(\frac{\pi_{ij}}{a_{i}b_{j}}\right)\pi_{ij}\right\\}$
(9)

where the notation $Q_{\mathbf{\Pi}}(\mathbf{w})$ in equation 10 denotes the
part of the objective function given fixed $\mathbf{\Pi}$:

$\displaystyle Q_{\mathbf{\Pi}}(\mathbf{w})$
$\displaystyle=\sum_{i=1}^{n}\sum_{j=1}^{n}\left\lVert
x_{i}(\mathbf{w})-y_{j}\right\rVert^{2}\pi_{ij}+\lambda\left\lVert\mathbf{w}-\mathbf{\bar{w}}\right\rVert^{2}$
(10a) $\displaystyle=\frac{1}{2}\underbrace{\sum_{i=1}^{n}\left\lVert
x_{i}(\mathbf{w})-y^{\prime}_{i}\right\rVert^{2}}_{K_{\mathbf{\Pi}}^{(1)}}+\frac{1}{2}\underbrace{\sum_{i=1}^{n}\left\lVert
x^{\prime}_{i}(\mathbf{w})-y_{i}\right\rVert^{2}}_{K_{\mathbf{\Pi}}^{(2)}}+\lambda\left\lVert\mathbf{w}-\mathbf{\bar{w}}\right\rVert^{2}$
(10b)

In $Q_{\mathbf{\Pi}}(\mathbf{w})$, for each index $i$, points $x^{\prime}_{i}$
and $y^{\prime}_{i}$ are chosen from the convex hulls formed by points in $x$
and $y$, as per the guidelines of Proposition 1. Now, contrasting this with
the LR model in equation 3, the objective $\bar{Q}(\mathbf{w})$ aims for
regression directly over the data points whereas every point from one
empirical set is matched for Euclidean distance computation to a point derived
from a convex combination of the other.

The infimum in equation 9 seeks the OT plan, $\mathbf{\Pi}$, that aligns the
empirical distributions $x$ and $y$ closely. In practical terms, for each data
point $x_{i}$, only a subset of $\\{y_{i}\\}_{i=1}^{n}$ will transport a
substantial mass, rather than the entire set. This behavior of $\mathbf{\Pi}$
effectively defines $n$ ”neighborhoods” for each data point $x_{i}$ within the
empirical distribution of $y$. Here, a ”neighborhood” refers to a group of
data points in $y$ that are proximate to a specific $x_{i}$ in the Euclidean
sense.

Neighborhood Size Control. A critical aspect of this formulation is the
entropic regularization term, which is used to modulate the size of these
neighborhoods. Specifically, increasing the value of $\varepsilon$ amplifies
the impact of the entropy term. This change broadens the neighborhoods,
drawing more data points into the fold of the associated convex hulls. An
illustrative edge case is when $\varepsilon=0$. Here, the OT does one-to-one
matching, implying that each data point $y_{i}$ primarily forms the convex
hull independently. On the contrary, when $\varepsilon\rightarrow\infty$ all
data points are equally weighted by $\mathbf{\Pi}$ and hence involved in
forming the convex hull as a neighborhood.

Capturing Covariance With Gradient Noise Reduction. For an arbitrary
$\mathbf{w}$, the EWR formulation essentially strikes a balance between
gradient noise reduction and covariance capturing. We show the analysis for
$K^{(1)}_{\mathbf{\Pi}}$ in equation 10, and
$K^{(2)}_{\mathbf{\Pi}}(\mathbf{w})$ follows the same logic exactly. Note that
$y^{\prime}_{i}=\sum_{j=1}^{n}\nu^{(i)}_{j}y_{j}=\sum_{j=1}^{n}\nu^{(i)}_{j}{{\nabla\bm{\ell}}}^{\top}_{j}\mathbf{\bar{w}}$,
where $\bm{\nu}^{(i)}$ are convex combination coefficients by Proposition 1.
Denote
${{\nabla\bm{\ell}}}^{\prime\top}_{i}=\sum_{j=1}^{n}\nu^{(i)}_{j}{{\nabla\bm{\ell}}}^{\top}_{j}$,
and
$\mathbf{G}^{\prime}=[\nabla\bm{\ell}^{\prime}_{1},\ldots,\nabla\bm{\ell}^{\prime}_{n}]^{\top}$.
The term $K^{(1)}_{\mathbf{\Pi}}$ expands as follows.

$\displaystyle K^{(1)}_{\mathbf{\Pi}}$
$\displaystyle=\sum_{i=1}^{n}({{\nabla\bm{\ell}}}_{i}^{\top}\mathbf{w}-{{{\nabla\bm{\ell}}}^{\prime}}^{\top}_{i}\mathbf{\bar{w}})^{\top}({{\nabla\bm{\ell}}}_{i}^{\top}\mathbf{w}-{{{{\nabla\bm{\ell}}}^{\prime}}}_{i}^{\top}\mathbf{\bar{w}})$
$\displaystyle=\sum_{i=1}^{n}(\mathbf{w}^{\top}{{\nabla\bm{\ell}}}_{i}{{\nabla\bm{\ell}}}_{i}^{\top}\mathbf{w}-\mathbf{w}^{\top}{{\nabla\bm{\ell}}}_{i}{\nabla\bm{\ell}}_{i}^{\prime}{}^{\top}\mathbf{\bar{w}}-\mathbf{\bar{w}}{}^{\top}{\nabla\bm{\ell}}_{i}^{\prime}{{\nabla\bm{\ell}}}_{i}{}^{\top}\mathbf{w}+\mathbf{\bar{w}}{}^{\top}{\nabla\bm{\ell}}_{i}^{\prime}{\nabla\bm{\ell}}_{i}^{\prime}{}^{\top}\mathbf{\bar{w}})$
(11)

Examining $K^{(1)}_{\mathbf{\Pi}}$ from equation 11, we see that it
effectively replaces half of ${{\nabla\bm{\ell}}}_{i}$ with
${{\nabla\bm{\ell}}}^{\prime}_{i}$, a version obtained through weighted
gradient averaging. Now let’s comparing the covariance between
$\nabla\bm{\ell}_{i}$ and $\nabla\bm{\ell}^{\prime}_{i}$. Assume that
${{\nabla\bm{\ell}}}_{i}$ ($1\leq i\leq n$) are i.i.d. with the same
covariance matrix $\Sigma$, then $\mathbf{G}^{\prime}$ is with equal or less
noise than $\mathbf{G}$. To show this, denote the covariance matrix of each
${{\nabla\bm{\ell}}}^{\prime}_{i}$ by

$\Sigma^{\prime}_{i}=\text{Cov}\bigg{[}\sum_{j=1}^{n}\nu^{(i)}_{j}{{\nabla\bm{\ell}}}_{j}\bigg{]}=\sum_{j=1}^{n}[\nu^{(i)}_{j}]^{2}\Sigma.$

The total variance of each gradient in $\mathbf{G}^{\prime}$ (i.e., the trace
of $\Sigma^{\prime}_{i}$) is then

$\text{trace}(\Sigma^{\prime}_{i})=\text{trace}\bigg{(}\sum_{j=1}^{n}\big{[}\nu^{(i)}_{j}\big{]}^{2}\Sigma\bigg{)}=\sum_{j=1}^{n}\big{[}\nu^{(i)}_{j}\big{]}^{2}\text{trace}(\Sigma)\leq\text{trace}(\Sigma).$

The last inequality follows from the fact that
$\sum_{j=1}^{n}[\nu^{(i)}_{j}]^{2}\leq 1$, which is a consequence of the
Cauchy-Schwarz inequality given that the coefficients $\nu^{(i)}_{j}$ form a
convex combination.

Originally, the covariance information of all data points is embedded in
${{\nabla\bm{\ell}}}_{i}{{\nabla\bm{\ell}}}_{i}^{\top}$ for $i=1,2,\ldots,n$.
An alternative representation is
${{\nabla\bm{\ell}}}^{\prime}_{i}{\nabla\bm{\ell}}_{i}^{\prime\top}$, which
prioritizes noise reduction, but sacrifices some covariance information. Both
${{\nabla\bm{\ell}}}^{\prime}_{i}{{\nabla\bm{\ell}}}_{i}^{\top}$ and
${{\nabla\bm{\ell}}}_{i}{\nabla\bm{\ell}}_{i}^{\prime\top}$ highlight a trade-
off. Notably, both the original covariance
${{\nabla\bm{\ell}}}_{i}{{\nabla\bm{\ell}}}_{i}^{\top}$ and its noise-reduced
counterpart ${\nabla\bm{\ell}}_{i}^{\prime}{\nabla\bm{\ell}}_{i}^{\prime\top}$
are retained in equation 11.

Difference From Averaging Prior to Optimization: Next, we show that such
gradient averaging differs from the averaging operation conducted prior to
optimization. Let
$\mathbf{G}^{\prime}=[{{\nabla\bm{\ell}}}^{\prime}_{1},{{\nabla\bm{\ell}}}^{\prime}_{2},\ldots,{{\nabla\bm{\ell}}}^{\prime}_{n}]^{\top}$
such that $\mathbf{G}^{\prime}$ represents the row-wise convex combination of
$\mathbf{G}$. Approximating the Hessian of the MIQP equation 2, two scenarios
emerge: using $\mathbf{G}$ that not performing averaging (case 1) and
$\mathbf{G}^{\prime}$ that performs averaging before optimization (case 2).

Case 1 is indeed the original LR formulation equation 3. Denote by $K$ its
term that is corresponding to $K^{(1)}_{\mathbf{\Pi}}$, we have

$\displaystyle K$
$\displaystyle=(\mathbf{w}-\mathbf{\bar{w}})^{\top}\mathbf{G}^{\top}\mathbf{G}(\mathbf{w}-\mathbf{\bar{w}})$
$\displaystyle=\sum_{i=1}^{n}(\mathbf{w}^{\top}{{\nabla\bm{\ell}}}_{i}{{{\nabla\bm{\ell}}}}_{i}^{\top}\mathbf{w}-\mathbf{w}^{\top}{{\nabla\bm{\ell}}}_{i}{{{\nabla\bm{\ell}}}}_{i}^{\top}\mathbf{\bar{w}}-\mathbf{\bar{w}}^{\top}{{\nabla\bm{\ell}}}_{i}{{{\nabla\bm{\ell}}}}_{i}^{\top}\mathbf{w}+\mathbf{\bar{w}}^{\top}{{\nabla\bm{\ell}}}_{i}{{{\nabla\bm{\ell}}}}_{i}^{\top}\mathbf{\bar{w}})$
(12)

Case 2 uses the less-noisy row-wise convex combination matrix
$\mathbf{G}^{\prime}$ instead of $\mathbf{G}$. Yet, the original covariance
${{\nabla\bm{\ell}}}_{i}{{\nabla\bm{\ell}}}_{i}^{\top}$ is lost: Denote by
$K^{\prime}$ the corresponding term, and we have

$K^{\prime}=\sum_{i=1}^{n}(\mathbf{w}^{\top}{\nabla\bm{\ell}}_{i}^{\prime}{{{{\nabla\bm{\ell}}}^{\prime}}}_{i}^{\top}\mathbf{w}-\mathbf{w}^{\top}{\nabla\bm{\ell}}_{i}^{\prime}{{{{\nabla\bm{\ell}}}^{\prime}}}_{i}^{\top}\mathbf{\bar{w}}-\mathbf{\bar{w}}^{\top}{\nabla\bm{\ell}}_{i}^{\prime}{{{{\nabla\bm{\ell}}}^{\prime}}}_{i}^{\top}\mathbf{w}+\mathbf{\bar{w}}^{\top}{\nabla\bm{\ell}}_{i}^{\prime}{{{{\nabla\bm{\ell}}}^{\prime}}}_{i}^{\top}\mathbf{\bar{w}})$
(13)

Inspecting the expressions, it can be observed that $K_{\mathbf{\Pi}}^{(1)}$
(also $K_{\mathbf{\Pi}}^{(2)}$) strikes a balance between $K$ and
$K^{\prime}$. There are two notable extreme cases for $\mathbf{\Pi}$:

  1. 1.

$\mathbf{\Pi}=\operatorname{diag}(1/n)$. This corresponds to the LR
formulation, as detailed in Section 2. A smaller value of $\varepsilon$ steers
the optimization in this direction.

  2. 2.

$\mathbf{\Pi}=(1/n^{2})\mathbf{1}\cdot\mathbf{1}^{\top}$. This arises when
$\varepsilon\rightarrow\infty$, meaning the entropy term holds sway in the OT
optimization. Here, mutual information is minimized to ensure an even
contribution from data points in the convex combination. Both $x^{\prime}_{i}$
and $y^{\prime}_{i}$ are the arithmetic means of their respective sets, and
all ${\nabla\bm{\ell}}_{i}^{\prime}$ are equivalent to the averaged gradient
over the $n$ points. Importantly, the original covariance remains intact even
in this edge case.

As $n$ grows indefinitely, the empirical OT formulation from equation 6
approaches its continuous counterpart given by equation 5. Intuitively, a
large dataset of high-quality training samples diminishes concerns over
gradient noise, making the empirical fisher a close approximation to the true
fisher. In such situations, $\varepsilon$ is set to zero. Brenier’s theorem
Peyré (2019) then suggests that the OT plan turns into a monotone map for
costs represented by squared Euclidean distances. This means $\mathbf{\Pi}$
tends towards $\operatorname{diag}(1/n)$. Consequently, the Wasserstein
distance formulation reduces to the Euclidean distance formulation, delivering
optimal performance with ample data.

An advantage of employing the EWR formulation is its inherent capability for
gradient averaging. This approach negates the need to manually determine the
convex combination coefficients or resort to density estimation to pinpoint
the nearest gradient neighbors for averaging. Importantly, this seamless
trade-off eludes the combination of the Euclidean distance with gradient
averaging. The reason is, the original covariance information will inevitably
be lost in the formulation equation 13, irrespective of the chosen averaging
method.

Sample Compexity of $W_{2}^{2}(\mu,\nu)$ is narrowed to $O(1/\sqrt{n})$ from
$O(1/n^{\frac{1}{4}})$ by the entropic regularization term. This ensures the
empirical distributions are close to the underlying distributions with a small
number of samples. Please see Appendix A.2 for details.

##  4 Algorithm Design

Algorithmic Framework. The algorithm addresses the network pruning problem
defined in equation 5. Drawing inspiration from Chen et al. (2022), the
algorithm incrementally adjusts the sparsity of the weights vector
$\mathbf{w}$ by using a descending sequence of non-zero elements
$k_{0},\ldots,k_{T}$. During each sparsity level, the weights $\mathbf{w}$ and
the transportation plan $\mathbf{\Pi}$ (can be obtained with efficient
algorithms, see Appendix A.3) are refined iteratively.

Algorithm 1 Entropic Wasserstein Regression Pruning

1:Pre-pruning weights $\mathbf{\bar{w}}$, sparsity $k$, regularization
parameter $\lambda$, $\varepsilon$, maximum rounds $T$, batches
$\mathcal{B}_{1},\ldots,\mathcal{B}_{T}$, optimization step size $\tau>0$

2:Post-pruning weights $\mathbf{w}$, satisfying
$\left\lVert\mathbf{w}\right\rVert\leq k$

3:Set $k_{0},k_{1},\ldots,k_{T}$ as a descending sequence, with $k_{0}<p$ and
$k_{T}=k$.

4:$\mathbf{w}^{(0)}\leftarrow\mathbf{\bar{w}}$

5:for $t\leftarrow 0,1,\ldots,T$ do

6: Compute
$\mathbf{G}=[{\nabla\bm{\ell}}_{1}(\mathbf{\bar{w}}),\ldots{\nabla\bm{\ell}}_{n}(\mathbf{\bar{w}})]^{\top}$
with batch $B_{t}$

7:
$\mathbf{x},\mathbf{y}\leftarrow\mathbf{G}\mathbf{w}^{(t)},\mathbf{G}\mathbf{\bar{w}}$

8: Compute the pairwise Euclidean distance matrix $\mathbf{C}$ between
$\mathbf{x}$ and $\mathbf{y}$

9: Compute OT planning $\mathbf{\Pi}^{(t)}$ (see Appendix A.3)

10: $\nabla
Q\leftarrow\mathbf{G}^{\top}(\mathbf{\Pi}(\mathbf{G}\mathbf{w}^{(t)}-\mathbf{G}\mathbf{\bar{w}}))+\lambda(\mathbf{w}^{(t)}-\mathbf{\bar{w}})$

11: $\mathbf{w}^{(t+\frac{1}{2})}\leftarrow\mathbf{w}^{(t)}-\tau\nabla Q$

12: $\mathbf{w}^{(t+1)}\leftarrow\text{Select from
$\mathbf{w}^{(t+\frac{1}{2})}$ $k_{t}$ components having largest absolute
values; Others zero }$

13: $\mathbf{\bar{w}}\leftarrow\mathbf{w}^{(t+1)}$

14:end for

Weights Optimization. The weights $\mathbf{w}$ are optimized using the
stochastic gradient descent (SGD) paired with the iterative hard threshold
(IHT) algorithm. We use $\nabla Q$ to represent the derivative of
$Q(\mathbf{w})$ for brevity, with its comprehensive derivation in the Appendix
A.4. The expression is

$\nabla
Q=\mathbf{G}^{\top}(\mathbf{\Pi}(\mathbf{G}\mathbf{w}-\mathbf{G}\mathbf{\bar{w}}))+\lambda(\mathbf{w}-\mathbf{\bar{w}}).$
(14)

Following the weight updates driven by SGD (as seen in line 11 of Algorithm
1), the IHT method is applied. Here, the $k_{t}$ components of $\mathbf{w}$
with the largest magnitudes are retained, while the remaining are set to zero,
ensuring adherence to the sparsity criteria.

A vital component of the optimization process is the choice of the stepsize
$\tau$ (referenced in line 11). Although a straightforward approach might be
to set $\tau=\frac{1}{L}$ (where $L$ denotes the Lipschitz constant of $Q$),
superior performance is achieved when the stepsize is ascertained using the
methodology delineated in (Chen et al., 2022, Algorithm 2). For the quadratic
function $Q$, the Lipschitz constant $L$ is given by $L=n\lambda+\left\lVert
G\right\rVert^{2}_{op}$, where $\left\lVert\cdot\right\rVert_{op}$ indicates
the foremost singular value.

Although there exist alternative strategies for refining the IHT
solution—including active set updates, coordinate descent, and the Woodbury
formula’s back-solving—a discussion on these falls outside the scope of this
paper. For an in-depth exploration, especially regarding the specialized case
described in equation 7, one can consult Bhatia et al. (2015); Chen & Mazumder
(2020); Hazimeh & Mazumder (2020); Benbaki et al. (2023).

##  5 Numerical Results

Table 1: Model Pruning Accuracy Benchmarking. Five runs are taken for LR and EWR, with the mean and 95% confidence interval (in the brackets) reported. The data of MP, WF, and CBS are copied from Yu et al. (2022). The superscript “$+\sigma^{2}$” indicates that 20% of data is with noise. Bold texts imply the best performance, with $0.1$ percentage as tolerance in difference. The sparsity column is the target sparsity.  Network | Sparsity | MP | WF | CBS | LR (EWRΠ=I/n) | EWR (proposed)
---|---|---|---|---|---|---
MLPNet on MNIST (93.97%) | 0.5 | 93.93 | 94.02 | 93.96 |  95.26 ($\pm$0.03) |  95.24 ($\pm$0.03)
0.6 | 93.78 | 93.82 | 93.96 |  95.13 ($\pm$0.02) |  95.13 ($\pm$0.01)
0.7 | 93.62 | 93.77 | 93.98 |  94.93 ($\pm$0.03) |  95.05 ($\pm$0.04)
0.8 | 92.89 | 93.57 | 93.90 |  94.82 ($\pm$0.04) |  94.84 ($\pm$0.03)
0.9 | 90.30 | 91.69 | 93.14 |  94.32 ($\pm$0.05) |  94.30 ($\pm$0.05)
0.95 | 83.64 | 85.54 | 88.92 |  92.82 ($\pm$0.06) |  92.86 ($\pm$0.05)
0.95${}^{+\sigma^{2}}$ | - | - | - | 90.11 ($\pm$0.08) |  90.50 ($\pm$0.07)
| 0.98 | 32.25 | 38.26 | 55.45 | 84.43 ($\pm$0.10) |  85.71 ($\pm$0.09)
| 0.98${}^{+\sigma^{2}}$ | - | - | - | 82.12 ($\pm$0.11) |  83.69 ($\pm$0.10)
ResNet20 on CIFAR10 (91.36%) | 0.5 | 88.44 | 90.23 | 90.58 |  92.06 ($\pm$0.04) |  92.04 ($\pm$0.03)
0.6 | 85.24 | 87.96 | 88.88 |  91.98 ($\pm$0.09) |  91.98 ($\pm$0.09)
0.7 | 78.79 | 81.05 | 81.84 | 91.09 ($\pm$0.10) |  91.89 ($\pm$0.10)
0.8 | 54.01 | 62.63 | 51.28 | 89.00 ($\pm$0.12) |  90.15 ($\pm$0.09)
0.9 | 11.79 | 11.49 | 13.68 | 87.63 ($\pm$0.11) |  88.82 ($\pm$0.10)
0.95 | - | - | - | 80.25 ($\pm$0.17) |  81.33 ($\pm$0.15)
0.95${}^{+\sigma^{2}}$ | - | - | - | 77.37 ($\pm$0.18) |  79.05 ($\pm$0.16)
| 0.98 | - | - | - | 68.15 ($\pm$0.27) |  69.21 ($\pm$0.24)
| 0.98${}^{+\sigma^{2}}$ | - | - | - | 65.04 ($\pm$0.27) |  68.01 ($\pm$0.25)
MobileNetV1 on ImageNet (71.95%) | 0.3 | 71.60 | 71.88 | 71.87 | 71.14 ($\pm$0.08) |  71.87 ($\pm$0.05)
0.4 | 69.16 | 71.15 | 71.45 | 71.12 ($\pm$0.10) |  71.44 ($\pm$0.07)
0.5 | 62.61 | 68.91 | 70.21 | 70.12 ($\pm$0.13) |  71.12 ($\pm$0.18)
0.6 | 41.94 | 60.90 | 66.37 | 70.05 ($\pm$0.22) |  70.92 ($\pm$0.18)
0.7 | 6.78 | 29.36 | 55.11 | 68.15 ($\pm$0.17) |  69.26 ($\pm$0.13)
0.8 | 0.11 | 0.24 | 16.38 | 65.72 ($\pm$0.19) |  66.82 ($\pm$0.14)
0.8${}^{+\sigma^{2}}$ | - | - | - | 60.29 ($\pm$0.18) |  63.62 ($\pm$0.15)
| 0.9 | - | - | - | 47.65 ($\pm$0.15) |  49.43 ($\pm$0.13)
| 0.9${}^{+\sigma^{2}}$ | - | - | - | 44.55 ($\pm$0.16) |  47.98 ($\pm$0.16)

Our method is compared with several existing SoTA methods including MP
(magnitude pruning Mozer & Smolensky (1989)), WF (WoodFisher Singh & Alistarh
(2020)), CBS (Combinatorial Brain Surgeon Yu et al. (2022)), and LR (i.e. the
sparse LR formulation adopted by Chen et al. (2022)). We refer to our proposed
method as EWR (i.e. sparse entropic Wasserstein regression). Remark that LR is
a special case of EWR, with $\mathbf{\Pi}=\operatorname{diag}(1/n)$. All the
methods are benchmarked on pre-trained neural networks: MLPNet (30k
parameters) trained on MNIST LeCun et al. (1998), ResNet20 He et al. (2016)
(200k parameters) trained on CIFAR10 Krizhevsky et al. (2009), and MobileNetV1
Howard et al. (2017) (4.2M parameters) trained on ImageNet Deng et al. (2009).
The experiment setup for reproducibility111The code is available on
https://anonymous.4open.science/r/Entropic-Wasserstein-Pruning-F222 is
detailed in AppendixA.5. We deliver more experiments results in Appendix
A.6-A.8.

Model Accuracy Performance Benchmarking. Table 1 compares the performance of
different networks across various sparsity levels, utilizing different
methods. MLPNet’s performance on MNIST is consistent across different sparsity
levels, with both LR and the proposed EWR method showing superior performance.
The advantages of EWR over the others are reflected by the two more
challenging tasks ResNet20 on CIFAR10 and MobileNetV1 on ImageNet, especially
in the presence of noisy gradients. In summary, the proposed EWR method
consistently outperforms or matches other methods. The LR method performs well
at lower sparsity levels but is surpassed otherwise.

Table 2: Comparison of testing loss values for LR and EWR for ResNet20. The
result is averaged over 25 runs. The 90% confidence interval is reported. The
target sparsity is set to be 0.95.

| 10% Noisy Gradient
---|---
Sparsity | Noise = $\sigma$ | Noise = $2\sigma$
| LR | EWR | Diff (%) | LR | EWR | Diff (%)
0.95 | 2.83 ($\pm$0.02) |  2.75 ($\pm$0.02) | 2.87% | 2.86 ($\pm$0.02) |  2.74 ($\pm$0.01) | 4.35%
0.84 | 1.58 ($\pm$0.01) |  1.54 ($\pm$0.01) | 2.73% | 1.62 ($\pm$0.03) |  1.54 ($\pm$0.02) | 5.15%
0.74 | 0.66 ($\pm$0.00) |  0.64 ($\pm$0.00) | 3.03% | 0.66 ($\pm$0.00) |  0.65 ($\pm$0.00) | 1.32%
0.63 | 0.35 ($\pm$0.00) | 0.35 ($\pm$0.00) | 0.00% | 0.35 ($\pm$0.00) | 0.35 ($\pm$0.00) | 0.85%

| 25% Noisy Gradient
---|---
Sparsity | Noise = $\sigma$ | Noise = $2\sigma$
| LR | EWR | Diff (%) | LR | EWR | Diff (%)
0.95 | 2.87 ($\pm$0.02) |  2.77 ($\pm$0.01) | 3.50% | 2.89 ($\pm$0.02) |  2.79 ($\pm$0.02) | 3.82%
0.84 | 1.72 ($\pm$0.02) |  1.65 ($\pm$0.02) | 4.07% | 1.76 ($\pm$0.02) |  1.69 ($\pm$0.02) | 4.05%
0.74 | 0.67 ($\pm$0.01) | 0.67 ($\pm$0.00) | 0.49% | 0.68 ($\pm$0.00) |  0.67 ($\pm$0.00) | 1.55%
0.63 | 0.36 ($\pm$0.00) | 0.36 ($\pm$0.00) | 0.00% | 0.35 ($\pm$0.00) | 0.35 ($\pm$0.00) | 0.00%

Table 3: Comparison of testing loss values for LR and EWR for MobileNetV1. The
result is averaged from 10 runs. The 90% confidence interval is reported. The
target sparsity is set to be 0.75.

| 10% Noisy Gradient
---|---
Sparsity | Noise = $\sigma$ | Noise = $2\sigma$
| LR | EWR | Diff (%) | LR | EWR | Diff (%)
0.75 | 4.54 ($\pm$0.06) | 4.34 ($\pm$0.06) | 4.41% | 4.62 ($\pm$0.07) | 4.40 ($\pm$0.06) | 5.02%
0.63 | 2.53 ($\pm$0.04) | 2.46 ($\pm$0.03) | 2.61% | 2.56 ($\pm$0.04) | 2.48 ($\pm$0.04) | 3.23%
0.53 | 1.64 ($\pm$0.02) | 1.56 ($\pm$0.02) | 5.16% | 1.66 ($\pm$0.03) | 1.56 ($\pm$0.02) | 6.01%
0.42 | 1.30 ($\pm$0.00) | 1.30 ($\pm$0.00) | 0.00% | 1.30 ($\pm$0.00) | 1.30 ($\pm$0.00) | 0.00%

| 25% Noisy Gradient
---|---
Sparsity | Noise = $\sigma$ | Noise = $2\sigma$
| LR | EWR | Diff (%) | LR | EWR | Diff (%)
0.75 | 4.93 ($\pm$0.06) | 4.53 ($\pm$0.06) | 8.13% | 5.02 ($\pm$0.06) | 4.66 ($\pm$0.05) | 7.92%
0.63 | 2.54 ($\pm$0.04) | 2.44 ($\pm$0.03) | 4.01% | 2.57 ($\pm$0.04) | 2.45 ($\pm$0.03) | 5.00%
0.53 | 1.66 ($\pm$0.02) | 1.57 ($\pm$0.02) | 5.16% | 1.66 ($\pm$0.02) | 1.55 ($\pm$0.02) | 6.63%
0.42 | 1.30 ($\pm$0.00) | 1.30 ($\pm$0.00) | 0.30% | 1.31 ($\pm$0.00) | 1.31 ($\pm$0.00) | 0.00%

Robustness with Noisy Gradient. By Section 3, EWR differs from LR in terms of
gradient noise reduction achieved by solving the OT problem to obtain a group
of non-trivial data pair weighting coefficients. Hence, LR that has the
transportation plan $\mathbf{\Pi}$ fixed to $\operatorname{diag}(1/n)$
naturally serves as a baseline for evaluating the effectiveness of such
optimization in terms of robustness against noise.

In two noisy scenarios, 10%, and 25%, we evaluate loss at noise levels of
$\sigma$ and $2\sigma$ across varying sparsity. Tables 2 and 3 contrast the
loss difference between LR and EWR. EWR consistently outperforms LR in both
ResNet20 and MobileNetV1, most notably in noisy conditions and at higher
sparsity. The peak performance difference is 8.13% favoring EWR on MobileNetV1
at 0.74 sparsity with 25% noise. Hence, EWR outperforms LR.

##  6 Conclusions and Future Impact

The paper offers a novel perspective on robust neural network pruning,
shedding light on more sophisticated techniques for optimization problems.
Specifically, a novel formulation based on EWR is proposed which strikes a
balance between covariance information preservation and noise reduction. The
work suggested promising avenues for applications in large-scale model
compression, though it may require further empirical validation and
exploration of practical implementations.

## References

  * Ang et al. (2020) Fan Ang, Li Chen, Nan Zhao, Yunfei Chen, Weidong Wang, and F Richard Yu.  Robust federated learning with noisy communication.  _IEEE Transactions on Communications_ , 68(6):3452–3464, 2020.
  * Benbaki et al. (2023) Riade Benbaki, Wenyu Chen, Xiang Meng, Hussein Hazimeh, Natalia Ponomareva, Zhe Zhao, and Rahul Mazumder.  Fast as chita: Neural network pruning with combinatorial optimization.  _arXiv preprint arXiv:2302.14623_ , 2023.
  * Bhatia et al. (2015) Kush Bhatia, Prateek Jain, and Purushottam Kar.  Robust regression via hard thresholding.  _Advances in neural information processing systems_ , 28, 2015.
  * Chen & Mazumder (2020) Wenyu Chen and Rahul Mazumder.  Multivariate convex regression at scale.  _arXiv preprint arXiv:2005.11588_ , 2020.
  * Chen et al. (2022) Wenyu Chen, Riade Benbaki, Xiang Meng, and Rahul Mazumder.  Network pruning at scale: A discrete optimization approach.  In _OPT 2022: Optimization for Machine Learning (NeurIPS 2022 Workshop)_ , 2022.
  * Deng et al. (2009) Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.  Imagenet: A large-scale hierarchical image database.  In _2009 IEEE conference on computer vision and pattern recognition_ , pp. 248–255. Ieee, 2009.
  * Garg et al. (2022) Sahaj Garg, Joe Lou, Anirudh Jain, Zhimu Guo, Bhavin J Shastri, and Mitchell Nahmias.  Dynamic precision analog computing for neural networks.  _IEEE Journal of Selected Topics in Quantum Electronics_ , 29(2: Optical Computing):1–12, 2022.
  * Genevay et al. (2019) Aude Genevay, Lénaic Chizat, Francis Bach, Marco Cuturi, and Gabriel Peyré.  Sample complexity of sinkhorn divergences.  In _The 22nd international conference on artificial intelligence and statistics_ , pp. 1574–1583. PMLR, 2019.
  * Gretton et al. (2006) Arthur Gretton, Karsten Borgwardt, Malte Rasch, Bernhard Schölkopf, and Alex Smola.  A kernel method for the two-sample-problem.  _Advances in neural information processing systems_ , 19, 2006.
  * Hassibi & Stork (1992) Babak Hassibi and David Stork.  Second order derivatives for network pruning: Optimal brain surgeon.  _Advances in neural information processing systems_ , 5, 1992.
  * Hazimeh & Mazumder (2020) Hussein Hazimeh and Rahul Mazumder.  Fast best subset selection: Coordinate descent and local combinatorial optimization algorithms.  _Operations Research_ , 68(5):1517–1537, 2020\.
  * He et al. (2016) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 770–778, 2016.
  * Howard et al. (2017) Andrew G Howard, Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang, Tobias Weyand, Marco Andreetto, and Hartwig Adam.  Mobilenets: Efficient convolutional neural networks for mobile vision applications.  _arXiv preprint arXiv:1704.04861_ , 2017.
  * Huang et al. (2022) Hong Huang, Lan Zhang, Chaoyue Sun, Ruogu Fang, Xiaoyong Yuan, and Dapeng Wu.  Fedtiny: Pruned federated learning towards specialized tiny models.  _arXiv preprint arXiv:2212.01977_ , 2022.
  * Isik et al. (2023) Berivan Isik, Kristy Choi, Xin Zheng, Tsachy Weissman, Stefano Ermon, H-S Philip Wong, and Armin Alaghi.  Neural network compression for noisy storage devices.  _ACM Transactions on Embedded Computing Systems_ , 22(3):1–29, 2023.
  * Janati et al. (2020) Hicham Janati, Boris Muzellec, Gabriel Peyré, and Marco Cuturi.  Entropic optimal transport between unbalanced gaussian measures has a closed form.  _Advances in neural information processing systems_ , 33:10468–10479, 2020.
  * Krizhevsky et al. (2009) Alex Krizhevsky, Geoffrey Hinton, et al.  Learning multiple layers of features from tiny images, 2009.
  * LeCun et al. (1989) Yann LeCun, John Denker, and Sara Solla.  Optimal brain damage.  _Advances in neural information processing systems_ , 2, 1989.
  * LeCun et al. (1998) Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner.  Gradient-based learning applied to document recognition.  _Proceedings of the IEEE_ , 86(11):2278–2324, 1998.
  * Li et al. (2020) Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith.  Federated learning: Challenges, methods, and future directions.  _IEEE signal processing magazine_ , 37(3):50–60, 2020.
  * Mahsereci et al. (2017) Maren Mahsereci, Lukas Balles, Christoph Lassner, and Philipp Hennig.  Early stopping without a validation set.  _arXiv preprint arXiv:1703.09580_ , 2017.
  * Mozer & Smolensky (1989) Michael C Mozer and Paul Smolensky.  Using relevance to reduce network size automatically.  _Connection Science_ , 1(1):3–16, 1989.
  * Nadjahi et al. (2020) Kimia Nadjahi, Alain Durmus, Lénaïc Chizat, Soheil Kolouri, Shahin Shahrampour, and Umut Simsekli.  Statistical and topological properties of sliced probability divergences.  _Advances in Neural Information Processing Systems_ , 33:20802–20812, 2020.
  * Nadjahi et al. (2021) Kimia Nadjahi, Alain Durmus, Pierre E Jacob, Roland Badeau, and Umut Simsekli.  Fast approximation of the sliced-wasserstein distance using concentration of random projections.  _Advances in Neural Information Processing Systems_ , 34:12411–12424, 2021.
  * Peyré (2019) Gabriel Peyré.  Course notes on computational optimal transport.  2019\.
  * Reeves (2017) Galen Reeves.  Conditional central limit theorems for gaussian projections.  In _2017 IEEE International Symposium on Information Theory (ISIT)_ , pp. 3045–3049. IEEE, 2017.
  * Siems et al. (2021) Julien Niklas Siems, Aaron Klein, Cedric Archambeau, and Maren Mahsereci.  Dynamic pruning of a neural network via gradient signal-to-noise ratio.  In _8th ICML Workshop on Automated Machine Learning (AutoML)_ , 2021\.
  * Singh & Alistarh (2020) Sidak Pal Singh and Dan Alistarh.  Woodfisher: Efficient second-order approximation for neural network compression.  _Advances in Neural Information Processing Systems_ , 33:18098–18109, 2020.
  * Steinhardt et al. (2017) Jacob Steinhardt, Pang Wei W Koh, and Percy S Liang.  Certified defenses for data poisoning attacks.  _Advances in neural information processing systems_ , 30, 2017.
  * Sudakov (1978) Vladimir Nikolaevich Sudakov.  Typical distributions of linear functionals in finite-dimensional spaces of higher dimension.  In _Doklady Akademii Nauk_ , volume 243, pp. 1402–1405. Russian Academy of Sciences, 1978.
  * Tuor et al. (2021) Tiffany Tuor, Shiqiang Wang, Bong Jun Ko, Changchang Liu, and Kin K Leung.  Overcoming noisy and irrelevant data in federated learning.  In _2020 25th International Conference on Pattern Recognition (ICPR)_ , pp. 5020–5027. IEEE, 2021.
  * Turan et al. (2022) Berkay Turan, César A Uribe, Hoi-To Wai, and Mahnoosh Alizadeh.  Robust distributed optimization with randomly corrupted gradients.  _IEEE Transactions on Signal Processing_ , 70:3484–3498, 2022.
  * Yang et al. (2020) Kai Yang, Tao Jiang, Yuanming Shi, and Zhi Ding.  Federated learning via over-the-air computation.  _IEEE transactions on wireless communications_ , 19(3):2022–2035, 2020.
  * Yu et al. (2022) Xin Yu, Thiago Serra, Srikumar Ramalingam, and Shandian Zhe.  The combinatorial brain surgeon: Pruning weights that cancel one another in neural networks.  In _International Conference on Machine Learning_ , pp. 25668–25683. PMLR, 2022.

##  Appendix A Appendix

###  A.1 Proof of Proposition 1

(Convex Hull Distance Equality) Consider a set $S$ and its convex hull
$\textit{Conv}(S)$ in a Euclidean space, and an arbitrary point ${x}$ in the
space. For any probability measure $\hat{\nu}$ on $S$, we can find a point
${y}^{\prime}$ in $\textit{Conv}(S)$ as
${y}^{\prime}=\int{y}\,\mathrm{d}\nu({y})$ such that
$\|{x}-{y}^{\prime}\|^{2}=\int\|{x}-{y}\|^{2}\,\mathrm{d}\hat{\nu}({y})$,
where $\nu$ is a measure on $\textit{Conv}(S)$.

###### Proof.

We define a function $f(\nu)=\left\lVert{x}-{y}^{\prime}\right\rVert^{2}$,
where ${y}^{\prime}=\int{y}\mathrm{d}\nu({y})$. This function takes the
empirical measure $\nu$ as input and computes the squared Euclidean distance
between ${x}$ and ${y}^{\prime}$. Similarly, we define a function
$g(\nu)=\int\left\lVert{x}-{y}\right\rVert^{2}\mathrm{d}\nu({y})$, which
computes the weighted average of squared Euclidean distances between ${x}$ and
the points in the set $S$ according to the probability measure $\nu$.

Without loss of generality, let’s assume that $S$ is contained within its
convex hull $\textit{Conv}(S)$. Then, the right-hand side of the equation to
be proved in the theorem takes the minimum and maximum values, respectively,
at points within $\textit{Conv}(S)$, i.e.,

$g_{\min}=\inf_{\nu}\int\|{x}-{y}\|^{2}\,\mathrm{d}\nu({y})$

and

$g_{\max}=\sup_{\nu}\int\|{x}-{y}\|^{2}\,\mathrm{d}\nu({y}).$

The function $f(\nu)$ takes its maximum value at ${z}$, where ${z}$ is the
farthest point inside $\textit{Conv}(S)$ from ${x}$, i.e.,

$f_{\max}=\sup_{\nu}\|{x}-{y}^{\prime}\|^{2}=\|{x}-{z}\|^{2}.$

Similarly, the minimum value is obtained at ${y}^{\prime}={z}^{\prime}$, where
${z}^{\prime}$ is the closest point inside $\textit{Conv}(S)$ to ${x}$. The
minimum value can reach zero if ${x}$ is inside $\textit{Conv}(S)$. Formally,

$f_{\min}=\inf_{\mu}\left\\{\|{x}-{z}\|^{2}\middle|~{}{z}\in\textit{Conv}(S)\right\\}.$

To establish that $f_{\max}\geq g_{\max}$, we consider the maximum values of
the two functions. The function $f(\nu)$ takes its maximum value at ${z}$,
which is the farthest point inside $\textit{Conv}(S)$ from ${x}$. This means
that $f_{\max}$ is the squared Euclidean distance between ${x}$ and ${z}$,
i.e., $f_{\max}=|{x}-{z}|^{2}$. On the other hand, the function $g(\nu)$
computes the weighted average of squared Euclidean distances between ${x}$ and
the points in $S$. The maximum value of $g(\nu)$, denoted as $g_{\max}$,
corresponds to the squared Euclidean distance between ${x}$ and the farthest
point in $S$. Since $\textit{Conv}(S)$ contains $S$, it follows that ${z}$ is
farther from ${x}$ than any point in $S$. Therefore, $f_{\max}\geq g_{\max}$.
Similarly, since $\textit{Conv}(S)$ contains $S$ and ${z}^{\prime}$ is the
closest point in $\textit{Conv}(S)$ to ${x}$, it follows that $f_{\min}\leq
g_{\min}$.

We apply the intermediate value theorem to finish the proof. For any measure
$\nu$, let $h(\mu)=f(\mu)-g(\nu)$. Once can find two measures $\mu_{1}$ and
$\mu_{2}$ such that $h(\mu_{1})\leq 0$ and $h(\mu_{2})\geq 0$, hence the zero
point of $h$ exists, followed by that there is a corresponding value of $\mu$
such that $f(\mu)=g(\nu)$.

Hence, the conclusion. ∎

###  A.2 Sample Complexity

The efficiency of an estimator is often measured by its ability to deliver
accurate estimates with fewer samples, a trait referred to as ’good sample
complexity’. The entropic Wasserstein distance formulation is posited to
enhance noise reduction by optimizing this sample complexity. When
$\varepsilon=0$, the sample complexity of $W_{2}^{2}(\mu,\nu)$ stands at
$O(1/n^{\frac{1}{4}})$. Specifically,

$\lim_{n\rightarrow\infty}\mathbb{E}[W_{2}^{2}(\mu_{n},\nu_{n})-W_{2}^{2}(\mu,\nu)]=O(1/n^{\frac{1}{4}})$

as per (Nadjahi et al., 2020, Corollary 2). As $\varepsilon$ approaches
infinity, $W_{2}^{2}(\mu,\nu)$ leverages the beneficial characteristics of the
maximum mean discrepancies (MMD) Gretton et al. (2006), which narrows the
sample complexity to $O(1/\sqrt{n})$ according to (Genevay et al., 2019,
Theorem 3). The coefficient $\varepsilon$ acts as a regulator, adjusting the
sample complexity within this specified range.

###  A.3 OT Plan Optimization

For an arbitrary $\mathbf{w}^{(t)}$, the transportation plan
$\mathbf{\Pi}^{(t)}$ is obtained by solving
$W_{2}^{2}(x(\mathbf{w}^{(t)}),y)$. The optimization of the transportation
plan $\mathbf{\Pi}$ in line 9 of Algorithm 1 is based on the Sinkhorn-Knopp
algorithm, shown in Algorithm 2 below.

Algorithm 2 Sinkhorn-Knopp Algorithm for Regularized OT

1:$\mathbf{C}$ ($n\times n$ Euclidean distance matrix),
$\mathbf{a},\mathbf{b}$ (probability mass of $\mathbf{x}$ and $\mathbf{y}$),
$\varepsilon$ (regularization parameter), $\epsilon$ (tolerance for stopping
criterion)

2:$\mathbf{\Pi}$

3:Initialize $\mathbf{K}=\exp(-\mathbf{C}/\varepsilon)$,
$\mathbf{u}=\mathbf{1}_{n}/n$, $\mathbf{v}=\mathbf{1}_{n}/n$

4:repeat

5: $\mathbf{u}=\nicefrac{{\mathbf{a}}}{{\mathbf{Kv}}}$

6: $\mathbf{v}=\nicefrac{{\mathbf{b}}}{{\mathbf{uK}}}$

7:until
$\sup\left\\{\left\lVert\mathbf{a}-\mathbf{Kvu}\right\rVert_{\infty},\left\lVert\mathbf{b}-\mathbf{uKv}\right\rVert_{\infty}\right\\}<\epsilon$

8:return $\mathbf{\Pi}=\mathbf{uKv}$

The Sinkhorn-Knopp algorithm is an iterative method used to solve regularized
OT problems, aiming to find a transportation plan that minimizes total cost
while adhering to specific source and target probability distributions. In the
algorithm, an initial matrix $\mathbf{K}$ is formed using the exponential of
the negative cost matrix divided by a regularization parameter, and two
vectors are initialized as uniform distributions. These vectors are then
iteratively updated using rules derived from the Kullback-Leibler divergence,
seeking to align the row and column sums of the resulting matrix with the
given source and target distributions. The process continues until the maximum
difference between the actual and desired row and column sums falls below a
specified tolerance, ensuring the solution is feasible.

Note that $x$ and $y$ are one-dimensional projections of high-dimensional
random variables $\nabla\ell$, of which the dimension is the number of
parameters of the neural network, a.k.a. $p$. Prior research has demonstrated
that when certain moderate criteria are met, the distribution of lower-
dimensional versions of high-dimensional random variables tends to closely
follow a Gaussian (or normal) distribution Sudakov (1978); Reeves (2017);
Nadjahi et al. (2021). In the research by Janati et al. (2020), it is
highlighted that given $x$ and $y$ obeying the Gaussian distribution,
$W_{2}^{2}(x,y)$ can be reduced to a concise closed-form solution. As
illustrated in Algorithm 3, the OT plan can be directly computed relying
solely on the statistics of the empirical distributions $\mathbf{x}$ and
$\mathbf{y}$.

Algorithm 3 Closed Form Algorithm for Regularized OT

1:$\psi=\sqrt{\frac{\varepsilon}{2}}$, $a=\text{mean}(\mathbf{x})$,
$b=\text{mean}(\mathbf{y})$, $\sigma^{2}_{a}=\text{var}(\mathbf{x})$,
$\sigma^{2}_{b}=\text{var}(\mathbf{y})$,
$d_{\psi}=(4\sigma_{a}\sigma_{b}^{2}\sigma_{a}+\psi^{4})^{\frac{1}{2}}$

2: $\mathbf{\Pi}\sim\mathcal{N}\left(\begin{bmatrix}a\\\
b\end{bmatrix},\begin{bmatrix}\sigma_{a}^{2}&\frac{1}{2}\sigma_{a}d_{\psi}\sigma_{a}^{-1}\\\
\frac{1}{2}\sigma_{a}d_{\psi}\sigma_{a}^{-1}&\sigma_{b}^{2}\end{bmatrix}\right)$

Specifically, the OT plan, represented as $\mathbf{\Pi}$, is governed by the
mean and variance of both $x$ and $y$, in tandem with the regularization
parameter $\varepsilon$. To derive the $n\times n$ matrix $\mathbf{\Pi}$,
$(x_{i},y_{i})$, where $i$ spans from 1 to $n$, are employed to extract $n$
samples from the distribution produced in Algorithm 3.

$0$$20$$40$$60$$80$$100$$1{,}000$$1{,}050$$1{,}100$$1{,}150$$10.24\%$$1.80\%$$\varepsilon$OT
Objective $W_{2}^{2}$(Without Regularization Term)Sinkhorn-Knopp vs Closed-
Form SolutionSinkhorn-KnoppClosed-Form Figure 1: Comparison between the
Sinkhorn-Knopp (i.e. Algorithm 2) and the closed-form solution (i.e. Algorithm
3). The plot is made based on the data of ResNet20 trained on Cifar10. The
relative difference is computed by (red - blue) / blue.

As depicted in Figure 1, a comparative evaluation of the two algorithms in
terms of attaining the OT objective, $W_{2}^{2}$, is presented. The entropic
regularization term’s value has been omitted considering its non-impact on the
pruning optimization of $\mathbf{w}$, given a constant $\mathbf{\Pi}$. It’s
notable that the disparity between the two algorithms in their objective
optimization magnifies as $\varepsilon$ increases. Typically, in real-world
applications, the value of $\varepsilon$ oscillates between 0 and 10 (and we
use $\varepsilon=1$ most frequently in our experiments). The variance in the
performance of the two algorithms concerning OT planning remains trivial,
echoing our practical observations during the algorithmic implementation in
this study.

###  A.4 Derivative of $Q(\mathbf{w})$

Let’s start by revisiting the function $Q(\mathbf{w})$:

$\displaystyle Q(\mathbf{w})$
$\displaystyle=\left\\{\sum_{i=1}^{n}\sum_{j=1}^{n}\left\lVert
x_{i}(\mathbf{w})-y_{j}\right\rVert^{2}\pi_{ij}\right\\}+\lambda\left\lVert\mathbf{w}-\mathbf{\bar{w}}\right\rVert^{2}$
(15)

Given:

$\displaystyle x_{i}(\mathbf{w})$
$\displaystyle={\nabla\bm{\ell}}_{i}^{\top}\mathbf{w}$ $\displaystyle y_{j}$
$\displaystyle={\nabla\bm{\ell}}_{j}^{\top}\mathbf{\bar{w}}$

Differentiating $Q(\mathbf{w})$ with respect to $\mathbf{w}$:

$\displaystyle\nabla Q(\mathbf{w})$
$\displaystyle=\nabla\left[\left\\{\sum_{i=1}^{n}\sum_{j=1}^{n}\left\lVert{\nabla\bm{\ell}}_{i}^{\top}\mathbf{w}-{\nabla\bm{\ell}}_{j}^{\top}\mathbf{\bar{w}}\right\rVert^{2}\pi_{ij}\right\\}+\lambda\left\lVert\mathbf{w}-\mathbf{\bar{w}}\right\rVert^{2}\right]$
$\displaystyle=2\lambda(\mathbf{w}-\mathbf{\bar{w}})+\nabla\left[\left\\{\sum_{i=1}^{n}\sum_{j=1}^{n}\left\lVert{\nabla\bm{\ell}}_{i}^{\top}\mathbf{w}-{\nabla\bm{\ell}}_{j}^{\top}\mathbf{\bar{w}}\right\rVert^{2}\pi_{ij}\right\\}\right]$
(16)

For the gradient of the inner term, consider:

$\displaystyle\nabla\left[\left\\{\sum_{i=1}^{n}\sum_{j=1}^{n}\left\lVert{\nabla\bm{\ell}}_{i}^{\top}\mathbf{w}-{\nabla\bm{\ell}}_{j}^{\top}\mathbf{\bar{w}}\right\rVert^{2}\pi_{ij}\right\\}\right]$
$\displaystyle=\sum_{i=1}^{n}\sum_{j=1}^{n}2\pi_{ij}({\nabla\bm{\ell}}_{i}^{\top}\mathbf{w}-{\nabla\bm{\ell}}_{j}^{\top}\mathbf{\bar{w}}){\nabla\bm{\ell}}_{i}$
(17)

Expressing in Matrix Form. Given matrices:

$\displaystyle\mathbf{G}$
$\displaystyle=\begin{bmatrix}{{\nabla\bm{\ell}}}_{1}\\\
{{\nabla\bm{\ell}}}_{2}\\\ \vdots\\\ {{\nabla\bm{\ell}}}_{n}\end{bmatrix}$
$\displaystyle\mathbf{\Pi}$ $\displaystyle=\text{Matrix with elements
}\pi_{ij}$

Consider the double summation:

$\sum_{i=1}^{n}\sum_{j=1}^{n}2\pi_{ij}({\nabla\bm{\ell}}_{i}^{\top}\mathbf{w}-{\nabla\bm{\ell}}_{j}^{\top}\mathbf{\bar{w}}){\nabla\bm{\ell}}_{i}$

For each $i$, the term ${\nabla\bm{\ell}}_{i}^{\top}\mathbf{w}$ projects
vector $\mathbf{w}$ onto ${\nabla\bm{\ell}}_{i}$. To compute this for every
$i$ simultaneously, it is $\mathbf{G}\mathbf{w}$. This results in an $n\times
1$ column vector. Similarly, for each $j$, it is $\mathbf{G}\mathbf{\bar{w}}$.
This also produces an $n\times 1$ column vector.

The difference between these two projections for each $i$ and $j$ is
$\mathbf{G}\mathbf{w}-\mathbf{G}\mathbf{\bar{w}}$. This results in an $n\times
1$ column vector. To incorporate the $\pi_{ij}$ weights, we have

$\mathbf{\Pi}\left(\mathbf{G}\mathbf{w}-\mathbf{G}\mathbf{\bar{w}}\right)$

This operation gives an $n\times 1$ column vector, where each element is a
summation over $j$ for the term
$\pi_{ij}({\nabla\bm{\ell}}_{i}^{\top}\mathbf{w}-{\nabla\bm{\ell}}_{j}^{\top}\mathbf{\bar{w}})$.

To finalize the summation, we multiply it by $\mathbf{G}^{\top}$, yielding

$\mathbf{G}^{\top}\left(\mathbf{\Pi}\left(\mathbf{G}\mathbf{w}-\mathbf{G}\mathbf{\bar{w}}\right)\right).$
(18)

Combining equation 16, equation 17, and equation 18, we get:

$\displaystyle\nabla Q(\mathbf{w})$
$\displaystyle=2\lambda(\mathbf{w}-\mathbf{\bar{w}})+2\mathbf{G}^{\top}(\mathbf{\Pi}(\mathbf{G}\mathbf{w}-\mathbf{G}\mathbf{\bar{w}}))$
$\displaystyle=2[\mathbf{G}^{\top}(\mathbf{\Pi}(\mathbf{G}\mathbf{w}-\mathbf{G}\mathbf{\bar{w}}))+\lambda(\mathbf{w}-\mathbf{\bar{w}})]$
(19)

###  A.5 Experiment Setup

The models MLPNet, ResNet20, and MobileNetV1 underwent a pre-training phase of
100 epochs utilizing 4 NVIDIA Tesla V100 32 GB GPUs connected with NVlink. The
training times were approximately 1 hour for MLPNet, 3 hours for ResNet20, and
1 day for MobileNetV1. Pre-pruning accuracy levels for these models are
detailed under their respective names in Table 1. For the pruning process, we
either utilized 2 NVIDIA Tesla V100 32 GB GPUs with NVlink or a single Tesla
A100 PCIE (available in 40 or 80 GB configurations). It’s worth emphasizing
the time-intensive nature of training and pruning the MobileNetV1 on ImageNet;
thus, harnessing multiple GPUs for concurrent training is highly recommended.

In Table 1, we set the pruning stage of LR and EWR to be $15$ for MLPNet and
ResNet20 and $10$ for MobileNetV1. The sparsity $k_{1},k_{2},\ldots k_{T}$ in
Algorithm 1 is arranged following an exponential gradual pruning schedule

$k_{t}=k_{T}+(k_{0}-k_{T})\left(1-\frac{t}{T}\right)^{3}$

with the initial sparsity $k_{0}$ set to zero. The fisher sample size setup
follows (Chen et al., 2022, Table 2), shown as Table 4 of this paper below.

Table 4: Comparison of Sample and Batch Sizes for Different Models Model | MLPNet | Resnet20 | MobileNet
---|---|---|---
| Sample | Batch | Sample | Batch | Sample | Batch
WF & CBS | 1000 | 1 | 1000 | 1 | 400 | 2400
LR & EWR | 1000 | 1 | 1000 | 1 | 1000 | 16

In Table 2, Table 3, and additional results provided in the Appendix, sparsity
is set using a linear gradual pruning strategy, progressing from 0 to 0.75 or
0.95 across ten distinct stages for MLPNet and ResNet, and from 0 to 0.75
across eight distinct stages for MobileNetV1. Notably, all recorded loss
values are captured immediately post-pruning, devoid of any subsequent fine-
tuning. This approach ensures that the loss values exclusively reflect the
impact of the pruning algorithms, without being clouded by external factors.
Contrasting with Table 1, where each row represents a full pruning cycle, the
loss values here are recorded across the ten incremental pruning stages. The
empirical fisher is computed based on 100 samples with batch size to be 1.

Throughout the paper, we set $\lambda$ in the optimization problem equation 6
to $0.01$. The regularization multiplier $\varepsilon$ is set to $1$, unless
specified otherwise. The noise level $\sigma$ is set to be the standard
deviation of the original gradients.

###  A.6 Ablation Study

Table 5: Comparison of loss values in terms of $\varepsilon$ for ResNet20. The results are obtained from 25 runs, with 10% Noisy gradient and noise level $\sigma$. The target sparsity is 0.95. Sparsity | LR Loss | EWR Loss
---|---|---
|  | $\varepsilon=0$ | $\varepsilon=1$ | $\varepsilon=2$ | $\varepsilon=10$ | $\varepsilon=\infty$
0.95 | 2.89 | 2.77 | 2.78 | 2.78 | 2.79 | 2.97
0.84 | 1.76 | 1.69 | 1.61 | 1.68 | 1.71 | 2.62
0.74 | 0.68 | 0.68 | 0.67 | 0.66 | 0.66 | 0.81
0.63 | 0.36 | 0.36 | 0.36 | 0.36 | 0.36 | 0.37
0.53 | 0.31 | 0.31 | 0.31 | 0.31 | 0.31 | 0.31
0.42 | 0.31 | 0.31 | 0.31 | 0.31 | 0.31 | 0.31
0.32 | 0.30 | 0.30 | 0.30 | 0.30 | 0.30 | 0.30
0.21 | 0.30 | 0.30 | 0.30 | 0.30 | 0.30 | 0.30
0.11 | 0.31 | 0.31 | 0.31 | 0.31 | 0.31 | 0.31
$\varepsilon=0$$\varepsilon=1$$\varepsilon=2$$\varepsilon=10$$0$$0.05$$0.1$$0.15$Loss
Reduction (LR - EWR)Sparsity = 0.74 Sparsity = 0.84 Sparsity = 0.95 Figure 2:
Loss reduction with different $\varepsilon$. The result is averaged over 25
runs for ResNet20, with 10% Noisy gradient and noise level $\sigma$. The error
bar shows 90% confidence interval. The target sparsity is 0.95.

The ablation study centered on ResNet20 gives insight into the influence of
the entropic regularization multiplier, $\varepsilon$, on pruning. The aim is
to understand how different values of this parameter affect the loss
$\mathcal{L}(\mathbf{w})$ during the pruning process.

Observations from Loss Values. From Table 5, we can glean the following
points. For higher sparsity levels (0.95, 0.84, 0.74), EWR Loss consistently
outperforms the LR Loss, except for $\varepsilon=\infty$. The optimum
performance of the EWR Loss, across varying sparsities, tends to occur when
$\varepsilon$ is set at values between 1 and 2. As $\varepsilon$ tends towards
infinity, the EWR loss exceeds the LR loss. This phenomenon aligns with the
neighborhood interpolation elaborated upon in Section 3. Specifically,
incorporating an excessive number of distant data points into the
interpolation detrimentally impacts performance. For lower sparsity levels
(from 0.63 downwards), the differences between LR Loss and EWR Loss across
different $\varepsilon$ values are minuscule.

Loss Reduction Insights. Referencing Figure 2, it’s evident that the reduction
in loss (difference between LR and EWR) is more pronounced at higher sparsity
levels. For a sparsity of 0.84, $\varepsilon=1$ demonstrates the most
significant loss reduction. The trend of EWR loss for different $\varepsilon$
values is consistent across varying sparsity levels. The impact $\varepsilon$
is clearly visible at higher sparsity levels. For mid to high sparsity levels,
lower values of $\varepsilon$ (specifically 1 and 2) seem to achieve the best
balance in terms of loss.

$0$$0.2$$0.4$$0.6$$0.8$$1$$1$$2$$3$SparsityEWR LossEWR Loss for Different
$\varepsilon$
Values$\varepsilon=0$$\varepsilon=1$$\varepsilon=2$$\varepsilon=10$ Figure 3:
EWR Loss in function of the sparsity. The result is obtained over 25 runs on
ResNet20, with 10% Noisy gradient and noise level $\sigma$. The target
sparsity is 0.95.

Figure 3 illustrates EWR Loss for various $\varepsilon$ values against
sparsity, there is a clear trend of decreasing loss as sparsity decreases,
consistent across all $\varepsilon$ values. Particularly at
$\varepsilon=\infty$, there’s a pronounced increase in loss at higher
sparsity, suggesting that extreme entropic regularization might hinder optimal
pruning. However, at low sparsity levels, the loss is consistent across all
$\varepsilon$, emphasizing the minimal impact of $\varepsilon$ in limited
pruning scenarios. This underscores the significance of choosing an
appropriate $\varepsilon$, balancing between regularization and pruning
efficiency.

The ablation study provides insights into the role of the parameter
$\varepsilon$ in pruning. Its influence diminishes at low sparsity levels but
becomes significant at extremely high sparsity. For ResNet20, an optimal range
for $\varepsilon$ appears to be between 1 and 2, ensuring effective pruning.
Generally, while the exact choice of $\varepsilon$ doesn’t drastically alter
the pruning outcome, exceedingly high values, such as $1\times 10^{8}$, might
lead to less than ideal pruning decisions in very sparse networks.

###  A.7 Analysis of Pruning Stage Versus Performance

Analysis of Loss. Figure 4 depicts the relationship between the sparsity and
the difference in loss (LR - EWR) under the influence of varying noise levels.

$0$$0.1$$0.2$$0.3$$0.4$$0.5$$0.6$$0.7$$0.8$$0.9$$1$$0$$0.02$$0.04$$0.06$$0.08$$0.1$$0.12$$0.14$4.35%5.15%1.32%0.85%2.87%2.73%3.03%SparsityDifference
in Loss (LR - EWR)10% Noisy GradientNoise = $2\sigma$Noise = $\sigma$

$0$$0.1$$0.2$$0.3$$0.4$$0.5$$0.6$$0.7$$0.8$$0.9$$1$$0$$0.02$$0.04$$0.06$$0.08$$0.1$$0.12$3.82%4.05%1.55%3.50%4.07%0.49%Sparsity25%
Noisy GradientNoise = $2\sigma$Noise = $\sigma$

Figure 4: Difference in loss between LR and EWR for ResNet20. The data is from
Table 2. The relative loss improvement of EWR over LR is reported. The target
sparsity is 0.95.

The left plot, representing data with a 10% noise level, depicts a prominent
decrease in the difference of loss as sparsity is reduced for both noise
levels $2\sigma$ and $\sigma$. Notably, in the $2\sigma$ noise scenario, there
is a sharp decline in loss difference when sparsity transitions from 0.95 to
0.53. Beyond this threshold, the loss difference stabilizes and remains near
zero. For the noise level $\sigma$, the decrease appears more gradual. It is
of interest to observe that the loss difference diminishes more swiftly for
$2\sigma$ compared to $\sigma$. The error bars offer insights into data
variability, showcasing broader intervals at elevated sparsity levels, which
suggests greater unpredictability at these levels, particularly in the
$2\sigma$ setting.

The right plot represents data with 25% noise. While the trends in loss
difference share similarities with the 10% noisy gradient, the exact values
differ slightly. In this 25% noise setting, the decline in loss difference of
the two noise levels are similar. The error bars, indicating confidence
intervals, highlight the increased variability at larger sparsity levels. This
variability is most noticeable for the $2\sigma$ setting at the top sparsity
levels.

Figure 5 depicts the difference in loss between LR and EWR algorithms applied
to the MobileNetV1. In both two cases 10% and 25% noisy gradient, as sparsity
increases, the loss difference diminishes, particularly when sparsity is
approximately 0.42 or less. The 10% noisy gradient reveals that the difference
in loss for noise level $2\sigma$ is marginally higher than that of $\sigma$
for most sparsity levels. In contrast, the 25% noisy gradient sometimes
exhibits a reversal in this trend, especially at the highest sparsity level of
0.74.

$0$$0.1$$0.2$$0.3$$0.4$$0.5$$0.6$$0.7$$0.8$$0$$0.05$$0.1$$0.15$$0.2$$0.25$5.02%3.23%6.01%4.41%2.61%5.16%SparsityDifference
in Loss (LR - EWR)10% Noisy GradientNoise = $2\sigma$Noise = $\sigma$

$0$$0.1$$0.2$$0.3$$0.4$$0.5$$0.6$$0.7$$0.8$$0$$0.1$$0.2$$0.3$$0.4$7.92%5.00%6.63%8.13%4.01%5.16%Sparsity25%
Noisy GradientNoise = $2\sigma$Noise = $\sigma$

Figure 5: Difference in loss between LR and EWR for MobileNetV1. The data is
from Table 3. The relative loss improvement of EWR over LR is reported. The
target sparsity is 0.75.

Confidence intervals provided at select data points underscore the reliability
of the data, with the 25% noisy gradient showing tighter intervals compared to
the 10% scenario. This infers a higher consistency in the measurements or a
minimized effect of outliers in the 25% noisy gradient. These plots accentuate
the interplay between sparsity, noise, and the performance difference of the
two algorithms, emphasizing the significance of noise levels in algorithmic
performance evaluations across various sparsity conditions. The underlying
rationale is that as the noise intensity is too high, the performance of both
LR and EWR tends to deteriorate. Consequently, their performances converge,
resulting in a diminished differential between the two.

$1$$2$$3$$4$$5$$6$$7$$8$$9$$0$$1$$2$$3$Pruning StageAccuracy Difference
(%)Pruning Stage vs. Accuracy (%)

$1$$2$$3$$4$$5$$6$$7$$8$$9$$0$$1$$2$$3$Pruning StageTop 5 Accuracy Difference
(%)Pruning Stage vs. Top 5 Accuracy (%)

Figure 6: Accuracy difference vs. pruning stage for ResNet20. The difference
is defined to be (EWR - LR) / LR. The data is obtained over 5-25 runs over
combinations of different setup: noise level $\sigma$ and $2\sigma$, noisy
gradient proportion 10% and 25%, $\varepsilon=1,2$. The target sparsity is
0.95.

Analysis of Accuracy. The box plots in Figure 6 show how the ResNet20 model
performs at different pruning stages in terms of top-1 accuracy (also referred
to as accuracy or the overall accuracy) and top-5 accuracy difference. Top-5
accuracy means any of our model’s top 5 highest probability answers match with
the expected answer. The difference is computed for EWR, using LR as the
baseline. When we look closely, we can see patterns that help us understand
how much pruning affects the model.

In the earlier pruning stages, both the overall accuracy and the top 5
accuracy differences for EWR are minimal, suggesting that EWR remains closely
aligned with the baseline LR in terms of performance. This minimal deviation
can be viewed as an advantage, as it implies that even with simplifications
brought by pruning, EWR retains its effectiveness compared to LR.

However, as pruning intensifies, the patterns begin to reveal more about EWR’s
relative strengths. Although the accuracy difference increases, this increase
in the context of the baseline suggests that EWR might be better at handling
intense pruning than LR. Notably, in the top 5 accuracy, the discrepancies
remain relatively low compared to the overall accuracy up until the more
aggressive pruning stages. This suggests that while the model’s primary
prediction confidence may decrease, the true class is still frequently among
its top 5 predictions. In essence, EWR seems to retain a broader spectrum of
potential correct classifications even when it’s uncertain about the primary
prediction.

$1$$2$$3$$4$$5$$6$$7$$0$$5$$10$$15$Pruning StageAccuracy Difference (%)Pruning
Stage vs. Accuracy (%)

$1$$2$$3$$4$$5$$6$$7$$0$$2$$4$$6$$8$$10$Pruning StageTop 5 Accuracy Difference
(%)Pruning Stage vs. Top 5 Accuracy (%)

Figure 7: Accuracy difference vs. pruning stage for MobileNetV1. The
difference is defined to be (EWR - LR) / LR. The data is obtained over 5-10
runs over combinations of different setup: noise level $\sigma$ and $2\sigma$,
noisy gradient proportion 10% and 25%, $\varepsilon=1,2$. The target sparsity
is 0.95.

Figure 7 elucidates the performance difference between EWR and LR using
MobileNetV1 across various pruning stages. In the first plot showcasing the
overall accuracy, there’s an evident upward trend in the median accuracy
difference as pruning intensifies. Initially, the difference is marginal,
which suggests that EWR’s performance closely mirrors LR during the early
pruning stages. However, as we progress to the 5th and 6th stages, the
accuracy difference widens considerably, indicating that EWR might be
outpacing LR. The 7th stage is particularly striking, with a median accuracy
difference surpassing 6%, pointing towards a potential superiority of EWR in
extreme pruning scenarios.

The second plot focuses on the top 5 accuracy differences, presenting a more
varied pattern. The early stages indicate a tight performance race between EWR
and LR. By the 2nd stage, there’s a minor dip, hinting at EWR’s possible
underperformance. This scenario changes by the 3rd stage as EWR regains
momentum. The later stages, especially the 5th and 6th, denote a significant
rise in the median difference in EWR’s favor. Much like the accuracy chart,
the 7th stage is distinct, with EWR showcasing a considerable advantage in the
top 5 accuracy over LR.

###  A.8 Optimal Transport Visualization

In this section, we showcase the optimized OT plan denoted as $\mathbf{\Pi}$.
This was derived from the pruning applied to ResNet20. Figure 8 displays two
data sets: 1) $\\{x_{i}\\}_{i=1}^{n}$: This is $\mathbf{G}\mathbf{w}$ where
$\mathbf{w}$ is the pruned model. 2) $\\{y_{i}\\}_{i=1}^{n}$: This represents
$\mathbf{G}\mathbf{\bar{w}}$ where $\mathbf{\bar{w}}$ is the original unpruned
model.

In Figure 9, the matrix $\mathbf{\Pi}$ is shown through vibrant heatmaps that
adjust with varying $\varepsilon$ values. For small $\varepsilon$, the
majority of the data remains near a diagonal. As $\varepsilon$ increases,
there’s a broader data distribution, notably for central data points.

$-60$$-50$$-40$$-30$$-20$$-10$$0$$10$$20$$30$$40$$50$$60$$0$$10$$20$$30$$40$$50$$60$Frequency$\\{y_{i}\\}_{i=1}^{n}$$\\{x_{i}\\}_{i=1}^{n}$

$-50$$-40$$-30$$-20$$-10$$0$$10$$20$$30$$40$$50$$0$$0.05$$0.1$$0.15$$0.2$$0.25$DensityKDE
$y$KDE $x$

Figure 8: The histogram and kde plot of the empirical distribution
$\\{x_{i}\\}_{i=1}^{n}$ and $\\{y_{i}\\}_{i=1}^{n}$.

(a) $\varepsilon=0$

(b) $\varepsilon=0.01$

(c) $\varepsilon=1$

(d) $\varepsilon=10$

Figure 9: Heatmap of the optimized OT plan $\mathbf{\Pi}$ across different
$\varepsilon$. Darker color indicate larger value in $\mathbf{\Pi}$.