summaries: The authors introduce a variational inference method that leverages Annealed Importance Sampling (AIS) and Langevin/HMC MCMC. In particular HMC dynamics are used to target a sequence of intermediate bridging densities in order to generate approximate samples from the target posterior distribution. Importantly, since the method is uncorrected (i.e. does without a accept/reject step) the resulting lower bound is fully differentiable. This makes various hard-to-set parameters (e.g. temperatures for bridging densities) learnable. In experiments the authors demonstrate that their method, Uncorrected Hamiltonian Annealing (UHA), leads to tight lower bounds on the log evidence.

summaries: The paper proposed to use the uncorrected HMC kernel in an AIS-type variational scheme. This enables optimizing the ELBO via reparameterization gradients for many algorithm tuning parameters, including step size, momentum covariance, and annealing schedule. The method is shown to outperform competitors in most cases.

summaries: The paper proposes UHA, an original and clever modification of Hamiltonian AIS. The idea is to change the AIS ratio such that it becomes tunable, but can still be computed in closed form (ratios of momentum distributions). This is achieved by dropping the accept-reject step in the HMC update such that a single UHA step only involves (re)sampling momenta and leapfrog integration. This modification allows the authors to tune parameters such as the proposal distribution, the step size and damping coefficients (used in momentum resampling). Many interesting ideas for future extensions are proposed. Experiments on various models demonstrate that UHA obtains a tighter ELBO than other combinations of VI and HMC.

summaries: The paper propose to use uncorrected HMC kernels (i.e. HMC kernels without accept-reject steps) in AIS to make the estimate of log normalisation constant differentiable w.r.t. to various parameters (e.g. step-size for numerical integration). The author(s) name(s) the method UHA. The promise is that one can then using reparameterisation gradient to tune these parameters. Through experiments on inference tasks, the benefit of tuning parameters in UHA is demonstrated against other methods that use HMC in AIS. For experiments on VAE training, UHA with tuned parameters has shown better test ELBO than IW.

summaries: This paper studies the annealed important sampling (AIS) using the Hamiltonian dynamics kernels in each interpolant. To circumvent the problem that the Hamiltonian dynamics involves the non-differentiable operations, this paper proposes to discard the accept-reject operations in the kernel, and derives the resulting density ratios in a simple form related to the momentum variables. The resulting approach, named Uncorrected Hamiltonian Annealing (UHA), is a fully differentiable AIS method. This paper validates the effectiveness of UHA through extensive experiments and highlights the importance of differentiability.

abstractText: Given an unnormalized target distribution we want to obtain approximate samples from it and a tight lower bound on its (log) normalization constant logZ. Annealed Importance Sampling (AIS) with Hamiltonian MCMC is a powerful method that can be used to do this. Its main drawback is that it uses non-differentiable transition kernels, which makes tuning its many parameters hard. We propose a framework to use an AIS-like procedure with Uncorrected Hamiltonian MCMC, called Uncorrected Hamiltonian Annealing. Our method leads to tight and differentiable lower bounds on logZ. We show empirically that our method yields better performances than other competing approaches, and that the ability to tune its parameters using reparameterization gradients may lead to large performance improvements.

authors:
Felix V Agakov
David Barber

title: An auxiliary variational method

venue: In International Conference on Neural Information Processing,

authors:
Michael Betancourt
Simon Byrne
Sam Livingstone
Mark Girolami

title: The geometric foundations of hamiltonian monte carlo

authors:
David M Blei
Alp Kucukelbir
Jon D McAuliffe

title: Variational inference: A review for statisticians

venue: Journal of the American Statistical Association,

authors:
James Bradbury
Roy Frostig
Peter Hawkins
Matthew James Johnson
Chris Leary
Dougal Maclaurin
George Necula
Adam Paszke
Jake VanderPlas
Skye Wanderman-Milne
Qiao Zhang

title: JAX: composable transformations of Python+NumPy programs, 2018

authors:
Yuri Burda
Roger Grosse
Ruslan Salakhutdinov

title: Importance weighted autoencoders

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

authors:
Anthony L Caterini
Arnaud Doucet
Dino Sejdinovic

title: Hamiltonian variational auto-encoder

venue: In Advances in Neural Information Processing Systems,

authors:
Tianqi Chen
Emily Fox
Carlos Guestrin

title: Stochastic gradient hamiltonian monte carlo

venue: In International conference on machine learning,

authors:
Xiang Cheng
Niladri S Chatterji
Peter L Bartlett
Michael I Jordan

title: Underdamped langevin mcmc: A non-asymptotic analysis

venue: In Conference on Learning Theory,

authors:
Tarin Clanuwat
Mikel Bober-Irizar
Asanobu Kitamoto
Alex Lamb
Kazuaki Yamamoto
David Ha

title: Deep learning for classical japanese literature

venue: arXiv preprint arXiv:1812.01718,

authors:
Gregory Cohen
Saeed Afshar
Jonathan Tapson
Andre Van Schaik

title: Emnist: Extending mnist to handwritten letters

venue: In 2017 International Joint Conference on Neural Networks (IJCNN),

title: Elements of information theory

authors:
Chris Cremer
Quaid Morris
David Duvenaud

title: Reinterpreting importance-weighted autoencoders

venue: arXiv preprint arXiv:1704.02916,

authors:
Xinqiang Ding
David J Freedman

title: Learning deep generative models with annealed importance sampling

venue: arXiv preprint arXiv:1906.04904,

authors:
Justin Domke
Daniel Sheldon

title: Importance weighting and variational inference

title: Divide and couple: Using monte carlo variational objectives for posterior approximation

authors:
Roger B Grosse
Siddharth Ancha
Daniel M Roy

title: Measuring the reliability of mcmc inference with bidirectional monte carlo

authors:
Roger B Grosse
Zoubin Ghahramani
Ryan P Adams

title: Sandwiching the marginal likelihood using bidirectional monte carlo

venue: arXiv preprint arXiv:1511.02543,

title: Learning deep latent gaussian models with markov chain monte carlo

authors:
Matthew D Hoffman
Andrew Gelman

title: The no-u-turn sampler: adaptively setting path lengths in hamiltonian monte carlo

venue: J. Mach. Learn. Res.,

authors:
Chin-Wei Huang
Shawn Tan
Alexandre Lacoste
Aaron Courville

title: Improving explorability in variational inference with annealed variational objectives

title: Equilibrium free-energy differences from nonequilibrium measurements: A master-equation approach

venue: Physical Review E,

authors:
Diederik P Kingma
Jimmy Ba

title: Adam: A method for stochastic optimization

authors:
Diederik P Kingma
Max Welling

title: Auto-encoding variational bayes

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

title: Gradient-based learning applied to document recognition

authors:
Yingzhen Li
Richard E Turner
Qiang Liu

title: Approximate inference with amortised mcmc

venue: arXiv preprint arXiv:1702.08343,

title: Annealed importance sampling

venue: Statistics and computing,

title: Hamiltonian importance sampling

venue: In talk presented at the Banff International Research Station (BIRS) workshop on Mathematical Issues in Molecular Dynamics,

title: Mcmc using hamiltonian dynamics. Handbook of markov chain monte carlo

authors:
Danilo Jimenez Rezende
Shakir Mohamed

title: Variational inference with normalizing flows

venue: In Proceedings of the 32nd International Conference on Machine Learning

authors:
Danilo Jimenez Rezende
Shakir Mohamed
Daan Wierstra

title: Stochastic backpropagation and approximate inference in deep generative models

venue: In Proceedings of the 31st International Conference on Machine Learning

authors:
Francisco Ruiz
Michalis Titsias

title: A contrastive divergence for combining variational inference and mcmc

venue: In International Conference on Machine Learning,

authors:
Ruslan Salakhutdinov
Iain Murray

title: On the quantitative analysis of deep belief networks

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

authors:
Tim Salimans
Diederik Kingma
Max Welling

title: Markov chain monte carlo and variational inference: Bridging the gap

authors:
Jascha Sohl-Dickstein
Benjamin J Culpepper

title: Hamiltonian annealed importance sampling for partition function estimation

venue: arXiv preprint arXiv:1205.1925,

authors:
Esteban G Tabak
Cristina V Turner

title: A family of nonparametric density estimation algorithms

venue: Communications on Pure and Applied Mathematics,

authors:
Achille Thin
Nikita Kotelevskii
Arnaud Doucet
Alain Durmus
Eric Moulines
Maxim Panov

title: Monte carlo variational auto-encoders

authors:
Michalis Titsias
Miguel Lázaro-Gredilla

title: Doubly stochastic variational bayes for nonconjugate inference

authors:
George Tucker
Dieterich Lawson
Shixiang Gu
Chris J Maddison

title: Doubly reparameterized gradient estimators for monte carlo objectives

authors:
Martin J Wainwright
Michael I Jordan

title: Graphical models, exponential families, and variational inference

venue: Foundations and Trends R © in Machine Learning,

authors:
Max Welling
Yee W Teh

title: Bayesian learning via stochastic gradient langevin dynamics

venue: In Proceedings of the 28th international conference on machine learning

authors:
Christopher Wolf
Maximilian Karl
Patrick van der Smagt

title: Variational inference with hamiltonian monte carlo

venue: arXiv preprint arXiv:1609.08203,

authors:
Yuhuai Wu
Yuri Burda
Ruslan Salakhutdinov
Roger Grosse

title: On the quantitative analysis of decoder-based generative models

authors:
Cheng Zhang
Judith Butepage
Hedvig Kjellstrom
Stephan Mandt

title: Advances in variational inference

venue: arXiv preprint arXiv:1711.05597,

heading: 1 Introduction

