Jump-Diffusion Langevin Dynamics for Multimodal Posterior Sampling

Jacopo Guidolin Vyacheslav Kungurtsev Ondřej Kuželka

University of Padua1 Czech Technical University2 Czech Technical University2

###### Abstract

Bayesian methods of sampling from a posterior distribution are becoming
increasingly popular due to their ability to precisely display the uncertainty
of a model fit. Classical methods based on iterative random sampling and
posterior evaluation such as Metropolis-Hastings are known to have desirable
long run mixing properties, however are slow to converge. Gradient based
methods, such as Langevin Dynamics (and its stochastic gradient counterpart)
exhibit favorable dimension-dependence and fast mixing times for log-concave,
and “close” to log-concave distributions, however also have long escape times
from local minimizers. Many contemporary applications such as Bayesian Neural
Networks are both high-dimensional and highly multimodal. In this paper we
investigate the performance of a hybrid Metropolis and Langevin sampling
method akin to Jump Diffusion on a range of synthetic and real data,
indicating that careful calibration of mixing sampling jumps with gradient
based chains significantly outperforms both pure gradient-based or sampling
based schemes.

##  1 INTRODUCTION

11footnotetext: The author acknowledges the support of the “Erasmus+”
Programme.22footnotetext: The authors acknowledge the support of the OP VVV
funded project CZ.02.1.01/0.0/0.0/16_019/0000765 “Research Center for
Informatics”. The access to the computational infrastructure of the OP VVV
funded project CZ.02.1.01/0.0/0.0/16_019/0000765 “Research Center for
Informatics” is also gratefully acknowledged.

One of the key tasks in machine learning is trying to estimate a complex and
unknown model with an approximated surrogate. Most of the time the latter one
will be a parametric model i.e. a model heavily dependent on a finite number
of learnable parameters.

In this work we will concentrate on the setting of Bayesian inference. In
particular our interest is to investigate effective techniques to productively
explore the posterior distribution on the space of the parameters of any given
parametric model; (Mackay, 1995), (Neal, 1996).

There are two main paradigms that are commonly followed in order to perform
this task. The first one being associated with techniques from the wide branch
of optimization and the second being sampling, specifically sampling by means
of Markov Chain Monte Carlo based methods (Barbu and Zhu, 2020).

Broadly speaking, sampling algorithms, which use a known proposal distribution
and careful probabilistic sample acceptance-rejection criteria, are known to
reliably mix to a wide range of posteriors. However, the mixing rate, or
convergence in terms of probability measures, is slow, and scales poorly with
the parameter dimension; (Mengersen and Tweedie, 1996), (Wang, 2020).

By contrast, optimization based, or gradient-informed methods, exhibit fast
ergodicity for log-concave potentials (see, e.g. Dalalyan (2014), Durmus and
Moulines (2016) and Cheng and Bartlett (2018)). In addition to that, while
they are long-run ergodic, practically they may be seen to exhibit “mode
seeking” behavior, wherein one mode in a multi-modal distribution is sampled
for a large portion of the chain, leaving the rest of the parameter space less
informed. This was more comprehensively studied by Ma et al. (2019), who
showed that for a specific class of nonconvex potentials, the dimensional
dependence of sampling is superior to that of gradient-based methods.
Recently, however, there has been interest in large scale multimodal posterior
models. In particular performing Bayesian inference on the parameters of a
Neural Network can provide desirable properties. Bayesian Neural Networks
(BNNs) can bring, along with reliable uncertainty estimates and an intrinsic
protection against overfitting, natural support for active learning, continual
learning, and decision-making under uncertainty (Izmailov et al., 2021). Yet,
the posteriors characterizing BNNs exhibit both significant multi-modality,
which should benefit sampling-based methods, while at the same time having a
large parameter dimension (and hence exhibit favorable performance for
gradient based samplers), and the modes have standalone significance for
inference (i.e., for MAP inference).

In this paper, we introduce a procedure based on performing Jump Diffusions on
a singular posterior, rather than explicitly jumping across components of a
mixture. The intention is to achieve the advantages of point-based and
gradient-based sampling simultaneously: periodic jumps ensure sufficient
exploration of the parameter space and subchains of Langevin inform the
precise details of the, generally, most significant modes. We shall
demonstrate that indeed, on both synthetic and real data sets, the combination
of the two procedures performs better than either point-based or gradient-
based sampling procedures alone.

##  2 BACKGROUND

In this Section we introduce two standard representatives of the two paradigms
of sampling as well as introduce our jump diffusion based procedure.

###  2.1 Metropolis Hastings Algorithm

Metropolis-Hastings (MH) is a reliable algorithm especially in the context of
unnormalized probabilities. It allows us to draw samples from a desired
probability distribution $P$ of density $p(x)$ if we have access to a function
$f(x)$ such that $f\propto p$, (Barbu and Zhu, 2020).

The main idea of the algorithm is to take any method that iteratively jumps
between different states and adding to the latter an additional _acceptance
step_ between the jumps. The purpose of the acceptance step is to “correct”
the initial algorithm in such a way that the Markov Chain generated by the
updated algorithm will satisfy the detailed balance equation. In particular,
given that $\mathcal{K}$ is the transition probability of the chain, that is:

$\mathcal{K}\left(X,Y\right)p(X)=\mathcal{K}\left(Y,X\right)p(Y).$

The standard way of proceeding is to start with an inexpensive proposal
distribution that allows us to conditionally sample from it rather
inexpensively, only then correct this sampling by applying MH. Each of these
proposed samples can then be either accepted or rejected according to a
particularly designed _acceptance probability_.

It has to be noted that the more we move into high dimensional settings, and
the more the geometry of the posterior distribution gets complex, the more an
accurate design for the proposal distribution plays a key role on yielding an
effective rate of convergence for MH. Formally, the mixing rate of MH depends
on the worst-case ratio of the desired and proposal distributions. This factor
typically scales unfavorably with the dimension; (Mengersen and Tweedie,
1996), (Wang, 2020).

We can see pseudocode for the Metropolis-Hastings with Independence Metropolis
Sampler design in Algorithm 1, (Barbu and Zhu, 2020).

Input: target distribution $p(\cdot)$, current state of the chain
$\theta_{t}$, a proposal probability $q(\cdot)$

Output: the next state of the chain $\theta_{t+1}$

1

2Generate a new proposal state $\bar{\theta}$ by sampling from $q$

3Calculate the acceptance probability:
$p_{a}(\theta,\bar{\theta})=\min\left(1,\frac{q(\theta)}{q(\bar{\theta})}\frac{p(\bar{\theta})}{p(\theta)}\right)$

Accept the new state with probability $p_{a}(\theta,\bar{\theta})$ and set
$\theta_{t+1}=\bar{\theta}$, otherwise $\theta_{t+1}=\theta_{t}$.

Metropolis-Hastings Algorithm, 1

###  2.2 Stochastic Gradient Langevin Dynamics

Let $\mathcal{M}$ be a parametric model with $\theta$ as vector of parameters,
$\theta\in\Theta$. Let $p\left(\cdot\right)$ be a prior distribution on
$\Theta$. Let $\mathcal{X}$ be the input space of our model and, given
$x\in\mathcal{X}$, let $p\left(x\mid\theta\right)$ be the density of the
quantity $x$ if our model has $\theta$ as parameter vector. Given a set of $n$
data items the update rule for the parameters is the following:

$\displaystyle\Delta\theta_{t}$
$\displaystyle=\frac{\varepsilon_{t}}{2}\Bigg{(}\underbrace{\nabla\log\left(p\left(\theta_{t}\right)\right)}_{\begin{subarray}{c}\text{priors
gradient}\end{subarray}}$
$\displaystyle\quad+\frac{N}{n}\sum_{i=1}^{n}\underbrace{\nabla\log\left(p(x_{i}\mid\theta_{t})\right)}_{\begin{subarray}{c}\text{likelihood
gradients}\end{subarray}}\Bigg{)}$
$\displaystyle\quad+\underbrace{\nu_{t}}_{\begin{subarray}{c}\text{noise}\end{subarray}}$

where $\nu_{t}\sim N\left(0,\varepsilon_{t}\right)$.

