summaries: This paper presents a method to speed up training of deep neural networks. The main contribution is a method to quickly identify winning lottery tickets (denoted early-bird, or EB by the authors), without running the model to convergence. The authors present interesting preliminary experiments that motivate their method, and show that it works on two image recognition datasets using two models.

summaries: The authors further study the lottery ticket hypothesis formulated by Frankle and Carbin. They demonstrate that the sparsity pattern corresponding to a lottery ticket for a given initialization can be uncovered via low-cost training. By doing so, they propose a method to: 1) first identify the lottery ticket efficiently and 2) exploit the sparsity of the resulting network to train it at a lower cost.

abstractText: Frankle & Carbin, 2019) shows that there exist winning tickets (small but critical subnetworks) for dense, randomly initialized networks, that can be trained alone to achieve a comparable accuracy to the latter in a similar number of iterations. However, the identification of these winning tickets still requires the costly train-prune-retrain process, limiting their practical benefits. In this paper, we discover for the first time that the winning tickets can be identified at a very early training stage, which we term as Early-Bird (EB) tickets, via lowcost training schemes (e.g., early stopping and low-precision training) at large learning rates. Our finding on the existence of EB tickets is consistent with recently reported observations that the key connectivity patterns of neural networks emerge early. Furthermore, we propose a mask distance metric that can be used to identify EB tickets with a low computational overhead, without needing to know the true winning tickets that emerge after the full training. Finally, we leverage the existence of EB tickets and the proposed mask distance to develop efficient training methods, which are achieved by first identifying EB tickets via low-cost schemes, and then continuing to train merely the EB tickets towards the target accuracy. Experiments based on various deep networks and datasets validate: 1) the existence of EB tickets and the effectiveness of mask distance in efficiently identifying them; and 2) that the proposed efficient training via EB tickets can achieve up to 5.8× ∼ 10.7× energy savings while maintaining comparable or even better accuracy as compared to the most competitive state-ofthe-art training methods, demonstrating a promising and easily adopted method for tackling the often cost-prohibitive deep network training. Codes available at https://github.com/RICE-EIC/Early-Bird-Tickets

authors:
Alessandro Achille
Matteo Rovere
Stefano Soatto

title: Critical learning periods in deep networks

venue: In International Conference on Learning Representations,

authors:
Ron Banner
Itay Hubara
Elad Hoffer
Daniel Soudry

title: Scalable methods for 8-bit training of neural networks

authors:
Y.H. Chen
T. Krishna
J.S. Emer
V. Sze

title: Eyeriss: An energy-efficient reconfigurable accelerator for deep convolutional neural networks

venue: IEEE Journal of Solid-State Circuits,

authors:
Emily L Denton
Wojciech Zaremba
Joan Bruna
Yann LeCun
Rob Fergus

title: Exploiting linear structure within convolutional networks for efficient evaluation

venue: Advances in Neural Information Processing Systems

authors:
Jonathan Frankle
Michael Carbin

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

authors:
Jonathan Frankle
David J. Schwab
Ari S. Morcos

title: The early phase of neural network training

authors:
Priya Goyal
Piotr Dollár
Ross B. Girshick
Pieter Noordhuis
Lukasz Wesolowski
Aapo Kyrola
Andrew Tulloch
Yangqing Jia
Kaiming He

title: Accurate, large minibatch SGD: training imagenet in 1 hour

authors:
Suyog Gupta
Ankur Agrawal
Kailash Gopalakrishnan
Pritish Narayanan

title: Deep learning with limited numerical precision

authors:
Song Han
Huizi Mao
William J Dally

title: Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding

venue: arXiv preprint arXiv:1510.00149,

authors:
Kaiming He
Xiangyu Zhang
Shaoqing Ren
Jian Sun

title: Deep residual learning for image recognition

title: Identity mappings in deep residual networks

venue: In European conference on computer vision,

authors:
Yang He
Guoliang Kang
Xuanyi Dong
Yanwei Fu
Yi Yang

title: Soft filter pruning for accelerating deep convolutional neural networks

venue: In International Joint Conference on Artificial Intelligence (IJCAI),

authors:
Itay Hubara
Matthieu Courbariaux
Daniel Soudry
Ran El-Yaniv
Yoshua Bengio

title: Quantized neural networks: Training neural networks with low precision weights and activations

authors:
Xianyan Jia
Shutao Song
Wei He
Yangzihao Wang
Haidong Rong
Feihu Zhou
Liqiang Xie
Zhenyu Guo
Yuanzhou Yang
Liwei Yu

title: Highly scalable deep learning training system with mixed-precision: Training imagenet in four minutes

venue: arXiv preprint arXiv:1807.11205,

authors:
Namhoon Lee
Thalaiyasingam Ajanthan
Philip HS Torr

title: Snip: Single-shot network pruning based on connection sensitivity

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

title: Pruning filters for efficient convnets

authors:
Yuanzhi Li
Colin Wei
Tengyu Ma

title: Towards explaining the regularization effect of initial large learning rate in training neural networks

authors:
Sicong Liu
Yingyan Lin
Zimu Zhou
Kaiming Nan
Hui Liu
Junzhao Du

title: On-demand deep model compression for mobile devices: A usage-driven model selection framework

venue: In Proceedings of the 16th Annual International Conference on Mobile Systems, Applications, and Services,

authors:
Zhuang Liu
Jianguo Li
Zhiqiang Shen
Gao Huang
Shoumeng Yan
Changshui Zhang

title: Learning efficient convolutional networks through network slimming

venue: In Proceedings of the IEEE International Conference on Computer Vision,

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

title: Rethinking the value of network pruning

venue: arXiv preprint arXiv:1810.05270,

authors:
Jian-Hao Luo
Jianxin Wu
Weiyao Lin

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

venue: In Proceedings of the IEEE international conference on computer vision,

authors:
Ari S. Morcos
Haonan Yu
Michela Paganini
Yuandong Tian

title: One ticket to win them all: generalizing lottery ticket initializations across datasets and optimizers

authors:
Nasim Rahaman
Aristide Baratin
Devansh Arpit
Felix Draxler
Min Lin
Fred Hamprecht
Yoshua Bengio
Aaron Courville

title: On the spectral bias of neural networks

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

authors:
Alex Renda
Jonathan Frankle
Michael Carbin

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

authors:
Jianghao Shen
Yonggan Fu
Yue Wang
Pengfei Xu
Zhangyang Wang
Yingyan Lin

title: Fractional Skipping: Towards Finer-Grained Dynamic Inference

venue: In The Thirty-Forth AAAI Conference on Artificial Intelligence,

authors:
Karen Simonyan
Andrew Zisserman

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

venue: arXiv preprint arXiv:1409.1556,

authors:
Nimit Sharad Sohoni
Christopher Richard Aberger
Megan Leszczynski
Jian Zhang
Christopher Ré

title: Low-memory neural network training: A technical report

authors:
Emma Strubell
Ananya Ganesh
Andrew McCallum

title: Energy and policy considerations for deep learning in NLP

authors:
Naigang Wang
Jungwook Choi
Daniel Brand
Chia-Yu Chen
Kailash Gopalakrishnan

title: Training deep neural networks with 8-bit floating point numbers

venue: In Advances in neural information processing systems,

authors:
Yue Wang
Tan Nguyen
Yang Zhao
Zhangyang Wang
Yingyan Lin
Richard Baraniuk

title: Energynet: Energy-efficient dynamic inference. 2018b

authors:
Yue Wang
Ziyu Jiang
Xiaohan Chen
Pengfei Xu
Yang Zhao
Yingyan Lin
Zhangyang Wang

title: E2-train: Training state-of-the-art cnns with over 80% energy savings

authors:
Yue Wang
Jianghao Shen
Ting-Kuei Hu
Pengfei Xu
Tan Nguyen
Richard Baraniuk
Zhangyang Wang
Yingyan Lin

title: Dual dynamic inference: Enabling more efficient, adaptive and controllable deep inference

venue: arXiv preprint arXiv:1907.04523,

authors:
Yunhe Wang
Chang Xu
XU Chunjing
Chao Xu
Dacheng Tao

title: Learning versatile filters for efficient convolutional neural networks

authors:
Yunhe Wang
Chang Xu
Chao Xu
Dacheng Tao

title: Packing convolutional neural networks in the frequency domain

venue: IEEE transactions on pattern analysis and machine intelligence,

authors:
Wei Wen
Chunpeng Wu
Yandan Wang
Yiran Chen
Hai Li

title: Learning structured sparsity in deep neural networks

authors:
Junru Wu
Yue Wang
Zhenyu Wu
Zhangyang Wang
Ashok Veeraraghavan
Yingyan Lin

title: Deep k-means: Re-training and parameter sharing with harder cluster assignments for compressing deep convolutions

venue: arXiv preprint arXiv:1806.09228,

authors:
Shuang Wu
Guoqi Li
Feng Chen
Luping Shi

title: Training and inference with integers in deep neural networks

authors:
Pengfei Xu
Xiaofan Zhang
Hao Cong
Yang Zhao
Yongan Zhang
Chaojian Li
Zetong Guan
Deming Chen
Yingyan Lin

title: AutoDNNchip: An automated dnn chip predictor and builderfor both fpgas and asics

venue: In 28th ACM/SIGDA International Symposium on Field-Programmable Gate Arrays,

authors:
Zhi-Qin John Xu
Yaoyu Zhang
Tao Luo
Yanyang Xiao
Zheng Ma

title: Frequency principle: Fourier analysis sheds light on deep neural networks

authors:
Guandao Yang
Tianyi Zhang
Polina Kirichenko
Junwen Bai
Andrew Gordon Wilson
Christopher De Sa

title: SWALP : Stochastic weight averaging in low precision training

authors:
Yukuan Yang
Shuang Wu
Lei Deng
Tianyi Yan
Yuan Xie
Guoqi Li

title: Training highperformance and large-scale deep neural networks with full 8-bit integers

venue: arXiv preprint arXiv:1909.02384,

authors:
Yang You
Zhao Zhang
Cho-Jui Hsieh
James Demmel
Kurt Keutzer

title: Imagenet training in minutes

venue: In Proceedings of the 47th International Conference on Parallel Processing,

authors:
Hattie Zhou
Janice Lan
Rosanne Liu
Jason Yosinski

title: Deconstructing lottery tickets: Zeros, signs, and the supermask

authors:
Barret Zoph
Quoc V. Le

title: Neural architecture search with reinforcement learning

heading: 1 INTRODUCTION