text:
Variational Inference (VI) [4, 41, 45] is a method to do approximate inference on a target distribution p(z) = p̄(z)/Z that is only known up to the normalization constant Z. The basic insights are, first, that the evidence lower bound (ELBO) Eq(z)[log p̄(z)− log q(z)] lower-bounds logZ and, second, that maximizing the ELBO is equivalent to minimizing the KL-divergence from q to p. The simplest VI method chooses a parameterized family for q and optimizes its parameters to maximize the ELBO.
A recent direction involves combining VI with Markov chain Monte Carlo (MCMC) [34, 43]. These methods can be seen as an instance of the auxiliary VI framework [2] – they create an augmented variational distribution that represents all intermediate random variables generated during the MCMC procedure. An augmented target distribution that attempts to capture the inverse MCMC dynamics is optimized jointly with this variational distribution. However, it has been observed that capturing inverse dynamics is challenging [43, §5.4] (further discussion in Section 4).
Annealed Importance Sampling (AIS) [22, 27] is a powerful technique used to build augmented distributions without the need of learning inverse dynamics. While it was originally proposed to estimate expectations using importance sampling, it can be easily used to build lower bounds on normalization constants of intractable densities [18, 44]. AIS creates a sequence of densities that bridge from a tractable initial approximation q to the target p̄. Then, the augmented variational distribution is given by a sequence of MCMC kernels targeting each bridging density, while the augmented target uses the reversals of those kernels. It turns out that the ratio of these augmented distributions can be computed using only evaluations of the bridging densities. Combining Hamiltonian MCMC kernels with AIS has been observed to produce strong lower bounds [35, 44].
However, these bounds are sensitive to numerous parameters, such as the initial distribution, bridging schedule, and parameters of the MCMC kernels. It would be desirable to optimize these parameters
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
to tighten the bound. Unfortunately, the presence of Metropolis-Hastings acceptance steps means that the the final estimator is non-differentiable, and thus reparameterization gradients cannot be used.
In this work, we propose Uncorrected Hamiltonian Annealing (UHA), a differentiable alternative to Hamiltonian AIS. We define an augmented variational distribution using Hamiltonian MCMC kernels, but dropping the accept-reject steps. This is motivated by the fact that Hamiltonian dynamics sometimes have high acceptance rates. Since these uncorrected MCMC kernels do not exactly hold the bridging densities invariant, an augmented target distribution cannot be defined in terms of reversals. Instead, we define our augmented target by deriving an algorithm for the exact reversal of the original (corrected) MCMC kernel and dropping the accept-reject step. Surprisingly, this yields a very simple expression for the resulting lower bound.
We use reparameterization gradients to tune various parameters involved in the lower bound produced by UHA, including the initial approximation q, parameters of the uncorrected MCMC kernel, and the bridging densities. Experimentally, tuning all these leads to large gains. For example, in several inference tasks we observe that tuning UHA with K = 64 bridging densities gives better results than traditional Hamiltonian AIS with K = 512.
Finally, we use UHA to train VAEs [24, 31]. In this case we observe that using UHA leads to higher ELBOs. In addition, we observe that increasing the number of bridging densities with UHA consistently leads to better results, and that for a large enough number of bridging densities the variational gap (difference between ELBO and true log-likelihood) becomes small, and models with higher log-likelihood are obtained.

heading: 2 Preliminaries

text:
Variational inference and augmentation. Suppose that p(z) = 1Z p̄(z) is some target density, where p̄ is unnormalized and Z = ∫ p̄(z)dz is the corresponding normalizer, and let
ELBO(q(z), p̄(z)) = E q(z)
log p̄(z)
q(z) (1)
be the "ELBO operator". Variational inference (VI) is based on the fact that for any q(z) we have [4]
logZ = ELBO(q(z), p̄(z)) + KL(q(z)‖p(z)). (2)
In VI, the parameters of q are tuned to maximize the "evidence lower bound" (ELBO). Since the KL-divergence is non-negative, this is always a lower bound on logZ. Also, maximizing the ELBO is equivalent to minimizing the KL-divergence from q to p.
To get tighter bounds and better approximations recent work has made use of augmented distributions [2, 21]. Let z1:M = (z1, · · · , zM ) and suppose that p̄(z1:M ) = p̄(zM )p(z1:M−1|zM ) augments the original target density while preserving its normalization constant. Then, for any q(z1:M ) we have
logZ = ELBO(q(z1:M ), p̄(z1:M )) + KL(q(z1:M )‖p(z1:M )). (3)
The first term is called the "augmented" ELBO and again lower bounds logZ. By the chain rule of KL-divergence [12], the KL-divergence from q to p over z1:M upper-bounds the KL-divergence over zM . This justifies using the marginal of q over zM to approximate the original target distribution.
Annealed Importance Sampling. A successful approach for creating augmented distributions is Annealed Importance Sampling (AIS) [27]. It creates an augmented proposal distribution q by applying a sequence of transition densities Tm(zm+1|zm), and an augmented target by defining transition densities Um(zm|zm+1). This gives the augmented densities
q(z1:M ) = q(z1) M−1∏ m=1 Tm(zm+1|zm) and p̄(z1:M ) = p̄(zM ) M−1∏ m=1 Um(zm|zm+1). (4)
Naively, the ratio of these densities is
p̄(z1:M ) q(z1:M ) = p̄(zM ) q(z1) M−1∏ m=1 Um(zm|zm+1) Tm(zm+1|zm) . (5)
To define the transitions Tm and Um, AIS creates a sequence of unnormalized densities π̄1, . . . , π̄M−1 that “bridge” from a starting distribution q to the target p̄, meaning that π̄1 is close to q and π̄M−1 is close to p̄. Then, for each intermediate distribution, Tm(zm+1|zm) is chosen to be a Markov kernel that holds πm invariant, and Um to be the reversal of Tm with respect to πm, defined as
Um(zm|zm+1) = T (zm+1|zm) πm(zm)
πm(zm+1) . (6)
This choice produces a simplification so that eq. 5 becomes
p̄(z1:M ) q(z1:M ) = p̄(zM ) q(z1) M−1∏ m=1 π̄m(zm) π̄m(zm+1) . (7)
This can be easily evaluated without needing to evaluate the transition densities. The ratio from eq. 7 can be used to get an expression for the lower bound ELBO(q(z1:M ), p̄(z1:M )). Research has shown that the AIS augmentation may lead to extremely tight lower bounds [18, 17, 35, 44].
Hamiltonian Dynamics. Many MCMC methods used to sample from p(z) are based on Hamiltonian dynamics [3, 8, 29, 42]. The idea is to create an augmented distribution p(z, ρ) = p(z)S(ρ), where S(ρ) is a distribution over a momentum variable ρ (e.g. a Multivariate Gaussian). Then, one can define numerical integration schemes where z and ρ evolve while nearly holding p(z, ρ) constant. When corrected by a Metropolis-Hastings acceptance step, this can be made to exactly hold p(z, ρ) invariant. This is alternated with a scheme that resamples the momentum ρ while holding S(ρ) invariant. When Hamiltonian dynamics work well, z can quickly move around, suppressing random-walk behavior.
There are a variety of different Hamiltonian MCMC methods, corresponding to different integration schemes, momentum distributions, and ways of resampling the momentum. For instance, HMC and Langevin dynamics use the leapfrog integrator, a Gaussian for the momentum variables and a full resampling of the momentum variables at each step [29, 42]. On the other hand, if the momentum variables are only partially resampled, the under-damped variants of HMC and Langevin dynamics are recovered [29]. It was observed that partial resampling may lead to improved perfomance [9].
It is easy to integrate Hamiltonian dynamics into AIS. First, define an augmented target p̄(z, ρ) = p̄(z)S(ρ) and an augmented starting distribution q(z, ρ) = q(z)S(ρ). Then, create a series of augmented densities π̄1(z, ρ), . . . , π̄M−1(z, ρ) bridging the two as π̄m(z, ρ) = π̄m(z)S(ρ). Finally, define the forward transition Tm(zm+1, ρm+1|zm, ρm) to be an iteration of a Hamiltonian MCMC method that leaves πm(z, ρ) invariant. We will describe a single transition Tm as a sequence of three steps: (1) resample the momentum; (2) simulate Hamiltonian dynamics and apply an accept-reject step; and (3) negate the momentum. The precise process that defines the transition is shown in Alg. 1. Note that this algorithm is quite general, and compatible with HMC, Langevin dynamics and their underdamped variants (by selecting an appropriate integrator and resampling method).
Algorithm 1 Corrected Tm(zm+1, ρm+1|zm, ρm) 1. Sample ρ′m from some s(ρ ′ m|ρm) that leaves S(ρ) invariant. Set z′m ← zm.
2. Simulate Hamiltonian dynamics as (z′′m, ρ′′m)← Tm(z′m, ρ′m). Calculate an acceptance probability α = min (1, π̄m(z′′m, ρ′′m)/π̄m(z′m, ρ′m)). With probability α, set (z′′′m , ρ′′′m)← (z′′m, ρ′′m). Otherwise, set (z′′′m , ρ′′′m)← (z′m, ρ′m). 3. Reverse the momentum as (zm+1, ρm+1)← (z′′′m ,−ρ′′′m). return (zm+1, ρm+1)
Representing Tm this way makes it easy to show it holds the density πm(z, ρ) invariant. The overall strategy is to show that each of the steps 1-3 holds πm invariant, and so does the composition of them [29, §3.2]. For steps 1 and 3 this is trivial, provided that S(ρ) = S(−ρ). For step 2, we require that the simulation Tm has unit Jacobian and satisfies T −1m = Tm. Then, Tm can be interpreted as a symmetric Metropolis-Hastings proposal, meaning the Metroplis-Hastings acceptance probability α is as given. A typical choice for Tm that satisfies these requirements is the leapfrog integrator with a momentum reversal at the end. (This reversal then gets "un-reversed" in step 3 for accepted moves.)
Since Tm holds πm invariant, we can define Um as the reversal of Tm wrt πm. Then, eq. 7 becomes
p̄(z1:M , ρ1:M ) q(z1:M , ρ1:M ) = p̄(zM , ρM ) q(z1, ρ1) M−1∏ m=1 π̄m(zm, ρm) π̄m(zm+1, ρm+1) . (8)
Using this ratio we get an expression for the lower bound ELBO(q(z1:M , ρ1:M ), p̄(z1:M , ρ1:M )) obtained with Hamiltonian AIS. While this method has been observed to yield strong lower bounds on logZ [35, 44] (see also Section 5.2), its performance depends on many parameters: initial distribution q(z), momentum distribution S, momentum resampling scheme, simulator Tm, and bridging densities. We would like to tune these parameters by maximizing the ELBO using reparameterization-based estimators. However, due to the accept-reject step required by the Hamiltonian MCMC transition, the resulting bound is not differentiable, and thus reparameterization gradients are not available.