It is important to note that the gradient step sizes and the standard
deviation of the injected noise are balanced in such a way that the variance
of the noise matches the variance of the posterior.

Naturally, discretization of Langevin dynamics, as necessary for practical
computation, introduces bias. Performing a Metropolis-Hastings acceptance-
rejection step, however, corrects this.

One other possible approach is to use a scheduled stepsize decay, as seen in
Welling and Teh (2011), which is also associated with a “stochastic gradient”
approach of using subsamples. This approach slows the per sample convergence
rate, however, and appears to be additionally mode-seeking.

Langevin dynamics has been studied extensively in regards to its mixing
properties. The associated SDE model appears in a number of physical phenomena
as well, and thus the Langevin type equation has received considerable
analysis Stroock and Varadhan (1997). As an algorithm for sampling, its
interest has increased due to the improved dimensional dependence of the
mixing rate. This is especially the case of log-concave potentials, for which
the performance is akin to gradient descent for convex functions, convergence
is fast. At the same time, however, the gains in local accuracy from using the
gradient are contrasted with the local inertia biasing against global
exploration of the parameter space. Namely, the wells corresponding to
attraction basins for the Langevin dynamics also exhibit long escape times,
and thus slow tendency to explore outside of a contemporaneous mode, (Chen et
al., 2020).

###  2.3 Jump-Diffusion

Our method is inspired by the concept of Jump-Diffusions, also known as Hybrid
Switching Diffusions in the literature (Yin and Zhu, 2010). Jump-Diffusions
are a particular class of diffusion processes involving multiple diffusions,
or SDEs with a drift and a Wiener process, together with a Poisson jump
process. The jump process randomly switches between the different diffusions.

Aside from describing some complex stochastic systems in physics, jump
diffusions have been applied to sample from mixtures. They were analyzed for
its statistical potential in (Green, 1995). Additional case studies appear in
(Grenander and Miller, 1994). Broadly speaking, the procedures have potential
for successfully sampling from posterior distributions with mixtures of
notably different distributions, with some notable success in applications
such as image segmentation (Han et al., 2004). Quality performance, however,
typically requires well informed data-driven priors.

The use of jump-diffusions for single potentials is, to the best of the
authors’ knowledge, unexplored, despite its natural intuitive potential. Below
we formally introduce our particular procedure inspired by the hybrid jump
diffusion literature.

##  3 JUMP-DIFFUSION LANGEVIN DYNAMICS

We now propose our jump-diffusion algorithm for single posterior sampling. The
method combines MH with adjusted Langevin. In particular, for some number of
samples, it performs a chain of adjusted Langevin. At the same time, after a
number of such samples, the number being itself sampled from a Poisson
distribution, an MH step occurs: a proposal distribution is sampled and then
accepted or rejected.

We define the constant stepsize of the Langevin sequence by $\varepsilon$ and
note that the frequency of jumps is regulated by a Poisson process,
$\tau\sim\text{Pois}(\lambda)$, where $\lambda$ can be considered a tunable
hyperparameter.

The underlying intuition is to prevent the Langevin diffusion from getting
stuck on certain modes by allowing the process to jump to a different state
without making use of first order information. We intend for the chain to, as
such, have high resolution for the most significant modes, while at the same
time ensuring that all nodes are explored, and the broad uncertainty as well
(but at coarser resolution). We shall see that this is precisely the behavior
we see in practice.

One can consider running parallel, or multiple sequential, chains with
different starting points. There are significant problems with this approach,
however. First, the lack of guarantees, since the jump process is not
Metropolis-regulated, as in our case. This would introduce bias. Second, this
would depend on the choice of the associated hyperparameters – there should be
a nice correspondence between the number of modes and the number of starting
points, otherwise we can’t expect such a method to perform well. By contrast,
with a large enough sample size, the procedure we introduce should effectively
characterize the posterior, with minimal hyperparameter dependence.

Input: target distribution $p(\cdot)$, current iteration $t$, current state of
the chain $\theta_{t}$, a proposal probability $q(\cdot)$, next jumping
iteration $\tilde{t}$, jump rate $\lambda$

Output: The next state of the chain $\theta_{t+1}$, the next jumping iteration
$\tilde{t}$

1

2if  _$t=\tilde{t}$_ then

3 Generate a new proposal state $\bar{\theta}$ by sampling from $q$

4 Calculate the acceptance probability:
$p_{a}(\theta_{t},\bar{\theta})=\min\left(1,\frac{q(\theta_{t})}{q(\bar{\theta})}\frac{p(\bar{\theta})}{p(\theta_{t})}\right)$

5 Accept the new state with probability $p_{a}(\theta_{t},\bar{\theta})$ and
set $\theta_{t+1}=\bar{\theta}$, otherwise $\theta_{t+1}=\theta_{t}$.

6 Sample $i\sim\text{Pois}(\lambda)$

7 $\tilde{t}=\tilde{t}+i$

8 else

9

10 Calculate the Langevin step update:  $\displaystyle\Delta\theta_{t}$
$\displaystyle=\frac{\varepsilon}{2}\Bigg{(}\nabla\log(p(\theta_{t}))$
$\displaystyle\quad+\sum_{i=1}^{N}\nabla\log(p(x_{i}\mid\theta_{t}))\Bigg{)}$
$\displaystyle\quad+\nu_{t}$

11 Set $\theta^{*}=\theta_{t}+\Delta\theta_{t}$.

12 Calculate the acceptance probability:
$p_{a}(\theta_{t},\theta^{*})=\min\left(1,\frac{q(\theta_{t})}{q(\theta^{*})}\frac{p(\theta^{*})}{p(\theta_{t})}\right)$

13 Accept the new state with probability $p_{a}(\theta_{t},\theta^{*})$ and
set $\theta_{t+1}=\theta^{*}$, otherwise $\theta_{t+1}=\theta_{t}$.

Our Algorithm, 2

The use of MH-regulated jumps is facilitated by a good choice of the proposal
probability. This can be done in a data-driven way, with some form of
pretraining. Or, alternatively, with sufficiently large variance, standard
Gaussians should work sufficiently well (Izmailov et al., 2021).

###  3.1 Theoretical Properties

The seminal reference on jump diffusions is the monograph Yin and Zhu (2010).
The canonical system is a set of stochastic diffusions and a Poisson jump
process that switches between them. Formally,

$dX(t)=b(X(t),\alpha(t))dt+\sigma(X(t),\alpha(t))dW$

where $\alpha(t)\in\mathcal{N}\subset\mathbb{N}$ is a set of integers
identifying distinct diffusion processes with (potentially state dependent)
transition matrix $q_{ij}(x)$.

Theorem 4.4 Yin and Zhu (2010) establishes weak convergence of the probability
density corresponding to the stochastic process to a stationary distribution
corresponding to a mixture of the stationary distributions for the set of
diffusions. Since in our case, the diffusions are the same, this corresponds
to the desired stationary distribution.

The use of Metropolis adjusted Langevin for the chains results in an unbiased
simulation of the diffusion, and as such the discretization maintains the
ergodic properties of the SDE. Exponential convergence under the log-concave
condition for the potentials is established in Roberts and Tweedie (1996).
Thus, during the procedure, any samples taken within a basin with a locally
strongly convex potential are converging geometrically in probability to the
distribution associated with that mode. Thus mixing is particularly fast in
the regions corresponding to potential modes, which is broadly desired as
accurate characterization of posterior modes is more significant for the user
than the rest of the phase space.

Taken together, we can conglomerate the conclusion of the theoretical results
as follows:

  * •

As Metropolis-adjusted discretized Langevin maintains the ergodicity of the
Langevin sampling procedure, and jump diffusions are known to be ergodic, we
expect the procedure to converge asymptotically, in standard probability
distances, to the desired posterior

  * •

When the chain is exploring a mode that is locally strongly log-concave, the
localized parameter region corresponding basin of attraction for that mode in
terms of gradient flows exhibits geometric ergodicity, i.e., exponential
convergence towards the mode.

With these in mind, we can see how depending on the landscape, we should see
how broadly this procedure should result in faster mixing than either Langevin
or Metropolis MCMC alone.