text:
The recent record-breaking predictive performance achieved by deep neural networks (DNNs) motivates a tremendously growing demand to bring DNN-powered intelligence into numerous applications (Xu et al., 2020). However, the excellent performance of modern DNNs comes at an often prohibitive training cost due to the required vast volume of training data and model parameters. As an illustrative example of the computational complexity of DNN training, one forward pass of the ResNet50 (He et al., 2016a) model requires 4 GFLOPs (FLOPs: floating point operations) of computations and training requires 1018 FLOPs, which takes 14 days on one state-of-the-art NVIDIA M40 GPU (You et al., 2018). As a result, training a state-of-the-art DNN model often demands considerable energy, along with the associated financial and environmental costs. For example, a recent report shows that training a single DNN can cost over $10K US dollars and emit as much carbon as five cars in their lifetimes (Strubell et al., 2019), limiting the rapid development of DNN innovations and raising various environmental concerns.
The recent trends of improving DNN efficiency mostly focus on compressing models and accelerating inference. An empirically adopted practice is the so-called progressive pruning and training
∗Correspondence should be addressed to: Zhangyang Wang and Yingyan Lin.
routine, i.e., training a large model fully, pruning it, and then retraining the pruned model to restore the performance (the process can be iterated several rounds). While this has been a standard practice for model compression (Han et al., 2015), some recent efforts start empirically linking it to the potential of more efficient training. Notably, the latest series of works (Frankle & Carbin, 2019; Liu et al., 2018b) reveals that dense, randomly-initialized networks contain small subnetworks which can match the test accuracy of original networks when trained alone themselves. These subnetworks are called winning tickets. Despite their insightful findings, there remains to be a major gap between the winning ticket observation and the goal of more efficient training, since winning tickets were only identified by pruning unimportant connections after fully training a dense network.
This paper closes this gap by demonstrating the Early-Bird (EB) tickets phenomenon: the winning tickets can be drawn very early in training and with aggressively low-cost training algorithms. Through a range of experiments on different DNNs and datasets, we observe the consistent existence of EB tickets and the cheap costs needed to reliably draw them, and develop a novel mask distance metric to detect their emergence. After being identified, re-training those EB tickets (using standard training) leads to comparable or even better final accuracies, compared to either standard training, or re-training the “ground-truth” winning tickets drawn after full training as in (Frankle & Carbin, 2019). Our observations seem to coincide with the recent findings by (Achille et al., 2019; Li et al., 2019) about the two-stage optimization trajectory in training. Taking advantage of EB tickets, we propose an efficient DNN training scheme termed EB Train. To our best knowledge, this is the first step taken towards exploiting winning tickets for a realistic efficient training goal.
Our contribution can be summarized as follow:
1. We discover the Early-Bird (EB) tickets, and show that they 1) consistently exist across DNN models and datasets; 2) can emerge very early in training; and 3) stay robust under (and sometimes even favor) various aggressive and low-cost training schemes (in addition to early stopping), including large learning rates and low-precision training.
2. We propose a practical, easy-to-compute mask distance as an indicator to draw EB tickets without accessing the “ground-truth” winning tickets (drawn after full training), fixing a major paradox for connecting winning tickets with the efficient training goal.
3. We design a novel efficient training framework based on EB tickets (EB Train). Experiments in state-of-the-art benchmarks and models show that EB Train can achieve up to 5.8× ∼ 10.7× energy savings, while maintaining the same or even better accuracy, compared to training with the original winning tickets.

heading: 2 RELATED WORKS

text:
Winning Ticket Hypothesis. The lottery ticket hypothesis (Frankle & Carbin, 2019) first points out that a small subnetwork, called the winning ticket, can be identified by pruning a fully trained dense network; when training it in isolation with the same weight initialization once assigned to the corresponding weights in the dense network, one can restore the comparable test accuracy to the dense network. However, finding winning tickets hinged on costly (iterative) pruning and retraining. (Morcos et al., 2019) studies the reuse of winning tickets, transferable across different datasets. (Zhou et al., 2019) discovers the existence of supermasks that can be applied to an untrained, randomly-initialized network. (Liu et al., 2018b) argues that the weight initialization might make less difference when trained with a large learning rate, while the searched connectivity is more of the winning ticket’s core value. It also explores the usage of both unstructured and (more hardware-friendly) structured pruning and shows that both lead to the emergence of winning tickets.
Another related work (Lee et al., 2019) prunes a network at single-shot with one mini-batch, in which the irrelevant connections are identified by a connection sensitivity criterion. Comparing to (Frankle & Carbin, 2019), the authors show their method to be more efficient in finding the good subnetwork (not the winning ticket), although its re-training accuracy/efficiency is found to be inferior, compared to training the “ground truth” winning ticket.
Other Relevant Observations in Training. (Rahaman et al., 2019; Xu et al., 2019) argue that deep networks will first learn low-complexity (lower-frequency) functional components, before absorbing high-frequency features: the former being more robust to perturbations. An important hint can be found in (Achille et al., 2019): the early stage of training seems to first discover the important connections and the connectivity patterns between layers, which becomes relatively fixed in the
later training stage. That seems to imply that the critical sub-network (connectivity) can be identified independent of, and seemingly also ahead of, the (final best) weights. Finally, Li et al. (2019) demonstrates that training a deep network with a large initial learning rate helps the model focus on memorizing easier-to-fit, more generalizable pattern faster and better – a direct inspiration for us to try drawing EB tickets using large learning rates.
Efficient Inference and Training. Model compression has been extensively studied for lighterweight inference. Popular means include pruning (Li et al., 2017; Liu et al., 2017; He et al., 2018; Wen et al., 2016; Luo et al., 2017; Liu et al., 2018a), weight factorization (Denton et al., 2014), weight sharing (Wu et al., 2018a), quantization (Hubara et al., 2017), dynamic inference (Wang et al., 2018b; 2019b; Shen et al., 2020), network architecture search (Zoph & Le, 2017), among many others (Wang et al., 2018c;d). On the other hand, the literature on efficient training appears to be much sparser. A handful of works (Goyal et al., 2017; Cho et al., 2017; You et al., 2018; Akiba et al., 2017; Jia et al., 2018; Gupta et al., 2015) focus on reducing the total training time in paralleled, communication-efficient distributed settings. In contrast, our goal is to shrink the total resource cost for in-situ, resource-constrained training, as (Wang et al., 2019a) advocated. (Banner et al., 2018; Wang et al., 2018a) presented low-precision training, which is aligned with our goal and can be incorporated into EB Train (see later).

heading: 3 DRAWING EARLY-BIRD TICKETS: HYPOTHESIS AND EXPERIMENTS

text: We hypothesize that the winning tickets can emerge at a very early training stage, which we term as an Early-Bird (EB) ticket. Consider a dense, randomly-initialized network f(x; θ), f reaches a minimum validation loss floss at the i-th iteration with a test accuracy facc, when optimized with stochastic gradient descent (SGD) on a training set. In addition, consider subnetworks f(x;m⊙ θ) with a mask m ∈ {0, 1} indicates the pruned and unpruned connections in f(x; θ). When being optimized with SGD on the same training set, f(x;m ⊙ θt) reach a minimum validation loss f ′loss with a test accuracy f ′acc, where θt denotes the weights at the t-th iteration of training. The EB tickets hypothesis articulates that there exists m such that f ′acc ≈ facc (even ≥), i.e., same or better generalization, with t≪ i (e.g., early stopping) and a sparse m (i.e., much reduced parameters). Section 3 addresses three key questions pertaining to the EB ticket hypothesis. We first show via an extensive set of experiments, that EB tickets can be observed across popular models and datasets (Section 3.1). We then try to be more aggressive to see if high-quality EB tickets still emerge under “cheaper” training (Section 3.2). We finally reveal that EB tickets can be identified using a novel mask distance between consecutive epochs, thus no full training needed (Section 3.3).

heading: 3.1 DO EARLY-BIRD TICKETS ALWAYS EXIST?

text:
We perform ablation simulations using two representative deep models: VGG16 (Simonyan & Zisserman, 2014) and pre-activation residual networks-101 (PreResNet101) (He et al., 2016b), on two popular datasets: CIFAR-10 and CIFAR-100. For drawing the tickets, we adopt a standard training protocol (Liu et al., 2018b) for both CIFAR-10 and CIFAR-100: the training takes 160 epochs in total and the batch size of 256; the initial learning rate is set to 0.1, and is divided by 10 at the 80th and 120th epochs, respectively; the SGD solver is adopted with a momentum of 0.9 and a weight decay of 10−4. For retraining the tickets, we keep the same setting by default.
We follow the main idea of (Frankle & Carbin, 2019), but instead prune networks trained at much earlier points (before the accuracies reach their final top values), to see if reliable tickets can still be drawn. We adopt the same channel pruning in (Liu et al., 2017) for all experiments since it is hardware friendly and aligns with our end goal of efficient training (Wen et al., 2016). Figure 1 reports the accuracies achieved by re-training the tickets drawn from different early epochs. All results consistently endorse that there exist high-quality tickets, at as early as the 20th epoch (w.r.t. a total of 160 epochs), that can achieve strong re-training accuracies. Comparing among different pruning ratios p, it is not too surprising to see over-pruning (e.g., p = 70%) makes drawing good tickets harder, indicating a balance that we need to calibrate between accuracy and training efficiency.
Two more striking observations from Figure 1: 1) there consistently exist EB tickets drawn at certain early epoch ranges, that outperform those drawn in a later stages, including the “ground-truth” winning tickets drawn at the 160th epoch. That intriguing phenomenon implies the possible “overcooking” when networks try to identify connectivity patterns at later stage (Achille et al., 2019),
that might hamper generalization; 2) some EB tickets are able to outperform even their unpruned, fully-trained models (e.g., the dashlines), potentially thanks to the sparse regularization learned by EB tickets.

heading: 3.2 DO EARLY-BIRD TICKETS STILL EMERGE UNDER LOWER-COST TRAINING?

text:
For EB tickets, only the important connections (connectivity) found by them matter, while the weights are to be re-trained anyway. One might hence come up with an idea, that more aggressive and “cheaper” training methods might be applicable to further shrink the cost of finding EB tickets (on top of the aforementioned early stopping of training thanks to the identification of EB tickets ), as long as the same significant connections still emerge. We experimentally investigate the impacts of two schemes (which are using appropriate large learning rates and training in lower precision): EB tickets are observed to survive well under both of them.
Appropriate large learning rates are important to the emergence of EB tickets. We first vary the learning rate schedule in the above experiments. The original schedule is denoted as [80LR→0.01, 120LR→0.001], i.e., starting from 0.1, decayed to 0.01 at the 80th epoch, and further decayed to 0.001 at the 120th epoch. In comparison, we test a new learning rate schedule [0LR→0.01, 100LR→0.001]: starting from 0.01 (the 0th epoch), and decayed to 0.001 at the 100th epoch. After drawing tickets, we re-train them all using the same learning rate schedule [80LR→0.01, 120LR→0.001] for sufficient training and a fair comparison. We can see from Table 1 that high-quality EB tickets always emerge earlier when searching with the schedule of a larger learning rates, i.e., [80LR→0.01, 120LR→0.001], whose final accuracies are also better. Note that the earlier emergence of good EB tickets contributes to lowering the training costs. It was observed that larger learning rates are beneficial for training the dense model fully to draw the winning tickets in (Frankle & Carbin, 2019; Liu et al., 2018b): here we show this benefit extends to EB tickets too. More experiments with larger learning rates can be found in Appendix A.2.
Low-precision training does not destroy EB tickets. We next examine the emergence of EB tickets within a state-of-the-art low-precision training algorithm (Wu et al., 2018b) , where all model weights, activations, gradients and errors are quantized to 8 bits throughout training. Note that here we only apply low-precision training to the stage of identifying EB tickets, and afterwards the tickets are trained in the same full-precision as in Section 3.1. Figure 2 shows the retraining accuracy and total number of FLOPs (EB ticket search (8 bits) + retraining (32 bits floating points)) for VGG16 and CIFAR-10/100. We can see that EB tickets still emerge when using aggressively low-precision for identifying EB tickets: the general trends resemble the full-precision training cases in Figure 1, except the emergence of good EB tickets seem to become even earlier up to initial epochs. In this
way, it will lead to cost savings in finding EB tickets, since low-precision updates can aggressively save energy compared to their full-precision baseline, as shown in the Table 2.