heading: 3 Uncorrected Hamiltonian Annealing

text:
The contribution of this paper is the development of uncorrected Hamiltonian Annealing (UHA). This method is similar to Hamiltonian AIS (eq. 8), but yields a differentiable lower bound. The main idea is simple. For any transitions Tm and Um, by the same logic as in eq. 5, we can define the ratio
p̄(z1:M , ρ1:M ) q(z1:M , ρ1:M ) = p̄(zM , ρM ) q(z1, ρ1) M−1∏ m=1 Um(zm, ρm|zm+1, ρm+1) Tm(zm+1, ρm+1|zm, ρm) . (9)
Hamiltonian AIS defines Tm as a Hamiltonian MCMC kernel that holds πm invariant, and Um as the reversal of Tm with respect to πm. While this leads to a nice simplification, there is no requirement that these choices be made. We can use any transitions as long as the ratio Um/Tm is tractable.
We propose to use the "uncorrected" versions of the transitions Tm and Um used by Hamiltonian AIS, obtained by dropping the accept-reject steps. To get an expression for the uncorrected Um we first derive the reversal Um used by Hamiltonian AIS (Alg. 2). These uncorrected transitions are no longer reversible with respect to the bridging densities πm(z, ρ), and thus we cannot use the simplification used by AIS to get eq. 8. Despite this, we show that the ratio Um/Tm for the uncorrected transitions can still be easily computed (Thm. 2). This produces a differentiable estimator, meaning the parameters can be tuned by stochastic gradient methods designed to maximize the ELBO.
We start by deriving the process that defines the transition Um used by Hamiltonian AIS. This is shown in Alg. 2. It can be observed that Um follows the same three steps of Tm (resample momentum, Hamiltonian simulation with accept-reject, momentum negation), but in reverse order.
Algorithm 2 Corrected Um(zm, ρm|zm+1, ρm+1) 1. Set (z′′′m , ρ′′′m)← (zm+1,−ρm+1). 2. Simulate Hamiltonian dynamics as (z′′m, ρ′′m)← Tm(z′′′m , ρ′′′m).
Calculate an acceptance probability α = min (1, π̄m(z′′m, ρ′′m)/π̄m(z′′′m , ρ′′′m)). With probability α, set (z′m, ρ′m)← (z′′m, ρ′′m). Otherwise, set (z′m, ρ′m)← (z′′′m , ρ′′′m).
3. Sample ρm from srev(ρm|ρ′m), the reversal of s(ρ′m|ρm) with respect to S(ρm). Set zm ← z′m. return (zm, ρm)
Lemma 1. The corrected Um (Alg. 2) is the reversal of the corrected Tm (Alg. 1) with respect to πm.
(Proof Sketch). First, we claim the general result that if T1, T2 and T3 have reversals U1, U2 and U3, respectively, then the composition T = T1 ◦T2 ◦T3 has reversal U = U3 ◦U2 ◦U1 (all reversals with respect to same density π). Then, we apply this to the corrected Tm and Um: Tm is the composition of three steps that hold πm invariant. Thus, its reversal Um is given by the composition of the reversals of those steps, applied in reversed order. A full proof is in Appendix F.
We now define the "uncorrected" transitions used by UHA, shown in Algs. 3 and 4. These are just the transitions used by Hamiltonian AIS but without the accept-reject steps. (If Hamiltonian dynamics are simulated exactly, the acceptance rate is one and the uncorrected and corrected transitions are equivalent.) We emphasize that, for the "uncorrected" transitions, Tm does not exactly hold πm invariant and Um is not the reversal of Tm. Thus, their ratio does not give a simple expression in terms of π̄m as in eq. 8. Nevertheless, the following result shows that their ratio has a simple form. Theorem 2. Let Tm and Um be the uncorrected transitions defined in Algs. 3 and 4, and let the dynamics simulator Tm(z, ρ) be volume preserving and self inverting. Then,
Um(zm, ρm|zm+1, ρm+1) Tm(zm+1, ρm+1|zm, ρm) = S(ρm) S(ρ′m) , (10)
Algorithm 3 Uncorrected Tm(zm+1, ρm+1|zm, ρm) 1. Sample ρ′m from some s(ρ ′ m|ρm) that leaves S(ρ) invariant. Set z′m ← zm.
2. Simulate Hamiltonian dynamics as (z′′m, ρ′′m)← Tm(z′m, ρ′m). 3. Reverse the momentum as (zm+1, ρm+1)← (z′′m,−ρ′′m). return (zm+1, ρm+1)
Algorithm 4 Uncorrected Um(zm, ρm|zm+1, ρm+1) 1. Set (z′′m, ρ′′m)← (zm+1,−ρm+1). 2. Simulate Hamiltonian dynamics as (z′m, ρ′m)← Tm(z′′m, ρ′′m). 3. Sample ρm from srev(ρm|ρ′m), the reversal of s(ρ′m|ρm) with respect to S(ρm). Set zm ← z′m. return (zm, ρm)
where ρ′m is the second component of Tm(zm+1,−ρm+1). (That is, ρ′m from Algs. 3 and 4.) (Proof Sketch.) We consider variants of Algs. 3 and 4 in which each time z is assigned we add Gaussian noise with some variance aI . We then derive the densities for Tm and Um using the rule for transformation of densities under invertible mappings, using that Tm is self-inverting and volume preserving. Taking the ratio gives eq. 10. Since this is true for arbitrary a, we take the stated result as the limit as a→ 0. A full proof is in Appendix G.
As an immediately corollary of eq. 9 and Theorem 2 we get that for UHA
p̄(z1:M , ρ1:M ) q(z1:M , ρ1:M ) = p̄(zM ) q(z1) M−1∏ m=1 S(ρm+1) S(ρ′m) . (11)
This ratio can be used to get an expression for the lower bound ELBO(q(z1:M , ρ1:M ), p̄(z1:M , ρ1:M )) obtained with UHA. As mentioned in Section 2, the parameters of the augmented distributions are tuned to maximize the ELBO, equivalent to minimizing the KL-divergence from q to p̄. While computing this ELBO exactly is typically intractable, an unbiased estimate can be obtained using a sample from q(z1:M , ρ1:M ) as shown in Alg. 5. If sampling is done using reparameterization, then unbiased reparameterization gradients may be used together with stochastic optimization algorithms to optimize the lower bound. In contrast, the variational lower bound obtained with Hamiltonian AIS (see Alg. 6 in Appendix A) does not allow the computation of unbiased reparameterization gradients.
Algorithm 5 Generating the (differentiable) uncorrected Hamiltonian annealing variational bound. Sample z1 ∼ q and ρ1 ∼ S. Initialize estimator as L ← − log q(z1). for m = 1, 2, · · · ,M − 1 do
Run uncorrected Tm (Alg. 3) on input (zm, ρm), storing ρ′m and the output (zm+1, ρm+1). Update estimator as L ← L+ log (S(ρm+1)/S(ρ′m)).
Update estimator as L ← L+ log p̄(zM ). return R

heading: 3.1 Algorithm Details

text:
Simulation of dynamics. We use the leapfrog operator with step-size to simulate Hamiltonian dynamics. This has unit Jacobian and satisfies Tm = T −1m (if the momentum is negated after the simulation), which are the properties required for eq. 11 to be correct (see Theorem 2).
Momentum distribution and resampling. We set the momentum distribution S(ρ) = N (ρ|0,Σ) to be a Gaussian with mean zero and covariance Σ. The resampling distribution s(ρ′|ρ) must hold this distribution invariant. As is common we use s(ρ′|ρ) = N (ρ′|ηρ, (1− η2)Σ), where η ∈ [0, 1) is the damping coefficient. If η = 0, the momentum is completely replaced with a new sample from S in each iteration (used by HMC and Langevin dynamics [29, 42]). For larger η, the momentum becomes correlated between iterations, which may help suppress random walk behavior and encourage faster mixing [9] (used by the underdamped variants of HMC and Langevin dynamics [29]).
Bridging densities. We set π̄m(z, ρ) = q(z, ρ)1−βm p̄(z, ρ)βm , where βm ∈ [0, 1] and βm < βm+1.
Computing gradients. We set the initial distribution q(z1) to be a Gaussian, and perform all sampling operations in Alg. 5 using reparameterization [24, 31, 39]. Thus, the whole procedure is differentiable and reparameterization-based gradients may be used to tune parameters by maximizing the ELBO. These parameters include the initial distribution q(z1), the covariance Σ of the momentum distribution, the step-size of the integrator, the damping coefficient η of the momentum resampling distribution, and the parameters of the bridging densities (including β), among others. As observed in Section 5.2.1 tuning all of these parameters may lead to considerable performance improvements.

heading: 4 Related Work

text:
UHA and slight variations have been proposed in concurrent work by Thin et al. [38], who use uncorrected Langevin dynamics together with the uncorrected reversal to build variational lower bounds, and by Zhang et al. [46], who proposed UHA with under-damped Langevin dynamics together with a convergence analysis for linear regression models.
There are three other lines of work that produce differentiable variational bounds integrating Monte Carlo methods. One is Hamiltonian VI (HVI) [34, 43]. It uses eq. 9 to build a lower bound on logZ, with Tm set to an uncorrected Hamiltonian transition (like UHA but without bridging densities) and Um set to conditional Gaussians parameterized by learnable functions. Typically, a single transition is used, and the parameters of the transitions are learned by maximizing the resulting ELBO.1
A second method is given by Hamiltonian VAE (HVAE) [7], based on Hamiltonian Importance sampling [28]. They augment the variational distribution with momentum variables, and use the leapfrog integrator to simulate Hamiltonian dynamics (a deterministic invertible transformation with unit Jacobian) with a tempering scheme as a target-informed flow [30, 37].
The third method is Importance Weighting (IW) [6, 13, 15]. Here, the idea is that ELBO(q(z), p̄(z)) ≤ E log 1K ∑ k p̄(zk)/q(zk), and that the latter bound can be optimized, rather than the traditional ELBO. More generally, other Monte-Carlo estimators can be used [16].
Some work defines novel contrastive-divergence-like objectives in terms of the final iteration of an MCMC chain [32, 26]. These do not provide an ELBO-like variational bound. While in some cases the initial distribution can be optimized to minimize the objective [32], gradients do not flow through the MCMC chains, meaning MCMC parameters cannot be optimized by gradient methods.
For latent variable models, Hoffman [19] suggested to run a few MCMC steps after sampling from the variational distribution before computing gradients with respect to the model parameters, which is expected to "debias" the gradient estimator to be closer to the true likelihood gradient. The variational distribution is simultaneously learned to optimize a standard ELBO. (AIS can also be used [14].)