##  4 EXPERIMENTS

Here we present the findings of our numerical experiments comparing the
sampling methods described above. In standard practice, we shall present
results with both low-dimensional models whose posteriors we can display
visually as well as more coarse comparisons for large-dimensional models.

###  4.1 Low-Dimensional Problems

####  4.1.1 One Dimensional Mixture

We begin with some standard one-dimensional examples. In particular, we
clearly show that if the distribution we aim to estimate is multimodal, it can
be difficult to estimate by means of pure Langevin Dynamics even in the one-
dimensional case.

In particular we assumed the posterior distribution to be of form
$p(x)=\frac{1}{Z}e^{-E(x)}$ where

$E(x)=x^{2}(\sin^{2}({2^{\alpha}x)}+\delta)$ (1)

with $\delta=2\cdot 10^{-2}$ and $\alpha\in\\{-2,-1,0,1\\}$. We tried to
estimate this hand-crafted distribution by means of Metropolis-Hastings,
Langevin Dynamics and Our method. All methods were run for 1M iterations
reaching the results shown in Figure 1. As we can see, the multimodality of
the distribution is barely captured when using pure Langevin Dynamics.

Figure 1: Target distribution and its estimations. Target, Metropolis-
Hastings, Langevin and Ours on the rows; $\alpha\in\\{-2,-1,0,1\\}$ on the
columns, respectively. Every method was run for 1M iterations.

####  4.1.2 Two Dimensional Mixture

Next we generated a two dimensional dataset in such a way that the resulting
posterior distribution was multimodal with distinct and slightly separated
modes. We suspect that these traits will prompt Stochastic Gradient Langevin
Dynamics to produce incorrectly estimated posterior distributions.

Figure 2: True posterior distribution (left) and posterior distribution
estimated using Metropolis-Hastings (right).

We let two parameters, $\theta_{0}$ and $\theta_{1}$, were sampled from a
Gaussian distribution: $\theta_{0},\theta_{1}\sim N\left(0,1\right)$. Then we
sampled 100 data points $\left\\{x_{1},...,x_{100}\right\\}$ with:

$x_{i}\sim\frac{1}{2}N\left(\theta_{0}^{2},1\right)+\frac{1}{2}N\left(\left(\theta_{0}+\theta_{1}\right)^{2},1\right)$

where $\theta_{0}=0$, $\theta_{1}=1$.

We know the true posterior distribution in the parameter space and so we are
able to plot it. In particular, working with log-probabilities:

$\displaystyle L\left(\theta_{0},\theta_{1}\right)$
$\displaystyle=\sum_{i=1}^{100}\log\left(\frac{1}{2}e^{\frac{\left(x_{i}-\theta_{0}^{2}\right)^{2}}{2}}+\frac{1}{2}e^{\frac{\left(x_{i}-(\theta_{0}+\theta_{1})^{2}\right)^{2}}{2}}\right)$
$\displaystyle\quad+\left(\frac{-\theta_{0}^{2}}{2}\right)+\left(\frac{-\theta_{1}^{2}}{2}\right)$

resulting in the plot seen in Figure 2.

In addition, in Figure 2, we show the empirical distribution generated by MH.
Even if 1M iteration can be said to be low, the resulting posterior
distribution is sound, the multimodality of the posterior is correctly
captured and the ratio between the modes seems to be preserved.

As described in Welling and Teh (2011), and due to the similar assumptions, we
ran the Stochastic Gradient Langevin process with a batch size of 1, for 10000
sweeps through the whole dataset (for a total of 1M iterations) and with
decaying learning rate from $10^{-2}$ to $10^{-4}$ regulated by
$\varepsilon_{t}=a\left(b+t\right)^{\gamma}$ where $a=0.2$, $b=231$ and
$\gamma=0.55$.

Figure 3: Examples of unsuccessfully estimated distributions using Stochastic
Gradient Langevin Dynamics.

By examining the posterior distribution estimated by SGLD (Figure 3), we see
the similar mode-seeking behavior. Due to the further separation of the modes
and the higher complexity of the posterior, the resulting estimated
distribution has notably truncated support. The method did not manage to visit
the parameter space extensively and the obtained distribution can be labeled
as unsatisfactory.

Next, in considering our Jump Langevin procedure, we extensively tested
multiple combinations of hyperparameter settings, involving varying stepsizes
and jump frequency $\lambda$, some of which are shown in Figure 4. We can say
that our method seems to be insensitive to the choice of the jump $\lambda$
parameter. On the stepsize does affect the results, but such is standard in
using Langevin discretizations.

Overall we can say that almost all the combinations of settings used in Figure
4 yielded satisfactory estimation of the posterior, managing to correctly
express its multimodality while producing a precise characterization of the
significant modes.

Figure 4: Distributions estimated with Our method with
$\lambda\in\\{50,100,200,1000\\}$ on the columns and $\varepsilon\in\\{5\cdot
10^{-5},1\cdot 10^{-5},5\cdot 10^{-6}\\}$ on the rows, respectively.

####  4.1.3 Synthetic Highly-Multimodal Posterior

Finally, among low dimensional examples, we now produce an even more
multimodal distribution to test the sampling algorithms under even more
challenging conditions.

For these experiments, we did not generate a dataset, instead we assumed
direct access to first and zeroth order information of the posterior
distribution.

In particular we assumed the posterior distribution to be of form
$p(x)=\frac{1}{Z}e^{-E(x)}$ where

$E(x)=(x^{2}+y^{2})(\sin^{2}{xy}+\delta)$ (2)

with $\delta=10^{-2}$. The obtained synthetic posterior proves to be highly
multimodal, with dense and separated modes as can be seen in Figure 5.

For a sound estimation of the posterior we generated a chain with 50M+ states
using Metropolis-Hastings and we removed the initial “burn-in” chunk,
obtaining a chain with 50M states.

Trying to properly estimate this distribution using pure Langevin Dynamics was
unachievable, the resulting distributions were almost exclusively unimodal and
heavily dependent on the starting point.

On the other hand, when trying to estimate the distribution with Our method,
we obtain the result showing in Figure 5. The presented distribution was
obtained after keeping the last 50M states of a chain with 50M+ states
generated with our method.

We also measure the “difference”, by means of the (W1) Wasserstein Distance,
between the distribution generated by the full MH chain and the distribution
generated with our method, as the number of chain-states grew. We show the
satisfactory results we obtain in Figure 6.

Figure 5: True posterior distribution (left) and posterior distribution
estimated with Our Method (right).

Figure 6: Plot of Wasserstein Distance for $\theta_{0}$ and $\theta_{1}$
distributions, respectively.

###  4.2 High-dimensional setting

For the experiments in a high dimensional setting we tested our method on
Bayesian Neural Networks.

In particular, we take into consideration a binary classification problem. The
dataset we used is the _fourclass_ dataset found in the LIBSVM library. For
this task we chose to use a fully connected neural network with 3 hidden
layers of 8, 8 and 4 neurons, respectively. The output layer had only 1 neuron
with a sigmoid as activation function.

We pretrained our network with SGD in order to gain some previous knowledge
about the parameters of the network, specifically we center the priors on the
optimal parameters, but use a relatively high standard deviation, in the
context of Neural Network parameters, of 0.5.

We run three separate chains using Metropolis-Hastings, Langevin Dynamics and
Our method for $3\cdot 10^{6}$ iterations, with a stepsize of $3\cdot 10^{6}$
for both gradient based methods and $\lambda=500$ for Our method.

To test the three methods we we took one sample 30k times from each chain. For
each sample of the parameter vector, we performed inference on the test set
and evaluated the performance by computing the accuracy of the predicted and
actual label values. From our findings, the Metropolis-Hasting results were
both noisy (i.e., exhibiting large variability across multiple runs) and
always significantly underperformed the other two methods. We can see a
comparison between pure Langevin Dynamics and Our method in Figure 7.

Figure 7: Distribution of accuracies for the sampled chain states. “Holes” on
the x axis are caused by the finite number of elements in the test set.

###  4.3 Additional Experiments