heading: 3.3 HOW TO IDENTIFY EARLY-BIRD TICKETS PRACTICALLY?

text:
Distance between Ticket Masks. For each time of pruning, we define a binary mask of the drawn ticket (pruned subnetwork) w.r.t. the full dense network. We follow (Liu et al., 2017) to consider the scaling factor r in batch normalization (BN) layers as indicators of the corresponding channels’ significance. Given a target pruning ratio p, we then prune the channels with top p-percentage smallest r values. Denote the pruned channels as 0 while the kept ones as 1, the original network can be mapped into a binary “ticket mask”. For any two sub-networks pruned from the same dense model, we calculate their mask distance via the Hamming distance between their two ticket masks.
Detecting EB Tickets via Mask Distances. We first visualize and observe the global behaviors of mask distances between consecutive epochs. Figure 3 plots the pairwise mask distance matrices (160 × 160) of the VGG16 and PreResNet101 experiments on CIFAR-100 (from Section 3.1), at different pruning ratios p, where (i, j)-th element in a matrix denotes the mask distance between subnetworks drawn from the i-th and j-th epochs in that corresponding experiment (160 epochs in total for all). For the ease of visualization, all elements in a matrix are linearly normalized between 0 and 1; Therefore, in the resulting matrix, a lower value (close to 0) indicates a smaller mask distance and is highlighted with a warmer color (same hereinafter).
Figure 3 demonstrates fairly consistent behaviors. Taking VGG16 on CIFAR-100 (p = 0.2) for an example: 1) at the very beginning, the mask distances change rapidly between epochs, manifested by the quickly “cooling” colors (yellow → green), from the diagonal line (lowest since that is comparing an epoch’s mask with itself), to off-diagonal (comparing different epochs; the more an element deviates from the diagonal, the more distant the two epochs are away from each other); 2) after ∼10 epochs, the off-diagonal elements become “yellow” too, and the color transition becomes smoother from diagonal to off-diagonal, indicating the masks change mildly only after passing this point; 3)
after∼80 epochs, the mask distances almost remain unchanged across epochs. Similar trends are observed in other plots too. It seems to concur the hypothesis in (Achille et al., 2019) that a network first learns important connectivity patterns and then fixs them and further tune their weights.
Our observation that the ticket masks quickly become stable and hardly changed in early training stages supports drawing EB tickets. We therefore measure the mask distance between the consecutive epochs, and draw EB tickets when such distance is smaller than a threshold . Practically, to improve the reliability of EB tickets, we will stop to draw EB tickets when the last five recorded mask distances are all smaller than , to avoid some irregular fluctuation in early training stages. In Figure 3, the red lines indicate the identification of EB tickets when is set to 0.1.

heading: 4 EFFICIENT TRAINING VIA EARLY BIRD TICKETS

text:
In this section, we present an efficient DNN training scheme, termed as EB Train, which leverages 1) the existence of EB tickets and 2) the proposed mask distance that can detect the emergence of EB tickets for time- and energy-efficient training. We will first provide a conceptual overview, and then describe the routine of EB Train, and finally show the evaluation performance of EB Train by benchmarking it with state-of-the-art methods of obtaining compressed DNN models based on representative datasets and DNNs.
Progressive Pruning and Training
Proposed: Efficient DNN Training via Early-Bird
Trained Model
…
EB Train (Proposed)
100% training (train 𝑵 epochs, e.g., 𝑵 = 𝟏𝟔𝟎)

heading: 4.1 WHY IS EB TRAIN MORE EFFICIENT?

text:
EB Train vs. Progressive Pruning and Training. Figure 4 illustrates an overview of our proposed EB Train and the progressive pruning and training scheme, e.g., as in (Frankle & Carbin, 2019). In particular, the progressive pruning and training scheme adopts a three-step routine of 1) training a large and dense model, 2) pruning it, and 3) then retraining the pruned model to restore performance, and these three steps can be iterated (Han et al., 2015). The first step often dominates (e.g., occupy 75% training FLOPs when using the PreResNet101 model and CIFAR-10 dataset) in terms of training energy and time costs as it involves training a large and dense model. The key feature of our EB Train scheme is that it replaces the aforementioned steps 1 and 2 with a lower-cost step of detecting the EB tickets, i.e., enabling early stopping during the time- and energy-dominant step of training the large and dense model, thus promising large savings in both training time and energy. For example, assuming the first step of the progressive pruning and training scheme requires N epochs to sufficiently train the model for achieving the target accuracy, the proposed EB Train needs only NEB epochs to identify the EB tickets by making use of the mask distance that can detect the emergence of EB tickets, where NEB ≪ N (e.g., NEB/N = 12.5% in the experiments summarized in Figure 1 and NEB/N = 6.25% in the experiment summarized in Table 1).
Algorithm 1: The Algorithm for Searching EB Tickets 1: Initialize the weights W , scaling factor r, pruning ratio p, and the FIFO queue Q with length l; 2: while t (epoch) < tmax do 3: Update W and r using SGD training; 4: Perform structured pruning based on rt towards the target ratio p; 5: Calculate the mask distance between the current and last subnetworks and add to Q. 6: t = t + 1 7: if Max(Q) < then 8: t∗ = t 9: Return f(x;mt∗ ⊙W ) (EB ticket); 10: end if 11: end while

heading: 4.2 HOW TO IMPLEMENT EB TRAIN?

text: From Figure 1, we can see that the EB Train scheme consists of a training (searching) step to identify EB tickets and a retraining step to retrain the EB tickets for achieving the target accuracy. Algorithm 1 describes the step of searching for EB tickets. Specifically, the EB searching process 1) first initializes the weights W and scaling factors r; 2) iterates the structured pruning process as in (Liu et al. (2017)) to calculate the mask distances between the consecutive resulting subnetworks, storing them into a first-in-first-out (FIFO) queue with a length of l = 5; and 3) exits when the maximum mask distance in the FIFO is smaller than a specified threshold (default 0.1 with the normalized distances of [0,1]). The output is the resulting EB ticket denoted as f(x;mt∗ ⊙W ), which will be retrained further to reach the target accuracy. Note that unlike the lottery ticket training in (Frankle & Carbin, 2019), EB Train inherits the unpruned weights from the drawn EB ticket instead of rewinding to the original initialization, as it has been shown that deeper networks are not robust to reinitializion with untrained weights (Frankle et al., 2020).

heading: 4.3 HOW DOES EB TRAIN PERFORM COMPARED TO STATE-OF-THE-ART TECHNIQUES?

text:
Experiment Setting. We consider training the VGG16 and PreResNet101 models on both CIFAR10/100 and ImageNet datasets following the basic setting of (Liu et al., 2018b) (i.e., training 160 epochs for CIFAR and 90 epochs for ImageNet). We measure the training energy costs from realdevice operations as the energy cost consists of both computational and data movement costs, the latter of which is often dominant but can not captured by the commonly used metrics, such as the number of FLOPs (Chen et al., 2017), we evaluate the proposed techniques against the baselines in terms of accuracy and real measured energy consumption. Specifically, all the energy consumption of full-precision models are obtained by training the corresponding models in an embedded GPU (NVIDIA JETSON TX2). The GPU measurement setup can be found in the Appendix. Note that the energy measurement results include the overhead of using the mask distance to detect the emergence of EB tickets, which is found negligible as compared to the total training cost. For example, for the VGG16 model, the overhead caused by computing mask distances is ≤ 0.029% in memory storage size, ≤ 0.026% in FLOPs, and ≤ 0.07% in real-measured energy costs. This is because 1) the mask distance evaluation only involves simple hamming distance calculations and 2) each epoch only calculates the distance once.
Results and Analysis. We first compare the computational savings with the baselines using a progressive pruning and training scheme (Liu et al., 2017) in terms of the total training FLOPs, when drawing subnetworks at different epochs. Figure 5 summarizes the results of the PreResNet101/VGG16 models and CIFAR-10/100 datasets, corresponding to the same set of experiments as in Figure 1. We see that EB Train can achieve 2.2 ∼ 2.4× FLOPs reduction over the baselines, while leading to comparable or even better accuracy (- 0.81% ∼ + 2.38% over the baseline). Table 2 compares the retraining accuracy and consumed FLOPs/energy of EB Train with four stateof-the-art progressive pruning and training techniques: the original lottery ticket (LT) training (Frankle & Carbin, 2019), network slimming (NS) (Liu et al., 2017), ThiNet (Luo et al., 2017) and SNIP (Lee et al., 2019). While all of them involve the process of training a dense network, pruning it, and retraining, they apply different pruning criteria, e.g., NS imposes L1-sparsity on channel-wise
scaling factors from BN layers, and ThiNet greedily prunes the channel that has the smallest effect on the next layer’s activation values. For EB Train, we by default follow (Liu et al., 2017) to inherit the same weights when re-training the searched ticket, adopting floating points for both the search and retrain stages (i.e., EB Train FF). We also notice existing debates (Liu et al., 2018b) on the initialization re-use, and thus also compare with a variant by re-training the ticket from a new random initialization, termed as EB Train (re-init). The comparisons in Table 2 further show that inheriting weights from the EB tickets favor the generalization of retraining as compared to both the random initialization and “over-cooked” weights, aligning well with the recent discussion between rewinding and fine-tuning (Renda et al., 2020). Furthermore, we apply the proposed EB Train on top of the low-precision training method (Yang et al., 2019a) and obtain experiment results of another two variants of EB Train: 1) EB Train with low-precision search and full-precision retrain (i.e., EB Train LF), and 2) EB Train with both low-precision search and retrain (i.e., EB Train LL).
Table 2 demonstrates that EB Train consistently outperforms all competitors in terms of saving training energy and computational costs, meanwhile improving the final accuracy in most cases. We use EB Train Improv. to record the performance margin (either accuracy or energy/computation) between EB Train and the strongest competitor among the four state-of-the-art baselines. Specifically, EB Train with full-precision floating point (FP32) search and retrain outperforms those pruning
methods by up to 1.2 ∼ 4.7× and 1.1 ∼ 4.5× in terms of the energy consumption and computational FLOPs, while always achieving comparable or even better accuracies, across three pruning ratios, two DNN models and two datasets. Moreover, EB Train with both low-precision (8 bits block floating point (FP8)) search and retrain outperforms the baselines by up to 5.8 ∼ 24.6× and 1.1 ∼ 5.0× in terms of energy consumption estimated using real-measured unit energy from (Yang et al., 2019b) and computational FLOPs. Note that FP8 and FP32 are counted as the same FLOPs count units (Sohoni et al., 2019). In addition, comparing with the re-init variant endorses the effectiveness of initialization inheritance in EB Train. As an additional highlight, EB Train naturally leads to more efficient inference of the pruned DNN models, unifying the boost of efficiency throughout the full learning lifecycle. ResNet18/50 on ImageNet. To study the performance of EB Train in a harder dataset, we conduct experiments using ResNet18/50 and ImageNet, and compare its resulting accuracy, total training FLOPs and energy cost with those of the method in (He et al., 2016a), NS (Liu et al., 2017), ThiNet (Luo et al., 2017) and SFP (He et al., 2018) as summarized in Table 3. We have three observations: 1) When training ResNet18 on ImageNet, EB Train achieves a better accuracy (+0.27%) over the unpruned one while introducing 10% channel sparsity; 2) EB Train outperforms other baselines for both ResNet18 and ResNet50 by reducing the total training costs while achieving a comparable or better accuracy. Specifically, EB Train achieves a reduced training FLOPs of 51.5% ∼ 56.1% and 48.6% ∼ 74.0% and a reduced training energy of 46.5% ∼ 53.2% and 49.1% ∼ 70.9%, while leading to a better top-1 accuracy of +0.19% ∼ +1.18% and +1.74% ∼ +2.34% for ResNet18 and ResNet50, respectively. For example, when training ResNet50 on ImageNet, EB Train with a pruning ratio of 70% achieves better (+1.74%) accuracy over ThiNet while saving 74% total training FLOPs and 71% training energy; 3) Different from the results on CIFAR-10/100 shown in Table 2, all methods performed on ImageNet (a harder dataset) start to yield accuracy reductions (-1.72% ∼ -3.55%) for ResNet18/50 with a pruning ratio of only 30%, compared with the unpruned one. Note that in Table 3, the unpruned results are based on the official implementation in (He et al., 2016a); The SFP results are obtained from their original paper (He et al., 2018), which does not provide results at a pruning ratio of 10%; all the ThiNet results under various pruning ratios are obtained from the original paper (Luo et al., 2017); and the NS results are obtained by conducting the experiments ourselves.