heading: 5 Experiments and Results

text:
This section presents results using UHA for Bayesian inference problems on several models of varying dimensionality and for VAE training. We compare against Hamiltonian AIS, IW, HVI and HVAE. We report the performance of each method for different values of K, the number of likelihood evaluations required to build the lower bound (e.g. number of samples used for IW, number of bridging densities plus one for UHA). Note that, for a fixed K, all methods have the same oracle complexity (i.e. number of target/target’s gradient evaluation), and that for K = 1 they all reduce to plain VI.
For UHA and Hamiltonian AIS we use under-damped Langevin dynamics, that is, we perform just one leapfrog step per transition and partially resample momentum. We implement all algorithms using Jax [5].
1The formulation of HVI allows the use of more than one transition. However, this leads to an increased number of reverse models that must be learned, and thus not typically used in practice. Indeed, experiments by Salimans et al. [34] use only one HMC step while varying the number of leapfrog integration steps, and results from Wolf et al. [43] show that increasing the number of transitions may actually yield worse bounds (they conjecture that this is due to the difficulty of learning inverse dynamics.).

heading: 5.1 Toy example

text:
This section compares results obtained with UHA and IW when the target is set to a factorized Studentt with mean zero, scale one, and three degrees of freedom. We tested three different dimensionalities: 20, 200 and 500. In all cases we have logZ = 0, so we can exactly analyze the tightness of the bounds obtained by the methods. We set the initial approximation to be a mean-field Gaussian, and optimize the objective using Adam [23] with a step-size of 0.001 for 5000 steps. For UHA we tune the initial approximation q(z), the integrator’s step-size and the damping coefficient η.
We ran UHA for K ∈ {4, 16, 64, 128} and IW for K ∈ {128, 1024}. Table 1 shows the results for the three dimensionalities considered. It can be observed that UHA performs significantly better than IW as the dimensionality increases; for the target with dimension 500, UHA with K = 16 yields better bounds than IW with K = 1024. On the other hand, the methods perform similarly for the low dimensional target. Finally, in this case both methods have similar time costs. For instance, for K = 128 UHA takes 14.2 seconds to optimize and IW takes 13.9.

heading: 5.2 Inference tasks

text:
This section shows results using UHA for Bayesian inference tasks. For this set of experiments, for UHA we tune the initial distribution q(z), the integrator’s step-size and the damping coefficient η. We include detailed results tuning more parameters in Section 5.2.1.
Models. We consider four models: Brownian motion (d = 32), which models a Brownian Motion process with a Gaussian observation model; Convection Lorenz bridge (d = 90), which models a nonlinear dynamical system for atmospheric convection; and Logistic regression with the a1a (d = 120) and madelon (d = 500) datasets. The first two obtained from the “Inference gym” [36].
Baselines. We compare UHA against IW, HVAE, a simple variant of HVI, and Hamiltonian AIS (HAIS). For all methods which rely on HMC (i.e. all except IW) we use a singe integration step-size common to all dimensions and fix the momentum distribution to a standard Gaussian. For HVI we learn the initial distribution q(z), integration step-size and the reverse dynamics Um (set to a factorized Gaussian with mean and variance given by affine functions), and for HVAE we learn q(z), and the tempering scheme (we use the quadratic scheme parameterized by a single parameter).
Training details. We set q(z) to be a mean-field Gaussian initialized to a maximizer of the ELBO, and tune the parameters of each method by running Adam for 5000 steps. We repeat all simulations for different step-sizes in {10−3, 10−4, 10−5}, and select the best one for each method. Since Hamiltonian AIS’ parameters cannot be tuned by gradient descent, we find a good pair ( , η) by grid search. We consider η ∈ {0.5, 0.9, 0.99} and three values of that correspond to three different rejection rates: 0.05, 0.25 and 0.5. We tested all 9 possible combinations and selected the best one.
Results are shown in Fig. 1. Our method yields better lower bounds than all other competing approaches for all models considered, and that increasing the number of bridging densities consistently leads to better results. The next best performing method is Hamiltonian AIS. IW also shows a good performance for the lower dimensional model Brownian motion. However, for models of higher dimensionality IW leads to bounds that are several nats worse than the ones achieved by UHA. Finally, HVI and HVAE yield bounds that are much worse than those achieved by the other three methods, and do not appear to improve consistently for larger K. For HVAE, these results are consistent with
the ones in the original paper [7, §4], in that higher K may sometimes hurt performance. For HVI, we believe this is related to the use of just one HMC step and suboptimal inverse dynamics.
Optimization times for Plain VI, IW and UHA (the latter two with K = 32) are 2.4, 3.4 and 4.4 seconds for the Brownian motion dataset, 2.5, 6.8 and 6.9 seconds for Lorenz convection, 2.8, 8.3 and 19.9 seconds for Logistic regression (A1A), and 4.6, 16.6 and 121.2 seconds for Logistic regression (Madelon). While IW and UHA have the same oracle complexity for the same K, we see that the difference between their time cost depends on the specific model under consideration. All other methods that use HMC have essentially the same time cost as UHA.

heading: 5.2.1 Tuning More Parameters with UHA

text:
A basic version of UHA involves fitting a variational distribution using plain VI, and then tuning the integration step-size and the damping coefficient η. However, more parameters could be tuned:
• Moment distribution cov Σ: We propose to learn a diagonal matrix instead of using the identity. • Bridging densities’ coefficients βm: Typically βm = m/M . We propose to learn the sequence β,
with the restrictions β0 = 0, βM = 1, βm < βm+1 and βm ∈ [0, 1]. • Initial distribution q(z): Instead of fixing q(z) to be a maximizer of the typical ELBO, we propose
to learn it to maximize the augmented ELBO obtained using UHA. • Integrator’s step-size : Instead of learning a unique step-size , we propose to learn a step-size
that is a function of β, i.e. (β). In our experiments we use an affine function. • Bridging densities parameters ψ: Instead of setting them-th bridging density to be q1−βmpβm , we
propose to set it to q1−βmψ(βm) p βm , where qψ(βm) is a mean-field Gaussian with a mean and diagonal covariance specified as affine functions of β.
We consider the four models described previously and compare three methods: UHA tuning all parameters described above, UHA tuning only the pair ( , η), and Hamiltonian AIS with parameters ( , η) obtained by grid-search. We perform the comparison forK ranging from 2 to 512. (ForK ≥ 64 we tune the UHA’s parameters using K = 64 and extrapolate them as explained in Appendix D.)
Results are shown in Fig. 2. It can be observed that tuning all parameters with UHA leads to significantly better lower bounds than those obtained by Hamiltonian AIS (or UHA tuning only and η). Indeed, for the Logistic regression models, UHA tuning all parameters for K = 64 leads to results comparable to the ones obtained by Hamiltonian AIS with K = 512.
To verify what parameters lead to larger performance improvements, we tested UHA with K = 64 tuning different subsets of { , η,Σ, β, q(z), (β), ψ(β)}. Fig. 3 shows the results. It can be observed
that tuning the bridging parameters β and the initial approximation q(z) leads to the largest gains in performance, and that tuning all parameters always outperforms tuning smaller subsets of parameters. We show a more thorough analysis, including more subsets and values of K in Appendix B.
Finally, Appendix E shows results comparing UHA (tuning several parameters) against HMC, mean field VI and IW in terms of the approximation accuracy achieved on a logistic regression model with a fixed computational budget.

heading: 5.3 VAE training

text:
Our method can be used to train latent variable models, such as Variational Auto-encoders (VAE) [24, 31]. In this case the initial approximation q(z|x) and the model p(x, z) are parameterized by two neural networks (encoder and decoder), whose parameters are trained by maximizing the ELBO. UHA can be used to train VAEs by augmenting these two distributions as described in Section 3.
Datasets. We use three datasets: mnist [25] (numbers 1-9), emnist-letters [11] (letters A-Z), and kmnist [10] (cursive Kuzushiji). All consist on greyscale images of 28×28 pixels. In all cases we use stochastic binarization [33] and a training set of 50000 samples, a validation set of 10000 samples, and a test set of 10000 samples. All datasets are available in tensorflow-datasets [1].
Baselines. We compare against Importance Weighted Auto-encoders [6] and plain VAE training [24].
Architecture details. We set q(z|x) to a diagonal Gaussian, p(z) to a standard Normal, and p(x|z) to a Bernoulli. We consider two architectures for the encoder and decoder: (1) Feed forward networks with one hidden layer of size 450 and Relu non-linearities, with a latent space dimensionality of 64; (2) Architecture used by Burda et al. [6], feed forward networks with two hidden layers of size 200 with tanh non-linearities, with a latent space dimensionality of 50.
Training details. In all cases the encoder and decoder are initialized to parameters that maximize the ELBO. For IW we tune the encoder and decoder parameters (using the doubly-reparameterized estimator [40]), and for UHA we tune the integration step-size , damping coefficient η, bridging parameters β, momentum covariance Σ (diagonal), and the decoder parameters. Following Caterini et al. [7] we constrain ∈ (0, 0.05) to avoid unstable behavior of the leapfrog discretization. We use Adam with a step-size of 10−4 to train for 100 epochs and use the validation set for early stopping. We repeated all simulations for three different random seeds. In all cases the standard deviation of the results was less than 0.1 nats (not shown in tables).
All methods achieved better results using the architecture with one hidden layer. These results are shown in Tables 2 and 3. The first one shows the ELBO on the test set achieved for different values of K, and the second one the log-likelihood on the test set estimated with AIS [44]. It can be observed that UHA leads to higher ELBOs, higher log-likelihoods, and smaller variational gaps (difference between ELBO and log-likelihood) than IW for all datasets, with the difference between both methods’ performance increasing for increasing K. Notably, for K = 64, the variational gap for UHA becomes quite small, ranging from 0.8 to 1.4 nats depending on the dataset.
Results for the architecture from Burda et al. [6] (two hidden layers) are shown in Tables 4 and 5 (Appendix C). Again, we observe that UHA consistently leads to higher ELBOs and the best test log-likelihood was consistently achieved by UHA with K = 64. However, for smaller K, IW sometimes had better log-likelihoods than UHA (despite worse ELBOs).