We want to provide a systematic test that highlights the comparison between
Langevin Dynamics and Jump Langevin (our method) on some increasingly
multimodal posteriors. To do that we synthetically generated four
unidimensional posterior distributions with, respectively, 2, 8, 16 and 32
equal and separated modes.

We then created a baseline in order to compare our methods. To do that, for
every one of the four posteriors, we generated a chain with 10M states using
the Metropolis-Hastings algorithm. We then generated, for each different
posterior, one chain with 10M states using Langevin Dynamics and one chain
with 10M states using Our method. To conclude we observed the evolution of the
measured Wasserstein Distance as the number of chain states increased.

Result show (Figure 8) the inability of correctly expressing the posteriors
when using pure Langevin Dynamics. Our method instead appears to effectively
converge to the appropriate stationary distribution.

Figure 8: Evolution of the Wasserstein Distance for eight different chains.
Langevin Dynamics on the top row, Jump Langevin / Our method on the bottom
row. The chains are estimating the posteriors with 2, 8, 16, 32 modes,
respectively, on the columns.

##  5 Conclusion

The two primary classes of methods for sampling, as based on solely zeroth-
order, or first-order posterior information, exhibit particular comparative
advantages and disadvantages that favor them for different problem classes. We
have shown that by carefully combining the two mechanisms of gradient seeking
together with proposal-defined jumps, we can sample from the challenging class
of high dimensional and multimodal distributions. It can be considered that
this resembles, e.g., trust region step acceptance rejection of gradient and
Hessian generated iterate estimates in standard nonconvex optimization
algorithms. As such, one could foresee a number of additional developments of
modularly considering different gradient-based chains (i.e., HMC) and jump
mechanisms (i.e., as including memory). This can open a broad research program
of investing more challenging posterior distributions as the accuracy and
robustness demands of contemporary AI systems advance.

#### References

##

  * Barbu and Zhu (2020) A. Barbu and S. Zhu.  _Monte Carlo Methods_.  Springer, 2020.  ISBN 9789811329708.
  * Chen et al. (2020) X. Chen, S. S. Du, and X. T. Tong.  On stationary-point hitting time and ergodicity of stochastic gradient langevin dynamics.  _Journal of Machine Learning Research_ , 2020.
  * Cheng and Bartlett (2018) X. Cheng and P. Bartlett.  Convergence of langevin mcmc in kl-divergence.  In F. Janoos, M. Mohri, and K. Sridharan, editors, _Proceedings of Algorithmic Learning Theory_ , volume 83 of _Proceedings of Machine Learning Research_ , pages 186–211. PMLR, 07–09 Apr 2018.
  * Dalalyan (2014) A. Dalalyan.  Theoretical guarantees for approximate sampling from smooth and log-concave densities.  _Journal of the Royal Statistical Society: Series B (Statistical Methodology)_ , 79, 12 2014.  doi: 10.1111/rssb.12183.
  * Durmus and Moulines (2016) A. Durmus and E. Moulines.  High-dimensional bayesian inference via the unadjusted langevin algorithm, 2016.
  * Green (1995) P. J. Green.  Reversible jump markov chain monte carlo computation and bayesian model determination.  _Biometrika_ , 82(4):711–732, 1995.
  * Grenander and Miller (1994) U. Grenander and M. I. Miller.  Representations of knowledge in complex systems.  _Journal of the Royal Statistical Society: Series B (Methodological)_ , 56(4):549–581, 1994.
  * Han et al. (2004) F. Han, Z. Tu, and S.-C. Zhu.  Range image segmentation by an effective jump-diffusion method.  _IEEE Transactions on pattern analysis and machine intelligence_ , 26(9):1138–1153, 2004.
  * Izmailov et al. (2021) P. Izmailov, S. Vikram, M. D. Hoffman, and A. G. G. Wilson.  What are bayesian neural network posteriors really like?  In M. Meila and T. Zhang, editors, _Proceedings of the 38th International Conference on Machine Learning_ , volume 139 of _Proceedings of Machine Learning Research_ , pages 4629–4640. PMLR, 18–24 Jul 2021.
  * Ma et al. (2019) Y.-A. Ma, Y. Chen, C. Jin, N. Flammarion, and M. I. Jordan.  Sampling can be faster than optimization.  _Proceedings of the National Academy of Sciences_ , 116(42):20881–20885, 2019.  doi: 10.1073/pnas.1820003116.
  * Mackay (1995) D. J. C. Mackay.  Probable networks and plausible predictions — a review of practical bayesian methods for supervised neural networks.  _Network: Computation in Neural Systems_ , 6(3):469–505, 1995.  doi: 10.1088/0954-898X_6_3_011.
  * Mengersen and Tweedie (1996) K. L. Mengersen and R. L. Tweedie.  Rates of convergence of the Hastings and Metropolis algorithms.  _The Annals of Statistics_ , 24(1):101 – 121, 1996.  doi: 10.1214/aos/1033066201.  URL https://doi.org/10.1214/aos/1033066201.
  * Neal (1996) R. M. Neal.  _Bayesian learning for neural networks_.  Springer Science & Business Media, 1996.
  * Roberts and Tweedie (1996) G. O. Roberts and R. L. Tweedie.  Exponential convergence of langevin distributions and their discrete approximations.  _Bernoulli_ , pages 341–363, 1996.
  * Stroock and Varadhan (1997) D. W. Stroock and S. S. Varadhan.  _Multidimensional diffusion processes_ , volume 233.  Springer Science & Business Media, 1997.
  * Wang (2020) G. Wang.  Exact convergence rate analysis of the independent metropolis-hastings algorithms, 2020.  URL https://arxiv.org/abs/2008.02455.
  * Welling and Teh (2011) M. Welling and Y. W. Teh.  Bayesian learning via stochastic gradient langevin dynamics.  In _Proceedings of the 28th International Conference on International Conference on Machine Learning_ , ICML’11, page 681–688, Madison, WI, USA, 2011. Omnipress.  ISBN 9781450306195.
  * Yin and Zhu (2010) G. Yin and C. Zhu.  _Hybrid switching diffusions. Properties and applications_ , volume 63.  01 2010.  ISBN 978-1-4419-1106-3.  doi: 10.1007/978-1-4419-1105-6.


# Compressing the Backward Pass of Large-Scale Neural Architectures by
Structured Activation Pruning

Daniel Barley and Holger Fröning  Computing Systems Group, Institute of
Computer Engineering
Heidelberg University, Germany
<EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

The rise of Deep Neural Networks (DNNs) has led to an increase in model size
and complexity, straining the memory capacity of GPUs. Sparsity in DNNs,
characterized as structural or ephemeral, has gained attention as a solution.
This work focuses on ephemeral sparsity, aiming to reduce memory consumption
during training. It emphasizes the significance of activations, an often
overlooked component, and their role in memory usage. This work employs
structured pruning in Block Sparse Compressed Row (BSR) format in combination
with a magnitude-based criterion to efficiently prune activations. We
furthermore introduce efficient block-sparse operators for GPUs and showcase
their effectiveness, as well as the superior compression offered by block
sparsity. We report the effectiveness of activation pruning by evaluating
training speed, accuracy, and memory usage of large-scale neural architectures
on the example of ResMLP on image classification tasks. As a result, we
observe a memory reduction of up to 32% while maintaining accuracy.
Ultimately, our approach aims to democratize large-scale model training,
reduce GPU requirements, and address ecological concerns.

##  I Introduction

Deep Neural Networks (DNNs) have grown exponentially in size and complexity in
recent years. Even though the success of deep learning is, in part, fueled by
advancements in Graphics Processor Unit (GPU) technology, this rapid increase
in model size is in direct conflict with the notoriously small memory capacity
of GPU architectures. Resource efficiency in general [1, 2] and sparsity [3]
in particular has therefore become an active field of research in deep
learning. There are many approaches to sparsity in DNNs. Fundamentally they
can be categorized into _structural_ , or _model_ sparsity, and _ephemeral_
sparsity. Structural sparsity is applied to parts of the network itself, this
may range from individual weights [4] to whole layers [5]. The sparse pattern
is then inherent to the network and applied equally to all inputs.
Structurally sparse approaches, like magnitude-based weight pruning, require
information on network parameters unknown prior to training. Ephemeral
sparsity, on the other hand, is applied dynamically per input. This allows for
pruning criteria independent of network parameters. Because network parameters
are subject to constant change during the training phase of a DNN, most
structural approaches, aside from simple static masks, disqualify for
training. An ephemeral pruning strategy is therefore pursued in this work.

