summaries: This paper proposes an efficient sparse training method to solves the problem of accelerating the training speed of deep neural networks and save memory usage.  The authors first formulate the training process as a continuous minimization problem under global sparsity constraints.  Then, separating the optimization process into two steps, corresponding to weight update and structure parameter update. A variance-reduced policy gradient estimator is proposed to update the structure parameter. Finally, the experimental results are solid and strong. The problem has lots of meanings for engineering implementation.

summaries: This paper addresses the topic of sparse training of neural networks. They target channel-level sparsity with the goal of realizing practical training speedups with existing software and hardware. The authors study the limitations of existing techniques in terms of their ability to exploit sparsity during training and show a wide array of empirical results with their proposed technique.

summaries: Paper proposes a variational based pruning method where, with 2 structurally sparse forward passes and 1 sparse backward, the training can be significantly faster. Description is clear and theoretically sounds, equations are easy to follow, code is provided, results are encouraging. The paper might be a great benchmark and starting point for faster pruning-while-training with efficient inference.

summaries: The authors propose a channel-level sparse training algorithm that can effectively utilize the channel sparsity to accelerate both forward and backward propagation process. They divide the sparse training process into weight updating step and mask updating step. And solve the problem separately. More specifically, they develop a variance-reduced policy gradient estimator to update the trainable pruning mask without the need for backpropagation, and hence a sparse computation on both forward and backward propagation.

summaries: Network sparsity has historically targeted only inference.  Recent explorations in training with sparsity fall short of their potential, though, by not fully exploiting sparsity in the forward and backward propagation steps.  The authors decompose training a channel-wise sparse network into two parts: training the weights themselves (as usual) and learning the sparse structure.  Rather than using a chain-rule based gradient step for the structure, a unique policy gradient estimator with reduced variance is used; it is composed of two forward passes, eliminating the need for the chain rule to be used to learn the sparse structure.  Experiments show competitive accuracy for saved parameters and FLOPs for a number of networks, as well as actual time saved during training with the proposed technique for VGG19 on CIFAR-10.

abstractText: Sparse training is a natural idea to accelerate the training speed of deep neural networks and save the memory usage, especially since large modern neural networks are significantly over-parameterized. However, most of the existing methods cannot achieve this goal in practice because the chain rule based gradient (w.r.t. structure parameters) estimators adopted by previous methods require dense computation at least in the backward propagation step. This paper solves this problem by proposing an efficient sparse training method with completely sparse forward and backward passes. We first formulate the training process as a continuous minimization problem under global sparsity constraint. We then separate the optimization process into two steps, corresponding to weight update and structure parameter update. For the former step, we use the conventional chain rule, which can be sparse via exploiting the sparse structure. For the latter step, instead of using the chain rule based gradient estimators as in existing methods, we propose a variance reduced policy gradient estimator, which only requires two forward passes without backward propagation, thus achieving completely sparse training. We prove that the variance of our gradient estimator is bounded. Extensive experimental results on real-world datasets demonstrate that compared to previous methods, our algorithm is much more effective in accelerating the training process, up to an order of magnitude faster.

authors:
M. Abadi
P. Barham
J. Chen
Z. Chen
A. Davis
J. Dean
M. Devin
S. Ghemawat
G. Irving
M. Isard

title: et al

venue: Tensorflow: A system for large-scale machine learning. In 12th {USENIX} symposium on operating systems design and implementation ({OSDI} 16), pages 265–283

authors:
Y. Bengio
N. Léonard
A. Courville

title: Estimating or propagating gradients through stochastic neurons for conditional computation

venue: arXiv preprint arXiv:1308.3432

authors:
J. Deng
W. Dong
R. Socher
L.-J. Li
K. Li
L. Fei-Fei

title: Imagenet: A large-scale hierarchical image database

venue: 2009 IEEE conference on computer vision and pattern recognition, pages 248–255. Ieee

authors:
T. Dettmers
L. Zettlemoyer

title: Sparse networks from scratch: Faster training without losing performance

venue: arXiv preprint arXiv:1907.04840

authors:
U. Evci
T. Gale
J. Menick
P.S. Castro
E. Elsen

title: Rigging the lottery: Making all tickets winners

venue: International Conference on Machine Learning, pages 2943–2952. PMLR

authors:
J. Frankle
G.K. Dziugaite
D.M. Roy
M. Carbin

title: Stabilizing the lottery ticket hypothesis

venue: arXiv preprint arXiv:1903.01611

authors:
Y. Gu
W. Zhang
C. Fang
J.D. Lee
T. Zhang

title: How to characterize the landscape of overparameterized convolutional neural networks

venue: Advances in Neural Information Processing Systems, volume 33, pages 3797–3807

authors:
Y. Guo
A. Yao
Y. Chen

title: Dynamic network surgery for efficient dnns

venue: Advances in neural information processing systems, pages 1379–1387

authors:
S. Han
H. Mao
W.J. Dally

title: Deep compression: Compressing deep neural networks with pruning

venue: trained quantization and huffman coding. International Conference on Learning Representations

authors:
S. Han
J. Pool
J. Tran
W. Dally

title: Learning both weights and connections for efficient neural network

venue: Advances in neural information processing systems, pages 1135–1143

authors:
K. He
X. Zhang
S. Ren
J. Sun

title: Deep residual learning for image recognition

venue: Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770–778

authors:
Y. He
G. Kang
X. Dong
Y. Fu
Y. Yang

title: Soft filter pruning for accelerating deep convolutional neural networks

venue: IJCAI International Joint Conference on Artificial Intelligence

authors:
Y. He
X. Zhang
J. Sun

title: Channel pruning for accelerating very deep neural networks

venue: Proceedings of the IEEE International Conference on Computer Vision, pages 1389–1397

authors:
A.G. Howard
M. Zhu
B. Chen
D. Kalenichenko
W. Wang
T. Weyand
M. Andreetto
H. Adam

title: Mobilenets: Efficient convolutional neural networks for mobile vision applications

venue: arXiv preprint arXiv:1704.04861

authors:
E. Jang
S. Gu
B. Poole

title: Categorical reparameterization with gumbel-softmax

venue: International Conference on Learning Representations

authors:
M. Kang
B. Han

title: Operation-aware soft channel pruning using differentiable masks

venue: International Conference on Machine Learning, pages 5122–5131. PMLR

title: Learning multiple layers of features from tiny images

venue: Master’s thesis, University of Tront

authors:
A. Kusupati
V. Ramanujan
R. Somani
M. Wortsman
P. Jain
S. Kakade
A. Farhadi

title: Soft threshold weight reparameterization for learnable sparsity

venue: In Proceedings of the International Conference on Machine Learning,

authors:
C. Lemaire
A. Achkar
P.-M. Jodoin

title: Structured pruning of neural networks with budget-aware regularization

venue: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 9108–9116

authors:
H. Li
A. Kadav
I. Durdanovic
H. Samet
H.P. Graf

title: Pruning filters for efficient convnets

authors:
L. Liebenwein
C. Baykal
H. Lang
D. Feldman
D. Rus

title: Provable filter pruning for efficient neural networks

authors:
B. Liu
M. Wang
H. Foroosh
M. Tappen
M. Pensky

title: Sparse convolutional neural networks

venue: Proceedings of the IEEE conference on computer vision and pattern recognition, pages 806–814

authors:
S. Liu
T. Chen
X. Chen
Z. Atashgahi
L. Yin
H. Kou
L. Shen
M. Pechenizkiy
Z. Wang
D.C. Mocanu

title: Sparse training via boosting pruning plasticity with neuroregeneration

venue: arXiv preprint arXiv:2106.10404

authors:
Z. Liu
J. Li
Z. Shen
G. Huang
S. Yan
C. Zhang

title: Learning efficient convolutional networks through network slimming

venue: Proceedings of the IEEE International Conference on Computer Vision, pages 2736–2744

authors:
Z. Liu
M. Sun
T. Zhou
G. Huang
T. Darrell

title: Rethinking the value of network pruning

authors:
C. Louizos
M. Welling
D.P. Kingma

title: Learning sparse neural networks through l_0 regularization

authors:
J.-H. Luo
J. Wu
W. Lin

title: Thinet: A filter level pruning method for deep neural network compression

venue: Proceedings of the IEEE international conference on computer vision, pages 5058–5066

authors:
S. Lym
E. Choukse
S. Zangeneh
W. Wen
S. Sanghavi
M. Erez

title: Prunetrain: fast neural network training by dynamic sparse model reconfiguration

venue: Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1–13

authors:
D.C. Mocanu
E. Mocanu
P. Stone
P.H. Nguyen
M. Gibescu
A. Liotta

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

venue: Nature communications, 9(1):1–12

authors:
H. Mostafa
X. Wang

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

venue: International Conference on Machine Learning, pages 4646–4655. PMLR

authors:
A. Paszke
S. Gross
F. Massa
A. Lerer
J. Bradbury
G. Chanan
T. Killeen
Z. Lin
N. Gimelshein
L. Antiga

title: et al

venue: Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32:8026–8037

authors:
A. Renda
J. Frankle
M. Carbin

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

authors:
D.J. Rezende
S. Mohamed
D. Wierstra

title: Stochastic backpropagation and approximate inference in deep generative models

venue: International conference on machine learning, pages 1278–1286. PMLR

authors:
P. Savarese
H. Silva
M. Maire

title: Winning the lottery with continuous sparsification

venue: Advances in Neural Information Processing Systems, 33

authors:
K. Simonyan
A. Zisserman

title: Very deep convolutional networks for large-scale image recognition

authors:
M. Song
A. Montanari
P. Nguyen

title: A mean field view of the landscape of two-layers neural networks

venue: Proceedings of the National Academy of Sciences, 115:E7665–E7671

authors:
S. Srinivas
A. Subramanya
R. Venkatesh Babu

title: Training sparse neural networks

venue: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition Workshops, pages 138–145

authors:
A. Vaswani
N. Shazeer
N. Parmar
J. Uszkoreit
L. Jones
A.N. Gomez

title: L

venue: u. Kaiser, and I. Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems

authors:
C. Wang
G. Zhang
R. Grosse

title: Picking winning tickets before training by preserving gradient flow