heading: 5 DISCUSSION AND CONCLUSION

text: We have demonstrated that winning tickets can be drawn at the very early training stage, i.e., EB tickets exist, in both the standard training and several lower-cost training schemes. That motivates a practical success of applying EB tickets to efficient training, whose results compare favorably against state-of-the-arts. Moreover, experiments show that EB Train with low-precision search and retraining achieve more efficient training. We believe many promising problems still remain open to be addressed. An immediate future work is to test low-precision EB Train algorithms on larger models/datasets. We are also curious whether more lower-cost training techniques could be associated with EB Train. Finally, sometimes high pruning ratios (e.g, p ≥ 0.7) can hurt the quality of EB tickets and the retrained networks. We look forward to automating the choice of p for different models/datasets, unleashing higher efficiency.

heading: 6 ACKNOWLEDGEMENT

text: The work is supported by the National Science Foundation (NSF) through the Real-Time Machine Learning (RTML) program (Award number: 1937592, 1937588).

heading: A APPENDIX

text:
A.1 MEASUREMENT OF ENERGY COST
Figure 6 shows our GPU measurement setup, in which the GPU board is connected to a laptop and a power meter. In particular, the training settings are downloaded from the laptop to the GPU board, and the real-measured energy consumption is obtained via the power meter and runtime measurement for the whole training process.
A.2 APPROPRIATE LARGE LEARNING RATE FAVORS THE EMERGENCE OF EB TICKETS
Here we show the effect of larger learning rates setting. The original schedule is denoted as [80LR→0.01, 120LR→0.001], i.e., starting from 0.1, decayed to 0.01 at the 80th epoch, and further decayed to 0.001 at the 120th epoch. In comparison, we test a larger initial learning rate, staring from 0.5 and 0.2 for VGG16 and PreResNet101 datasets, respectively. After drawing tickets, we retrain them all using the same learning rate schedule with their searching phase for sufficient training and a fair comparison. We see from Table 4 that high-quality EB tickets also emerge early when searching with the larger initial learning rate, whose final accuracies are even better.

title: FICIENT TRAINING OF DEEP NETWORKS


summaries: The paper proposed a adversarial learning framework that tries to align the hidden features of data with simple prior distributions.  A training strategy similar to GAN  was exploited. The proposed framework was argued that it can well deal with the adversarial perturbations. Some experiments were conducted, verifying that the proposed algorithm seems useful and robust.

summaries: This paper proposes a new regularization technique called Embedding Regularization to improve the adversarial robustness.  The idea is to use generative adversarial networks (GAN) to perform inference on the latent space by matching the aggregated posterior of the hidden space vector with a prior distribution. The proposed strategy could be combined with adversarial training to achieve state-of-the-art adversarial accuracy on several benchmark datasets.

abstractText: Recent studies have demonstrated the vulnerability of deep convolutional neural networks against adversarial examples. Inspired by the observation that the intrinsic dimension of image data is much smaller than its pixel space dimension and the vulnerability of neural networks grows with the input dimension, we propose to embed high-dimensional input images into a low-dimensional space to perform classification. However, arbitrarily projecting the input images to a lowdimensional space without regularization will not improve the robustness of deep neural networks. We propose a new framework, Embedding Regularized Classifier (ER-Classifier), which improves the adversarial robustness of the classifier through embedding regularization. Experimental results on several benchmark datasets show that, our proposed framework achieves state-of-the-art performance against strong adversarial attack methods.

authors:
Anish Athalye
Nicholas Carlini
David Wagner

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

venue: arXiv preprint arXiv:1802.00420,

authors:
Nicholas Carlini
David Wagner

title: Adversarial examples are not easily detected: Bypassing ten detection methods

venue: In Proceedings of the 10th ACM Workshop on Artificial Intelligence and Security, AISec

title: Towards evaluating the robustness of neural networks

venue: In Security and Privacy (SP),

authors:
Pin-Yu Chen
Huan Zhang
Yash Sharma
Jinfeng Yi
Cho-Jui Hsieh

title: Zoo: Zeroth order optimization based black-box attacks to deep neural networks without training substitute models

authors:
Pin-Yu Chen
Yash Sharma
Huan Zhang
Jinfeng Yi
Cho-Jui Hsieh

title: Ead: elastic-net attacks to deep neural networks via adversarial examples

venue: In AAAI,

authors:
Adam Coates
Andrew Ng
Honglak Lee

title: An analysis of single-layer networks in unsupervised feature learning

venue: In Proceedings of the fourteenth international conference on artificial intelligence and statistics,

authors:
Jia Deng
Wei Dong
Richard Socher
Li-Jia Li
Kai Li
Li Fei-Fei

title: Imagenet: A large-scale hierarchical image database

venue: In Computer Vision and Pattern Recognition,

authors:
Gavin Weiguang Ding
Kry Yik Chau Lui
Xiaomeng Jin
Luyu Wang
Ruitong Huang

title: On the sensitivity of adversarial robustness to input data distributions

authors:
Ian Goodfellow
Jean Pouget-Abadie
Mehdi Mirza
Bing Xu
David Warde-Farley
Sherjil Ozair
Aaron Courville
Yoshua Bengio

title: Generative adversarial nets

authors:
Ian J Goodfellow
Jonathon Shlens
Christian Szegedy

title: Explaining and harnessing adversarial examples (2014). arXiv preprint arXiv:1412.6572

authors:
Andrew Ilyas
Ajil Jalal
Eirini Asteri
Constantinos Daskalakis
Alexandros G Dimakis

title: The robust manifold defense: Adversarial training using generative models

venue: arXiv preprint arXiv:1712.09196,

authors:
Andrew Ilyas
Shibani Santurkar
Dimitris Tsipras
Logan Engstrom
Brandon Tran
Aleksander Madry

title: Adversarial examples are not bugs, they are features

authors:
Alex Krizhevsky
Geoffrey Hinton

title: Learning multiple layers of features from tiny images

venue: Technical report, Citeseer,

authors:
Alexey Kurakin
Ian Goodfellow
Samy Bengio

title: Adversarial examples in the physical world

venue: arXiv preprint arXiv:1607.02533,

title: The mnist database of handwritten digits. http://yann

venue: lecun. com/exdb/mnist/,

authors:
Elizaveta Levina
Peter J Bickel

title: Maximum likelihood estimation of intrinsic dimension

authors:
Yandong Li
Lijun Li
Liqiang Wang
Tong Zhang
Boqing Gong

title: Nattack: Learning the distributions of adversarial examples for an improved black-box attack on deep neural networks

venue: In ICML,

authors:
Xuanqing Liu
Minhao Cheng
Huan Zhang
Cho-Jui Hsieh

title: Towards robust neural networks via random self-ensemble

venue: arXiv preprint arXiv:1712.00673,

authors:
Laurens van der Maaten
Geoffrey Hinton

title: Visualizing data using t-sne

venue: Journal of machine learning research,

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:
Dongyu Meng
Hao Chen

title: Magnet: a two-pronged defense against adversarial examples

venue: In Proceedings of the 2017 ACM SIGSAC Conference on Computer and Communications Security,

authors:
Takeru Miyato
Shin-ichi Maeda
Masanori Koyama
Shin Ishii

title: Virtual adversarial training: a regularization method for supervised and semi-supervised learning

authors:
Seyed-Mohsen Moosavi-Dezfooli
Alhussein Fawzi
Pascal Frossard

title: Deepfool: a simple and accurate method to fool deep neural networks

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

authors:
Seyed-Mohsen Moosavi-Dezfooli
Alhussein Fawzi
Omar Fawzi
Pascal Frossard

title: Universal adversarial perturbations

venue: arXiv preprint,

authors:
Taesik Na
Jong Hwan Ko
Saibal Mukhopadhyay

title: Cascade adversarial machine learning regularized with a unified embedding

venue: arXiv preprint arXiv:1708.02582,

authors:
Nicolas Papernot
Patrick McDaniel
Xi Wu
Somesh Jha
Ananthram Swami

title: Distillation as a defense to adversarial perturbations against deep neural networks

authors:
Nicolas Papernot
Patrick McDaniel
Ian Goodfellow
Somesh Jha
Z Berkay Celik
Ananthram Swami

title: Practical black-box attacks against machine learning

venue: In Proceedings of the 2017 ACM on Asia Conference on Computer and Communications Security,

authors:
Pouya Samangouei
Maya Kabkab
Rama Chellappa

title: Defense-gan: Protecting classifiers against adversarial attacks using generative models

venue: arXiv preprint arXiv:1805.06605,

authors:
Carl-Johann Simon-Gabriel
Yann Ollivier
Bernhard Schölkopf
Léon Bottou
David LopezPaz

title: Adversarial vulnerability of neural networks increases with input dimension

venue: arXiv preprint arXiv:1802.01421,

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,

title: Optimal transport: old and new, volume 338

venue: Springer Science & Business Media,

authors:
Cihang Xie
Jianyu Wang
Zhishuai Zhang
Zhou Ren
Alan Yuille

title: Mitigating adversarial effects through randomization

venue: arXiv preprint arXiv:1711.01991,

authors:
Shufei Zhang
Kaizhu Huang
Jianke Zhu
Yang Liu

title: Manifold adversarial learning, 2018

heading: 1 INTRODUCTION