heading: 6 Discussion

text:
Since UHA yields a differentiable lower bound, one could tune other parameters not considered in this work. For instance, a different momentum distribution per bridging density could be used, that is, π̄m(z, ρ) = π̄m(z)Sm(ρ). We believe additions such as this may yield further gains. Also, our method can be used to get tight and differentiable upper bounds on logZ using the reversed AIS procedure described by Grosse et al. [18].
Finally, removing accept-reject steps might sometimes lead to instabilities during optimization if the step-size becomes large. We observed this effect when training VAEs on some datasets for the larger values of K. We solved this by constraining the range of (previously done by Caterini et al. [7]). While this simple solution works well, we believe that other approaches (e.g. regularization, automatic adaptation) could work even better. We leave the study of such alternatives for future work.
Acknowledgments and Disclosure of Funding
This material is based upon work supported in part by the National Science Foundation under Grant No. 1908577.

title: MCMC Variational Inference via Uncorrected Hamiltonian Annealing


summaries: This paper studies certified robustness from the perspective of computing Lipschitzness bounds. Prior work compute global bounds on the Lipschitz constant of the network; this paper takes an alternative approach which computes local Lipschitz bounds around a specific input. These local Lipschitz bounds are based on interval bounds for the hidden layers on each data point, and computing a tighter Lipschitz constant based on eliminating rows/columns of the weight matrices when the interval bounds indicate that a neuron is inactive.   To train certifiably robust models, the worst-case loss is optimized based on these Lipschitz bounds. There are several other ingredients for the proposed method: a modified ReLU activation which truncates when the output is a large positive value, and a sparsity-inducing regularization meant to limit the number of neurons where the upper and lower activation bounds differ. The authors compare certified robustness w.r.t. $\ell_2$ perturbations against other Lipschitzness-based baselines and demonstrate improvements.

summaries: This paper provides a trainable local Lipschitz upper-bound for a neural network that is tighter than the global Lipschitz upper-bound computed via multiplying the Lipschitz constant of each individual layer. By making use of the interaction between the weight matrices and the piece-wise linear activation functions, the authors demonstrate that the global Lipschitz bound could be improved when the activation functions are constant locally. In addition, the authors proposed a variant of ReLU that is clipped at a certain upper threshold to increase the constant region to help achieve a tighter Lipschitz upper-bound. Empirically, the authors showed that optimizing the local Lipschitz bound, along with using the modified ReLU, improved the certified robust performance of existing algorithms on several standard image classification tasks.

summaries: Per data-example, a local Lipschitz bound on the neural network output is computed and then used to certify robustness. This bound is made tighter by exploiting invariances of the ReLU activations throughout the network under small perturbations around the given data point. The paper further increases invariances by clipping ReLU and training the clipping threshold.

summaries: This paper proposes a new method for certified robustness that tightly bounds the local Lipschitz constants for models by exploiting zero-activation of ReLU. For each data point, the proposed method first calculates the positions where feature maps become zero by ReLU. Next, since the row vector of the weight matrix before the ReLU outputs zero does not affect the output, the proposed method eliminates such row vectors and computes the Lipschitz constants of the modified weight matrix. Experiments demonstrate that the proposed method outperforms previous methods in terms of clean accuracy, robust accuracy against PGD, and certified robust accuracy.

summaries: This paper proposes using local Lipschitz bound to obtain certified robust neural networks. The use of local Lipschitz constant strictly results in a provable tighter bound than the commonly used global Lipschitz bound. To incorporate the proposed certification method into training, the authors further proposed a modified ReLU function with a learnable upper threshold, and hinge loss to encourage the pre-activation states to be constant past ReLU function.  The proposed method is sound, numerical experiments show promising improvement over the existing certified robust training methods.

abstractText: Certified robustness is a desirable property for deep neural networks in safetycritical applications, and popular training algorithms can certify robustness of a neural network by computing a global bound on its Lipschitz constant. However, such a bound is often loose: it tends to over-regularize the neural network and degrade its natural accuracy. A tighter Lipschitz bound may provide a better tradeoff between natural and certified accuracy, but is generally hard to compute exactly due to non-convexity of the network. In this work, we propose an efficient and trainable local Lipschitz upper bound by considering the interactions between activation functions (e.g. ReLU) and weight matrices. Specifically, when computing the induced norm of a weight matrix, we eliminate the corresponding rows and columns where the activation function is guaranteed to be a constant in the neighborhood of each given data point, which provides a provably tighter bound than the global Lipschitz constant of the neural network. Our method can be used as a plug-in module to tighten the Lipschitz bound in many certifiable training algorithms. Furthermore, we propose to clip activation functions (e.g., ReLU and MaxMin) with a learnable upper threshold and a sparsity loss to assist the network to achieve an even tighter local Lipschitz bound. Experimentally, we show that our method consistently outperforms state-of-the-art methods in both clean and certified accuracy on MNIST, CIFAR-10 and TinyImageNet datasets with various network architectures.

authors:
Anish Athalye
Nicholas Carlini
David Wagner

title: Obfuscated gradients give a false sense of security: Circumventing defenses to adversarial examples

authors:
Moustapha Cisse
Piotr Bojanowski
Edouard Grave
Yann Dauphin
Nicolas Usunier

title: Parseval networks: Improving robustness to adversarial examples

authors:
Haifeng Qian
Mark N Wegman

title: L2-nonexpansive neural networks

venue: arXiv preprint arXiv:1802.07896,

authors:
Klas Leino
Zifan Wang
Matt Fredrikson

title: Globally-robust neural networks

venue: arXiv preprint arXiv:2102.08452,

authors:
Todd Huster
Cho-Yu Jason Chiang
Ritu Chadha

title: Limitations of the lipschitz constant as a defense against adversarial examples

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

authors:
Aleksander Madry
Aleksandar Makelov
Ludwig Schmidt
Dimitris Tsipras
Adrian Vladu

title: Towards deep learning models resistant to adversarial attacks

venue: arXiv preprint arXiv:1706.06083,

authors:
Matthias Hein
Maksym Andriushchenko

title: Formal guarantees on the robustness of a classifier against adversarial manipulation

venue: Advances in Neural Information Processing Systems,

authors:
Huan Zhang
Pengchuan Zhang
Cho-Jui Hsieh

title: Recurjac: An efficient recursive algorithm for bounding jacobian matrix of neural networks and its applications

venue: In Proceedings of the AAAI Conference on Artificial Intelligence,

authors:
Guy Katz
Clark Barrett
David L Dill
Kyle Julian
Mykel J Kochenderfer

title: Reluplex: An efficient smt solver for verifying deep neural networks

venue: In International Conference on Computer Aided Verification,

authors:
Mahyar Fazlyab
Alexander Robey
Hamed Hassani
Manfred Morari
George J Pappas

title: Efficient and accurate estimation of lipschitz constants for deep neural networks

authors:
Matt Jordan
Alexandros G Dimakis

title: Exactly computing the local lipschitz constant of relu networks

venue: arXiv preprint arXiv:2003.01219,

authors:
Sven Gowal
Krishnamurthy Dvijotham
Robert Stanforth
Rudy Bunel
Chongli Qin
Jonathan Uesato
Timothy Mann
Pushmeet Kohli

title: On the effectiveness of interval bound propagation for training verifiably robust models

venue: arXiv preprint arXiv:1810.12715,

authors:
Matthew Mirman
Timon Gehr
Martin Vechev

title: Differentiable abstract interpretation for provably robust neural networks

authors:
Huan Zhang
Hongge Chen
Chaowei Xiao
Bo Li
Duane Boning
Cho-Jui Hsieh

title: Towards stable and efficient training of verifiably robust neural networks

venue: In International Conference on Learning Representations,

authors:
Cem Anil
James Lucas
Roger Grosse

title: Sorting out lipschitz function approximation

authors:
Kevin Scaman
Aladin Virmaux

title: Lipschitz regularity of deep neural networks: analysis and efficient estimation

venue: arXiv preprint arXiv:1805.10965,

authors:
Lily Weng
Huan Zhang
Hongge Chen
Zhao Song
Cho-Jui Hsieh
Luca Daniel
Duane Boning
Inderjit Dhillon

title: Towards fast computation of certified robustness for relu networks

authors:
Christian Szegedy
Wojciech Zaremba
Ilya Sutskever
Joan Bruna
Dumitru Erhan
Ian Goodfellow
Rob Fergus

title: Intriguing properties of neural networks

venue: arXiv preprint arXiv:1312.6199,

authors:
Yusuke Tsuzuku
Issei Sato
Masashi Sugiyama

title: Lipschitz-margin training: Scalable certification of perturbation invariance for deep neural networks

venue: arXiv preprint arXiv:1802.04034,

authors:
Sungyoon Lee
Jaewook Lee
Saerom Park

title: Lipschitz-certifiable training with a tight outer bound

authors:
Eric Wong
J Zico Kolter

title: Provable defenses against adversarial examples via the convex outer adversarial polytope

venue: In ICML,

authors:
Eric Wong
Frank R Schmidt
Jan Hendrik Metzen
J Zico Kolter

title: Scaling provable adversarial defenses

venue: In NeurIPS,

authors:
Zhouxing Shi
Yihan Wang
Huan Zhang
Jinfeng Yi
Cho-Jui Hsieh

title: Fast certified robust training via better initialization and shorter warmup

venue: Advances in Neural Information Processing Systems (NeurIPS),

authors:
Jeremy M Cohen
Elan Rosenfeld
J Zico Kolter

title: Certified adversarial robustness via randomized smoothing

authors:
Bai Li
Changyou Chen
Wenlin Wang
Lawrence Carin

title: Certified adversarial robustness with additive noise

venue: In Advances in Neural Information Processing Systems,

authors:
Mathias Lecuyer
Vaggelis Atlidakis
Roxana Geambasu
Daniel Hsu
Suman Jana