authors:
Z. Wang
J. Wohlwend
T. Lei

title: Structured pruning of large language models

venue: Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 6151–6162

authors:
W. Wen
C. Wu
Y. Wang
Y. Chen
H. Li

title: Learning structured sparsity in deep neural networks

venue: Advances in neural information processing systems, pages 2074–2082

authors:
X. Xiao
Z. Wang
S. Rajasekaran

title: Autoprune: Automatic network pruning by regularizing auxiliary parameters

venue: Advances in Neural Information Processing Systems, pages 13681–13691

authors:
M. Ye
C. Gong
L. Nie
D. Zhou
A. Klivans
Q. Liu

title: Good subnetworks provably exist: Pruning via greedy forward selection

venue: International Conference on Machine Learning, pages 10820–10830. PMLR

authors:
X. Yuan
P.H.P. Savarese
M. Maire

title: Growing efficient deep networks by structured continuous sparsification

authors:
S. Zagoruyko
N. Komodakis

title: Wide residual networks

venue: British Machine Vision Conference 2016. British Machine Vision Association

authors:
W. Zeng
R. Urtasun

title: MLPrune: Multi-layer pruning for automated neural network compression, 2019

authors:
X. Zhou
W. Zhang
H. Xu
T. Zhang

title: Effective sparsification of neural networks with global sparsity constraint

venue: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 3599–3608

authors:
M. Zhu
S. Gupta

title: To prune

venue: or not to prune: exploring the efficacy of pruning for model compression. arXiv preprint arXiv:1710.01878

heading: 1 Introduction

text:
In the last decade, deep neural networks (DNNs) [35, 11, 38] have proved their outstanding performance in various fields such as computer vision and natural language processing. However, training such large-sized networks is still very challenging, requiring huge computational power and storage. This hinders us from exploring larger networks, which are likely to have better performance. Moreover, it is a widely-recognized property that modern neural networks are significantly overparameterized, which means that a fully trained network can always be sparsified dramatically by network pruning techniques [9, 8, 25, 46, 20] into a small sub-network with negligible degradation in accuracy. After pruning, the inference efficiency can be greatly improved. Therefore, a natural question is can we exploit this sparsity to improve the training efficiency?
The emerging technique called sparse network training [10] is closely related with our question, which can obtain sparse networks by training from scratch. We can divide existing methods into two categories, i.e., parametric and non-parametric, based on whether they explicitly parameterize network structures with trainable variables (termed structure parameters). Empirical results [24, 34, 44, 23] demonstrate that the sparse networks they obtain have comparable accuracy with those ∗Equal contribution †Jointly with Google Research
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
obtained from network pruning. However, most of them narrowly aim at finding a sparse subnetwork instead of simultaneously sparsifying the computation of training by exploiting the sparse structure. As a consequence, it is hard for them to effectively accelerate the training process in practice on general platforms, e.g., Tensorflow [1] and Pytorch [31]. Detailed reasons are discussed below:
• Non-parametric methods find the sparse network by repeating a two-stage procedure that alternates between weight optimization and pruning [10, 6], or by adding a proper sparsityinducing regularizer on the weights to the objective [22, 41]. The two-stage methods prune the networks in weight space and usually require retraining the obtained subnetwork from scratch every time when new weights are pruned, which makes training process even more time-consuming. Moreover, the computation of regularized methods is dense since the gradients of a zero-valued weights/filters are still nonzero.
• All the parametric approaches estimate the gradients based on chain rule. The gradient w.r.t. the structure parameters can be nonzero even when the corresponding channel/weight is pruned. Thus, to calculate the gradient via backward propagation, the error has to be propagated through all the neurons/channels. This means that the computation of backward propagation has to be dense. Concrete analysis can be found in Section 3.
We notice that some existing methods [4, 28] can achieve training speedup by careful implementation. For example, the dense to sparse algorithm [28] removes some channels if the corresponding weights are quite small for a long time. However, these methods always need to work with a large model at the beginning epochs and consume huge memory and heavy computation in the early stage. Therefore, even with such careful implementations, the speedups they can achieve are still limited.
In this paper, we propose an efficient channel-level parametric sparse neural network training method, which is comprised of completely sparse (See Remark 1) forward and backward propagation. We adopt channel-level sparsity since such sparsity can be efficiently implemented on the current training platforms to save the computational cost. In our method, we first parameterize the network structure by associating each filter with a binary mask modeled as an independent Bernoulli random variable, which can be continuously parameterized by the probability. Next, inspired by the recent work [47], we globally control the network size during the whole training process by controlling the sum of the Bernoulli distribution parameters. Thus, we can formulate the sparse network training problem into a constrained minimization problem on both the weights and structure parameters (i.e., the probability). The main novelty and contribution of this paper lies in our efficient training method called completely sparse neural network training for solving the minimization problem. Specifically, to fully exploit the sparse structure, we separate training iteration into two parts, i.e., weight update and structure parameter update. For weight update, the conventional backward propagation is used to calculate the gradient, which can be sparsified completely because the gradients of the filters with zero valued masks are also zero. For structure parameter update, we develop a new variance reduced policy gradient estimator (VR-PGE). Unlike the conventional chain rule based gradient estimators (e.g., straight through[2]), VR-PGE estimates the gradient via two forward propagations, which is completely sparse because of the sparse subnetwork. Finally, extensive empirical results demonstrate that our method can significantly accelerate the training process of neural networks.
The main contributions of this paper can be summarized as follows:
• We develop an efficient sparse neural network training algorithm with the following three appealing features:
– In our algorithm, the computation in both forward and backward propagations is completely sparse, i.e., they do not need to go through any pruned channels, making the computational complexity significantly lower than that in standard training.
– During the whole training procedure, our algorithm works on small sub-networks with the target sparsity instead of follows a dense-to-sparse scheme.
– Our algorithm can be implemented easily on widely-used platforms, e.g., Pytorch and Tensorflow, to achieve practical speedup.
• We develop a variance reduced policy gradient estimator VR-PGE specifically for sparse neural network training, and prove that its variance is bounded.
• Experimental results demonstrate that our methods can achieve significant speed-up in training sparse neural networks. This implies that our method can enable us to explore larger-sized neural networks in the future.
Remark 1. We call a sparse training algorithm completely sparse if both its forward and backward propagation do not need to go through any pruned channels. For such algorithms, the computational cost in forward and backward propagation cost can be roughly reduced to ρ2 ∗ 100%, with ρ being the ratio of remaining unpruned channels.

heading: 2 Related Work

text: In this section, we briefly review the studies on neural network pruning, which refers to the algorithms that prune DNNs after fully trained, and the recent works on sparse neural network training.

heading: 2.1 Neural Network Pruning

text:
Network Pruning [10] is a promising technique for reducing the model size and inference time of DNNs. The key idea of existing methods [10, 8, 46, 20, 27, 13, 48, 40, 43, 32, 16] is to develop effective criteria (e.g, weight magnitude) to identify and remove the massive unimportant weights contained in networks after training. To achieve practical speedup on general devices, some of them prune networks in a structured manner, i.e., remove the weights in a certain group (e.g., filter) together, while others prune the weights individually. It has been reported in the literature [8, 25, 46, 20] that they can improve inference efficiency and reduce memory usage of DNNs by orders of magnitudes with minor loss in accuracy, which enables the deployment of DNNs on low-power devices.
We notice that although some pruning methods can be easily extended to train sparse networks, they cannot accelerate or could even slow down the training process. One reason is they are developed in the scenario that a fully trained dense network is given, and cannot work well on the models learned in the early stage of training. Another reason is after each pruning iteration, one has to fine tune or even retrain the network for lots of epoch to compensate the caused accuracy degradation.

heading: 2.2 Sparse Neural Network Training

text:
The research on sparse neural network training has emerged in the recent years. Different from the pruning methods, they can find sparse networks without pre-training a dense one. Existing works can be divided into four categories based on their granularity in pruning and whether the network structures are explicitly parameterized. To the best of our knowledge, no significant training speedups achieved in practice are reported in the literature. Table 1 summarizes some representative works.
Weight-level non-parametric methods, e.g., [6, 10, 48, 29, 30], always adopt a two-stage training procedure that alternates between weight optimization and pruning. They differ in the schedules of tuning the prune ratio over training and layers. [10] prunes the weights with the magnitude below a certain threshold and [48, 6] gradually increase the pruning rate during training. [30, 5] automatically reallocate parameters across layers during training via controlling the global sparsity.
Channel-level non-parametric methods [12, 41] are proposed to achieve a practical acceleration in inference. [41] is a structured sparse learning method, which adds a group Lasso regularization into the objective function of DNNs with each group comprised of the weights in a filter. [12] proposes a soft filter pruning method. It zeroizes instead of hard pruning the filters with small `2 norm, after which these filters are treated the same with other filters in training. It is obvious that these methods cannot achieve significant speedup in training since they need to calculate the full gradient in backward propagation although the forward propagation could be sparsified if implemented carefully.
Parametric methods multiply each weight/channel with a binary [47, 44, 37, 42] or continuous [24, 26, 19, 18] mask, which can be either deterministic [24, 42] or stochastic [47, 44, 26, 37, 19, 18]. The mask is always parameterized via a continuous trainable variable, i.e., structure parameter. The
sparsity is achieved by adding sparsity-inducing regularizers on the masks. The novelties of these methods lie in estimating the gradients w.r.t structure parameters in training. To be precise,
• Deterministic Binary Mask. [42] parameterizes its deterministic binary mask as a simple step function and estimates the gradients via sigmoid straight through estimator (STE) [2].
• Deterministic Continuous Mask. [24] uses the linear coefficients of batch normalization (BN) as a continuous mask and enforces most of them to 0 by penalizing the objective with `1 norm of the coefficients. [18] defines the mask as a soft threshold function with learnable threshold. These methods can estimate the gradients via standard backward propagation.
• Stochastic Binary Mask. [44, 37] model the mask as a bernoulli random variable and the gradients w.r.t. the parameters of bernoulli distributions are estimated via STE. [47] estimates the gradients via Gumbel-Softmax trick [15], which is more accurate than STE.
• Stochastic Continuous Mask. [26, 19] parameterize the mask as a continuous function g(c, ), which is differentiable w.r.t. c, and is a parameter free noise, e.g., Gaussian noise N (0, 1). In this way, the gradients can be calcuated via conventional backward propagation.
Therefore, we can see that all of these parametric methods estimate the gradients of the structure parameters based on the chain rule in backward propagation. This makes the training iteration cannot be sparsified by exploiting the sparse network structure. For the details, please refer to Section 3.