With the type of sparsity in place, the sparse component remains to be
determined. One can observe that the state of neural architectures is composed
of weights, biases, gradients, and activations. Activations are often overseen
as they seem to exist only temporarily during forward inference. During
training, however, the backpropagation algorithm requires them to be stored
until the backward path is executed. Moreover, activations make up the vast
majority of the memory footprint during training as they scale with an
additional mini-batch dimension. For the targeted ResMLP architecture [6] we
observe that 90% of the total memory consumption is due to activations. This
percentage is even larger for convolutional architectures, as weight sharing
reduces memory pressure with regard to parameters, effectively increasing the
percentual contribution made by activations to the overall state. Considering
Amdahl’s Law, we identify activations to offer the largest potential for
reducing state.

TABLE I: ResMLP-S12 Memory Footprint by Component Component | Input | Model | Optimizer | Activations | Total
---|---|---|---|---|---
Memory [$\mathrm{MiB}$] | 18.5 | 59.6 | 121.1 | 1599.2 | 1798.4
% | 1.0 | 3.3 | 6.7 | 89.0 | 100.0

While pruning has become a widely used method, it is equally well known that
most processor architectures do not handle unstructured sparsity well. This is
especially true for GPUs, where irregular memory access patterns lead to a
rapid degradation of memory bandwidth. Thus, the present work considers
pruning in a structured manner, so that the number of activations required to
be stored is reduced substantially, while the resulting tensor shapes are
inline with an efficient execution on GPUs. The Block Sparse Compressed Row
(BSR) format we opted for in this work has the additional benefit of a better
compression ratio for increasing block sizes.

Furthermore, we use a simple magnitude-based criterion to select activations
during the forward pass and prune them only after the computation is
completed. This ensures an accurate loss computation and avoids problems like
vanishing activations. The $l2$-norm is used to compare blocks of activations,
which are then ranked using an efficient top-k approach, pruning the lowest
scoring blocks.

One of the challenges faced exploiting structured sparse activations for
training is the lack of suitable block-sparse operators in publicly available
libraries. We therefore present a set of suitable operators with support for
BSR-compressed activations. Using these, we can assess the effectiveness of
block-sparse activation pruning in terms of model accuracy, training time, and
reduction of state, the latter being the main focus of this work.

In summary, this work makes the following contributions:

  1. 1.

We introduce a set of efficient block-sparse operators for GPU architectures,
which combine comparable performance with dense linear algebra libraries but
support as low as 30% of sparsity efficiently.

  2. 2.

We propose a simple yet effective method to prune activations in a structured
manner, creating sparsity patterns inline with the previous operators.

  3. 3.

We evaluate training speed, accuracy, and memory footprint of large-scale
neural architectures for image classification in comparison to baseline
implementations, highlighting the effectiveness of the activation pruning
concept.

As we will see in the following, the magnitude-based activation pruning method
is simple yet surprisingly effective. Overall, we see this work as a first
step in democratizing the training of large-scale models in a way that it
becomes more accessible for researchers without access to large-scale training
clusters. Furthermore, requiring less GPUs for a given training task is
reducing the associated carbon footprint, which is of increasing importance
given the trend towards larger and larger models while ecological implications
are becoming more severe.

##  II Background and Related Work

###  II-A Sparse data structures

Sparse linear algebra has a long tradition in many disciplines like scientific
computing, data science, and, of course more recently, machine learning. As a
result, there exists a plethora of sparse matrix formats and libraries,
tailored to the specific needs of the application. Compressing sparse matrices
is achieved by storing the non-zero elements only, alongside metadata to
reconstruct the original matrix. Well known formats like Compressed Sparse
Row/Column (CSR/CSC) formats use lists of corresponding row and column indices
of non-zero values to that end. To be viable, the ratio of zero elements to
the total number of values, or sparsity, has to be sufficiently large. If
sparsity is very low, the additional encoding overhead may even cause memory
consumption to be larger than saving the original matrix. For CSR, which is
natively supported by PyTorch, sparsity has to be larger than 50% to
effectively achieve compression, otherwise the coding overhead will dominate.
As will be discussed later in this work, for sparse training we can only
expect sparsity up to around 80%.

Another format supported by PyTorch is the Block Sparse Compressed Row (BSR)
format. BSR works analogously to CSR, the difference being that BSR does not
encode individual values, but rather blocks of values. Like with CSR, an array
of compressed row indices (crow) is used to encode the cummulative number of
non-zero blocks per row of the original matrix. The last element in crow is
then the total number of non-zero blocks. To get the number of non-zero blocks
for a specific row $n$ we subtract the $n$-th crow entry from the $n+1$-th. To
then get the block’s position within a row the column index (col) array is
accessed at the value of $\texttt{crow}[n]$. In case of multiple blocks per
row, following col entries encode their positions. The larger the BSR block
size, the fewer indices are required, improving compression. Table II shows
the additional encoding overhead of BSR for various block sizes $b$ at
increasing sparsity $s$. In addition to better compression, block-sparse
formats complement the warp execution paradigm of GPUs very well. For
sufficiently large blocks load and store operations can be coalesced and a
high throughput thereby maintained.

While PyTorch offers a BSR tensor construct, most operators do not yet support
it, hence there is a need for custom compute kernels.

TABLE II: BSR Memory Overhead Over Ideal Compression $1-s$ $b$ | 1 | 4 | 8 | 16 | 32 | 64 | 128 | 384
---|---|---|---|---|---|---|---|---
$s$ [%] | Additional Encoding Overhead [%] |
0 | 100.26 | 25.26 | 12.76 | 6.51 | 3.39 | 1.82 | 1.04 | 0.52
20 | 80.26 | 20.26 | 10.26 | 5.26 | 2.78 | 1.53 | 0.82 | 0.57
40 | 60.26 | 15.26 | 7.76 | 4.00 | 2.13 | 1.23 | 0.76 | 0.62
60 | 40.26 | 10.26 | 5.26 | 2.77 | 1.52 | 0.85 | 0.54 | 0.16
80 | 20.26 | 5.26 | 2.77 | 1.52 | 0.87 | 0.56 | 0.49 | 0.21
100 | 0.26 | 0.26 | 0.26 | 0.26 | 0.26 | 0.26 | 0.26 | 0.26

###  II-B Neural architectures for computer vision

The Vision Transformer (ViT) [7] is a state-of-the-art deep learning model
that uses the Transformer architecture to analyze images by breaking them into
patches and processing them as sequences. Its self-attention mechanism makes
it very effective for tasks like image classification. However, assessing the
effectiveness of activation pruning on ViT would require substantial
engineering effort due to various operators found in this neural architecture,
and, as previously reported, block-sparse operators are not supported natively
by standard tool stacks.

Thus, the ResMLP architecture proposed by Touvron et al. [6] was chosen for
experimentation for two main reasons: (1) It is a rather simple architecture
consisting of linear layers only, which lowers the engineering overhead for
custom operators. (2) As a descendant of the Vision Transformer, it employs a
similar embedding scheme, processing patches of the input image which
naturally defines a hardware friendly structure for pruning.

The input image is split into rectangular patches which are then linearized.
In case of color images, color channels are interleaved. The resulting row
vector is then scaled to a hidden dimension consistent throughout the network.
For the full image we obtain an embedding matrix with rows corresponding to
patches, as shown in Fig. 1. From this embedding matrix we can easily prune
rows, or at least parts thereof. Because rows are contiguous in memory, no
rearrangement of data is needed to encode the pruned embedding matrix to BSR.
The underlying intuition is that we only want to group values that have a
physical relationship.

Figure 1: Patch embeddings used in ResMLP and other vision MLPs. The input
image ($(N,C,H,W)$ tensor) is divided into patches $P_{i,j}$ which are then
linearized and color-channel interleaved. Linearized patches are finally
scaled to a configurable hidden dimension $D$ via a linear layer.