title: Certified robustness to adversarial examples with differential privacy

venue: IEEE Symposium on Security and Privacy (SP),

authors:
Hadi Salman
Jerry Li
Ilya Razenshteyn
Pengchuan Zhang
Huan Zhang
Sebastien Bubeck
Greg Yang

title: Provably robust deep learning via adversarially trained smoothed classifiers

authors:
Huan Zhang
Tsui-Wei Weng
Pin-Yu Chen
Cho-Jui Hsieh
Luca Daniel

title: Efficient neural network robustness certification with general activation functions

authors:
Andrew G Howard
Menglong Zhu
Bo Chen
Dmitry Kalenichenko
Weijun Wang
Tobias Weyand
Marco Andreetto
Hartwig Adam

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

venue: arXiv preprint arXiv:1704.04861,

authors:
Asher Trockman
J Zico Kolter

title: Orthogonalizing convolutional layers with the cayley transform

venue: arXiv preprint arXiv:2104.07167,

authors:
Alex Krizhevsky
Geoffrey Hinton

title: Learning multiple layers of features from tiny images

authors:
Kai Y Xiao
Vincent Tjeng
Nur Muhammad Shafiullah
Aleksander Madry

title: Training for faster adversarial robustness verification via inducing relu stability

venue: arXiv preprint arXiv:1809.03008,

authors:
Diederik P Kingma
Jimmy Ba

title: Adam: A method for stochastic optimization

venue: arXiv preprint arXiv:1412.6980,

authors:
Kaidi Xu
Huan Zhang
Shiqi Wang
Yihan Wang
Suman Jana
Xue Lin
Cho-Jui Hsieh

title: Fast and Complete: Enabling complete neural network verification with rapid and massively parallel incomplete verifiers

authors:
Shiqi Wang
Huan Zhang
Kaidi Xu
Xue Lin
Suman Jana
Cho-Jui Hsieh
J Zico Kolter

title: Beta-CROWN: Efficient bound propagation with per-neuron split constraints for complete and incomplete neural network verification

venue: arXiv preprint arXiv:2103.06624,

authors:
Stanley Bak
Changliu Liu
Taylor Johnson

title: The second international verification of neural networks competition (vnn-comp 2021): Summary and results

venue: arXiv preprint arXiv:2109.00498,

authors:
Hongyang Zhang
Yaodong Yu
Jiantao Jiao
Eric Xing
Laurent El Ghaoui
Michael Jordan

title: Theoretically principled trade-off between robustness and accuracy

heading: 1 Introduction