heading: 3 Why Existing Parameteric Methods Cannot Achieve Practical Speedup?

text:
In this section, we reformulate existing parametric channel-level methods into a unified framework to explain why they cannot accelerate the training process in practice.
Notice that convolutional layer can be viewed as a generalized fully connected layer, i.e., viewing the channels as neurons and convolution of two matrices as a generalized multiplication (see [7]). Hence, for simplicity, we consider the fully connected network in Figure 1. Moreover, since the channels in CNNs are corresponding to the neurons in fully connected networks, we consider neuron-level instead of weight-level sparse training in our example.
As discussed in Section 2, existing methods parameterize the 4 kinds of mask in the following ways:
(i): mi = φ(si); (ii): mi = ψ(si); (iii): mi = g(si, ), ∼ N (0, 1); (iv): m ∼ Bern(pi(s)), where the function φ(si) is binary, e.g., step function; ψ(si) is a continuous function; g(si, ) is differentiable w.r.t. si. All existing methods estimate the gradient of the loss `(ŷ, y) w.r.t. si based on chain rule, which can be formulated into a unified form below.
Specifically, we take the pruned neuron x3 in Figure 1 as an example, the gradient is calculated as
∇s3`(ŷ, y) = ∂`(ŷ, y)
∂x̂3︸ ︷︷ ︸ a
( w>:,3xin )︸ ︷︷ ︸ forward ∂m3 ∂s3 . (1)
Existing parametric methods developed different ways to estimate ∂m3∂s3 . Actually, for cases (ii) and (iii), the gradients are well-defined and thus can be calculated directly. STE is used to estimate the gradient in case (i) [42]. For cases (iv), [44, 37, 47] adopt STE and Gumbel-Softmax.
In Eqn.(1), the term (a) is always nonzero especially when x̂3 is followed by BN. Hence, we can see that even for the pruned neuron x3, the gradient ∂m3∂s3 can be nonzero in all four cases. This means the backward propagation has to go though all the neurons/channels, leading to dense computation.
At last, we can know from Eqn.(1) that forward propagation in existing methods cannot be completely sparse. Although w>:,3xin can be computed sparsely as in general models xin could be a sparse tensor of a layer with some channels being pruned, we need to calculate it for each neuron via forward propagation to calculate RHS of Eqn.(1). Thus, even if carefully implemented, the computational cost of forward propagation can only be reduced to ρ ∗ 100% instead of ρ2 ∗ 100% as in inference. That’s why we argue that existing methods need dense computation at least in backward propagation. So they cannot speed up the training process effectively in practice. Remark 2. The authors of GrowEfficient [44] confirmed that actually they also calculated the gradient of qc w.r.t, sc in their Eqn.(6) via STE even if qc = 0. Thus need dense backward propagation.

heading: 4 Channel-level Completely Sparse Neural Network Training

text: Below, we present our sparse neural network training framework and the efficient training algorithm.

heading: 4.1 Framework of Channel-level Sparse Training

text:
Given a convolutional network f(x;w), let {Fc : c ∈ C} be the set of filters with C being the set of indices of all the channels. To parameterize the network structure, we associate each Fc with a binary mask mc, which is an independent Bernoulli random variable. Thus, each channel is computed as
xout, c = xin ∗ (Fcmc) ,
with ∗ being the convolution operation. Inspired by [47], to avoid the problems, e.g., gradient vanishing, we parameterize mc directly on the probability sc, i.e., mc equals to 1 and 0 with the probabilities sc and 1 − sc, respectively. Thus, we can control the channel size by the sum of sc. Following [47], we can formulate channel-level sparse network training into the following framework:
min w,s
Ep(m|s) L(w,m) := 1
N N∑ i=1 ` (f (xi;w,m) ,yi) (2)
s.t. w ∈ Rn, s ∈ S := {s ∈ [0, 1]|C| : 1>s ≤ K},
where {(xi,yi)}Ni=1 is the training dataset, w is the weights of the original network, f (·; ·, ·) is the pruned network, and `(·, ·) is the loss function, e.g, cross entropy loss. K = ρ|C| controls the remaining channel size with ρ being the remaining ratio of the channels.
Discussion. We’d like to point out that although our framework is inspired by [47], our main contribution is the efficient solver comprised of completely sparse forward/backward propagation for Problem (2). Moreover, our framework can prune the weights in fully connected layers together, since we can associate each weight with an independent mask.

heading: 4.2 Completely Sparse Training with Variance Reduced Policy Gradient

text: Now we present our completely sparse training method, which can solve Problem (2) via completely sparse forward and backward propagation. The key idea is to separate the training iteration into filter update and structure parameter update so that the sparsity can be fully exploited.

heading: 4.2.1 Filter Update via Completely Sparse Computation

text:
It is easy to see that the computation of the gradient w.r.t. the filters can be sparsified completely. To prove this point, we just need to clarify the following two things:
• We do not need to update the filters corresponding to the pruned channels. Consider a pruned channel c, i.e., mc = 0, then due to the chain rule, we can have
∂` (f (xi;w,m)) ∂Fc = ∂` (f (xi;w,m)) ∂xout,c ∂xout,c ∂Fc ≡ 0,
the last equation holds since xout,c ≡ 0. This indicates that the gradient w.r.t the pruned filter Fc is always 0, and thus Fc does not need to be updated.
• The error cannot pass the pruned channels via backward propagation. Consider a pruned channel c, we denote its output before masking as x̂out,c = xin ∗ Fc, then the error propagating through this channel can be computed as
∂` (f (xi;w,m)) ∂x̂out,c = ∂` (f (xi;w,m)) ∂xout,c ∂xout,c x̂out,c ≡ 0.
This demonstrates that to calculate the gradient w.r.t. the unpruned filters, the backward propagation does not need to go through any pruned channels.
Therefore, the filters can be updated via completely sparse backward propagation.

heading: 4.2.2 Structure Parameter Update via Variance Reduced Policy Gradient

text:
We notice that policy gradient estimator (PGE) can estimate the gradient via forward propagation, avoiding the pathology of chain rule based estimators as dicussed in Section 3. For abbreviation, we denote L(w,m) as L(m) since w can be viewed as a constant here. The objective can be written as
Φ(s) = Ep(m|s) L(m),
which can be optimized using gradient descent:
s← s− η∇Φ(s).
with learning rate η. One can obtain a stochastic unbiased estimate of the gradient∇Φ(s) using PGE:
∇Φ(s) = Ep(m|s) L(m)∇s ln p(m|s), (PGE)
leading to Policy Gradient method, which may be regarded as a stochastic gradient descent algorithm:
s← s− ηL(m)∇s ln p(m|s). (3)
In Eqn.(3), L(m) can be computed via completely sparse forward propagation and the computational cost of ∇s ln p(m|s) = m−ss(1−s) is negligible, therefore PGE is computationally efficient.
However, in accordance with the empirical results reported in [33, 15], we found that standard PGE suffers from high variance and does not work in practice. Below we will develop a Variance Reduced Policy Gradient Estimator (VR-PGE) starting from theoretically analyzing the variance of PGE.
Firstly, we know that this variance of PGE is
Ep(m|s) L2(m)‖∇s ln p(m|s)‖22 − ‖∇Φ(s)‖22,
which can be large because L(m) is large. Mean Field theory [36] indicates that, while L(m) can be large, the term L(m)− L(m′) is small when m and m′ are two independent masks sampled from a same distributionp(m|s) (see the appendix for the details). This means that we may consider the following variance reduced preconditioned policy gradient estimator:
Em′∼p(m′|s)Em∼p(m|s) (L(m)− L(m′))Hα(s)∇s ln p(m|s), (VR-PGE)
where Hα(s) is a specific diagonal preconditioning matrix
Hα(s) = diag (s ◦ (1− s))α , (4)
with α ∈ (0, 1) and ◦ being the element-wise product. It plays a role as adaptive step size and it is shown that this term can reduce the variance of the stochastic PGE term∇s ln p(m|s). The details can be found in the appendix. Thus Φ(s) can be optimized via:
s← s− η (L (m)− L (m′))Hα(s)∇s ln p(m|s). (5)
In our experiments, we set α to be 12 for our estimator VR-PGE. The theorem below demonstrates that VR-PGE can have bounded variance.
Algorithm 1 Completely Sparse Neural Network Training Input: target remaining ratio ρ, a dense network w, the step size η, and parameter α in (4) . 1: Initialize w, let s = ρ1. 2: for training epoch t = 1, 2 . . . T do 3: for each training iteration do 4: Sample mini batch of data B = {(x1,y1) , . . . , (xB ,yB)}. 5: Sample m(i) from p(m|s), i = 1, 2. 6: Update s and w
s← projS(z) with z = s− η ( LB(w,m(1))− LB(w,m(2)) ) Hα(s)m (1)−s s(1−s) ,
w ← w − η∇wLB ( w,m(1) ) 7: end for 8: end for 9: return A pruned network w ◦m by sampling a mask m from the distribution p(m|s).
Theorem 1. Suppose m and m′ are two independent masks sampled from the Bernoulli distribution p(m|s), then for any α ∈ [ 12 , 1) and s ∈ (0, 1) |C|, the variance is bounded for
(L(m)− L(m′))Hα(s)∇s ln p(m|s)
Finally, we provide a complete view of our sparse training algorithm in Algorithm 1, which is essentially a projected stochastic gradient descent equipped with our efficient gradient estimators above. The projection operator in Algorithm 1 can be computed efficiently using Theorem 1 of [47].
Discussion. In our algorithm, benefited from our constraint on s, the channel size of the neural network during training can be strictly controlled. This is in contrast with GrowEfficient [44], which ultilizes regularizer term to control the model size and has situations where model size largely drift away from desired. This will have larger demand for the GPU memory storage and have more risk that memory usage may explode, especially when we utilize sparse learning to explore larger models. Moreover, our forward and backward propagations are completely sparse, i.e., they do not need to go through any pruned channels. Therefore, the computational cost of each training iteration can be roughly reduced to ρ2 ∗ 100% of the dense network.

heading: 5 Experiments

text: In this section, we conduct a series of experiments to demonstrate the outstanding performance of our method. We divide the experiments into five parts. In part one, we compare our method with several state-of-the-art methods on CIFAR-10 [17] using VGG-16 [35], ResNet-20 [11] and WideResNet-28-10 [45] to directly showcase the superiority of our method. In part two, we directly compare with state-of-the-art method GrowEfficient [44] especially on extremely sparse regions, and on two high capacity networks VGG-19 [35] and ResNet-32 [11] on CIFAR-10/100 [17]. In part three, we conduct experiments on a large-scale dataset ImageNet [3] with ResNet-50 [11] and MobileNetV1 [14] and compare with GrowEfficient [44] across a wide sparsity region. In part four, we present the train-computational time as a supplementary to the conceptual train-cost savings to justify the applicability of sparse training method into practice. In part five, we present further analysis on epoch-wise train-cost dynamics and experimental justification of variance reduction of VR-PGE. Due to the space limitation, we postpone the experimental configurations, calculation schemes on train-cost savings and train-computational time and additional experiments into appendix.

heading: 5.1 VGG-16, ResNet-20 and WideResNet-28-10 on CIFAR-10

text:
Table 2 presents Top-1 validation accuracy, parameters, FLOPs and train-cost savings comparisons with channel pruning methods L1-Pruning [20], SoftNet [12], ThiNet [27], Provable [21] and sparse training method GrowEfficient [44]. SoftNet can train from scratch but requires completely dense computation. Other pruning methods all require pretraining of dense model and multiple rounds of pruning and finetuning, which makes them slower than vanilla dense model training. Therefore the train-cost savings of these methods are below 1× and thus shown as ("-") in Table 2.
GrowEfficient [44] is a recently proposed state-of-the-art channel-level sparse training method showing train-cost savings compared with dense training. As described in Section 3, GrowEfficient features completely dense backward and partially sparse forward pass, making its train-cost saving limited by 32 . By contrast, the train-cost savings of our method is not limited by any constraint. The details of how train-cost savings are computed can be found in appendix.
Table 2 shows that our method generally exhibits better performance in terms of validation accuracy, parameters and particularly FLOPs. In terms of train-cost savings, our method shows at least 1.85× speed-up against GrowEfficient [44] and up to 9.39× speed-up against dense training.

heading: 5.2 Wider Range of Sparsity on CIFAR-10/100 on VGG-19 and ResNet-32

text:
In this section, we explore sparser regions of training efficiency to present a broader comparision with state-of-the-art channel sparse training method GrowEfficient [44].
We plot eight figures demonstrating the relationships between the Top-1 validation accuracy, FLOPs and train-cost savings. We find that our method generally achieves higher accuracy under same FLOPs
settings. To be noted, the train-cost savings of our method is drastically higher than GrowEfficient [44], reaching up to 58.8× when sparisty approches 1.56% on ResNet-32 on CIFAR-100, while the speed-up of GrowEfficient is limited by 32 .

heading: 5.3 ResNet-50 and MobileNetV1 on ImageNet-1K

text: In this section, we present the performance boost obtained by our method on ResNet-50 and MobileNetV1 on ImageNet-1K [3]. Our method searches a model with 76.0% Top-1 accuracy, 48.2% parameters and 46.8% FLOPs beating all compared state-of-the-art methods. The train-cost saving comes up to 1.60× and is not prominent due to the accuracy constraint to match up with compared methods. Therefore we give a harder limit to the channel size and present sparser results on the same Table 3, reaching up to 7.36× speed-up while still preserving 69.3% Top-1 accuracy. For the already compact model MobileNetV1, we plot two figures in Figure 3 comparing with GrowEfficient [44]. We find that our method is much stabler in sparse regions and obtains much higher train-cost savings.

heading: 5.4 Actual Training Computational Time Testing

text:
In this section, we provide actual training computational time on VGG-19 and CIFAR-10. The GPU in test is RTX 2080 Ti and the deep learning framework is Pytorch [31]. The intent of this section is to justify the feasibility of our method in reducing actual computational time cost, rather than staying in conceptual training FLOPs reduction. The computational time cost is measured by wall clock time, focusing on forward and backward propagation. We present training computational time in Table 4 with varying sparsity as in Figure 2. It shows that the computational time savings increases steadily with the sparisty. We also notice the gap between the savings in FLOPS and computational time. The gap comes from the difference between FLOPs and actual forward/backward time. More specifically, forward/backward time is slowed down by data-loading processes and generally affected by hardware latency and throughput, network architecture, etc. At extremely sparse regions, the pure computational time of sparse networks only occupies little of the forward/backward time and the cost of data management and hardware latency dominates the wall-clock time. Despite this gap, it can be expected that our train-cost savings can be better translated into real speed-up in exploring large
models where the pure computational time dominates the forward/backward time, which promises a bright future for making training infeasibly large models into practice.

heading: 5.5 Further Analysis

text: [Epoch-wise Train-cost Dynamics of Sparse Training Process] We plot the train-cost dynamics in Figure 3. The vertical label is the ratio of train-cost to dense training, the inverse of train-cost savings. This demonstrates huge difference between our method and GrowEfficient [44]. The model searched by our method exhibits 92.73% Top-1 accuracy, 16.68% parameters, 14.28% FLOPs with 5.28× train-cost savings, while the model searched by GrowEfficient exhibits 92.47% Top-1 accuracy, 18.08% parameters, 22.74% FLOPs with 1.21× train-cost savings. [Experimental Verification of Variance Reduction of VR-PGE against PGE] We plot the mean of variance of gradients of channels from different layers. The model checkpoint and input data are selected randomly. The gradients are calculated in two approaches, VR-PGE and PGE. From the rightmost graph of Figure 3, we find that the VR-PGE reduces variance significantly, up to 3 orders of magnitude.

heading: 6 Conclusion

text:
This paper proposes an efficient sparse neural network training method with completely sparse forward and backward passes. A novel gradient estimator named VR-PGE is developed for updating structure parameters, which estimates the gradient via two sparse forward propagation. We theoretically proved that VR-PGE has bounded variance. In this way, we can separate the weight and structure update in training and making the whole training process completely sparse. Emprical results demonstrate that the proposed method can significantly accelerate the training process of DNNs in practice. This enables us to explore larger-sized neural networks in the future.
Acknowledgments and Disclosure of Funding
This work is supported by GRF 16201320.

title: Efficient Neural Network Training via Forward and Backward Propagation Sparsification


summaries: Based on Non-Equilibrium IS (NEIS), the authors propose the new importance sampler NEO-IS. The main advantage of NEO-IS over NEIS seems to be that NEO-IS provides unbiased estimates of the normalization constant. They prove the unbiasedness and also propose an MCMC extension for sampling (which was not provided by NEIS). They empirically show that NEO-IS (and its MCMC variant) can be better that other recent methods for tasks of  (1) finding the marginal likelihood, (2) Sampling from posterior distribution, (3) Inpaiting.

summaries: The paper proposes a new family of importance samplers based on iterating an invertible map $T$ forwards and backwards, thus generating orbits that serve as importance samples after proper weighting. If normalized importance weights can be computed, the resulting IS estimate is unbiased. Also a self-normalized version is proposed, which is biased but still converges against the correct expectation for $N\to\infty$. Since the map $T$ is not assumed to leave the target invariant, the proposed algorithm is called non-equilibrium orbit importance sampling (NEO-IS).   The authors propose to use conformal Hamiltonian dynamics as a mapping (standard Hamilton dynamics based on the negative log target plus a friction term). The dissipative friction term allows the authors to control how exhaustively phase space is explored.  The second main contribution is an MCMC algorithm based on the ideas used in NEO-IS. NEO-MCMC applies iterative Sampling Importance Resampling (SIR) within the NEO-IS framework. NEO-MCMC jumps from one point on an orbit to a point on another orbit. Orbits are sampled according to their importance weight.  The algorithms are benchmarked on various targets (25 Gaussian mixture, Funnel) and applied to training a VAE.

summaries: The paper puts forward a new family of algorithms, NEO, for computing normalizing constants and sampling from complex distributions. NEO-IS is an importance sampling estimator of the normalizing constant. NEO-MCMC combines NEO-IS unbiased estimator of the normalizing constant with iterated sampling- importance resampling methods.

summaries: The paper introduce two sampling approaches named NEO-IS and NEO-MCMC respectively. The former method aims to estimate normalizing constant for an unnormalized target distribution, and the latter method can be executed to do sampling from unnormalized distribution. Empirical experiments reveal impressive performance for both the approaches on Normalizing Constant Estimation and Sampling tasks.  Exact theoretical analysis, especially convergence guarantees are also given in this paper.

summaries: This paper proposes NEO which extends Non-Equilibrium Importance Sampling (NEIS) to use discrete orbit and extra weights. The probability distribution and sampling process is derived based on the Jacobian determinant. This method could be applied to estimate the partition function and a biased estimation called self-normalized IS is proposed based on estimated partition function. This paper also derive upper bound for the bias and MSE. A damped Hamiltonian system is adopted to equip the original system with extra momentum and provide orbit trajectory for NEO. NEO is also combined to Sampling Importance Resampling to derive the NEO-MCMC procedure. The authors present reasonable experiments to evaluate their methods.

abstractText: Sampling from a complex distribution π and approximating its intractable normalizing constant Z are challenging problems. In this paper, a novel family of importance samplers (IS) and Markov chain Monte Carlo (MCMC) samplers is derived. Given an invertible map T, these schemes combine (with weights) elements from the forward and backward Orbits through points sampled from a proposal distribution ρ. The map T does not leave the target π invariant, hence the name NEO, standing for Non-Equilibrium Orbits. NEO-IS provides unbiased estimators of the normalizing constant and self-normalized IS estimators of expectations under π while NEO-MCMC combines multiple NEO-IS estimates of the normalizing constant and an iterated sampling-importance resampling mechanism to sample from π. For T chosen as a discrete-time integrator of a conformal Hamiltonian system, NEO-IS achieves state-of-the art performance on difficult benchmarks and NEO-MCMC is able to explore highly multimodal targets. Additionally, we provide detailed theoretical results for both methods. In particular, we show that NEO-MCMC is uniformly geometrically ergodic and establish explicit mixing time estimates under mild conditions. 35th Conference on Neural Information Processing Systems (NeurIPS 2021).

authors:
S. Agapiou
O. Papaspiliopoulos
D. Sanz-Alonso
A. Stuart

title: Importance sampling: Intrinsic dimension and computational cost

venue: Statistical Science,

authors:
Ö.D. Akyildiz
J. Míguez

title: Convergence rates for optimised adaptive importance samplers

venue: Statistics and Computing,

authors:
C. Andrieu
A. Doucet
R. Holenstein

title: Particle Markov chain Monte Carlo methods

venue: Journal of the Royal Statistical Society: Series B,

authors:
E. Bingham
J.P. Chen
M. Jankowiak
F. Obermeyer
N. Pradhan
T. Karaletsos
R. Singh
P. Szerlip
P. Horsfall
N.D. Goodman

title: Pyro: Deep universal probabilistic programming

venue: Journal of Machine Learning Research,

authors:
Y. Burda
R. Grosse
R. Salakhutdinov

title: Importance weighted autoencoders

venue: In The 4th International Conference on Learning Representations (ICLR),

authors:
T. Che
R. Zhang
J. Sohl-Dickstein
H. Larochelle
L. Paull
Y. Cao
Y. Bengio

title: Your GAN is secretly an energy-based model and you should use discriminator driven latent sampling

venue: arXiv preprint arXiv:2003.06060,

authors:
R.V. Craiu
C. Lemieux

title: Acceleration of the multiple-try Metropolis algorithm using antithetic and stratified sampling

authors:
P. Del Moral
A. Doucet
A. Jasra

title: Sequential Monte Carlo samplers

authors:
X. Ding
D.J. Freedman

title: Learning deep generative models with annealed importance sampling

venue: arXiv preprint arXiv:1906.04904,

authors:
T.A. El Moselhy
Y.M. Marzouk

title: Bayesian inference with optimal maps

venue: Journal of Computational Physics,

authors:
G. Franca
J. Sulam
D.P. Robinson
R. Vidal

title: Conformal symplectic and relativistic optimization

venue: Journal of Statistical Mechanics: Theory and Experiment,

authors:
H. Jia
U. Seljak

title: Normalizing constant estimation with Gaussianized bridge sampling

venue: In Symposium on Advances in Approximate Bayesian Inference,

authors:
D. Kingma
M. Welling

title: Auto-encoding variational Bayes

venue: In International Conference on Learning Representations,

authors:
D. Levy
M.D. Hoffman
J. Sohl-Dickstein

title: Generalizing Hamiltonian Monte Carlo with neural networks

authors:
Z. Liu
P. Luo
X. Wang
X. Tang

title: Large-scale celebfaces attributes (celeba) dataset

venue: Retrieved August,

authors:
C.J. Maddison
D. Paulin
Y.W. Teh
B. O’Donoghue
A. Doucet

title: Hamiltonian descent methods

venue: arXiv preprint arXiv:1809.05042,

title: Annealed importance sampling

authors:
A. Owen
Y. Zhou

title: Safe and effective importance sampling

venue: Journal of the American Statistical Association,

authors:
G. Papamakarios
E. Nalisnick
D.J. Rezende
S. Mohamed
B. Lakshminarayanan

title: Normalizing flows for probabilistic modeling and inference

title: Distilling importance sampling

venue: arXiv preprint arXiv:1910.03632,

authors:
G. Rotskoff
E. Vanden-Eijnden

title: Dynamical computation of the density of states and Bayes factors using nonequilibrium importance sampling

venue: Physical Review Letters,

title: Comment: A noniterative sampling/importance resampling alternative to the data augmentation algorithm for creating a few imputations when fractions of missing information are modest: The SIR algorithm

authors:
F.J. Ruiz
M.K. Titsias
T. Cemgil
A. Doucet

title: Unbiased gradient estimation for variational auto-encoders using coupled Markov chains

venue: In Uncertainty in Artificial Intelligence,

authors:
A.Y. Shestopaloff
Neal
R. M

title: Sampling latent states for high-dimensional non-linear state space models with the embedded hmm method

venue: Bayesian Analysis,

authors:
Ø. Skare
E. Bølviken
L. Holden

title: Improved sampling-importance resampling and reduced bias importance sampling

venue: Scandinavian Journal of Statistics,

title: Nested sampling for general Bayesian computation

title: Bayesian analysis of nonlinear and non-Gaussian state space models via multiple-try sampling methods

authors:
S.T. Tokdar
R.E. Kass

title: Importance sampling: a review

venue: Wiley Interdisciplinary Reviews: Computational Statistics,

authors:
R. Turner
J. Hung
E. Frank
Y. Saatchi
J. Yosinski

title: Metropolis–Hastings generative adversarial networks

venue: In International Conference on Machine Learning,

authors:
P. Wirnsberger
A.J. Ballard
G. Papamakarios
S. Abercrombie
S. Racanière
A. Pritzel
D. Jimenez Rezende
C. Blundell

title: Targeted free energy estimation via learned mappings

venue: The Journal of Chemical Physics,

authors:
H. Wu
J. Köhler
F. Noe

title: Stochastic normalizing flows

venue: In Advances in Neural Information Processing Systems,

authors:
Y. Wu
Y. Burda
R. Salakhutdinov
R. Grosse

title: On the quantitative analysis of decoderbased generative models

venue: arXiv preprint arXiv:1611.04273,

heading: 1 Introduction

text:
Consider a target distribution of the form π(x) ∝ ρ(x)L(x) where ρ is a probability density function (pdf) on Rd and L is a nonnegative function. Typically, in a Bayesian setting, π is a posterior distribution associated with a prior distribution ρ and a likelihood function L. Another situation of interest is generative modeling where π is the distribution implicitly defined by a Generative Adversarial Networks (GAN) discriminator-generator pair where ρ is the distribution of the generator and L is derived from the discriminator [31, 6]. In a Variational Auto Encoder (VAE) context [14, 5], π could be the true posterior distribution, ρ the approximate posterior distribution output by the encoder and L an importance weight between the true posterior and approximate posterior distributions. We are interested in this paper in sampling from π and approximating its intractable normalizing constant Z = ∫ ρ(x)L(x)dx. These problems arise in many applications in statistics, molecular dynamics or machine learning, and remain challenging.
Many approaches to compute normalizing constants are based on Importance Sampling (IS) - see [1, 2] and the references therein - and its variations, among others, Annealed Importance Sampling (AIS) [19, 34, 10] and Sequential Monte Carlo (SMC) [9]. More recently, Neural IS has also become very popular in machine learning; see e.g. [11, 18, 21, 22, 32, 33]. Neural IS is an adaptive IS which relies on an importance function obtained by applying a normalizing flow to a reference distribution. The parameters of this normalizing flow are chosen by minimizing a divergence between the proposal and the target (such as the Kullback–Leibler [18] or the χ2-divergence [1]). Recent work on the subject proposes to add stochastic moves in order to enhance the performance of the normalizing flows [33].
More recently, the Non-Equilibrium IS (NEIS) method has been introduced by [23] as an alternative to these approaches. Similar to Neural IS, NEIS consists in transporting samples {Xi}Ni=1 from a reference distribution using a family of deterministic mappings. For NEIS, this family is chosen to be an homogeneous differential flow (φt)t∈R. In contrast to Neural IS, for any i ∈ [N ], the sample Xi is propagated both forward and backward in time along the orbits associated with (φt)t∈R until stopping conditions are met. Moreover, the resulting estimator of the normalizing constant is obtained by computing weighted averages of the whole orbit (φt(Xi))t∈[τ+,i,τ−,i], where τ+,i, τ−,i are the resulting stopping times, and not only the endpoints φτ+,i(X i), φτ−,i(X i). In [23], the authors provide an application of NEIS with (φt)t∈R associated with a conformal Hamiltonian dynamics, and reports impressive numerical results on difficult normalizing constants estimation problems, in particular for high-dimensional multimodal distributions.
We propose in this work NEO-IS which alleviates the shortcomings of NEIS. Similar to NEIS, samples are drawn from a reference distribution, typically set to ρ, and are propagated under the forward and backward orbits of a discrete-time dynamical system associated with an invertible transform T. An estimator of the normalizing constant is obtained by reweighting all the points on the whole orbits using the IS rule. Contrary to NEIS, the NEO-IS estimator of Z is unbiased under assumptions that are mild and easy to verify. It is more flexible than NEIS because it does not rely on the accuracy of the discretization of a continuous-time dynamical system.
We then show how it is possible to leverage the unbiased estimator of Z defined by NEO-IS to obtain NEO-MCMC, a novel massively parallel MCMC algorithm to sample from π. In a nutshell, NEOMCMC relies on parallel walkers which each estimates the normalizing constant but are allowed to interact through a resampling mechanism.
Our contributions can be summarized as follows.
(i) We present a novel class of IS estimators of the normalizing constant Z referred to as NEO-IS. More broadly, a small modification of this algorithm also allows us to estimate integrals with respect to π. Both finite sample and asymptotic guarantees are provided for these two methodologies.
(ii) We develop a new massively parallel MCMC method, NEO-MCMC. NEO-MCMC combines NEO-IS unbiased estimator of the normalizing constant with iterated samplingimportance resampling methods. We prove that it is π-reversible and ergodic under very general conditions. We derive also conditions which imply that NEO-MCMC is uniformly geometrically ergodic (with an explicit expression of the mixing time).
(iii) We illustrate our findings using numerical benchmarks which show that both NEO-IS and NEO-MCMC outperform state-of-the-art (SOTA) methods in difficult settings.

heading: 2 NEO-IS algorithm

text:
In this section, we derive the NEO-IS algorithm. The two key ingredients for this algorithm are (1) the reference distribution ρ and (2) a transformation T assumed to be a C1-diffeomorphism with inverse T−1. Write, for k ∈ N∗ = N\{0}, Tk = T ◦Tk−1, T0 = Idd and similarly T−k = T−1 ◦T−(k−1). For any k ∈ Z, denote by ρk : Rd → R+ the pushforward of ρ by Tk, defined for x ∈ Rd by ρk(x) = ρ(T
−k(x))JT−k(x), where JΦ(x) ∈ R+ is the absolute value of the Jacobian determinant of Φ : Rd → Rd evaluated at x. In line with multiple importance sampling à la Owen and Zhou [20], we introduce the proposal density
ρT(x) = Ω −1 ∑
k∈Z $kρk(x) , (1) where {$k}k∈Z is a nonnegative sequence and Ω = ∑ k∈Z$k. Note that we assume in the sequel that the support of the weight sequence defined as {k ∈ Z : $k 6= 0} is finite. Thus, the mixture distribution in (1) is a finite mixture. Given x ∈ Rd, ρT(x) is a function of the forward and backward orbit of T through x.
For any nonnegative function f , the definition of ρT implies that ∫ f(y)ρT(y)dy =
Ω−1 ∫ ∑
k∈Z$kf(T k(x))ρ(x)dx. Assuming that $0 > 0, the ratio ρ(x)/ρT(x) ≤ $−10 Ω <∞ is
bounded. We can therefore apply the IS principle which allows to write the identity∫ f(x)ρ(x)dx = ∫ ( f(y) ρ(y)
ρT(y)
) ρT(y)dy = ∫ ∑ k∈Z f(Tk(x))wk(x)ρ(x)dx , (2)
where the weights are given by (see Appendix A.2 for a detailed derivation),
wk(x) = $kρ(T k(x))/{ΩρT(Tk(x))} = $kρ−k(x) /∑ i∈Z $k+iρi(x) . (3)
We assume in the sequel that $0 > 0. In particular, note that under this condition, the weights wk are also upper bounded uniformly in x: for any x ∈ Rd, wk(x) ≤ $k/$0. Equations (2) and (3) suggest to estimate the integral ∫ f(x)ρ(x)dx by INEO$,N (f) = N −1∑N i=1 ∑ k∈Z wk(X i)f(Tk(Xi)) where {Xi}Ni=1 are i.i.d. samples from the proposal ρ, which is denoted by X1:N iid∼ ρ.
Algorithm 1 NEO-IS Sampler
1. Sample X1:N iid∼ ρ for i ∈ [N ]. 2. For i ∈ [N ], compute the path (Tj(Xi))j∈Z and weights (wj(Xi))j∈Z.
3. INEO$,N (f) = N −1∑N i=1 ∑ k∈Z wk(X i)f(Tk(Xi)).
This estimator is obtained by a weighted combination of the elements of the independent forward and backward orbits {Tk(Xi)}k∈Z with X1:N
iid∼ ρ. This estimator is referred to as NEO-IS. Choosing f ≡ L provides the NEO-IS estimator of the normalizing constant of π:
Ẑ $ Xi = ∑ k∈Z L(T k(Xi))wk(X i) , Ẑ $ X1:N = N −1∑N i=1 Ẑ $ Xi . (4)
We now study the performance of the NEO-IS estimator. The following two quantities play a fundamental role in the analysis:
E$T = EX∼ρ [(∑ k∈Z wk(X)L(T k(X))/Z )2] ,M$T = supx∈Rd ∑ k∈Z wk(x)L(T k(x))/Z . (5)
Theorem 1. Ẑ $
X1:N is an unbiased estimator of Z. If E $ T < ∞, then, E[|Ẑ
$
X1:N /Z−1|2] = N−1(E$T − 1). If M$T <∞, then, for any δ ∈ (0, 1), with probability 1− δ, √ N ∣∣∣Ẑ$X1:N /Z−1∣∣∣ ≤
M$T √ log(2/δ)/2.
The (elementary) proof is postponed to Appendix A.3. E$T plays the role of the second-order moment of the importance weights EX∼ρ[L2(X)] which is key to the performance of IS algorithms
[1, 2]. In addition, since the NEO-IS estimator Ẑ $
X1:N is unbiased, the Cauchy–Schwarz inequality implies that EX∼ρ [(∑ k∈Z wk(X)L(T k(X)))2 ] ≥ Z2 and hence that E$T ≥ 1. Note that if ‖L‖∞ = supx∈Rd L(x) < ∞, then since the weights are uniformly bounded by Ω$−10 , we have M$T ≤ ‖L‖∞Ω$ −1 0 /Z.
Using the NEO-IS estimate Ẑ $
X1:N of the normalizing constant, we can construct a self-normalized IS estimate of ∫ f(x)π(x)dx:
JNEO$,N (f) = N −1 N∑ i=1 Ẑ $ Xi Ẑ $ X1:N ∑ k∈Z L(Tk(Xi))wk(X i) Ẑ $ Xi f(Tk(Xi)) , (6)
referred to as NEO-SNIS estimator. This expression may seem unnecessarily complicated but highlights the hierarchical structure of the estimator. We combine estimators (Ẑ $ Xi) −1∑ k∈Z L(T k(Xi))wk(X
i)f(Tk(Xi)) evaluated on the forward and backward orbits through the points {Xi}Ni=1 using the normalized weights {Ẑ $ Xi/Ẑ $
X1:N }Ni=1. Although the NEO-IS estimator is unbiased, the NEO-SNIS is in general biased. However, for bounded functions, both the bias and the variance of the NEO-SNIS estimator are O(N−1), with constants proportional to E$T . For g a π-integrable function, we set π(g) = ∫ g(x)π(x)dx.
Theorem 2. Assume that E$T <∞. Then, for any function g satisfying supx∈Rd |g(x)| ≤ 1 on Rd, and N ∈ N
E X1:N iid∼ρ
[ |JNEO$,N (g)− π(g)|2 ] ≤ 4 ·N−1E$T , (7)∣∣∣E
X1:N iid∼ρ
[ JNEO$,N (g)− π(g) ]∣∣∣ ≤ 2 ·N−1E$T . (8) If M$T <∞, then for δ ∈ (0, 1], with probability at least 1− δ,√
N |JNEO$,N (g)− π(g)| ≤ ‖g‖∞M$T √ 32 log(4/δ) . (9)
The proof is postponed to Appendix A.4. These results extend to NEO-SNIS estimators the results known for self-normalized IS estimators; see e.g., [1, 2] and the references therein. The upper bounds stated in this result suggest it is good practice to keep E$T /N small in order to obtain sensible approximations. For two pdfs p and q on Rd, denote by Dχ2(p, q) = ∫ {p(x)/q(x)− 1}2q(x)dx the χ2-divergence between p and q. Lemma 3. For any nonnegative sequence ($k)k∈Z, we have E$T ≤ Dχ2(π‖ρT) + 1.
The proof is postponed to Appendix A.5. Lemma 3 suggests that accurate sampling requires N to scale linearly with the χ2-divergence between the target π and the extended proposal ρT.
Remark 1. We can extend NEO to non homogeneous flows, replacing the family {Tk : k ∈ Z} with a collection of mappings {Tk : k ∈ Z}. This would allow us to consider further flexible classes of transformations such as normalizing flows; see e.g. [21]. The χ2-divergence Dχ2(π‖ρT) provides a natural criterion for learning the transformation. We leave this extension to future work.
Conformal Hamiltonian transform The efficiency of NEO relies heavily on the choice of T. Intuitively, a sensible choice of T requires that (i) E$T is small, i.e. ρT should be close to π by Lemma 3 (see (5)), (ii) the inverse T−1 and the Jacobian of T are easy to compute. Following [23], we use for T a discretization of a conformal Hamiltonian dynamics. Assume that U(·) = − log π(·) is continuously differentiable. We consider the augmented distribution π̃(q, p) ∝ exp{−U(q)−K(p)} on R2d, where q is the position, p is the momentum, and K(p) = pTM−1p/2 is the kinetic energy, with M a positive definite mass matrix. By construction, the marginal distribution of the momentum under π̃ is the target pdf π(q) = ∫ π̃(q, p)dp. The conformal Hamiltonian ODE associated with π̃ is defined by
dqt/dt = ∇pH(qt, pt) = M−1pt , (10) dpt/dt = −∇qH(qt, pt)− γpt = −∇U(qt)− γpt ,
where H(q, p) = U(q) +K(p), and γ > 0 is a damping constant. Any solution (qt, pt)t≥0 of (10) satisfies setting Ht = H(qt, pt), dHt/dt = −γpTt M−1pt ≤ 0. Hence, all orbits converge to fixed points that satisfy∇U(q) = 0 and p = 0; see e.g. [12, 17].
In the applications below, we consider the conformal version of the symplectic Euler (SE) method of (10), see [12]. This integrator can be constructed as a splitting of the two conformal and conservative parts of the system (10). When composing a dissipative with a symplectic operator, we set for all (q, p) ∈ R2d, Th(q, p) = (q + hM−1{e−hγp − h∇U(q)}, e−hγp − h∇U(q)), where h > 0 is a discretization stepsize. This transformation can be connected with classical momentum optimization schemes, see [12, Section 4]. By [12, Section 3], for any h > 0 Th is a C1-diffeomorphism on R2d with Jacobian given by JTh(q, p) = e−γhd. In addition, its inverse is T −1 h (q, p) = (q − hM−1p, eγh{p+ h∇U(q − hM−1p)}). Therefore, the weight (3) of the NEO estimator is given by
wk(q, p) = $kρ̃(T
k h(q, p))e −γkhd∑ j∈Z$k+j ρ̃(T −j h (q, p))e γjhd ,
where ρ̃(q, p) ∝ ρ(q)e−K(p). Figure 1 displays for different values of γ on a log-scale the bound E 1[0:K]
Th − 1 appearing in Theorem 1 as a function of K, here we use the sequence of weights ($k)k∈Z = (1[0:K](k))k∈Z (i.e. only the K + 1 first elements of the forward orbits are used and are equally weighted). For comparison, we also present on the same plot the bounds achieved by averaging K + 1 independent IS estimates, EIS(K)− 1 = (K + 1)−1EX∼ρ[L(X)2]. Interestingly, Figure 1 shows that there is a trade-off in the choice of γ which controls the exploration of the state space by the Hamiltonian dynamics since the higher γ, the faster the orbits converge towards the modes. This fast convergence prevents a “good” exploration of the space; e.g. E 1[0:K]
Th is smaller for
γ = 1.0 than for γ = 2.0 when K > 7.
3 NEO-MCMC algorithm
We now derive an MCMC method to sample from π based on the NEO-IS estimator. A natural idea consists in adapting the Sampling Importance Resampling procedure (SIR) (see for example [24, 27]) to the NEO framework.
The SIR method to sample JNEO$,N (see (6)) consists of 4 steps. (SIR-1) Draw independently X1:N iid∼ ρ and compute the associated forward and backward orbits {Tk(Xi)}k∈Z of the point. (SIR-2) Compute the normalizing constants associated with each orbit {Ẑ $
Xi}Ni=1. (SIR-3) Sample an orbit index IN ∈ [N ] with probability {Ẑ $ Xi/ ∑N j=1 Ẑ $
Xj}Ni=1. (SIR-4) Draw the iteration index KN on the IN -th orbit with probability {L(Tk(XIN ))wk(XI N )/Ẑ $ XIN }k∈Z.
The resulting draw is denoted by UN = TK N (XI N
). By construction, for any bounded function f , we get that E [ f(UN ) ∣∣X1:N , IN ] = {Ẑ$XIN }−1∑k∈Z wk(XIN )L(Tk(XIN )) which implies E [ f(UN )
∣∣X1:N ] = JNEO$,N (f) (see (6)). Using Theorem 2, we therefore obtain |E[f(UN )] −∫ f(z)π(z)dz| ≤ 101/2‖f‖∞E$T N−1, showing that the law of the random variable µN = Law(UN ) converges in total variation to π as N →∞,
‖µN − π‖TV = sup ‖f‖∞≤1 |µN (f)− π(f)| ≤ 101/2E$T N−1 . (11)
Algorithm 2 NEO-MCMC Sampler At step n ∈ N∗, given the conditioning orbit point Yn−1. Step 1: Update the conditioning point
1. Set X1n = Yn−1 and for any i ∈ {2, . . . , N}, sample Xin
iid∼ ρ. 2. Sample the orbit index In with probability
proportional to (Ẑ $
Xin )i∈[N ], (4).
3. Set Yn = XInn . Step 2: Output a sample
4. Sample index Kn with probability proportional to {wk(Yn)L(Tk(Yn))/Ẑ $
Yn}k∈Z 5. Output Un = TKn(Yn).
Based on [3], we now derive the NEOMCMC procedure, which in a nutshell consists in iterating the SIR procedure while keeping a conditioning point (or equivalently, orbit); see Appendix C. The convergence of NEO-MCMC does not rely on letting N →∞: the NEO-MCMC works as soon as N ≥ 2, although as we will see below the mixing time decreases as N increases.
This procedure is summarized in Algorithm 2. The NEO-MCMC procedure is an iterated algorithm which produces a sequence {(Yn, Un)}n∈N of points in Rd. The n-th iteration of the NEO-MCMC algorithm consists in two main steps: 1) updating the conditioning point Yn−1 → Yn 2) sampling Un by selecting a point in the orbit {Tk(Yn)}k∈Z of the conditioning point. Compared to SIR, only the generation of the points (step (SIR-1)) is modified: we set X1n = Yn−1 (the conditioning point), and then draw X2:Nn
iid∼ ρ. The sequence {Yn}n∈N defined by Algorithm 2 is a Markov chain: P (Yn ∈ A |Y0:n−1) = P (Yn ∈ A |Yn−1) = P (Yn,A) where
P (y,A) = ∫ δy(dx 1) N∏ j=2 ρ(xj)dxj N∑ i=1
Ẑ $ xi∑N j=1 Ẑ $ xj 1A(x i) , y ∈ Rd ,A ∈ B(Rd) . (12)
Note that this Markov kernel describes the way, at stage n + 1, the conditioning point Yn+1 is selected given Yn, which depends only on the estimator of the normalizing constants associated with each orbit, but not on the sample Un selected on the conditioning orbit. In addition, given the conditioning point Yn at the n-th iteration, the conditional distribution of the output sample Un is P ( Un ∈ B | In, X1:Nn ) = P (Un ∈ B |Yn) = Q(Yn,B) where
Q(y,B) = ∑ k∈Z wk(y)L(T k(y)) Ẑ $ y 1B(T k(y)) , y ∈ Rd ,B ∈ B(Rd) . (13)
With these notations, if the Markov chain is started at Y0 = y, then for any n ∈ N, the law of the n-th conditioning point is P (Yn ∈ A |Y0 = y) = Pn(y,A) and the law of the n-th sample is P (Un ∈ B |Y0) = PnQ(y,B). Define π̃ the pdf given for y ∈ Rd by
π̃(y) = ρ(y)
Z ∑ k∈Z wk(y)L(T k(y)) = ρ(y)Ẑ $ y Z . (14)
The following theorem shows that, for any initial condition y ∈ Rd, the distribution of the variable Yn converges in total variation to π̃ and that the distribution of Un converges to π. Theorem 4. The Markov kernel P is reversible with respect to the distribution π̃, ergodic and Harris positive, i.e., for all y ∈ Rd, limn→∞ ‖Pn(y, ·) − π̃‖TV = 0. In addition, π = π̃Q and limn→∞ ‖PnQ(y, ·) − π‖TV = 0. Moreover, for any bounded function g and any y ∈ Rd, limn→∞ n −1∑n−1 i=0 g(Ui) = π(g), P-almost surely, where {Ui}i∈N is defined in Algorithm 2 with Y0 = y.
The proof is postponed to Appendix A.6. Remark 2. We may provide another sampling procedure of {Yn}n∈N. Define the pdf on the extended space [N ]× RdN by π̌(i, x1:N ) = N−1π̃(xi) ∏N j=1,j 6=i ρ(x j). Consider a Gibbs sampler
targeting π̌ consisting in (a) sampling X1:N\{In−1}n |(In−1, Xn−1) ∼ ∏ j 6=In−1 ρ(x j), (b) sampling In|X1:Nn ∼ Cat({Ẑ $ Xin / ∑N j=1 Ẑ $ Xjn }Ni=1 and (c) set Yn = XInn . This algorithm is a Gibbs sampler on π̌ and we easily verify that the distribution of {Yn}n∈N is the same as Algorithm 2.
The next theorem provides non asymptotic quantitative bounds on the convergence in total variation. The main interest of NEO-MCMC algorithm is motivated empirically from observed behaviour: the mixing time of the corresponding Markov chain improves as N increases. This behaviour is quantified theoretically in the next theorem. Moreover, this improvement is obtained with little extra computational overhead, since sampling N points from the proposal distribution ρ, computing the forward and backward orbits of the points and evaluating the normalizing constants {Ẑ $
Xin }Ni=1 can
be performed in parallel. Theorem 5. Assume thatM$T <∞, see (5). Set N = (N − 1)/(2M$T +N − 2) and κN = 1− N . Then, for any y ∈ Rd and k ∈ N, ‖P k(y, ·)− π̃‖TV ≤ κkN and ‖P kQ(y, ·)− π‖TV ≤ κkN .
Instead of sampling the new points X2:Nn independently from ρ (Step 1 in Algorithm 2), it is possible to draw the proposals X1:Nn conditional to the current point Yn−1; see [29, 8, 26, 25] for related works. Following [25], we use a reversible Markov kernel with respect to the proposal ρ, i.e., such that ρ(x)m(x, x′) = ρ(x′)m(x′, x), assuming for simplicity that this kernel has density m(x, x′). If ρ = N(0, σ2 Idd) , an appropriate choice is an autoregressive kernel m(x, x′) = N(x′;αx, σ2(1− α2) Idd). More generally, we can use a Metropolis–Hastings kernel with invariant distribution ρ. In this case, r1(x1, x1:N\{1}) = ∏N j=2m(x j−1, xj) and for each i ∈ [2 : N ],
ri(x i, x1:N\{i}) = i−1∏ j=1 m(xj+1, xj) N∏ j=i+1 m(xj−1, xj) . (15)
Since m is reversible with respect to ρ, for all i, j ∈ [N ], ρ(xi)ri(xi, x1:N\{i}) = ρ(xj)rj(x j , x1:N\{j}) where ri(xi;x1:N\{i}) defines the the conditional distribution of X1:N\{i} given Xi = xi. The only modification in Algorithm 2 is Step 1, which is replaced by: Draw Un ∈ [N ] uniformly, set XUnn = Yn−1 and sample X 1:N\{Un} n ∼ rUn(XUnn , ·). The validity of this procedure is established in Appendix A.6.
4 Continuous-time version of NEO and NEIS
The NEO framework can be thought of as an extension of NEIS introduced in [23]. NEIS focuses on normalizing constant estimation and should be therefore compared with NEO-IS. In [23], the authors do not consider possible extensions of these ideas to sampling problems. We consider here how NEO could be adapted to continuous-time dynamical system. Proofs of the statements and detailed technical conditions are postponed to Appendix B.
Consider the Ordinary Differential Equation (ODE) ẋt = b(xt) , where b : Rd → Rd is a smooth vector field. Denote by (φt)t∈R the flow of this ODE (assumed to be well-behaved). Under appropriate regularity condition Jφt(x) = exp( ∫ t 0 ∇ · b(φs(x))ds); see Lemma S5. Let $ : R → R+ be a
nonnegative smooth function with finite support, with Ωc = ∫∞ −∞$(t)dt. The continuous-time counterpart of the proposal distribution (1) is ρcT(x) = (Ω c)−1 ∫∞ −∞$(t)ρ(φ−t(x))Jφ−t(x)dt, which is a continuous mixture of the pushforward of the proposal ρ by the flow of (φs)s∈R. Assuming for simplicity that ρ(x) > 0 for all x ∈ Rd, then ρcT(x) > 0 for all x ∈ Rd, and using again the IS formula, for any nonnegative function f ,∫
f(x)ρ(x)dx = ∫ f(x) ρ(x)
ρcT(x) ρcT(x)dx = ∫ [∫ ∞ −∞ wct (x)f(φt(x))dt ] ρ(x)dx, (16)
wct (x) = $(t)ρ(φt(x))Jφt(x) /∫ ∞ −∞ $(s+ t)ρ(φs(x))Jφs(x)ds . (17)
These relations are the continuous-time counterparts of (2). Eqs. (16)-(17) define a version of NEIS [23], with a finite support weight function $; see Appendices B.2 and B.3 for weight functions with infinite support. This identity is of theoretical interest but must be discretized to obtain a computationally tractable estimator. For h > 0, denote by Th an integrator with stepsize h > 0 of the ODE ẋ = b(x). We may construct NEO-IS and NEO-SNIS estimators based on the transform T ← Th and weights $k ← $(kh). We might show that for any bounded function f and for any x ∈ Rd, limh↓0 ∑ k∈Z wk(x)f(T k h(x)) = ∫∞ −∞ w c t (x)f(φt(x))dt, where we omitted here the
dependency in h of wk. Therefore, taking h ↓ 0+, the NEO-IS converges to the continuous-version (16)-(17). There is however an important difference between NEO and the NEIS method in [23] which stems from the way (16)-(17) are discretized. Compared to NEIS, NEO-IS using T← Th and weights $k ← $(kh) is unbiased for any stepsize h > 0. NEIS uses an approach inspired by the nested-sampling approach, which amounts to discretizing the integral in (16) also in the state-variable x; see [28, 7]. This discretization is biased which prevents the use of this approach to develop MCMC sampling algorithm; see Appendix B.

heading: 5 Experiments and Applications

text:
Normalizing constant estimation The performance of NEO-IS is assessed on different normalizing constant estimation benchmarks; see [13]. We focus on two challenging examples. Additional experiments and discussion on hyperparameter choice are given in the supplementary material, see Appendix D.1.
(1) Mixture of Gaussian (MG25): π(x) = P−1 ∑P i=1 N(x;µi,j , Dd), where d ∈ {10, 20, 45}, Dd = diag(0.01, 0.01, 0.1, . . . , 0.1) and µi,j = [i, j, 0, . . . , 0]T with i, j ∈ {−2, . . . , 2}.
(2) Funnel distribution (Fun) π(x) = N(x1; 0, a2) ∏d i=1 N(xi; 0, e
2bx1) with d ∈ {10, 20, 45}, a = 1, and b = 0.5. In both case, the proposal is ρ = N(0, σ2ρ Idd) with σ 2 ρ = 5.
The NEO-IS estimator is compared with (i) the IS estimator using the proposal ρ, (ii) the Adaptive Importance Sampling (AIS) estimator of [30], (iii) Stochastic Normalizing Flows (SNF)1 and (iv) the Neural Importance Sampling (NIS)2. For NEO-IS, we use $k = 1[K](k) with K = 10 (ten steps on the forward orbit), and conformal Hamiltonian dynamics γ = 1, M = 5 · Idd for dimensions d = {10, 20}, and γ = 2.5 for d = 45 (where γ is the damping factor, M the mass matrix, h is the stepsize of the integrator). The parameters of AIS are set to obtain a complexity comparable to NEO-IS; see Appendix D.1. For NIS, we use the default parameters and for SNF we used the same architectures as in [33]. In Fun, we set γ = 0.2, K = 10, M = 5 · Idd, and h = 0.3. The IS estimator was based on 5 · 105 samples, and NIS, NEO-IS and AIS were computed with 5 · 104 samples. Figure 2 shows that NEO-IS consistently outperforms the competing methods.
1Implementation available at https://github.com/noegroup/stochastic_normalizing_flows. 2Implementation available at https://github.com/ndeutschmann/zunis.
Sampling NEO-MCMC is assessed for the distributions (MG25) (d = 40) and Fun (d = 20). NEO-MCMC sampler is compared with (i) the No-U-Turn Sampler - Pyro library [4] - and (ii) i-SIR algorithm [25]. The proposal distribution is ρ = N(0, σ2ρ Idd) with σ 2 ρ = 5. Dependent proposals are used (see (15)) with m(x, x′) = N(x′;αx, σ2ρ(1− α2) Idd) with α = 0.99. For NUTS, the default parameters are used. For i-SIR, we use the same number of proposals N = 10, proposal distribution and dependent proposal as for NEO-MCMC. To perform a fair comparison, we use the same clock time for all three algorithms. The number of iterations for correlated i-SIR, NEO-MCMC, and NUTS are n = 4 · 106, n = 4 · 105, and n = 5 · 105, respectively. Figure 3 displays the empirical
two-dimensional histograms of the two first coordinates of samples from the ground truth, i-SIR, NUTS and NEO-MCMC sampler. It is worthwhile to note that NEO-MCMC algorithm performs much better for MG25 which is a very challenging distribution, even for SOTA algorithm such as NUTS, which struggles to cross energy barriers between modes. For Fun, NEO-MCMC performs favourably with respect to NUTS, which is well adapted for this type of distributions.
Block Gibbs Inpainting with Deep Generative models and NEO-MCMC We apply NEOMCMC to the task of sampling the posterior of a deep latent variable model. To be consistent with the rest of the paper, we use non-standard notation here with x being the latent variable and z the observation. More precisely, we assume that x ∼ N(0, Idd) and a conditional distribution p(z | x) which generates an image z = (z1, . . . , zD) ∈ RD. Given a family of parametric decoders {x 7→ pθ(z | x), θ ∈ Θ}, and a training set D = {zi}Mi=1, training involves finding the MLE θ∗ = arg maxθ∈Θ pθ(D). As pθ(z) = ∫ pθ(z | x)p(x)dx, the likelihood is intractable and to alleviate this problem, [14] proposed to train jointly an approximate posterior qφ(x|z) that maximizes a tractable lower-bound on the log-likelihood: ELBO(z, θ, φ) = EX∼qφ(·|z)[log pθ(z,X)/qφ(X|z)] ≤ pθ(z), where qφ(x | z) is a tractable conditional distribution with parameters φ ∈ Φ. It is assumed in the sequel that conditional to the latent variable x, the coordinates are independent, i.e. pθ(z | x) = ∏D i=1 pθ(z
i|x). Note that it is possible to train VAE with the NEO algorithm, using the unbiased estimate of the normalizing constant to construct an ELBO. This approach is described in the supplement Appendix E. We do not focus on this approach here and assume that the VAE has been trained and we are only interested in the sampling problem. In our experiment, we use a VAE trained on CelebA dataset 3 [16]. We consider the Block Gibbs inpainting task introduced in [15, Section 5.2.2]. Given an image z, denote by [zt, zb] the top and the bottom half pixels. Assume only zt? is observed, then we are interested in in-painting the bottom of an image by the posterior distribution of zb given zt?. This is achieved using Block Gibbs sampling. A two-stage Gibbs sampler amounts to (a) sampling pθ∗(x|zt, zb) and (b) sampling pθ∗(zb|x, zt) = pθ∗(zb|x) (since zb and zt are independent conditional on x). Given zk = (zt?, z b k), we sample at each step xk ∼ pθ∗(x | zk) and then
3Publicly available online, see https://github.com/YannDubs/disentangling-vae/tree/master/ results/betaH_celeba
zbk+1 ∼ pθ∗(zb | xk). We then set zk+1 = (zt∗, zbk+1). Stage (b) is elementary but stage (a) is challenging. We use an MCMC-within-Gibbs scheme using different samplers. We use the following decomposition of pθ∗(x | z) ∝ ρ(x)L(x) for ρ(x) ∝ qβφ∗(x | z) and L(x) = pθ∗(x, z)/q β φ∗(x | z) with β ∈ (0, 1). It is possible to sample from ρ(x) as qφ∗(x | z) is Gaussian. In our experiments with CelebA and the chosen trained VAE, we have x ∈ R10 (recall that x is our latent variable here), z ∈ R12288, and use β = 0.1. We then compare i-SIR, HMC and NEO-MCMC sampler in stage (a), with the same computational complexity (N = 10, K = 12, γ = 0.2 for NEO-MCMC, N = 120 for i-SIR, and HMC is run with K = 20 leap-frog steps). Again, NEO-MCMC and i-SIR use dependent proposals, with m a Random Walk Metropolis kernel with stepsize 0.1. For each algorithm, 10 steps are performed. Figure 4 displays the evolution of the resulting Markov chains. The samples clearly illustrate that NEO-MCMC mixes better than i-SIR and HMC. More details and examples are presented in the supplementary.

heading: 6 Conclusion

text:
In this paper, we have proposed a new family of algorithms, NEO, for computing normalizing constants and sampling from complex distributions. This methodology comes with asymptotic and non-asymptotic convergence guarantees. For normalizing constant estimation, NEO-IS compares favorably to state-of-the-art algorithms on difficult benchmarks. NEO-MCMC is also able to sample some complex distributions: it is particularly well-adapted to sampling multimodal distributions, thanks to its proposal mechanism which avoids being trapped in local modes. There are numerous potential extensions to this work. For example, it would be interesting to consider deterministic transformations other than conformal Hamiltonian dynamics integrators. These transformations could be trained, as for Neural IS, using a variation lower bound. It would also be interesting to further investigate the influence of the mixture weights {$k}k∈Z on the efficiency of NEO. Broader impact: Sampling from complex target distributions and computing their normalizing constants has numerous applications. Our work proposes novel methods to address such problems and has thus potential applications in many areas. This work does not present any foreseeable societal consequence.
Acknowledgments and Disclosure of Funding
Arnaud Doucet is partly supported by the EPSRC grant EP/R034710/1. He also acknowledges support of the UK Defence Science and Technology Laboratory (DSTL) and EPSRC under grant EP/R013616/1. This is part of the collaboration between US DOD, UK MOD and UK EPSRC under the Multidisciplinary University Research Initiative. Alain Durmus and Eric Moulines acknowledge support of the Lagrange Mathematical and Computing Research Center.

title: NEO: Non Equilibrium Sampling on the Orbit of a Deterministic Transform