text:
Deep neural networks (DNNs) have been widely used for tackling numerous machine learning problems that were once believed to be challenging. With their remarkable ability of fitting training data, DNNs have achieved revolutionary successes in many fields such as computer vision, natural language progressing, and robotics. However, they were shown to be vulnerable to adversarial examples that are generated by adding carefully crafted perturbations to original images. The adversarial perturbations can arbitrarily change the network’s prediction but often too small to affect human recognition (Szegedy et al., 2013; Kurakin et al., 2016). This phenomenon brings out security concerns for practical applications of deep learning.
Two main types of attack settings have been considered in recent research (Goodfellow et al.; Carlini & Wagner, 2017a; Chen et al., 2017; Papernot et al., 2017): black-box and white-box settings. In the black-box setting, the attacker can provide any inputs and receive the corresponding predictions. However, the attacker cannot get access to the gradients or model parameters under this setting; whereas in the white-box setting, the attacker is allowed to analytically compute the model’s gradients, and have full access to the model architecture and weights. In this paper, we focus on defending against the white-box attack which is the harder task.
Recent work (Simon-Gabriel et al., 2018) presented both theoretical arguments and an empirical one-to-one relationship between input dimension and adversarial vulnerability, showing that the vulnerability of neural networks grows with the input dimension. Therefore, reducing the data dimension may help improve the robustness of neural networks. Furthermore, a consensus in the highdimensional data analysis community is that, a method working well on the high-dimensional data is because the data is not really of high-dimension (Levina & Bickel, 2005). These high-dimensional data, such as images, are actually embedded in a low dimensional space. Hence, carefully reducing the input dimension may improve the robustness of the model without sacrificing performance.
Inspired by the observation that the intrinsic dimension of image data is actually much smaller than its pixel space dimension (Levina & Bickel, 2005) and the vulnerability of a model grows with its input dimension (Simon-Gabriel et al., 2018), we propose a defense framework that embeds input images into a low-dimensional space using a deep encoder and performs classification based on the latent embedding with a classifier network. However, an arbitrary projection does not guarantee improving the robustness of the model, because there are a lot of mapping functions including non-robust ones from the raw input space to the low-dimensional space capable of minimizing the classification loss. To constrain the mapping function, we employ distribution regularization in the
embedding space leveraging optimal transport theory. We call our new classification framework Embedding Regularized Classifier (ER-Classifier). To be more specific, we introduce a discriminator in the latent space which tries to separate the generated code vectors from the encoder network and the ideal code vectors sampled from a prior distribution, i.e., a standard Gaussian distribution. Employing a similar powerful competitive mechanism as demonstrated by Generative Adversarial Networks (Goodfellow et al., 2014), the discriminator enforces the embedding space of the model to follow the prior distribution.
In our ER-Classifier framework, the encoder and discriminator structures together project the input data to a low-dimensional space with a nice shape, then the classifier performs prediction based on the lowdimensional embedding. Based on the optimal transport theory, the proposed ER-Classifier minimizes the discrepancy between the distribution of the true label and the distribution of the framework output, thus only retaining important features for classification in the embedding space. With a small embedding dimension,
the effect of the adversarial perturbation is largely diminished through the projection process.
We compare ER-Classifier with other state-of-the-art defense methods on MNIST, CIFAR10, STL10 and Tiny Imagenet. Experimental results demonstrate that our proposed ER-Classifier outperforms other methods by a large margin. To sum up, this paper makes the following three main contributions:
• A novel unified end-to-end robust deep neural network framework against adversarial attacks is proposed, where the input image is first projected to a low-dimensional space and then classified. • An objective is induced to minimize the optimal transport cost between the true class distribution and the framework output distribution, guiding the encoder and discriminator to project the input image to a low-dimensional space without losing important features for classification. • Extensive experiments demonstrate the robustness of our proposed ER-Classifier framework under the white-box attacks, and show that ER-Classifier outperforms other state-ofthe-art approaches on several benchmark image datasets.
As far as we know, our approach is the first that applies optimal transport theory, i.e., a Wasserstein distance regularization, to a bottleneck embedding layer of a deep neural network in a purely supervised learning setting without considering any reconstruction loss, although optimal transport theory or a discriminator loss has been applied to generative models in an unsupervised learning setting (Makhzani et al., 2015; Tolstikhin et al., 2017); (2) Our method is also the first that establishes the connection between a Wasserstein distance regularization and the robustness of deep neural networks for defending against adversarial examples.

heading: 2 RELATED WORK

text: In this section, we summarize related work into three categories: attack methods, defense mechanisms and optimal transport theory. We first discuss different white-box attack methods, followed by a description of different defense mechanisms against, and finally optimal transport theory.

heading: 2.1 ATTACK METHODS

text:
Under the white-box setting, attackers have all information about the targeted neural network, including network structure and gradients. Most white-box attacks generate adversarial examples based on the gradient of loss function with respect to the input. An algorithm called fast gradient sign method (FGSM) was proposed in (Goodfellow et al.) which generates adversarial examples based on the sign of gradient. Many other white-box attack methods have been proposed
recently (Moosavi-Dezfooli et al., 2016; Chen et al., 2018; Madry et al., 2017; Carlini & Wagner, 2017b), and among them C&W and PGD attacks have been widely used to test the robustness of machine learning models.
C&W attack: The adversarial attack method proposed by Carlini and Wagner (Carlini & Wagner, 2017b) is one of the strongest white-box attack methods. They formulate the adversarial example generating process as an optimization problem. The proposed objective function aims at increasing the probability of the target class and minimizing the distance between the adversarial example and the original input image. Therefore, C&W attack can be viewed as a gradient-descent based adversarial attack.
PGD attack: The projected gradient descent attack is proposed by (Madry et al., 2017), which finds adversarial examples in an -ball of the image. The PGD attack updates in the direction that decreases the probability of the original class most, then projects the result back to the -ball of the input. An advantage of PGD attack over C&W attack is that it allows direct control of distortion level by changing , while for C&W attack, one can only do so indirectly via hyper-parameter tuning.
Both C&W attack and PGD attack have been frequently used to benchmark the defense algorithms due to their effectiveness (Athalye et al., 2018). In this paper, we mainly use l∞-PGD untargeted attack to evaluate the effectiveness of the defense method under white-box setting.
Instead of crafting different adversarial perturbation for different input image, an algorithm was proposed by (Moosavi-Dezfooli et al., 2017) to construct a universal perturbation that causes natural images to be misclassified. However, since this universal perturbation is image-agnostic, it is usually larger than the image-specific perturbation generated by PGD and C&W.

heading: 2.2 DEFENSE MECHANISMS

text:
Many works have been done to improve the robustness of deep neural networks. To defend against adversarial examples, defenses that aim to increase model robustness fall into three main categories: i) augmenting the training data with adversarial examples to enhance the existing classifiers (Madry et al., 2017; Na et al., 2017; Goodfellow et al.); ii) leveraging model-specific strategies to enforce model properties such as smoothness (Papernot et al., 2016); and, iii) trying to remove adversarial perturbations from the inputs (Xie et al., 2017; Samangouei et al., 2018; Meng & Chen, 2017). We select three representative methods that are effective under white-box setting.
Adversarial training: Augmenting the training data with adversarial examples can increase the robustness of the deep neural network. Madry et al. (Madry et al., 2017) recently introduced a minmax formulation against adversarial attacks. The proposed model is not only trained on the original dataset but also adversarial example in the -ball of each input image.
Random Self-Ensemble: Another effective defense method under white-box setting is RSE (Liu et al., 2017). The authors proposed a “noise layer”, which fuses output of each layer with Gaussian noise. They empirically show that the noise layer can help improve the robustness of deep neural networks. The noise layer is applied in both training and testing phases, so the prediction accuracy will not be largely affected.
Defense-GAN: Defense-GAN (Samangouei et al., 2018) leverages the expressive capability of GANs to defend deep neural networks against adversarial examples. It is trained to project input images onto the range of the GAN’s generator to remove the effect of the adversarial perturbation. Another defense method that uses the generative model to filter out noise is MagNet proposed by (Meng & Chen, 2017). However, the differences between ER-Classifier and the two methods are obvious. ER-Classifier focuses on reducing the dimension, and performing classification based on the low-dimensional embedding, while Defense-GAN and MagNet mainly apply the generative model to filter out the adversarial noise, and both Defense-GAN and MagNet perform classification on the original dimension space. (Samangouei et al., 2018) showed that Defense-GAN is more robust than MagNet, so we only compare with Defense-GAN in the experiments.
Other Related Methods: Zhang et al. (2018) regularizes the latent space with Gaussian Mixture Model and applies KL-divergence to do the optimization. However, our method employs a simple but nice-shaped Gaussian prior for Wasserstein distance minimization to constrain the global shape of the latent embeddings, while permitting high freedom for the shapes of individual class distributions of latent embeddings. We want the classifier to decide the optimal class-specific distributions of
latent embeddings. Miyato et al. (2018) shares a similar idea to adversarial learning, but it employs virtual labels generated by a current classifier to identify search directions that can smooth the output label distribution of the classifier and is best suitable for semi-supervised learning. Please note that both methods in (Zhang et al., 2018; Miyato et al., 2018) are designed for improving generalization performance but not for defending against adversarial examples. A recent paper (Ilyas et al., 2019) shows that adversarial examples are purely human phenomenon and models tend to learn features that are not robust yet generalize well. We show that our Wasserstein distance regularization helps to identify robust features, which will be discussed later.
Notations In this paper, we use l∞ and l2 distortion metrics to measure similarity. We report l∞ distance in the normalized [0, 1] space, so that a distortion of 0.031 corresponds to 8/256, and l2 distance as the total root-mean-square distortion normalized by the total number of pixels.
We use calligraphic letters for sets (i.e., X ), capital letters for random variables (i.e., X), and lower case letters for their values (i.e., x). The probability distributions are denoted with capital letters (i.e., PX ) and corresponding densities with lower case letters (i.e., pX ).

heading: 3 PROPOSED FRAMEWORK: EMBEDDING REGULARIZED CLASSIFIER

heading: 3.1 FRAMEWORK DETAILS