###  II-C Related work

Reviewing other ephemeral approaches to sparsity helped with the development
tremendously. One of the most well known examples is _dropout_ [8]. Originally
developed as a regularizer, dropout randomly forces activations to zero,
sparsifying the connections in a layer. This aims to prevent co-adaption,
where feature detectors strongly correlate with other feature detectors,
making them useless on their own, which also prevents overfitting. In its
original form, dropout induces unstructured sparsity, which is, as
established, not suited for compression. There exists a DropBlock [9] variant,
which drops spatially contiguous blocks. However, leaving the forward pass
unaltered was one of the goals of this work.

Many alternative approaches seek to prune gradients. Sparse gradients are
especially beneficial in the context of distributed training, where they have
to be broadcasted among compute ranks. Most approaches also prune based on
magnitude, using either fixed thresholds or adaptive ones obtained via
reduction. Zhang et al. propose a technique called Memorized Sparse
Backpropagation [10]. The network essentially estimates a gradient that is
empirically shown to be sufficiently similar to the real gradient to
effectively train. Unpropagated gradients are still stored for future updates
to counteract information loss. While this reduces communication, memory is
not saved in this approach. We largely adapt the top-k pruning strategy from
this work, however.

A more sophisticated pruning criterion is presented by Liu et al. [11]. Their
Dynamic Sparse Training (DST) strategy turns pruning thresholds into learnable
parameters. Because the model has to experiment with different configurations,
DST uses soft pruning, meaning pruned values are not actually discarded, but
masked instead for future use. Combined with the fact that DST focuses on
weights this makes this strategy inapplicable to our work.

The closest to our work is Sparse Weight Activation Training (SWAT) [12]. SWAT
uses a very similar approach for pruning activations in the context of CNNs,
but not with the goal of compression, but faster execution on sparse processor
architectures instead. In addition to activation pruning SWAT also adaptively
masks weights during the forward computation, similar to dropout. While
structured pruning is mentioned, the effects of different granularities of
said structure are not explored. As mentioned SWAT focuses on special sparse
accelerators in the context of latency based on hardware simulations, while we
focus on memory compression and efficient execution on GPUs.

##  III Block Sparse Operators

In this section we discuss the design and implementation of the block-sparse
linear operator developed in the context of this work. While a linear operator
is presented here, the technique used can be extended to any operator in
principle. Additionally, we report performance for possible configurations
using input shapes characteristic to vision MLPs. The benchmarks are performed
on an Nvidia A30 GPU1113574 CUDA cores at
$10.3\text{\,}\mathrm{T}\mathrm{F}\mathrm{/}\mathrm{s}$ peak FP32,
$24\text{\,}\mathrm{GB}$ HBM2 at
$933\text{\,}\mathrm{GB}\text{\,}{\mathrm{s}}^{-1}$, $165\text{\,}\mathrm{W}$
TDP, PCIe 4.0 x16, https://www.nvidia.com/en-us/data-center/products/a30-gpu.

The block-sparse linear operator is designed to be a plug-in replacement for
PyTorch’s
nn.linear222https://pytorch.org/docs/stable/generated/torch.nn.Linear.html
layer with two additional parameters: the level of sparsity, and the block
size of the sparse structure. As mentioned in section I the forward
computation is kept dense to ensure an accurate loss computation. After
completion of the forward computation, the smallest blocks of input
activations are selected using a magnitude based criterion described in
further detail in section IV. The selected blocks are forced to zero and the
activation tensor is converted to the BSR format. Fig. 2 shows an overview.

Figure 2: High-level overview of the block-sparse linear operator. After the
dense forward computation is completed, blocks of activations are pruned based
on magnitude (top-k pruning). The pruned activations are then saved for the
backward pass. During the backward pass they are used to calculate the weight
gradient via block-sparse matrix multiplication (BSpMM)

A custom CUDA block-sparse matrix multiply (BSpMM) kernel is used to calculate
the weight gradient during the backward pass. Because the input and bias
gradients do not depend on activations, they can be computed normally using
native PyTorch operations. The BSpMM implementation is heavily inspired by
CUTLASS333https://github.com/nvidia/cutlass. The kernel can be divided into a
load and a compute phase. We use vectorized loads and transposition of
fragments in shared memory during the load phase to ensure high throughput. To
parse the activation’s sparse structure, threads are grouped into cooperative
groups with each group being responsible for a BSR block. Each group checks if
their respective block is non-zero. If it is zero, the group will skip the
load and mark a shared memory internal predicate bit. During the compute phase
threads will skip the corresponding dot product calculation, if the predicate
vector indicates a zero block. See Fig. 3 for a visualization of the loading
process.

(a) Load BSR indices from global memory. (b) On matching index load non-zero
values to registers.

(c) Cache in shared memory

Figure 3: Cooperative thread groups loading a $4\times 8$ BSR encoded tensor
with block size $b=4$ to local shared memory. (a): First parse compressed row
indices to get the number of non-zero blocks (nnz) in the respective row. Use
nnz to access the column index array. (b): If row and column indices found in
(a) match the group load corresponding block from global memory to registers.
(c): Store loaded values to shared memory for following dot product
calculation. We use a bit vector to mark zero blocks in shared memory, so they
are skipped during the compute phase Because the number of non-zero blocks is,
a priori, unknown enough shared memory has to be allocated to fit a fully
dense block in the worst case.

The BSR block size has large implications on the performance of the kernel.
Smaller blocks lead to more indices loaded from global memory, lowering the
kernel’s FLOP per byte ratio. Also, due to the warp execution style of the
GPU, thread divergence may occur at group sizes smaller than the canonical
warp size of 32. Another, less obvious, limitation occurs at large BSR block
sizes. Because of the limit of 1024 threads per CUDA thread block it becomes
necessary to split the logical cooperative groups among multiple thread
blocks. Because communication among thread blocks is not allowed, indices for
the same logical group will have to be fetched by every participating thread
block individually from global memory, negating the benefit of fewer total
indices. Generally, increasing sparsity decreases latency. However,
performance gains diminish for larger block sizes. At block size 4, the
overhead due to decoding makes up a large percentage of the overall execution
time. Therefore, lowering this overhead via increased sparsity has a greater
effect and vice versa for larger block sizes. Fig. 4 shows the operator’s
latency at different block sizes for increasing levels of sparsity.

Measuring throughput at block size 64 for input shapes corresponding to
ImageNet, as shown in Figure 5, shows slight speedups over the dense operator
for sufficiently large inputs at moderate levels of sparsity (around 50%). For
the smaller embedding dimension $D=384$, speedups are achieved even at 10%
sparsity. At 90% sparsity we achieve around 140% peak floating point 32
performance. Note that we, of course, calculate throughput in regard to the
dense number of floating point operations. While saving memory is the main
focus of this work, faster execution is a welcome side effect.

Figure 4: Latency of the BSpMM operator for inputs of shape $(64,196,384)$,
different block sizes $b$, and increasing levels of sparsity. The red line
marks the latency of PyTorch’s dense torch.matmul routine.  Figure 5: BSpMM
throughput at block size 64 compared to dense torch.matmul for common input
shapes encountered in ResMLP. Note that FLOPs refers to the dense number of
operations in this context. The peak FP32 performance of the A30 GPU used is
shown in red.

##  IV Activation Pruning

In this work we use a simple magnitude-based top-k pruning approach to
sparsify activations. To force blocks of activations to zero, their vector or
matrix norms are compared and ranked. We then proceed to zero out the $k$
smallest blocks, where $k$ is determined by multiplying the total number of
blocks $N$ by the sparsity parameter $s$. The $l2$-norm is chosen to weigh
larger individual values stronger and retain a larger total norm of the
activation tensor, see Fig. 6.

Blocks are only compared locally, not among other activations in the mini-
batch, because there is no direct relationship. For example, one could imagine
a darker image being pruned completely otherwise. This makes the comparison
computationally inexpensive. Additionally, this also causes all activations in
a mini-batch to have the same amount of zero and non-zero blocks, reducing the
BSR encoding overhead.