text:
With the ever-growing deployment of deep neural networks, formal robustness guarantees are needed in many safety-critical applications. Strategies to improve robustness such as adversarial training only provide empirical robustness, without formal guarantees, and many existing adversarial defenses have been successfully broken using stronger attacks [1]. In contrast, certified defenses give formal robustness guarantees that any norm-bounded adversary cannot alter the prediction of a given network.
Bounding the global Lipschitz constant of a neural network is a computationally efficient and scalable approach to provide certifiable robustness guarantees [2–4]. The global Lipschitz bound is typically computed as the product of the spectral norm of each layer. However, this bound can be quite loose because it needs to hold for all points from the input domain, including those inputs that are far away from each other. Training a network while constraining this loose bound often imposes to high a degree of regularization and reduces network capacity. It leads to considerably lower clean accuracy in certified training compared to standard and adversarial training [5, 6].
A local Lipschitz constant, on the other hand, bounds the norm of output perturbation only for inputs from a small region, usually selected as a neighborhood around each data point. It produces a
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
tighter bound by considering the geometry in a local region and often yields much better robustness certification [7, 8]. Unfortunately, computing the exact local Lipschitz constant is NP-complete [9]. Obtaining reasonably tight local Lipschitz bounds via semidefinite programming [10] or mixed integer programming [11] is typically only applicable to small, previously trained networks since it is difficult to parallelize the optimization solver and make it differentiable for training. On the other hand, many existing certified defense methods have achieved success by using a training-based approach with a relatively weak but efficient bound [12–14]. Therefore, to incorporate local Lipschitz bound in training, a computationally efficient and training-friendly method must be developed.
Our contributions: We propose an efficient method to incorporate a local Lipschitz bound in training deep networks, by considering the interactions between an activation layer such as a Rectified Linear Unit (ReLU) layer and a linear (or convolution) layer. Our bound is computed for each data point and this translates to different types of outputs from the activation function: constant or varying under input perturbations. If the outputs of some activation neurons are constant under local perturbation, we eliminate the corresponding rows in the previous weight layer and the corresponding columns in the next weight layer, and then compute the spectral norm of the reduced matrix.
Our main insight is to use training to make the proposed local Lipschitz bound tight. This is different from existing works that find local Lipschitz bound for a fixed network [8, 10, 11]. Instead, we aim to enable a network to learn to tighten our proposed local Lipschitz bound during training. To achieve this, we propose to clip the activation function with an individually learnable threshold ✓. Take ReLU for example, the output of a “clipped” ReLU becomes a constant when input is greater than this threshold (see Figure 1). Once the input of the ReLU is greater than the threshold or less than 0, then this ReLU neuron does not contribute to the local Lipschitz constant, and thus the corresponding row or column of weight matrices can be removed. We also apply this method to non-ReLU activation functions such as MaxMin [15] to create constant output regions. Additionally, we also use a hinge loss function to encourage more neurons to have constant outputs. Our method can be used as a plug-in module in existing certifiable training algorithms that involve computing Lipschitz bound. Our contributions can be summarized as:
• To the best of our knowledge, we are the first to incorporate a local Lipschitz bound during training for certified robustness. Our bound is provably tighter than the global Lipschitz bound and is also computationally efficient for training.
• We propose to use activation functions with learnable threshold to encourage more fixed neurons during training, which assists the network to learn to tighten our bound. We show that more than 45% rows and columns can be removed from weight matrices.
• We consistently outperform state-of-the-art Lipschitz based certified defense methods for `2 norm robustness. On CIFAR-10 with perturbation ✏ = 36255 , we obtain 54.3% verified accuracy with ReLU activation function and 60.7% accuracy with MaxMin [15], outperforming the SOTA baselines, and also achieve better clean accuracy. Our code is available at https://github. com/yjhuangcd/local-lipschitz.

heading: 2 Related work

text:
Bounds on Local Lipschitz Constant A sound upper bound of local Lipschitz constant (simply referred to as “local Lispschitz bound” in our paper) is a crucial property to determine the robustness
of a classifier. Finding an exact local Lipschitz constant for a neural network is generally NP hard [16], so most works focus on finding a sound upper bound. Hein and Andriushchenko [7] derived an analytical bound for 2 layer neural networks and found that local Lipschitz bounds could be much tighter than the global one and give better robustness certificates. RecurJac [8] is a recursive algorithm that analyzes the local Lipschitz constant in a neural network using a bound propagation [14] based approach. FastLip [17] is a special and weaker form of RecurJac. Fazlyab et al. [10] used a stronger semidefinite relaxation to compute a tighter bound of local Lipschitz constant. Jordan and Dimakis [11] formulated the computation of Lipschitz as an mixed integer linear programming (MILP) problem and they were able to solve the exact local Lipschitz constant. Although these approaches can obtain reasonably tight and sound local Lipschitz constants, none of them have been demonstrated effective for training a certifiably robust network, where high efficiency and scalability are required. Note that although an empirical estimate of local Lipschitz constant can be easily found via gradient ascent (e.g., the local lower bounds reported in [4]), it is not a sound bound and does not provide certifiable robustness guarantees.
Certifiably Robust Training using Lipschitz constants The Lipschitz constant plays a central role in many works on training a certifiably robust neural network, especially for `2 norm robustness. Since naturally trained networks usually have very large global Lipschitz constant bounds [18], most existing works train the network to encourage small a Lipschitz bound. Cisse et al. [2] designed networks with orthogonal weights, whose Lipschitz constants are exactly 1. As this can be too restrictive, later works mostly use power iteration to obtain per-layer induced norms, whose product is a Lipschitz constant. Lipschitz Margin Training (LMT) [19] and Globally-Robust Neural Networks (Gloro) [4] both upper bound the worst margin via global Lipschitz constant with different loss functions. LMT constructs a new logit by adding the worst margin to all its entries except the ground truth class. Gloro construct a new logit with one more class than the original logit vector, determines whether the input sample can be certified. However, these approaches did not exploit the available local information to tighten Lipschitz bound and improve certified robustness. Box constrained propagation (BCP) [20] achieves a tighter outer bound than global Lipschitz based outer bound, by taking local information into consideration via interval bound (box) propagation. They compute the worst case logit based on the intersection of a (global) ball and a (local) box. Although box propagation considers local information, the ball propagation still uses global Lipschitz constant, and its improvement is still limited with low clean accuracy.
Other Certified Defenses Besides using Lipschitz constants, one of the most popular certifiable defense against `1 norm bounded inputs is via the convex outer adversarial polytope [21, 22]. [13] takes a similar approach via abstract interpretation. These methods uses linear relaxations of neural networks to compute an outer bound at the final layer. However, because the convex relaxations employed are relatively expensive, these methods are typically slow to train. A simple and fast certifiable defense for `1 norm bounded inputs is interval bound propagation (IBP) [12, 13]. Since the IBP bound can be quite loose for general networks, its good performance relies on appropriate hyper-parameters. CROWN-IBP [14] outperforms previous methods by combining IBP bound in a forward bounding pass and a tighter linear relaxation bound in a backward bound pass. [23] improved IBP with better initialization to accelerate training. Additionally, randomized smoothing [24–26] is a probabilistic method to certify `2 norm robustness with arbitrarily high confidence. The prediction of a randomized smooth classifier is the most likely prediction returned by the base classifier that is fed by samples from a Gaussian distribution. Salman et al. [27] further improves the performance of randomized smoothing via adversarial training.

heading: 3 Method

text: We begin with notations and background for Lipschitz bound. We introduce our method for local Lipschitz bound computation in Section 3.2. Then we introduce how to incorporate our efficient local Lipschitz bound in robust training (Section 3.3).

heading: 3.1 Notation and Background

text:
Notations. We denote the Euclidean norm of a vector x as kxk and kAk is the spectral norm of matrix A. Subscript of vector x denotes element, i.e., xi is the i-th element of x. We use LBl and UBl to denote lower bounds and upper bounds of pre-ReLU activation values for layer l. Definition 1. The Lipschitz constant of a function f : Rd ! Rm over an open set X is defined as,
L(f,X ) := sup x,y2X ,x 6=y ||f(y) f(x)|| ||x y|| ,
If L(f,X ) exists and is finite, we say that f is Lipschitz continuous over X . Suppose X = dom(f), L(f,X ) is the global Lipschitz constant of f ; if X is defined as the ✏-ball at point x, i.e., X := {x0|||x x0||  ✏}, then L(f,X ) is the local Lipschitz constant of f at x.
Global Lipschitz bound in existing works Consider a L-layer ReLU neural network which maps input x to output zL+1 = F (x;W ) using the following architecture, for l = 1, ..., L 1
z1 = x; zl+1 = (W lzl); zL+1 = WLzL (1) where W = {W 1:L} are the parameters, and (·) = max(·, 0) is the element-wise ReLU activation functions. Here we consider the bias parameters to be zero because they do not contribute to the Lipschitz bound. Since the Lipschitz constant of ReLU activation (·) is equal to 1, a global Lipschitz bound of F is, Lglob  ||WL|| · ||WL 1|| · · · ||W 1|| (2) where ||W l|| equals the spectral norm (maximum singular value) of the weight matrix W l. However, the global Lipchitz bound ignores the highly nonlinear property of deep neural networks. In what follows, we introduce our method that considers the interaction between ReLU and linear layer to obtain a tighter local Lipchitz bound in a computationally efficient way, that allows us to train a certifiably robust network using local Lipschitz bounds.

heading: 3.2 Our Approach for Efficient Local Lipschitz Bound

text:
In this section, we use ReLU as an example to describe how we compute our efficient local Lipschitz bound. We will discuss how to apply our method on other types of activation functions in Section 3.3. To exploit the piece-wise linear properties of ReLU neurons, we discuss the outputs of ReLU case by case. Intuitively, if the input of a ReLU neuron is always less or equal to zero, its output will always be zero, which is a constant and not contributing to Lipschitz bound. If the input of a ReLU’s can sometimes be greater than zero, the ReLU output will vary based on the input.
We define diagonal indicator matrices I lV(zl) to represent the entries where the ReLU outputs are varying and I lC(z
l) for entries where the ReLU outputs are constant under perturbation. Here zl 2 Rdl denotes the feature map of input x at layer l. Throughout this paper, unless otherwise mentioned, the indicator matrix is a function of the feature value zl, evaluated at a given input x.
Given an input perturbation kx0 xk  ✏, suppose zl(x0) is bounded element-wise as LBl  zl(x0)  UBl, we define diagonal matrix I lV and I lC as:
I lV(i, i) = ⇢ 1 if UBli > 0 0 otherwise , I lC(i, i) = ⇢ 1 if UBli  0 0 otherwise
(3)
By this definition, the ReLU output can either be constant or vary with respect to input perturbation. Hence we have I lV + I lC = I , where I is the identity matrix. LB and UB can be obtained cheaply from interval bound propagation [12] or other bound propagation mechanisms [21, 28].
A crucial observation is that to compute the local Lipschitz bound, we only need to consider the ReLU neurons which are non-fixed. The fixed ReLU neurons are always zero (locally, in the prescribed neighborhood around x) and thus have no impact to final outcome. We define a diagonal matrix DV to represent the ReLU outputs that are varying. Then, a neural network function (denoted as F (x;W )) can be rewritten as:
F (x;W ) = WLDL 1V W L 1 · · ·D1VW 1x
= (WLIL 1V )D L 1 V (I L 1 V W L 1IL 2V ) · · ·D1V(I1VW 1)x (4)
where
DlV(i, i) =
⇢ (ReLU(zli) > 0) if I lV(i, i) = 1
0 if I lV(i, i) = 0 , (5)
where denotes an indicator function.
Note here that we ignore bias terms for simplicity. Based on (4), an important insight used in our approach is that by combining the ReLU function with weight matrix, we have the opportunity to tighten Lipschitz bound by considering I lVW lI l 1 V as a whole based on whether there are ReLU outputs stay at constant under perturbation. Importantly, since DlV depends on UB, (4) only holds in a local region of x, which leads to a local Lipschitz constant bound at input x:
Llocal(x)  kWLIL 1V kkI L 1 V W L 1IL 2V k · · · kI1VW 1k (6)
The following theorem states that the local Lipchitz bound calculated via (6) is always tighter than the global Lipchitz bound in Eq (2), for all inputs. Theorem 1 (Tighter Lipchitz Bound). For any input x 2 Rn and L-layer ReLU neural network F (x;W ), the local Lipchitz bound calculated via (6) in any neighborhood of x is no larger than the global Lipchitz bound in Eq (2), i.e., 8x, Llocal(x)  Lglob.
The proof of Theorem 1 leverages the following proposition. Proposition 1. If a column and/or row is added to a matrix, then the matrix spectral norm (maximum singular value) will be no less than the spectral norm of the original matrix. That is, given matrix A 2 Rm⇥n, and y 2 Rm, z 2 Rn, then
max([A|y]) max(A) , max(  A z ) max(A). (7)
Proof of Proposition 1. Let A0 =  A z . The singular value of A0 is defined as the square roots of the
eigenvalues of A0TA0, where A0TA0 = ATA+zT z ATA, that simply imply ||A0|| = max(A0) max(A) = ||A||. Similar holds for A0 = [A|y].
By Proposition 1, it is straightforward to show that kIL 1V WL 1I L 2 V k  kWL 1k since the left hand side is the spectral norm of the reduced matrix, after removing corresponding rows/columns in WL 1 where the neuron output under local perturbation is constant. Therefore, the product of spectral norm of the reduced matrices is no larger than the product of the spectral norm of raw weight matrices, which leads to 8x, Llocal(x)  Lglob.

heading: 3.3 Training for Tight Local Lipschitz

text:
To encourage the network to learn which rows and columns need to be eliminated to make local Lipschitz bound tighter, we combine our local Lipschitz bound computation with certifiably robust training. This is different from existing works leveraging optimization tools to find local Lipschitz bound for a fixed network [10, 11]. By training with the proposed local Lipschitz bound, we can achieve good certified robustness on large neural networks.
More precisely, using our local Lipschitz bound, we can obtain the worst case logit z⇤ that is used to form a robust loss for training: E(x,y)⇠DL(z⇤(x), y), where L is the cross entropy loss function, (x, y) is the image and label pair from the training datasets. A simple way to compute the worst logit is z⇤i = zi + p 2✏LLocal for i 6= y, z⇤y = zy (see [19]). Our approach is also compatible with tighter bounds on the worst case logit, such as the one used in BCP [20]. To give the network more capability to learn to tighten our proposed local Lipschitz bound, we propose the following approaches:
Allowing More Eliminated Rows via ReLU✓ The key to tighten our local Lipschitz bound is to delete rows and columns in weight matrices that align with singular vectors corresponding to the largest singular value. To encourage more rows and columns to be deleted, we need to have more ReLU outputs to be at constant under perturbation. Standard ReLU is only lower bounded by zero, but is not upper bounded. If we can set an “upper bound” of ReLU output, we can have more neurons
have fixed outputs at this upper bound. An upper bounded ReLU unit called ReLU6 is proposed in [29], where the maximum output is set to a constant 6. Different from ReLU6 that sets a constant maximum output threshold, we make the threshold to be a learnable parameter. We name the new type of activation function ReLU✓ , which is defined as:
ReLU✓ (zi; ✓i) =
8 <
: 0, if zi <= 0 zi, if 0 < zi < ✓i ✓i, if zi >= ✓i
(8)
where ✓i is a learnable upper bound of the ReLU output.
Similar to (3), the indicator matrices for the varying outputs of ReLU✓ are
I lV(i, i) =
⇢ 1 if UBli > 0 and LB l i < ✓i
0 otherwise (9)
Depending on the ReLU✓ activation status, the output of a ReLU✓ neural network is,
F (x;W, ✓) = WL(DL 1V W L 1 · · · (D1VW 1x+D1✓) · · ·+DL 1✓ ) , (10)
where Dl✓ denotes the ReLU output fixed at the maximum output value,
Dl✓(i, i) = ⇢ ✓i if I l✓(i, i) = 1 0 if I l✓(i, i) = 0 , I l✓(i, i) = ⇢ 1 if LBli ✓i 0 otherwise
(11)
The local Lipchitz bound is still calculated as (6). However, the bound can be potentially learned tighter because we encourage ReLU outputs to be constant in both directions, and there could be less varying outputs in Eq (9) than in Eq (3).
Extension to non-ReLU activation functions Our local Lipschitz bound can be applied on nonReLU activation functions. The key is to create constant output regions for the activation function and delete the corresponding rows or columns in the weight matrices. Since the MaxMin activation function [15] has been shown to outperform ReLU on certified robustness [4, 30], we take MaxMin as an example to explain how to apply our local Lipschitz bound. Let x1 and x2 be two groups of the input, the output of MaxMin is max(x1, x2),min(x1, x2). To exploit local Lipschitz, we created a clipped version of MaxMin, similar to ReLU✓. The output of the clipped MaxMin is min(max(x1, x2), a),max(min(x1, x2), b), where a is a learnable upper threshold for the max output in MaxMin and b is a learnable lower threshold for the min output in MaxMin. Box propagation rule through MaxMin is straightforward to derive, so we can get the box bound on each entry after MaxMin. If the lower bounds of the Max entries are bigger than the upper threshold a, or the upper bounds of the Min entries are smaller than the lower threshold b, we can delete the corresponding columns in the successive matrix (similar to the procedure for ReLU networks).
Encouraging Fixed Neurons via a Sparsity Loss To encourage more rows and columns to be deleted in the weight matrices, We design a sparsity loss to regularize the neural network towards this goal. For a ReLU neural network, assuming that the i-th entry of the feature map at layer l is bounded by LBli  zli  UBli, we hope the neural network can learn to make as many UBi to be smaller than zero and LBi to be larger than ✓i without sacrificing too much of classification accuracy. For a MaxMin neural network, let LBmax be the lower bounds of the Max entries, and UBmin be the upper bounds of the Min entries. We hope the neural network can learn to make as many LBmax to be larger than the upper threshold a and UBmin to be smaller than the lower threshold b. The sparsity losses for ReLU and MaxMin networks are as follows:
LReLUsparsity = max(0,UBli)+max(0, ✓i LBli) , LMaxMinsparsity = max(0,UBmin b)+max(0, a LBmax) (12) Finally, our full training procedure is presented in Algorithm 1.
Computational cost To obtain the local Lipschitz bound, we must perform power iterations to compute the spectral norm of reduced weight matrices for every input and its feature maps at each layer. Compared to other methods that use optimization tools (e.g., SDP, MILP) to bound local Lipschitz, our method is computationally efficient since only matrix vector multiplication is used.
Algorithm 1: Local Lipchitz based Certifiably Robust Training Input : Training data (x, y) ⇠ D, perturbation size ✏, number of iterations for power method n, a neural network with L layers. repeat
Compute the box outer bound [LBl,UBl] for layers 1 to L ; Compute indicator matrix IV using Eq (9) ; // Compute local Lipcthiz bound L(x) for every input x using Eq (6) for layer from 1 to L do
// Power method Initialize ul with the updated ul from the previous training episode ; for i < n do
If layer is conv: v I lVconv(W l, I l 1V ul)/kI lVconv(W l, I l 1 V u
l)k u I l 1V conv|(W l, I lVv)/kI l 1 V conv
|(W l, I lVv)k If layer is linear: v I lVW lI l 1V ul/kI lVW lI l 1 V u
lk u I l 1V W l | I lVv/kI l 1V W l | I lVvk
end If layer is conv: l(x) vlI lVconv(W l, I l 1V ul) If layer is linear: l(x) vlI lVW lI l 1V ul
end Compute the worst logits using our local Lipschitz bound ; Update model parameters based on some loss functions (e.g., Cross-entropy loss).
until training ends; return Parameters of a robust neural network
During training, compared to methods that only uses global Lipschitz bound, the local Lipcshitz bound varies based on the inputs. For global Lipschitz bound, a common practice is to keep track of the iterate vector u in power iteration, and use it to initialize the power iteration in the next training batch. With this initialization, only a few number of iterations is performed during training (typically the number of iterations is between 1 to 10 [20, 4]). To extend the initialization strategy for u to compute local Lipschitz, we need to keep track of the iterate vectors based on every input and its feature maps for every layer. Fortunately, this vector only provides an initialization for power iteration so it does not need to be stored accurately, and can be stored using low precision tensors. Further extensions could use dimension reduction or compression methods to store these vectors, or learn a network along the way to predict a good initializer for power iteration. An alternative approach is to random initialize u in every mini-batch, but we find that more power iterations need to be performed during training for this approach to have comparable performance as the using saved u.
During evaluation, to minimize the extra computational cost, we can avoid bound computation for two types of inputs: inputs that can already be certified using global Lipschitz bound or can be attacked by adversaries (e.g., a 100-step PGD attack). In practice, this typically rules out more than 80% samples on CIFAR-10 or larger datasets, and greatly reduce computational cost required to compute local Lipschitz constants. In Section 4 we will show more empirical results on this aspect.

heading: 4 Experiment

text:
In this section, we first show our method achieves tighter Lipschitz bounds in neural networks. When combined with certifiable training algorithms such as BCP [20] and Gloro [4], as well as training algorithms using orthogonal convolution and MaxMin activation function [30], our method achieves both higher clean and certified accuracy.
Experiment setup We train with our method to certify robustness within a `2 ball of radius 1.58 on MNIST [31] and 36/255 on CIFAR-10 [32] and Tiny-Imagenet 1 on various network architectures.
1https://tiny-imagenet.herokuapp.com
We denote neural network architecture by the number of convolutional layers and the number of fully-connected layers. For instance, 6C2F indicates that there are 6 convolution layers and 2 fullyconnected layers in the neural network. Networks have ReLU activation function unless mentioned otherwise. For more details and hyper-parameters in training, please refer to Appendix C. Our code is available at https://github.com/yjhuangcd/local-lipschitz.
Tighter Lipschitz bound We compared the training process of our method and BCP [20] in Figure 2 (a-c). During training, our method uses local Lipschitz bound while BCP uses global Lipschitz bound for robust loss. We also tracked the global Lipschitz bound during our training and the average local Lipschitz bound (computed by our method) during BCP training for comparison. We can see from Figure 2 (a) that our local Lipschitz bound is always tighter the global Lipschitz bound. Furthermore, it is crucial to incorporate our bound in training to enable the network to learn to tighten the bound. We can see that if we directly apply our method to a BCP trained network after training, the local Lipschitz bound has much less improvement over global Lipschitz bound. In addition, a tighter local Lipschitz bound by our method allows the neural networks to have larger global Lipschitz bound in the beginning of training. This potentially provides larger model capacity and eases the training in the early stage. As a consequence, we see a large improvement of both clean loss (Figure 2 (b)) and also robust loss (Figure 2 (c)) throughout the training.
Sparsity of varying ReLU outputs We examine our 6C2F model trained on CIFAR-10 and report the proportion of varying (non-constant) ReLU outputs at all layers except the last fully-connected layer (Figure 3). We compared the proportion with that of a standard CNN (trained with only cross entropy loss on unperturbed inputs) and a robust CNN trained by BCP. As we can see, the standard neural networks has the most varying ReLU neurons, indicating that dense varying ReLU outputs may provide larger model capacity for clean accuracy but reduce robustness. Our method has more varying ReLU outputs than BCP, while achieving tighter local Lipschitz bound than BCP. This indicates that our training method encourages the neural network to learn to delete rows and columns that contribute most to local Lipschitz constant during training, while keeping ReLUs for other rows or columns varying to obtain better natural accuracy.
Certified robustness Our method can be used as a plug-in module in certifiable training algorithms that involves using `2 Lipschitz constant such as BCP [20] and Gloro [4]. We use Local-Lip-G and
Local-Lip-B to denote our method trained with the Gloro loss and BCP loss respectively (detailed formulation for each loss can be found in Appendix B). We compare the performance of our method against competitive baselines on both ReLU neural networks [20, 4, 19, 12, 22, 33] and MaxMin neural networks [4, 30]. For each method, we report the clean accuracy (accuracy on non-perturbed inputs), the PGD accuracy (accuracy on adversarial inputs generated by PGD attack [6]), and the certified accuracy (the proportion of inputs that can be correctly classified and certified within ✏-ball). For PGD attack, we use 100 steps with step size of ✏/4. In our experiments, we use box propagation (as done in BCP) to obtain the lower bound and upper bound of every neuron. With our tighter Lipchitz bound, we further improve clean, PGD, and certified accuracy upon BCP and Gloro, and achieve the state-of-the-art performance on certified robustness (Table 2). On CIFAR-10 with ReLU activation function, we improved certified accuracy from 51.3% (SOTA) to 54.3%. When MaxMin activation function is used, our local Lipschitz training approach also consistently improves clean, PGD and verified accuracy over baselines on both CIFAR-10 and TinyImageNet datasets.
To demonstrate the effectiveness of incorporating our bound in training, we also use our local bound to directly compute certified accuracy for pretrained models using BCP. Our bound improves the certified accuracy of a pretrained BCP model from 51.3% (reported in Table 2) to 51.8%. The improvement is less than training with our bound (54.3% in Table 2). Therefore, it is crucial to incorporate our bound in training to gain non-trivial robustness improvements.
Initialization strategy for power method We used two initialization strategies for singular vectors u in power method during training. One option is to store u for all the inputs and feature maps and initialize u in the current training epoch with u from the previous epoch, which requires storage. The other option is to random initialize u. Table 1 shows the performance of these two approaches on CIFAR-10 with the 6C2F architecture. The number of power iterations during training is listed in the bracket. Although random initialization is memory-efficient, it needs more power iterations during training to achieve comparable performance compared to the approach of storing u. Too few iterations tend to cause inaccurate singular value and overfitting, resulting in lower certified accuracy.
Computational cost during evaluation time Since local Lipschitz bound needs to be evaluated for every input and global Lipschitz bound does not depend on the input, our method involves additional computation cost during certification. Let u(t) be the singular vector computed by power iteration at iteration t, we stop power iteration when ||u(t + 1) u(t)||  1e 3. To analyze the computational cost during evaluation time, we plot the histogram of number of iterations for power method to converge for the second last layer in the 6C2F model in Figure 4. The average number of iterations for convergence is 137. The histograms for other layers are in Appendix C. To reduce the computational cost, we only need to compute local Lipschitz bound for samples that cannot be certified by global Lipschitz bound or cannot be attacked by adversaries. The proportion of those samples is (100% PGD Err Global Certified Acc). For the 6C2F model on CIFAR-10, the proportion of samples that can be certified using global Lipschitz bound is 51.0%, and the error under PGD attack is 35.2%. Hence we only need to evaluate local Lipschitz bounds on the remaining 13.8% samples, which greatly reduces the overhead of computing the local Lipschitz bounds.

heading: 5 Conclusion

text: In this work, we propose an efficient way to incorporate local Lipschitz bound for training certifiably robust neural networks. We classify the outputs of activation function as being constant or varying under input perturbations and consider them separately. Specifically, we remove the redundant rows and columns corresponding to the constant activation outputs in the weight matrix to get a tighter local Lipschitz bound. We propose a learnable bounded activation and the use of a sparsity encouraging loss during training to assist the neural network to learn to tighten our local Lipschitz bound. Our method consistently outperforms state-of-the-art Lipschitz based certified defend methods for `2 norm robustness. We see no immediate negative societal impact in the proposed approach.

heading: Acknowledgements

text: Y. Huang is supported by DARPA LwLL grants. A. Anandkumar is supported in part by Bren endowed chair, DARPA LwLL grants, Microsoft, Google, Adobe faculty fellowships, and DE Logi grant. Huan Zhang is supported by funding from the Bosch Center for Artificial Intelligence.

title: Training Certifiably Robust Neural Networks with Efficient Local Lipschitz Bounds