text:
We propose a novel defense framework, ER-Classifier, which aims at projecting the image data to a low-dimensional space to remove noise and stabilize the classification model by minimizing the optimal transport cost between the true label distribution PY and the distribution of the ER-Classifier output (PC). An overview of the framework is shown in Figure 1. The encoder and discriminator structures together help diminish the effect of the adversarial perturbation by projecting input data to a space of lower dimension, then the classifier part performs classification based on the lowdimensional embedding.
Mathematically, input images X ∈ X = Rd are projected to a low-dimensional embedding vector Z ∈ Z = Rk through the encoder Qφ. The discriminator Dγ discriminates between the generated code Z̃ ∼ Qφ(Z|X) and the ideal code Z ∼ PZ . The classifier Cτ performs classification based on the generated code Z̃, producing output U ∈ U = Rm, where m is the number of classes. The label of X is denoted as Y ∈ U . The Kantorovich’s distance induced by the optimal transport problem is given by
Wc(PY , PC) := inf Γ∈P(Y∼PY ,U∼PC) E(Y,U)∼Γ {c(Y,U)} ,
whereP(Y ∼ PY , U ∼ PC) is the set of all joint distributions of (Y,U) with marginals PY and PC , and c(y, u) : U × U 7→ R+ is any measurable cost function. Wc(PY , PC) measures the divergence between probability distributions PY and PC . When the probability measures are on a metric space, the p-th root of Wc is called the p-Wasserstein distance.
To minimize the Wasserstein distance between the distribution of the true label (PY ) and the distribution of the ER-Classifier output (PC), we can prove that it is sufficient to find a conditional distribution Q(Z|X) such that its marginal distribution QZ is identical to a prior distribution PZ . The theorem and the proof are deferred to the Appendix. In this paper, we apply standard Gaussian as our prior distribution PZ , but other priors may be used for different cases. The final objective of ER-Classifier is:
inf Q(Z|X)∈Q EPXEQ(Z|X) {`(f(X),C(Z))}+ λD(QZ , PZ), (1)
where Q can be a deterministic encoder as focused by this paper due to its simplicity or stochastic encoder as the one in a standard Variational Autoencoder, λ > 0 is a hyper-parameter and D is an arbitrary divergence between QZ and PZ .
To estimate the divergences between QZ and PZ , we apply a GAN-based framework, fitting a discriminator to minimize the 1-Wasserstein distance between QZ and PZ :
W (QZ , PZ) = inf Γ∈P(Z̃∼QZ ,Z∼PZ) E(Z̃,Z)∼Γ‖Z̃ − Z‖.
We have also tried the Jensen-Shannon divergence, but as expected, Wasserstein distance provides more stable training and better results. When training the framework, the weight clipping method proposed in Wasserstein GAN (Arjovsky et al., 2017) is applied to help stabilize the training of discriminator Dγ . The training algorithm is summarized in Algorithm 1.
At training stage, the encoder Qφ first maps the input x to a low-dimensional space, resulting in generated code (z̃). Another ideal code (z) is sampled from the prior distribution, and the discriminator Dγ discriminates between the ideal code (positive data) and the generated code (negative data). The classifier (Cτ ) predicts the image label based on the generated code (z̃).
Algorithm 1 Training ER-Classifier 1: Input: Regularization coefficient λ > 0, encoder Qφ, discriminator Dγ , and classifier Cτ . 2: Note: ` stands for the cross-entropy loss. 3: while (φ, γ, τ) not converged do 4: Sample {(x1, y1), ..., (xn, yn)} from the training set 5: Sample {z1, ..., zn} from the prior PZ 6: Directly obtain or sample z̃i from Qφ(Z|xi) for i = 1, ..., n 7: Update Dγ by ascending the following objective by 1-step Adam:
λ
n n∑ i=1 Dγ(zi)−Dγ(z̃i)
8: Update Qφ and Cτ by descending the following objective by 1-step Adam:
1
n n∑ i=1 `(Cτ (Qφ(xi)), yi)
9: Update Qφ by ascending the following objective by 1-step Adam:
λ
n n∑ i=1 Dγ(Qφ(xi))
10: end while
At inference time, only the encoder Qφ and the classifier Cτ are used. The input image x is first mapped to a low-dimensional space by the encoder (z̃ = Qφ(x)), then the latent code z̃ is fed into the classifier to obtain the predicted label.
The main goal of ER-Classifier is leveraging input space dimension reduction to remove adversarial perturbations. Therefore, other defense methods can also benefit from this property. Our framework is trained with min-max robust optimization (Madry et al., 2017).

heading: 3.2 JUSTIFICATIONS OF OUR FRAMEWORK

text:
There are two Wasserstein distances (W-distances) in our framework. One is the W-distance between the aggregated latent embedding distribution Q(z) and the prior distribution PZ , and the other one is the W-distance between the true label distribution PY and the distribution of the ER-Classifier output (PC). In Algorithm 1, we are minimizing the first one. The theorem in the Appendix shows that minimizing the first W-distance in combination with minimizing a standard cross-entropy loss as done in Algorithm 1 is equivalent to minimizing the second W-distance, which guarantees that the training process is not distracted from the main goal of the framework, classification. That is to say, Algorithm 1 will result in a classifier with the following property: the global output distribution of the classifier will match the global ground-truth label distribution in the data no matter whether the encoder Qφ is deterministic or stochastic (the second W-distance is automatically minimized).
It’s hard to analyze the importance of the theorem in the Appendix if we just look at a deterministic encoder. Let’s convert this deterministic encoder to a stochastic encoder that outputs a Gaussian z with a fixed variance and the mean being the same as its corresponding deterministic version. The theory tells us that, by minimizing the first W-distance over all sampled z’s from this stochastic encoder and the standard cross-entropy loss, we will automatically minimize the second W-distance
and preserve the global label frequency in the dataset, even though these z’s are only -close to the deterministic encoding features of training data.
Moreover, we find that minimizing the W-distance helps the encoder identify some robust features instead of non-robust features (Ilyas et al., 2017), because our proposed regularization constrains the -ball around each Qφ(Z|X) to contribute to preserving the global label distribution in the data, even with X integrated out. From this perspective, we can view our proposed framework as “a supervised variant” of Generative Adversarial Network or Wasserstein Autoencoder in which the Generator or Decoder is replaced by a Classifier that generates labels from low-dimensional latent embeddings preserving global label frequency in the training dataset. Replacing W-distance with KL divergence loses all these nice properties.
In our framework, we use a simple but nice-shaped Gaussian prior PZ for W-distance minimization to constrain the global shape of the latent embeddings, while permitting high freedom for the shapes of individual class distributions of latent embeddings. We want the classifier to decide the optimal class-specific distributions of latent embeddings. In addition, it is interesting to explore how to set -ball to make sure the stochastic encoder to best align the latent embedding z to human-perceived robust features, which will be left as future work.

heading: 4 EXPERIMENTS

text:
In this section, we compare the performance of our proposed algorithm (ER-Classifier) with other state-of-the-art defense methods on several benchmark datasets:
• MNIST (LeCun, 1998): handwritten digit dataset, which consists of 60, 000 training images and 10, 000 testing images. Theses are 28×28 black and white images in ten different classes. • CIFAR10 (Krizhevsky & Hinton, 2009): natural image dataset, which contains 50, 000 training images and 10, 000 testing images in ten different classes. These are low resolution 32× 32 color images. • STL10 (Coates et al., 2011): color image dataset similar to CIFAR10, but contains only 5, 000 training images and 8, 000 testing images in ten different classes. The images are of higher resolution 96× 96. • Tiny Imagenet (Deng et al., 2009): a subset of Imagenet dataset. Tiny Imagenet has 200 classes, and each class has 500 training images, 50 testing images, making it a challenging benchmark for defense task. The resolution of the images is 64× 64.
Various defense methods have been proposed to improve the robustness of deep neural networks. Here we compare our algorithm with state-of-the-art methods that are robust in white-box setting. Madry’s adversarial training (Madry’s Adv) has been recognized as one of the most successful defense method in white-box setting, as shown in (Athalye et al., 2018).
Random Self-Ensemble (RSE) method introduced by (Liu et al., 2017) adds stochastic components in the neural network, achieving similar performance to Madry’s adversarial training algorithm.
Another method we would like to compare with is Defense-GAN (Samangouei et al., 2018). It first trains a generative adversarial network to model the distribution of the training data. At inference time, it finds a close output to the input image and feed that output into the classifier. This process “projects” input images onto the range of GAN’s generator, which helps remove the effect of adversarial perturbations. In (Samangouei et al., 2018), the author demonstrated the performance of Defense-GAN on MNIST and Fashion-MNIST, so we will compare our method with Defense-GAN on MNIST.
Since the main goal of ER-Classifier is using dimension reduction to improve adversarial robustness, other defense methods can also benefit from this property. The proposed ER-Classifier is trained with min-max robust optimization (Madry et al., 2017). To demonstrate the dimension reduction ability of ER-Classifier, we include a variant ER-Classifier− which trains ER-Classifier without min-max robust optimization. 4.1 EVALUATE MODELS UNDER WHITE-BOX l∞-PGD ATTACK
In this section, we evaluate the defense methods against l∞-PGD untargeted attack, which is one of the strongest white-box attack methods. Models are evaluated under different distortion level ( ),
and the larger the distortion the stronger the attack. Depending on the image scale and type, different datasets are sensitive to different strength of attack.
Models on MNIST are evaluated under distortion level from 0 to 0.4 by 0.025. Models on CIFAR10 and STL10 are evaluated under ∈ [0, 0.06, 0.005]. Models on Tiny Imagenet are evaluated under ∈ [0, 0.02, 0.002]. As mentioned in the notation part, all the distortion levels are reported in the normalized [0, 1] space. The experimental results are shown in Figure 2. To demonstrate the results more clearly, we show part of the results in Table 1.
Based on Figure 2 and Table 1, we can see that ER-Classifier is the most robust one on a variety of datasets. ER-Classifier without min-max robust optimization can also improve the robustness of deep neural network. Compare the performance of ER-Classifier− with the performance of model without defense method (No Defense), we can see that ERClassifier− is much more robust than the model with no defense method on all benchmark datasets. Besides, when the distortion level ( ) is large, ER-Classifier− tends to perform better than other state-of-the-art defense methods on MNIST, CIFAR10 and Tiny Imagenet. This phenomenon is obvious on CIFAR10 and it even performs better than ER-Classifier when
the attack strength is strong. The reason might be that without min-max robust optimization, it is easier to regularize the embedding space.
We also compare Defense-GAN with our method ERClassifier on MNIST. Although Defense-GAN was shown to be partly broken by (Athalye et al., 2018; Ilyas et al., 2017), both ER-Classifier and DefenseGAN share the similar idea of projecting the input to a learned manifold, and comparing to Defense-GAN is important to demonstrate the advantage of our novel Wassserstein distance regularization. Please note that Defense-GAN is not our major comparison baseline in
this paper. Both ER-Classifier and Defense-GAN are evaluated against the l2-C&W untargeted attack, one of the strongest white-box attack proposed in (Carlini & Wagner, 2017b). Defense-GAN is evaluated using the method proposed in (Athalye et al., 2018), and the code is available on github1. ER-Classifier is evaluated against l2-C&W untargeted attack with the same hyper-parameter values as those used in the evaluation of Defense-GAN. The results under l2 ≤ 0.005 threshold are shown in Table 2. Based on Table 2, ER-Classifier is much more robust than Defense-GAN un-
1Publicly available at https://github.com/anishathalye/obfuscated-gradients/ tree/master/defensegan
der the l2 ≤ 0.005 threshold. Since (Samangouei et al., 2018) did not evaluate Defense-GAN on CIFAR10, STL10 and Tiny Imagenet, without details of GAN structure, we can not compare with Defense-GAN on these datasets.

heading: 4.2 EVALUATE MODELS UNDER BLACK-BOX ATTACK

text: We evaluate Madry’s adversarial training, ER-Classifier, and ER-Classifier−, against a recently proposed black-box attack method called Nattack (Li et al., 2019)2 on CIFAR10. Nattack is only performed on the first 100 images of CIFAR10 since the attack process takes a long time. We report the accuracy = number of correctly classified / number of attacked images (exactly 100). The accuracy of Madry’s adv, ER-Classifier, and ER-Classifier− is, respectively, 38%, 43%, and 32%. ER-Classifier still outperforms Madry’s adv.

heading: 4.3 EVALUATE THE EFFECT OF DISCRIMINATOR

text:
ER-Classifier framework consists of three parts, and the classification task is done by the encoder Qφ and classifier Cτ . Without the discriminator part, the encoder can also project the input images to a low-dimensional space. However, arbitrarily projecting the images to a low-dimensional space with only the encoder part cannot improve the robustness of the model. In contrast, sometimes it even decreases the robustness of the model.
To show that arbitrarily projecting the input images to a low-dimensional space can not improve the robustness, we fit a framework with only the encoder and classifier part (E-CLA), where the encoder and classifier have the same structures as in ER-Classifier, and compare E-CLA with the ER-Classifier framework. For a fair comparison, both structures are trained without min-max robust optimization. The results are shown in Figure 3.
Based on Figure 3, we can observe that ER-Classifier is much more robust than just the encoder and classifier structure on MNIST, CIFAR10 and Tiny Imagenet. It is also more robust on STL10 but not that much. The reason might be that there are only 5, 000 training images in STL10 and the resolution is 96× 96. Therefore, it is harder to learn a good embedding with limited amount of images. However, even when the number of training images is limited, ER-Classifier is still much more robust than the E-CLA structure. This observation demonstrates that regularization on the embedding space helps improve the adversarial robustness. Notice that the performance of E-CLA structure is similar to the performance of model without defense method on CIFAR10, STL10 and Tiny Imagenet, and worse on MNIST, which means the robustness of ER-Classifier does not come from the structure design.
Variational auto-encoder can project the images to low-dimensional space and use Kullback–Leibler divergence loss to regularize the embedding distribution, which does not need discriminator structure. Therefore, we also tried VAE-CLA, which applies Variational auto-encoder structure to do the projection and regularization. The experimental results in Figure 3 show that VAE-CLA does not perform as well as ER-Classifier. Based on the observation of the Kullback–Leibler loss and classification loss during the training process, it seems difficult for VAE-CLA to balance between the two tasks. The reason might be that Kullback–Leibler distances are not sensible cost functions when learning distributions supported by low dimensional manifolds (Arjovsky et al., 2017).
2https://github.com/Cold-Winter/Nattack

heading: 4.4 PRIOR SELECTION

text:
ER-Classifier does not have restrictions on the choice of prior. However, the selection of prior is important as it imposes different restrictions on the embedding space.
Three different prior distributions are tested on MNIST and CIFAR10 datasets. They are standard Gaussian, Uniform(−3, 3) and Cauchy(0, 1), where Cauchy(0, 1) has the same support as standard Gaussian but is heavy tailed and 99.7% of the standard Gaussian points lies within [−3, 3]. All the models are trained without min-max robust optimization, and the experimental results are shown in Figure 4. Based on the results, all three priors work well, but standard Gaussian performs best on both datasets.
Ding et al. (Ding et al., 2019) prove that adversarial robustness is sensitive to the input data distribution, and if the data is uniformly distributed in the input space, no algorithm can achieve good robustness. They also empirically show that cornered/concentrated data distributions tend to achieve better robustness. This helps explain why regularizing the embedding space can help improve robustness. Though the projection process reduces the input dimension, the embedding space is still large. Prior distribution helps push the embedding space to be more concentrated, reducing the valid perturbation space.
Details of hyper-parameter selection, model structure and code are included in the supplementary part. Embedding space visualization can also be found in the supplementary material.

heading: 5 CONCLUSION

text: In this paper, we propose a new defense framework, ER-Classifier, which projects the input images to a low-dimensional space to remove adversarial perturbation and stabilize the model through minimizing the discrepancy between the true label distribution and the framework output distribution. We empirically show that ER-Classifier is much more robust than other state-of-the-art defense methods on several benchmark datasets. Future work will include further exploration of the low-dimensional space to improve the robustness of deep neural network.

heading: APPENDIX

heading: THEORETICAL ANALYSIS OF OUR FRAMEWORK

text:
Mathematically, input images X ∈ X = Rd are projected to a low-dimensional embedding vector Z ∈ Z = Rk through the encoder Qφ. The discriminator Dγ discriminates between the generated code Z̃ ∼ Qφ(Z|X) and the ideal code Z ∼ PZ . The classifier Cτ performs classification based on the generated code Z̃, producing output U ∈ U = Rm, where m is the number of classes. The label of X is denoted as Y ∈ U . The ER-Classifier framework embeds important classification features by minimizing the discrepancy between the distribution of the true label (PY ) and the distribution of the framework output (PC). In the framework, the classifier (PC(U |Z)) maps a latent code Z sampled from a fixed distribution in a latent space Z , to the output U ∈ U = Rm. The density of ER-Classifier output is defined as follow:
pC(u) := ∫ Z pC(u|z)pZ(z)dz, ∀u ∈ U . (2)
In this paper we apply standard Gaussian as our prior distribution PZ , but other priors may be used for different cases. Assume there is an oracle f : X 7→ U assigning the image data (X ∈ X ) its true label (Y ∈ U). We want to minimize the optimal transport cost between the distribution of the true label (PY ) and the distribution of the ER-Classifier output (PC).

heading: OPTIMAL TRANSPORT THEORY

text:
There are various ways to define the distance or divergence between the target distribution and the model distribution. In this paper, we turn to the optimal transport theory (Villani, 2008), which provides a much weaker topology than many others. In real applications, data is usually embedded in a space of a much lower dimension, such as a non-linear manifold. Kullback-Leibler divergence, Jensen-Shannon divergence and Total Variation distance are not sensible cost functions when learning distributions supported by lower dimensional manifolds (Arjovsky et al., 2017). In contrast, the optimal transport cost is more sensible in this setting. Kantorovich’s distance induced by the optimal transport problem is given by
Wc(PY , PC) := inf Γ∈P(Y∼PY ,U∼PC) E(Y,U)∼Γ {c(Y,U)} ,
whereP(Y ∼ PY , U ∼ PC) is the set of all joint distributions of (Y, U) with marginals PY and PC , and c(y, u) : U × U 7→ R+ is any measurable cost function. Wc(PY , PC) measures the divergence between probability distributions PY and PC . When the probability measures are on a metric space, the p-th root of Wc is called the p-Wasserstein distance.
To minimize the optimal transport cost between the distribution of the true label (PY ) and the distribution of the ER-Classifier output (PC), it is sufficient to find a conditional distribution Q(Z|X) such that its marginal distribution QZ is identical to the prior distribution PZ .
Theorem 1 For PC as defined above with a deterministic PC(U |Z) and any function C : Z 7→ U
inf Γ∈P(Y∼PY ,U∼PC)
E(Y,U)∼Γ {`(Y,U)}
= inf Q:QZ=PZ EPXEQ(Z|X) {`(f(X),C(Z))} ,
where Γ ∈ P(Y ∼ PY , U ∼ PC) is the set of all joint distributions of (Y, U) with marginals PY and PC , and `(y, u) : U × U 7→ R+ is any measurable cost function. QZ is the marginal distribution of Z when X ∼ PX and Z ∼ Q(Z|X). (The proof is presented later.)
Therefore, optimizing over the objective on the r.h.s is equivalent to minimizing the discrepancy between the true label distribution (PY ) and the output distribution PC , thus the important classification features are embedded in the low-dimensional space. This is the core idea of the paper,
summarizing the high-dimensional data in a space of much lower dimension without losing important features for classification. To implement the r.h.s objective, the constraint on QZ can be relaxed by adding a penalty term. The final objective of ER-Classifier is:
inf Q(Z|X)∈Q EPXEQ(Z|X) {`(f(X),C(Z))}+ λD(QZ , PZ), (3)
whereQ is any nonparametric set of probabilistic encoders, λ > 0 is a hyper-parameter and D is an arbitrary divergence between QZ and PZ .
To estimate the divergences between QZ and PZ , we apply a GAN-based framework, fitting a discriminator to minimize the 1-Wasserstein distance between QZ and PZ :
W (QZ , PZ) = inf Γ∈P(Z̃∼QZ ,Z∼PZ) E(Z̃,Z)∼Γ‖Z̃ − Z‖.
We have also tried the Jensen-Shannon divergence, but as expected, Wasserstein distance provides more stable training and better results. When training the framework, the weight clipping method proposed in Wasserstein GAN (Arjovsky et al., 2017) is applied to help stabilize the training of discriminator Dγ .

heading: PROOF OF THEOREM 1

text:
The proof of Theorem 1 is adapted from the proof of Theorem 1 in (Tolstikhin et al., 2017). Consider certain sets of joint probability distributions of three random variables (X,U,Z) ∈ X × U × Z . X can be taken as the input images, U as the output of the framework, and Z as the latent codes. PC,Z(U,Z) represents a joint distribution of a variable pair (U,Z), where Z is first sampled from PZ and then U from PC(U |Z). PC defined in (2) is the marginal distribution of U when (U,Z) ∼ PC,Z .
The joint distributions Γ(X,U) or couplings between values of X and U can be written as Γ(X,U) = Γ(U |X)PX(X) due to the marginal constraint. Γ(U |X) can be decomposed into an encoding distribution Q(Z|X) and the generating distribution PC(U |Z), and Theorem 1 mainly shows how to factor it through Z.
In the first part, we will show that if PC(U |Z) are Dirac measures, we have
inf Γ∈P(X∼PX ,U∼PC)
E(X,U)∼Γ {`(f(X), U)}
= inf Γ∈PX,U
E(X,U)∼Γ {`(f(X), U)} , (4)
where P(X ∼ PX , U ∼ PC) denotes the set of all joint distributions of (X,U) with marginals PX , PC , and likewise for P(X ∼ PX , Z ∼ PZ). The set of all joint distributions of (X,U,Z) such that X ∼ PX , (U,Z) ∼ PC,Z , and (U ⊥ X)|Z are denoted by PX,U,Z . PX,U and PX,Z denote the sets of marginals on (X,U) and (X,Z) induced by PX,U,Z . From the definition, it is clear that PX,U ⊆ P(PX , PC). Therefore, we have
inf Γ∈P(X∼PX ,U∼PC)
E(X,U)∼Γ {`(f(X), U)}
≤ inf Γ∈PX,U E(X,U)∼Γ {`(f(X), U)} , (5)
The identity is satisfied if PC(U |Z) are Dirac measures, such as U = C(Z). This is proved by the following Lemma in (Tolstikhin et al., 2017). -5pt
Lemma 1 PX,U ⊆ P(PX , PC) with identity if PC(U |Z = z) are Dirac for all z ∈ Z . (see details in (Tolstikhin et al., 2017).)
In the following part, we show that
inf Γ∈PX,U
E(X,U)∼Γ {`(f(X), U)}
= inf Q:QZ=PZ EPXEQ(Z|X) {`(f(X),C(Z))} . (6)
Based on the definition, P(PX , PC), PX,U,Z and PX,U depend on the choice of conditional distributions PC(U |Z), but PX,Z does not. It is also easy to check that PX,Z = P(X ∼ PX , Z ∼ PZ). The tower rule of expectation, and the conditional independence property of PX,U,Z implies
inf Γ∈PX,U
E(X,U)∼Γ {`(f(X), U)}
= inf Γ∈PX,U,Z
E(X,U,Z)∼Γ {`(f(X), U)}
= inf Γ∈PX,U,Z EPZEX∼P (X|Z)EU∼P (U |Z) {`(f(X), U)}
= inf Γ∈PX,U,Z EPZEX∼P (X|Z) {`(f(X),C(Z))}
= inf Γ∈PX,Z
E(X,Z)∼Γ {`(f(X),C(Z))}
= inf Q:QZ=PZ EPXEQ(Z|X) {`(f(X),C(Z))} (7)
Finally, since Y = f(X), it is easy to get
inf Γ∈P(Y∼PY ,U∼PC)
E(Y,U)∼Γ {`(Y,U)}
= inf Γ∈P(X∼PX ,U∼PC)
E(X,U)∼Γ {`(f(X), U)} (8)
Now (4), (6) and (8) are proved and the three together prove Theorem 1.
Our proposed framework readily applies to non-deterministic case. If the classifier part is nondeterministic, Lemma 1 provides only the inclusion of sets PX,U ⊆ P(PX , PU ), and we can get an upper bound on the Wasserstein distance between the ground-truth and predicted label distributions:
inf Γ∈P(X∼PX ,U∼PC) E(X,U)∼Γ {`(f(X), U)} ≤ inf Γ∈PX,U E(X,U)∼Γ {`(f(X), U)}
≤ d∑ i=1 σ2i + inf Γ∈PX∼PX,Z∼PZ E(X,Z)∼Γ { ‖f(X)−C(Z)‖2 } , (9)
where we assume the conditional distributions PC(U |Z = z) have mean values C(z) ∈ Rd and marginal variances σ21 , ..., σ 2 d ≥ 0 for all z ∈ Z , where C : Z → X , and `(y, u) = ‖y − u‖2. The above upper bound is derived by:
inf Γ∈PX,U
E(X,U)∼Γ { ‖f(X)− U‖2 } = inf
Γ∈PX,U,Z EPZEX∼P (X|Z)EU∼P (U |Z){‖f(X)− U‖2}
(10)
and
EU∼P (U |Z){‖f(X)− U‖2} = EU∼P (U |Z){‖f(X)−C(Z) + C(Z)− U‖2} = ‖f(X)−C(Z)‖2 + EU∼P (U |Z){< f(X)−C(Z),C(Z)− U >}+ EU∼P (U |Z){‖C(Z)| − U‖2}
= ‖f(X)−C(Z)‖2 + d∑ i=1 σ2i . (11)
In equation 11, the second term of the second last row becomes 0 since the optimization will drive f(X)−C(Z) to zero.

heading: HYPER-PARAMETER SELECTION

heading: DIMENSION OF EMBEDDING SPACE

text:
One important hyper-parameter for the ER-Classifier is the dimension of the embedding space. If the dimension is too small, important features are “collapsed” onto the same dimension, and if the dimension is too large, the projection will not extract useful information, which results in too much noise and instability. The maximum likelihood estimation of intrinsic dimension proposed in
(Levina & Bickel, 2005)3 is used to calculate the intrinsic dimension of each image dataset, serving as a guide for selecting the embedding dimension. The sample size used in calculating the intrinsic dimension is 1, 000, and changing the sample size does not influence the results much. Based on the intrinsic dimension calculated by (Levina & Bickel, 2005), we test several different values around the suggested intrinsic dimension and evaluate the models against l∞-PGD attack. All models are trained without min-max robust optimization, and the experimental results are shown in Figure 5.
The final embedding dimension is chosen based on robustness, number of parameters, and testing accuracy when there is no attack. The final embedding dimensions and suggested intrinsic dimensions are shown in Table 3.
Based on Figure 5, the embedding dimension close to the calculated intrinsic dimension usually offers better results except on MNIST. One explanation may be that MNIST is a simple handwritten digit dataset, so performing classification on MNIST may not require that many dimensions.

heading: EPSILON SELECTION

text:
Epsilon ( ) is an important hyper-parameter for adversarial training. When doing Madry’s adversarial training, we test the model robustness with different and choose the best one. The experiment results are shown in Figure 6.
Based on Figure 6, we use = 0.3, 0.03, 0.03 in Madry’s adversarial training on MNIST, CIFAR10 and STL10 respectively. For Tiny Imagenet, we use = 0.01. To make a fair comparison, we use the same when training ER-Classifier.
3Code publicly available at https://github.com/OFAI/hub-toolbox-python3
EMBEDDING VISUALIZATION
In this section, we compare the embedding learned by Encoder+Classifier structure (E-CLA) and the embedding learned by ER-Classifier on several datasets without min-max robust optimization. We first generate embedding of testing data using the encoder (z̃ = Qφ(x)), then project the embedding points (z̃) to 2-D space by tSNE(Maaten & Hinton, 2008). Then we generate adversarial images (xadv) against E-CLA and ER-Classifier using l∞-PGD attack. The adversarial embedding is generated by feeding the adversarial images into the encoder (z̃adv = Qφ(xadv)). Finally, we project the adversarial embedding points (z̃adv) to 2-D space. The results are shown in Figure 7. The plots in the first and second rows are embedding visualization plots for E-CLA, and the plots in the third and last rows are the embedding visualization plots for ER-Classifier. In adversarial embedding visualization plots, the misclassified point is marked as “down triangle”, which means the PGD attack successfully changed the prediction, and the correctly classified point is marked as “point”, which means the attack fails.
Based on Figure 7, we can see that E-CLA can learn a good embedding on legitimate images of MNIST. Embedding points for different classes are separated on the 2D space, but under adversarial attack, some embedding points of different classes are mixed together. However, ER-Classifier can generate good separated embeddings on both legitimate and adversarial images. On CIFAR10, the E-CLA can not generate good separated embeddings on either legitimate images or adversarial images, while ER-Classifier can generate good separated embeddings for both.

heading: PSEUDO-CODE

text: Code for reproduction will be made available online at github later. The pseudocode for training ER-Classifier is shown in Listing 1.

heading: MODEL STRUCTURE

text:
MNIST, STL10 and TinyImagenet classifier structures used for baseline methods are shown in Figure 8. We use VGG19 for the baseline methods on CIFAR10. Details of ER-Classifier structures on the four benchmark datasets are shown in Figure 9-10.
Original Embedding of E-CLA on MNIST
Adversarial Embedding of E-CLA on MNIST Classified Right Classified Wrong
Original Embedding of E-CLA on CIFAR10
Adversarial Embedding of E-CLA on CIFAR10 Classified Right Classified Wrong
Original Embedding of ER-classifier on MNIST
Adversarial Embedding of ER-classifier on MNIST Classified Right Classified Wrong
Original Embedding of ER-classifier on CIFAR10
Adversarial Embedding of ER-classifier on CIFAR10 Classified Right Classified Wrong
1 d e f t r a i n e r c l a ( t r a i n l o a d e r , t e s t l o a d e r , encoder , d i s c r i m i n a t o r , 2 c l a s s i f i e r , o t h e r h y p e r p a r a m e t e r s ) : 3 c r i t e r i o n = nn . C r o s s E n t r o p y L o s s ( ) 4 e n c o d e r . t r a i n ( ) 5 d i s c r i m i n a t o r . t r a i n ( ) 6 c l a s s i f i e r . t r a i n ( ) 7 # O p t i m i z e r s 8 e n c o p t i m = opt im . Adam( e n c o d e r . p a r a m e t e r s ( ) , l r = l r ) 9 d i s o p t i m = opt im . Adam( d i s c r i m i n a t o r . p a r a m e t e r s ( ) , l r = 0 . 5 ∗ l r )
10 c l a o p t i m = opt im . Adam( c l a s s i f i e r . p a r a m e t e r s ( ) , l r = 0 . 0 5 ∗ l r ) 11 e n c s c h e d u l e r = StepLR ( enc op t im , s t e p s i z e =30 , gamma = 0 . 5 ) 12 d i s s c h e d u l e r = StepLR ( d i s o p t i m , s t e p s i z e =30 , gamma = 0 . 5 ) 13 c l a s c h e d u l e r = StepLR ( c l a o p t i m , s t e p s i z e =30 , gamma = 0 . 5 ) 14 one = t o r c h . Tensor ( [ 1 ] ) 15 mone = one ∗ −1 16 f o r epoch i n r a n g e ( num epoch ) : 17 s t e p = 0 18 f o r images , l a b e l s i n tqdm ( t r a i n l o a d e r ) : 19 e n c o d e r . z e r o g r a d ( ) 20 d i s c r i m i n a t o r . z e r o g r a d ( ) 21 c l a s s i f i e r . z e r o g r a d ( ) 22 # ======== Min−Max Robus t O p t i m i z a t i o n ======== # 23 images = a d v g e t ( images , c l a s s i f i e r , e n c o d e r ) 24 # ======== T r a i n D i s c r i m i n a t o r ======== # 25 f r o z e n p a r a m s ( e n c o d e r ) 26 f r o z e n p a r a m s ( c l a s s i f i e r ) 27 f r e e p a r a m s ( d i s c r i m i n a t o r ) 28 z f a k e = s a m p l e z ( p r i o r , n z , b a t c h s i z e , s igma ) 29 d f a k e = d i s c r i m i n a t o r ( t o v a r ( z f a k e ) ) 30 z r e a l = e n c o d e r ( images ) 31 d r e a l = d i s c r i m i n a t o r ( t o v a r ( z r e a l ) ) 32 d i s c f a k e = LAMBDA ∗ d f a k e . mean ( ) 33 d i s c r e a l = LAMBDA ∗ d r e a l . mean ( ) 34 d i s c f a k e . backward ( one ) 35 d i s c r e a l . backward ( mone ) 36 d i s s l o s s = d i s c f a k e − d i s c r e a l 37 d i s o p t i m . s t e p ( ) 38 c l i p p a r a m s ( d i s c r i m i n a t o r ) 39 # ======== T r a i n C l a s s i f i e r and Encoder ======== # 40 f r e e p a r a m s ( e n c o d e r ) 41 f r e e p a r a m s ( c l a s s i f i e r ) 42 f r o z e n p a r a m s ( d i s c r i m i n a t o r ) 43 p r e d l a b e l s = c l a s s i f i e r ( e n c o d e r ( t o v a r ( images ) ) ) 44 c l a s s l o s s = LAMBDA0 ∗ c r i t e r i o n ( p r e d l a b e l s , l a b e l s ) 45 c l a s s l o s s . backward ( ) 46 c l a o p t i m . s t e p ( ) 47 e n c o p t i m . s t e p ( ) 48 # ======== T r a i n Encoder ======== # 49 f r e e p a r a m s ( e n c o d e r ) 50 f r o z e n p a r a m s ( c l a s s i f i e r ) 51 f r o z e n p a r a m s ( d i s c r i m i n a t o r ) 52 z r e a l = e n c o d e r ( images ) 53 d r e a l = d i s c r i m i n a t o r ( e n c o d e r ( V a r i a b l e ( images . d a t a ) ) ) 54 d l o s s = LAMBDA1 ∗ ( d r e a l . mean ( ) ) 55 d l o s s . backward ( one ) 56 e n c o p t i m . s t e p ( ) 57 s t e p += 1 58 s a v e f i l e ( f i l e n a m e , encoder , d i s c r i m i n a t o r , c l a s s i f i e r , d a t a s e t = d a t a s e t ) 59 r e t u r n c l a s s i f i e r , e n c o d e r
Listing 1: Pseudocode for training ER-Classifier