At this time, pruning is performed equally among all layers. For instance,
setting sparsity to 50% will force all linear layers in the network to prune
50% of their activations. As other work suggests that different layers have
different sensitivity to perturbations [13, 14], future work will analyze each
layer’s sensitivity to pruning during training.

Figure 6: Example top-k block pruning step with block size $3$ at 50%
sparsity. We calculate the $l2$-norm per block and select the lowest scoring
blocks using top-k. Blocks are only compared within a single activation, not
among others in the mini-batch.

##  V Experiments

We evaluate the feasibility of block-sparse training on the ResMLP-S12
architecture in the context of image classification. Due to time and compute
resource constrains, a full parameter grid search over the ImageNet dataset is
not viable. We therefore gauge the network’s sensitivity to block-sparsity
using the smaller CIFAR10 and CIFAR100 datasets first. Based on the results,
we selectively test our method on ImageNet.

###  V-A Grid Search CIFAR10/100

Fig. 8 shows the top-1 accuracy achieved relative to a densely trained network
for various configurations with regard to amount of sparsity and block size.
Besides these two parameters, all experiments share the same
hyperparameters444Hyperparameters are adapted from the original paper [6].

As to be expected, higher levels of sparsity lead to a decrease in accuracy.
This effect is stronger at larger block sizes. Intuitively, pruning larger
blocks will lead to ”important” values being pruned more often. Interestingly,
low levels of sparsity can actually improve accuracy, possibly due to pruning
noise acting as a regularizer: as smaller components will no longer
contribute, parameter updates become smaller in general.

Similarly, the effects of the block size become more pronounced at higher
levels of sparsity. Therefore, smaller block sizes are preferred. For low
levels of sparsity, there is no discernible trend. This preference towards
smaller block sizes is in direct opposition to our goal of compression.
However, viable configurations for $b\geq 16$ can be found up to $70\%$ to
$80\%$ of sparsity.

There is some statistical variation in the results. Some configurations
perform inexplicably well, like $(s=80\%,b=1)$ for CIFAR10, or badly, like
$(s=30\%,b=32)$ for CIFAR100. For reference, running the dense baseline
configuration with different weight initialization seeds shows a standard
deviation of circa $\pm 1\%$ in top-1 accuracy. Future work should therefore
encompass multiple runs per configuration for statistical analysis.

Observing the loss curves for different configurations shows the training to
remain stable in terms of spread of the loss curve, see Fig. 7. $b=384$ is
chosen as it produces the largest variation in results. Due to the dense
forward computation an accurate loss value is obtained. However, because
gradient calculation is lossy, parameter updates become less accurate,
affecting the rate of convergence. We can observe the loss curves diverging
within the first 10 epochs and running almost in parallel afterward.

Figure 7: Comparing loss curves at increasing levels of sparsity for
ResMLP-S12/CIFAR10. $b=384$ was chosen, as it produces the most extreme split.
Because of the dense forward computation, the loss does not become more noisy,
only the rate of convergence of the loss curve changes.

(a)  (b)

Figure 8: Heatmap showing the top-1 accuracy relative to dense training for
ResMLP-S12 on CIFAR10 (a) and CIFAR100 (b). Values above baseline accuracy are
shown in green. Based on the observed standard deviation for dense training we
define an arbitrary level of 1.5% below baseline as acceptable. Acceptable
values are shown in orange, while accuracies lower than 1.5% below baseline
are shown in red.

###  V-B Early Scaling Experiments

Using the results shown in Fig. 8 as a heuristic, we evaluate block sparse
training performance on the ImageNet dataset next. In Fig. 9 we inspect the
loss curves, yielding results similar to CIFAR10 and CIFAR100, with a more
drastic spread, however. Translating loss to accuracy, $s=60\%$ causes a
reduction by $-5\%$ in accuracy, $s=70\%$ causes $-7\%$, and $s=80\%$ leads to
$-9\%$ top-1 validation accuracy. However, based on the trends of the loss
curves we expect slightly longer training with an adapted learning rate
schedule to improve accuracy to near baseline. Further parameter studies are
needed of course to conclude this expectation.

To quantify the actual memory saved we additionally log memory allocations via
a custom allocator. Table III shows the absolute and relative memory consumed
by activations for ImageNet training at mini-batch size 32 for selected
potential configurations.

TABLE III: Memory Saved by Pruning (ResMLP-S12/ImageNet, Mini-Batch Size of 32) Sparsity | Block Size | Activations | $\delta$ | $\delta$
---|---|---|---|---
$[\%]$ | – | [$\mathrm{MiB}$] | [$\mathrm{MiB}$] | $[\%]$
Dense | n.a. | 1599.2 | – | –
60 | 16 | 1234.0 | -365.2 | -22.8
60 | 32 | 1215.0 | -384.2 | -24.0
60 | 64 | 1205.3 | -393.9 | -24.6
70 | 16 | 1161.5 | -437.7 | -27.4
70 | 32 | 1142.9 | -456.3 | -28.5
70 | 64 | 1137.2 | -462.0 | -29.0
80 | 16 | 1082.8 | -516.4 | -32.3
80 | 32 | 1075.1 | -524.1 | -32.8
80 | 64 | 1071.0 | -528.2 | -33.0
Figure 9: ResMLP-S12/ImageNet loss for $b=16$. A behavior similar to Fig. 7
can be observed. Again, only the rate of convergence is impacted by sparsity.
The effect is more severe, however.

##  VI Discussion

Mini-batch size: A seemingly straight forward alternative to pruning is
reducing the mini-batch size. While directly reducing the memory footprint,
doing so comes with performance and, moreover, prediction accuracy
implications. Fig. 10 shows the latency of training ResMLP-S12/CIFAR10 for one
epoch at various mini-batch sizes. For small inputs there is not enough work
for the throughput oriented GPU to hide latency. Consequentially, training
becomes slower. At larger mini-batch sizes the opposite holds and the training
time decreases.

Comparing the training behavior of sparsely and densely trained network at
different mini-batch sizes, shows the sparse training to be much closer to
dense training at the same mini-batch size. This is shown in Fig. 11, where we
compare a densely trained ResMLP-S12 on CIFAR10 to training with 50% sparsity.
This comparison is not entirely fair, as 50% sparsity does not translate to
50% less total memory, but based on the observations from Fig. 8 there is
potential to move to higher levels of sparsity without sacrificing accuracy.
In this regard, Table IV compares the achieved top-1 validation accuracies for
different mini-batch sizes. We can see that, on top of training significantly
slower, the network loses 7% accuracy going from mini-batch size 32 to 16.
Meanwhile, the accuracy for the sparse network at mini-batch size 32 does not
differ from the dense network. This could be a valuable alternative in
resource-constrained systems, where a viable mini-batch could otherwise not be
used due to memory constraints.

TABLE IV: Comparing ResMLP-S12/CIFAR10 Accuracies at Different Mini-Batch Sizes to Sparse Training Mini-Batch Size | 16 | 32 | 64 | 128
---|---|---|---|---
Model | Top-1 Accuracy [%]
dense | 82 | 89 | 91 | 90
50% sparse | 82 | 89 | 90 | 90
Figure 10: ResMLP-S12/CIFAR10 time per epoch versus mini-batch size. At lower
mini-batch sizes the GPU is not utilized optimally, increasing the latency.
Figure 11: ResMLP-S12/CIFAR10 loss at various mini-batch sizes for dense and
sparse training at $s=50\%$. The sparse loss follows the corresponding dense
curve very closely performing better than the next smaller mini-batch size.

Overparameterization: Although reasons for this are not yet fully understood,
overparameterization of networks, in combination with regularization, seems to
lead to better generalization [15, 16, 17]. Our method can be used in this
context to increase the model’s complexity to leverage this observation,
counteracting the normally associated increase in memory consumption to a
degree.

While block sparse training applies to CIFAR10 and CIFAR100 without
repercussions in terms of accuracy and without any changes to the training
configuration, it seems that, in this configuration, further hyperparameter
tuning is required for ImageNet. It is however noteworthy, that at 15 million
parameters ResMLP-S12 is a relatively small architecture in the context of the
ImageNet dataset. Overparameterization is therefore expected to be low, thus a
larger variant is expected to show better results.

Operator implementation: We are also currently working on porting the operator
implementation from CUDA to Triton[18] for better portability and
extensibility. A block-sparse version of the affine scaling layer found in
ResMLP is now also available, potentially allowing for further memory savings,
and a convolution operator is planned. A block-sparse convolution operator
allows for interesting fundamental comparisons between vision MLPs and CNNs,
as the former emulate some of the properties found in the latter. The cross-
channel layer mimics the cross-channel property found in plain convolutions.
Similarly, the cross-patch layer can be seen as a 1D convolution, mixing
information among patches. Effects of differences, like weight sharing, on
robustness to block sparsity remain to be observed. CNNs also make for an
interesting candidate, as hardware friendly block sizes are relatively small.
The most typical kernel sizes in computer vision are $3\times 3$ or $5\times
5$. Due to the multitude of parameters, implementing an efficient block sparse
convolutional operator is challenging, however. Nvidia’s cuDNN library, for
example, employs different algorithms depending on the
parameters555https://docs.nvidia.com/deeplearning/performance/dl-performance-
convolutional/index.html, Section 1 via autotuning. Triton might also help in
this regard, reducing development time and increasing code readability.

Stochasticity: Related work [19] on gradient compression shows stochastic
components to compression to be beneficial. Therefore, future work will also
encompass random components in the pruning process. One could imagine randomly
swapping blocks near the top-k threshold, resulting in some blocks above the
threshold being pruned anyway, and vice-versa below threshold.

##  VII Summary

We have shown block sparse activation pruning based on a simple top-k
criterion to be a viable method for compressing the memory footprint of neural
architectures at the example of ResMLP. For CIFAR10/100 we can save around 32%
of the memory consumed by activations, which translates to 28% of the total
memory footprint, without affecting accuracy significantly. We have
additionally demonstrated the scalability of our block-sparse linear operator,
maintaining decent latency and throughput for as low as 10% sparsity for
adequate block sizes.

Future work will encompass the directions mentioned in the discussion, with an
emphasis on larger neural architectures, larger data sets, as well as advanced
pruning methods including stochasticity, second-order pruning methods, and
more sophisticated pruning schedules.

## Acknowledgements

This work is part of the ”Model-Based AI” project, which is funded by the Carl
Zeiss Foundation.

## References

  * [1] J. Lee, L. Mukhanov, A. S. Molahosseini, U. Minhas, Y. Hua, J. Martinez del Rincon, K. Dichev, C.-H. Hong, and H. Vandierendonck, “Resource-efficient convolutional networks: A survey on model-, arithmetic-, and implementation-level techniques,” _ACM Comput. Surv._ , vol. 55, no. 13s, Jul 2023. [Online]. Available: https://doi.org/10.1145/3587095
  * [2] W. Roth, G. Schindler, M. Zöhrer, L. Pfeifenberger, R. Peharz, S. Tschiatschek, H. Fröning, F. Pernkopf, and Z. Ghahramani, “Resource-efficient neural networks for embedded systems,” _CoRR_ , vol. abs/2001.03048, 2020. [Online]. Available: http://arxiv.org/abs/2001.03048
  * [3] Z. Liu, M. Sun, T. Zhou, G. Huang, and T. Darrell, “Rethinking the value of network pruning,” _CoRR_ , vol. abs/1810.05270, 2018. [Online]. Available: http://arxiv.org/abs/1810.05270
  * [4] L. Prechelt, “Connection pruning with static and adaptive pruning schedules,” vol. 16, no. 1, pp. 49–61. [Online]. Available: https://www.sciencedirect.com/science/article/pii/S0925231296000549
  * [5] G. Schindler, W. Roth, F. Pernkopf, and H. Froening, “Parameterized Structured Pruning for Deep Neural Networks,” _CoRR_ , vol. abs/1906.05180. [Online]. Available: http://arxiv.org/abs/1906.05180
  * [6] H. Touvron, P. Bojanowski, M. Caron, M. Cord, A. El-Nouby, E. Grave, G. Izacard, A. Joulin, G. Synnaeve, J. Verbeek, and H. Jégou, “ResMLP: Feedforward networks for image classification with data-efficient training,” _CoRR_ , vol. abs/2105.03404. [Online]. Available: http://arxiv.org/abs/2105.03404
  * [7] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, J. Uszkoreit, and N. Houlsby, “An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale,” _CoRR_ , vol. abs/2010.11929. [Online]. Available: http://arxiv.org/abs/2010.11929
  * [8] G. E. Hinton, N. Srivastava, A. Krizhevsky, I. Sutskever, and R. R. Salakhutdinov, “Improving neural networks by preventing co-adaptation of feature detectors,” _CoRR_ , vol. abs/1207.0580. [Online]. Available: http://arxiv.org/abs/1207.0580
  * [9] G. Ghiasi, T.-Y. Lin, and Q. V. Le, “DropBlock: A regularization method for convolutional networks,” _CoRR_ , vol. abs/1810.12890. [Online]. Available: http://arxiv.org/abs/1810.12890
  * [10] Z. Zhang, P. Yang, X. Ren, Q. Su, and X. Sun, “Memorized sparse backpropagation,” vol. 415, pp. 397–407. [Online]. Available: https://www.sciencedirect.com/science/article/pii/S0925231220313357
  * [11] J. Liu, Z. Xu, R. Shi, R. C. C. Cheung, and H. K. H. So, “Dynamic Sparse Training: Find Efficient Sparse Network From Scratch With Trainable Masked Layers,” _CoRR_ , vol. abs/2005.06870. [Online]. Available: http://arxiv.org/abs/2005.06870
  * [12] M. A. Raihan and T. Aamodt, “Sparse Weight Activation Training,” in _Advances in Neural Information Processing Systems_ , vol. 33. Curran Associates, Inc., pp. 15 625–15 638. [Online]. Available: https://proceedings.neurips.cc/paper/2020/hash/b44182379bf9fae976e6ae5996e13cd8-Abstract.html
  * [13] H. Borras, B. Klein, and H. Fröning, “Walking noise: Understanding implications of noisy computations on classification tasks,” _CoRR_ , vol. abs/2212.10430, Dec 2022. [Online]. Available: https://arxiv.org/abs/2212.10430
  * [14] T. Krieger, B. Klein, and H. Fröning, “Towards hardware-specific automatic compression of neural networks,” _CoRR_ , vol. abs/2212.07818, Dec 2022. [Online]. Available: https://arxiv.org/abs/2212.07818
  * [15] C. Zhang, S. Bengio, M. Hardt, B. Recht, and O. Vinyals, “Understanding deep learning (still) requires rethinking generalization,” _Communications of the ACM_ , vol. 64, no. 3, pp. 107–115, Feb. 2021. [Online]. Available: https://dl.acm.org/doi/10.1145/3446776
  * [16] J. Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, S. Gray, A. Radford, J. Wu, and D. Amodei, “Scaling Laws for Neural Language Models,” _CoRR_ , vol. abs/2001.08361. [Online]. Available: http://arxiv.org/abs/2001.08361
  * [17] Z. Li, E. Wallace, S. Shen, K. Lin, K. Keutzer, D. Klein, and J. Gonzalez, “Train Big, Then Compress: Rethinking Model Size for Efficient Training and Inference of Transformers,” in _Proceedings of the 37th International Conference on Machine Learning_. PMLR, pp. 5958–5968. [Online]. Available: https://proceedings.mlr.press/v119/li20m.html
  * [18] P. Tillet, H. T. Kung, and D. Cox, “Triton: An intermediate language and compiler for tiled neural network computations,” in _Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages - MAPL 2019_. ACM Press, pp. 10–19. [Online]. Available: http://dl.acm.org/citation.cfm?doid=3315508.3329973
  * [19] D. Alistarh, D. Grubic, J. Li, R. Tomioka, and M. Vojnovic, “Qsgd: Communication-efficient sgd via gradient quantization and encoding,” _CoRR_ , vol. abs/1610.02132, 2017. [Online]. Available: http://arxiv.org/abs/1610.02132