summaries: This paper introduces a optimisation for BERT models based on using block matrices for the attention layers. This allows to reduce the memory footprint and the processing  time during training while reaching state-of-the-art results on 5 datasets. An interesting study on memory consumption in BERT is conducted. No results are given at test time : is there also a memory and processing time reduction ?

summaries: The paper propose to sparsify the attention matrix to decrease memory usage and to speed up training. The authors experiment the model on multiple tasks. The model gains ~20% efficiency with ~20% decrease in memory use while maintaining comparable performance to the state of the art model. To keep the performance comparable, the authors use the same training corpus. The authors also discuss how block size could change the performance of the model. The paper is clear and well organized with good experiment results.

abstractText: We present BlockBERT, a lightweight and efficient BERT model that is designed to better modeling long-distance dependencies. Our model extends BERT by introducing sparse block structures into the attention matrix to reduce both memory consumption and training time, which also enables attention heads to capture either shortor long-range contextual information. We conduct experiments on several benchmark question answering datasets with various paragraph lengths. Results show that BlockBERT uses 18.7-36.1% less memory and reduces the training time by 12.0-25.1%, while having comparable and sometimes better prediction accuracy, compared to an advanced BERT-based model, RoBERTa.

authors:
Tianqi Chen
Bing Xu
Chiyuan Zhang
Carlos Guestrin

title: Training deep nets with sublinear memory cost

venue: arXiv preprint arXiv:1604.06174,

authors:
Rewon Child
Scott Gray
Alec Radford
Ilya Sutskever

title: Generating long sequences with sparse transformers

venue: arXiv preprint arXiv:1904.10509,

authors:
Gonçalo M Correia
Vlad Niculae
André FT Martins

title: Adaptively sparse transformers

venue: arXiv preprint arXiv:1909.00015,

authors:
Zihang Dai
Zhilin Yang
Yiming Yang
William W Cohen
Jaime Carbonell
Quoc V Le
Ruslan Salakhutdinov

title: Transformer-xl: Attentive language models beyond a fixed-length context

authors:
Jacob Devlin
Ming-Wei Chang
Kenton Lee
Kristina Toutanova

title: Bert: Pre-training of deep bidirectional transformers for language understanding

venue: In NAACL-HLT’

authors:
Matthew Dunn
Levent Sagun
Mike Higgins
V Ugur Guney
Volkan Cirik
Kyunghyun Cho

title: Searchqa: A new q&a dataset augmented with context from a search engine

venue: arXiv preprint arXiv:1704.05179,

authors:
Jonathan Frankle
Michael Carbin

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

venue: arXiv preprint arXiv:1803.03635,

authors:
Scott Gray
Alec Radford
Diederik P Kingma

title: Gpu kernels for block-sparse weights

venue: arXiv preprint arXiv:1711.09224,

authors:
Qipeng Guo
Xipeng Qiu
Pengfei Liu
Yunfan Shao
Xiangyang Xue
Zheng Zhang

title: Startransformer. In NAACL-HLT

authors:
Geoffrey Hinton
Oriol Vinyals
Jeff Dean

title: Distilling the knowledge in a neural network

venue: arXiv preprint arXiv:1503.02531,

authors:
Sepp Hochreiter
Jürgen Schmidhuber

title: Long short-term memory

venue: Neural computation,

authors:
Yanping Huang
Yonglong Cheng
Dehao Chen
HyoukJoong Lee
Jiquan Ngiam
Quoc V Le
Zhifeng Chen

title: Gpipe: Efficient training of giant neural networks using pipeline parallelism

venue: arXiv preprint arXiv:1811.06965,

authors:
Mandar Joshi
Eunsol Choi
Daniel Weld
Luke Zettlemoyer

title: Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension

venue: In ACL’

authors:
Mandar Joshi
Danqi Chen
Yinhan Liu
Daniel S Weld
Luke Zettlemoyer
Omer Levy

title: Spanbert: Improving pre-training by representing and predicting spans

venue: arXiv preprint arXiv:1907.10529,

authors:
Mandar Joshi
Omer Levy
Daniel S Weld
Luke Zettlemoyer

title: Bert for coreference resolution: Baselines and analysis

venue: arXiv preprint arXiv:1908.09091,

venue: arXiv preprint arXiv:1412.6980,

authors:
Tom Kwiatkowski
Jennimaria Palomaki
Olivia Redfield
Michael Collins
Ankur Parikh
Chris Alberti
Danielle Epstein
Illia Polosukhin
Jacob Devlin
Kenton Lee

title: Natural questions: a benchmark for question answering research

venue: Transactions of the Association for Computational Linguistics,

authors:
Kenton Lee
Ming-Wei Chang
Kristina Toutanova

title: Latent retrieval for weakly supervised open domain question answering

venue: arXiv preprint arXiv:1906.00300,

authors:
Shiyang Li
Xiaoyong Jin
Yao Xuan
Xiyou Zhou
Wenhu Chen
Yu-Xiang Wang
Xifeng Yan

title: Enhancing the locality and breaking the memory bottleneck of transformer on time series forecasting

authors:
Yinhan Liu
Myle Ott
Naman Goyal
Jingfei Du
Mandar Joshi
Danqi Chen
Omer Levy
Mike Lewis
Luke Zettlemoyer
Veselin Stoyanov

title: Roberta: A robustly optimized bert pretraining approach

authors:
Paul Michel
Omer Levy
Graham Neubig

title: Are sixteen heads really better than one

venue: arXiv preprint arXiv:1905.10650,

authors:
Paulius Micikevicius
Sharan Narang
Jonah Alben
Gregory Diamos
Erich Elsen
David Garcia
Boris Ginsburg
Michael Houston
Oleksii Kuchaiev
Ganesh Venkatesh

title: Mixed precision training

venue: arXiv preprint arXiv:1710.03740,

authors:
Lesly Miculicich
Dhananjay Ram
Nikolaos Pappas
James Henderson

title: Document-level neural machine translation with hierarchical attention networks

venue: In EMNLP’

authors:
Rodrigo Nogueira
Kyunghyun Cho

title: Passage re-ranking with bert

venue: arXiv preprint arXiv:1901.04085,

authors:
Myle Ott
Sergey Edunov
Alexei Baevski
Angela Fan
Sam Gross
Nathan Ng
David Grangier
Michael Auli

title: fairseq: A fast, extensible toolkit for sequence modeling

authors:
Matthew Peters
Mark Neumann
Mohit Iyyer
Matt Gardner
Christopher Clark
Kenton Lee
Luke Zettlemoyer

title: Deep contextualized word representations

venue: In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies,

authors:
Alec Radford
Jeff Wu
Rewon Child
David Luan
Dario Amodei
Ilya Sutskever

title: Language models are unsupervised multitask learners

authors:
Pranav Rajpurkar
Robin Jia
Percy Liang

title: Know what you don’t know: Unanswerable questions for squad

venue: arXiv preprint arXiv:1806.03822,

authors:
Alexander Rives
Siddharth Goyal
Joshua Meier
Demi Guo
Myle Ott
C Lawrence Zitnick
Jerry Ma
Rob Fergus

title: Biological structure and function emerge from scaling unsupervised learning to 250 million protein

venue: sequences. bioRxiv,

authors:
Livio Baldini Soares
Nicholas FitzGerald
Jeffrey Ling
Tom Kwiatkowski

title: Matching the blanks: Distributional similarity for relation learning

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

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

authors:
Sainbayar Sukhbaatar
Edouard Grave
Piotr Bojanowski
Armand Joulin

title: Adaptive attention span in transformers

venue: arXiv preprint arXiv:1905.07799,

authors:
Raphael Tang
Yao Lu
Linqing Liu
Lili Mou
Olga Vechtomova
Jimmy Lin

title: Distilling taskspecific knowledge from bert into simple neural networks

authors:
Adam Trischler
Tong Wang
Xingdi Yuan
Justin Harris
Alessandro Sordoni
Philip Bachman
Kaheer Suleman

title: Newsqa: A machine comprehension dataset

venue: In Proceedings of the 2nd Workshop on Representation Learning for NLP,

authors:
Ashish Vaswani
Noam Shazeer
Niki Parmar
Jakob Uszkoreit
Llion Jones
Aidan N Gomez
Łukasz Kaiser
Illia Polosukhin

title: Attention is all you need

venue: In Advances in neural information processing systems,

authors:
Elena Voita
David Talbot
Fedor Moiseev
Rico Sennrich
Ivan Titov

title: Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned

authors:
Zhilin Yang
Peng Qi
Saizheng Zhang
Yoshua Bengio
William W. Cohen
Ruslan Salakhutdinov
Christopher D. Manning

title: HotpotQA: A dataset for diverse, explainable multi-hop question answering

authors:
Zhilin Yang
Zihang Dai
Yiming Yang
Jaime Carbonell
Ruslan Salakhutdinov
Quoc V Le

title: Xlnet: Generalized autoregressive pretraining for language understanding

heading: 1 INTRODUCTION

text:
Recent emergence of the pre-training and fine-tuning paradigm, exemplified by methods like ELMo (Peters et al., 2018), GPT-2 (Radford et al., 2019), BERT (Devlin et al., 2019), XLNet (Yang et al., 2019) and RoBERTa (Liu et al., 2019), has drastically reshaped the landscape of the natural language processing research. These methods first pre-train a deep model with language model objectives using a large corpus and then fine-tune the model using in-domain supervised data for target applications. Despite its conceptual simplicity, this paradigm has reestablished the new state-ofthe-art baselines across various tasks, such as question answering (Devlin et al., 2019), coreference resolution (Joshi et al., 2019b), relation extraction (Soares et al., 2019) and text retrieval (Lee et al., 2019; Nogueira & Cho, 2019), to name a few.
Building such models in practice, however, is an extremely resource-intensive process. For instance, the training of BERT-family models is notoriously expensive. Devlin et al. (2019) report that it takes four days for pre-training BERT-Base/BERT-Large on 4/16 Cloud TPUs, respectively. In order to reduce the pre-training time of RoBERTa to 1 day, Liu et al. (2019) use 1,024 V100 GPUs. One crucial factor that contributes to the long training time is the memory consumption of these deep models, as it directly affects the batch size. Although the fine-tuning stage is relatively inexpensive, the memory issue still restricts the scenarios in which BERT can be used. For instance, “it is currently not possible to re-produce most of the BERT-Large results on the paper using a GPU with 12GB-16GB of RAM, because the maximum batch size that can fit in memory is too small.1”
Although one may think that model size is the main contributor to the large memory consumption, our analysis (Section 2.1) shows that one of the main bottlenecks is actually dot-product selfattention, operated in multiple layers of Transformers (Vaswani et al., 2017), the building block of BERT. As the attention operation is quadratic to the sequence length, this fundamentally limits the maximum length of the input sequence, and thus restricts the model capacity in terms of capturing long-distance dependencies. As a result, downstream tasks have to either truncate their sequences to leading tokens (Nogueira & Cho, 2019) or split their sequences with a sliding window (Joshi et al., 2019a;b). Ad-hoc handling of long sequences is also required in the pre-training stage, such as updating the model using only short sequences in the early stage (Devlin et al., 2019).
Common strategies for reducing memory consumption, unfortunately, do not work. For instance, shrinking the model by lowering the number of layers L, attention heads A, or hidden units H leads to significant performance degradation (Vaswani et al., 2017; Devlin et al., 2019) and does not address the long sequence issue. Alternatively, general low-memory training techniques, such as
1https://github.com/google-research/bert
microbatching (Huang et al., 2018) and gradient checkpointing (Chen et al., 2016) essentially trade off training time for memory consumption, prolongs the already lengthy training process.
In this work, we explore a different strategy, sparsifying the attention layers, intending to design a lightweight and effective BERT that can model long sequences in a memory-efficient way. Our BlockBERT extends BERT by introducing sparse block substructures into the attention matrix to reduce both memory consumption and the number of floating point operations (FLOPs), which also enables attention heads to capture either short- or long-range contextual information. Compared to the previous method that also enforces sparsity (e.g., Child et al. (2019)), our approach is much simpler mathematically and very easy to implement. More importantly, the results of experiments conducted on several benchmark question answering datasets with various paragraph lengths show that BlockBERT performs comparably or even better than the original BERT-family models, while enjoying an 18.7-36.1% reduction in memory usage and 12.0-25.1% reduction in training time.
The rest of the paper is organized as follows. Section 2 gives a brief introduction of the BERT model, along with an in-depth analysis of its memory usage during training time. We describe our proposed model in Section 3 and contrast it with existing methods that aim for creating a lighter model. Section 4 presents the experimental results and ablation studies, followed by a short survey of other related work in Section 5 and the conclusion in Section 6.

heading: 2 BACKGROUND: MEMORY BOTTLENECK IN TRAINING BERT

text: We briefly review BERT and introduce its memory profiling in this section. Following the paradigm of language model pre-training and down-stream task fine-tuning, BERT (Devlin et al., 2019) consists of multiple layers of bidirectional Transformers (Vaswani et al., 2017), where each Transformer encoder has a multi-head self-attention layer and a position-wise feed-forward layer. Using the same notation as in (Devlin et al., 2019), we denote the number of Transformer layers by L, the number of hidden units by H , the number of attention heads by A, the sequence length by N and the batch size by B. We also assume the feed-forward hidden unit size to be 4H .2

heading: 2.1 MEMORY PROFILING

text:
Training BERT is a memory-intensive process. In order to identify the bottleneck, we follow the memory model proposed by Sohoni et al. (2019), where the memory usage throughout neural network training is categorized into three main types: (1) Model Memory is used to store model parameters; (2) Optimizer Memory is the additional memory used by the specific learning algorithm during the process; (3) Activation Memory consists of the outputs of each layer, which are cached for reuse in backpropagation to compute gradients.
Take BERT-Base training as an example. The model has 110M parameters, so the model memory uses 0.2 GB if stored in half-precision floating-point format (FP16). For Adam (Kingma & Ba, 2014), the optimizer needs additional memory to store the gradients, first moments, and second moments of model parameters. If stored using the same precision, the optimizer memory should be three times of model memory.3 To calculate the exact size of activation memory is not trivial because it depends heavily on the implementation of the toolkit. Instead, we measure it empirically by training BERT-Base using Adam with a memory profiler (more details are provided in Appendix A.2).
We use 32 NVIDIA V100 GPUs for training. Each single GPU thus consumes a mini-batch of size b = B/32 = 8. Figure 1a shows the profiling result for a single GPU, where the model/optimizer/activation memory consumes 0.21/1.03/8.49 GB, resp. We can see that activation memory accounts for the vast majority of the total GPU memory (87.6%) and is clearly the bottleneck. Notice that although our analysis is done on BERT-Base, it can be easily generalized to BERT-Large and other models such as RoBERTa (Liu et al., 2019) and XLNet (Yang et al., 2019).
2The default parameter settings for BERT-Base and BERT-Large can be found in Table 5 in Appendix A.1. 3In the current PyTorch Adam implementation, the first and second moments are stored in single precision.
Consequently, BERT’s optimizer memory (1 GB) is five times of model memory (0.2 GB).

heading: 2.2 A REGRESSION ANALYSIS ON ACTIVATION MEMORY

text: For BERT, or more specifically, Transformer, the activation memory corresponds to intermediate results of different layers It grows linearly in all the model hyper-parameters, except the sequence length N , due to the attention layers. To quantify more clearly the O(N) and O(N2) components in the activation memory, we conduct a regression analysis as follows. Assume that the activation memory (in each GPU) is a polynomial a2bN2 + a1bN + a0, where b is the batch size in each GPU. If we fix the total number of tokens in a GPU, i.e., b×N , to be constant (in our case, 4096), we should have a linear function w.r.t. N , i.e., 4096a2N + 4096a1 + a0. We enumerate N from {128, 256, 512, 1024} in our experiments, and plot the corresponding profiled activation memory in Figure 1b. Using ordinary least squares (OLS), with b × N = 4096, the estimated linear function for activation memory is 0.00715 × N + 4.83, where the first term is responsible for the O(N2) component. When N = 512, we can see that for BERT-Base, the O(N2) component accounts for 3.66 GB and O(N) accounts for 4.83 GB. When the sequence length N increases to 1024, however, the O(N2) component increases to 7.32 GB, while O(N) is unchanged.

heading: 2.3 GENERAL TECHNIQUES FOR REDUCING MEMORY USAGE IN MODEL TRAINING

text:
Observing that activation memory is the bottleneck, we discuss the effectiveness of common memory reduction techniques for BERT training below.
Low Precision (Micikevicius et al., 2017): Low precision is to use half-precision or mixed-precision for training neural networks. This technique has been widely used in Transformer training (Ott et al., 2019; Liu et al., 2019). In this work, we already assume to use mixed-precision training by default, as indicated in the aforementioned analysis.
Microbatching (Huang et al., 2018): Microbatching is to split a batch into small microbatches (which can be fit into memory), and then run forward and backward passes on them separately with gradients for each micro-batch accumulated. Because it runs forward/backward pass multiple times for a single batch, it trades off time for memory.
Gradient Checkpointing (Chen et al., 2016): Gradient checkpointing saves memory by only caching activations of a subset of layers. The un-cached activations will be recomputed during backpropagation from the latest checkpoint. This strategy trades off time for memory by repeating computations that require large memory and will obviously extend the model training time.
Knowledge Distillation (Hinton et al., 2015): Knowledge distillation aims to compress and transfer knowledge from a teacher model to a simpler student model. However, knowledge distillation relies on a teacher model (which is still expensive in training time) and usually suffers from a certain degree of performance degradation.
As common techniques are limited in reducing both the training time and memory usage, we investigate how to optimize the dot-product attention layers and introduce our approach next.

heading: 3 MODEL: BLOCKBERT

text:
Following (Vaswani et al., 2017), the dot-product attention in Transformer is defined as: Attention(Q,K,V ) = softmax ( QK>√
d
) V .
where Q,K,V ∈ RN×d with N to be the sequence length and d to be a hidden dimension. As we can see, the inner product between Q and K consumes O(N2) memory. One simple way to reduce memory consumption of attention is to sparsify the attention matrix. Suppose we have a masking matrix M ∈ {0, 1}N×N . We define a masked version of attention as follows:
Attention(Q,K,V ,M) = softmax ( QK>√
d M
) V , (1)
with operator defined by
(A M)ij = { Aij if Mij = 1 −∞ if Mij = 0 .
In this work, we design M to be a sparse block matrix, which not only reduces memory and the number of floating point operations (FLOPs) but also benefits from efficient dense matrix support from deep learning frameworks, such as PyTorch and Tensorflow. More formally, we split the length-N input sequence into n partitions, with each partition of length Nn .
4 The N × N attention matrix is then partitioned into n × n blocks, where each block matrix is of size Nn × N n . A sparse block matrix M can be defined by a permutation π of {1, 2, · · · , n}:
Mij =
{ 1 if π ( b (i−1)n
N + 1c
) = b (j−1)n
N + 1c
0 otherwise. (2)
By writing Q,K,V as be block matrices, such that Q = [Q>1 · · · Q>n ] > ,K = [K>1 · · · K>n ] > and V = [V >1 · · · V >n ] > and pluging them into Equation 1, we can formally define Blockwise Attention as follows:
Blockwise-Attention(Q,K,V ,M) =  softmax ( Q1K > π(1)√ d ) Vπ(1) ...
softmax
( QnK
> π(n)√ d ) Vπ(n)
 . (3)
As a result, it only needs to compute and store QiK>π(i) (i = 1, · · ·n), each of which has size N n × N n . In other words, BlockBERT reduces the corresponding O(N
2) memory consumption and FLOPs by a factor of n, since Nn × N n × n = N×N n .

heading: 3.1 BLOCKWISE MULTI-HEAD ATTENTION

text:
Analogous to Multi-head Attention (Vaswani et al., 2017), we allow queries, keys, and values to be projected multiple times and perform blockwise attentions in parallel. Moreover, different blockwise attention heads can use different masking matrices. The outputs of multiple heads are then concatenated and aggregated with another linear projection. Let A be the number of attention heads and H the number of hidden units. Blockwise multi-head attention is formally defined by:
Blockwise-Multi-head-Attention(Q,K,V ) = Concat(head1, · · · headA)WO
where headi = Blockwise-Attention(QWQi ,KW K i ,V W V i ,Mi),
with d = HA ,W Q i ,W K i ,W V i ∈ RH×d and the projection matrix WO ∈ RH×H . Each masking matrix Mi is determined by permutation πi according to Equation 2. In particular, we choose π from permutations generated by shifting one position: σ = (2, 3, · · · , n, 1), i.e., we select π ∈ {σ, σ2, · · · , σn}. For example, with 12 attention heads (A = 12) and 2 blocks (n = 2), one
4We assume N can be divided by n. If not, we pad the input sequence to make N divisible.
configuration can be assigning 10 heads to permutation (1, 2) and the other 2 heads to permutation (2, 1). Figure 2 illustrates the blockwise multi-head attention with the block numbers n ∈ {2, 3}. Blockwise sparsity captures both local and long-distance dependencies in a memory-efficiency way, which is crucial for long-document understanding tasks. For instance, the identity permutation, i.e., (1, 2, · · · , n), enables each token to attend its nearby tokens in self-attention. Tokens within the same local group attend a long-distance group of tokens together in other permutations. Our proposed BlockBERT essentially replaces the multi-head attention layers in Transformer/BERT with blockwise multi-head attention.

heading: 3.2 ANALYSIS OF MEMORY USAGE REDUCTION

text: To validate our claim that BlockBERT with n×n blocks can reduce theO(N2) memory use by a factor of n, we perform the same memory profiling as described in sections 2.1 and 2.2. Again, We fix the number of tokens in each GPU (b×N = 4096) and chooseN from {128, 256, 512, 1024, 2048}.5 As we can see from Figure 3 and Table 1, the empirical results align well with the theoretical values. When we set block size to be 2 and 3 for BlockBERT, their estimated O(N2) activation memory decreases to 1/2 and 1/3 of BERT’s O(N2) activation memory, resp. As shown in Table 2, for the sequence length N = 512, BlockBERT with 2 / 3 blocks saves 18.7% / 23.8% overall memory, resp. The saving is more significant for longer sequences. When N = 1024, the overall memory reduction of BlockBERT with 2 / 3 blocks is 27.3% / 36.1%, resp.

heading: 4 EXPERIMENTS

text:
We evaluate the pre-training and fine-tuning performance of BlockBERT. In particular, when n = 2, we denote 10:2 to be the configuration which distributes 10 heads to permutation (1, 2) and 2 to permutation (2, 1); when n = 3, we denote 8:2:2 to be the configuration which assigns 8, 2, 2 heads
5We use GPUs of 16 GB memory for profiling. BERT with N = 2048 fails due to an out-of-memory error.
to permutation (1, 2, 3), (2, 3, 1), and (3, 1, 2), resp. We compare BlockBERT with the following baselines:
Google BERT The pre-trained base model from Devlin et al. (2019).
RoBERTa-2seq and RoBERTa-1seq We compare with two versions of RoBERTa (Liu et al., 2019). RoBERTa-2seq is trained with both masked language model (MLM) task and next sentence prediction (NSP) task, while RoBERTa-1seq refers to the pre-training model with only MLM task.
SparseBERT We pre-train BERT models with its Transformer encoder replaced by a Sparse Transformer encoder (Child et al., 2019). We set its sparsity hyper-parameters stride ` = 128 and expressivity c = 32. The attention masks used for Sparse Transformer encoder are illustrated in Figure 5.

heading: 4.1 PRE-TRAINING

text: All the models follow the base setting, i.e., L = 12, H = 768, A = 12 and are trained on the same corpus — BooksCorpus and English Wikipedia with uncased word piece tokens. We fix the number of tokens per batch B ×N = 131, 072, i.e., if sequence length N = 512 then batch size B = 256, if sequence length N = 1024 then batch size B = 128. The detailed pre-training configuration is listed in Table 6 in Appendix A.1. Moreover, the pre-training of SparseBERT and BlockBERT follows the RoBERTa-1seq setting, i.e., we drop the NSP (Next Sentence Prediction) task, and an input sequence is up to N tokens until it reaches a document boundary. A summary of the pretraining performance comparison between BlockBERT and RoBERTa-1seq is shown in Table 2. Besides memory saving, we also achieve a significant speedup. For example, when N = 1024, BlockBERT (n = 2) reduces the training time from RoBERTa’s 9.7 days to 7.5 days.

heading: 4.2 FINE-TUNING TASKS

text:
We evaluate BlockBERT on several question answering tasks, including SQuAD 1.1/2.0 (Rajpurkar et al., 2018) and five other tasks from the MrQA shared task6 — HotpotQA (Yang et al., 2018), NewsQA (Trischler et al., 2017), SearchQA (Dunn et al., 2017), TriviaQA (Joshi et al., 2017) and NaturalQA (Kwiatkowski et al., 2019). Since MrQA does not have an official test set, we follow Joshi et al. (2019a) who split the development set evenly to build a new development set and test set.
These QA datasets have different paragraph length distribution patterns and are thus ideal for testing the effectiveness of BlockBERT. For example, SQuAD, NaturalQA, and HotpotQA consist of mostly short paragraphs (shorter than 512), while paragraphs in SearchQA (average length 1,004) and TriviaQA (average length 934) have around 1,000 tokens. This means that for SearchQA and TriviaQA, a BERT model with sequence length N = 512 can only capture half of the context. The detailed paragraph length distributions can be found in Figure 6.
For all the pre-trained models, we adopt the same fine-tuning QA setup from Devlin et al. (2019). The tokenized paragraph (p1, · · · , ps) and question (q1, · · · , qt) are concatenated to be a sequence [CLS]q1 · · · qt[SEP]p1 · · · ps[SEP]. The sequence is then fed into the pre-trained model with two extra linear layers for predicting the start and end positions of the answer spans. The detailed fine-tuning setting is listed in Appendix A.4. Table 3 and Table 4 report the experimental results.
6https://mrqa.github.io
BlockBERT (n=2) v.s. RoBERTa-1seq Comparing BlockBERT (n = 2) with RoBERTa-1seq on pre-trained model withN = 512, we observe an absolute F1 difference from 0.04 (in NaturalQA) to 1.18 (in NewsQA), with average difference to be 0.55. For N = 1024, BlockBERT achieves more comparable or even better performance (in SearchQA, NewsQA, and HotpotQA) to RoBERTa-1seq. The average difference on F1 reduces to 0.27.
BlockBERT v.s. SparseBERT ForN = 512, it is interesting that BlockBERT with 3 blocks (density 33.33%) performs better then SparseBERT (density 44.20%) in both SQuAD and MrQA tasks. Similar patterns can be observed for N = 1024. These results show that off-diagonal masking matrices, e.g., the masking matrix defined by permutation (2, 1), play crucial roles in BlockBERT.
Effect of Long Sequence Pre-training Our observations are twofold. (1) Long sequence pretraining benefits long sequence fine-tuning. In TriviaQA and SearchQA, of which paragraph lengths are around 1024, pre-training models with N = 1024 achieve significantly better performance. (2) The heterogeneity of pre-training and fine-tuning sequence length may hurt performance. For example, in SQuAD, we do not see significant performance gain by using pre-trained models with N = 1024; in HotpotQA and NewsQA, longer sequence pre-training even hurts performance.
Effect of #Blocks It is not surprising that BlockBERT with 2 blocks (n = 2) performs better than that with 3 blocks (n = 3), because it keeps more attention matrix entries. The biggest difference is in SQuAD 2.0 and NewsQA with N = 1024, where we observe an absolute loss of 1.6 F1 by increasing block number from 2 to 3.
In summary, not only BlockBERT saves training time and memory, but it also has competitive and sometimes better performance, especially for tasks with longer sequences. This demonstrates the effectiveness of our blockwise multi-head attention approach.

heading: 4.3 ABLATION STUDY

text:
We fix the assignment of attention heads in above experiments. For example, BlockBERT with sequence length N = 512 and 2 blocks is trained with ten heads using permutation (1, 2) and the other two using permutation (2, 1). However, we know that there are other ways to partition twelve attention heads, e.g., seven heads for permutation (1, 2) and the other five for permutation (2, 1). It would be interesting to see how the assignment of heads affects model performance. In this section, we grid search attention head assignments and plot their best validation performance in 1.2M training steps. The results are shown in Figure 4.
Our observations are threefold: (1) Identity permutations, i.e., (1, 2) and (1, 2, 3), are important. As shown in Figure 4, all optimal solutions assign considerable attention heads to block diagonal matrices, since those matrices enable each token to attend to its nearby tokens; (2) Non-identity permutations follow the rule of “vital few and trivial many.” Although identity permutations are important, assigning all attention heads to them (corresponding to 12:0 and 12:0:0 in Figure 4) significantly hurts performance, since the model can not learn long-term dependencies with only identity permutation; (3) Pre-training performance and fine-tuning performance are correlated but not always consistent. When n = 3, pre-training performance suggests 10:1:1 to be the best head assignment — ten heads for permutation (1, 2, 3), one head for (2, 3, 1) and one head for (3, 1, 2), but we observe that the configuration of 8:2:2 achieves better performance in fine-tuning tasks.

heading: 5 RELATED WORK

text:
In this section, we review the related work of memory optimization for neural network training and recent efforts to simplify Transformer and BERT. In recent years, there is an increasing interest in training neural networks with low-memory (Sohoni et al., 2019). Mainstream techniques include low-precision training (Micikevicius et al., 2017), microbatching (Huang et al., 2018), gradient checkpointing (Chen et al., 2016). Another line of researches studies this problem from a theoretical perspective, including the recently proposed lottery ticket hypothesis (Frankle & Carbin, 2018). Since the invention of Transformer (Vaswani et al., 2017; Dai et al., 2019) and its successful application on language model pre-training (Devlin et al., 2019; Radford et al., 2019; Yang et al., 2019; Liu et al., 2019), there have been several studies attempted to simplify it from different perspectives. The first line of research focuses on attention matrix sparsification, such as Star Transformer (Guo et al., 2019), Sparse Transformer (Child et al., 2019), Adaptive Sparse Transformer (Correia et al., 2019; Sukhbaatar et al., 2019), Log-Sparse Transformer (Li et al., 2019), etc. However, due to limited support for sparse tensor from current deep learning platforms, most of studies have to represent a sparse matrix using a dense matrix with a binary mask or rely on customized CUDA kernels (Gray et al., 2017). The second line of research attempts to prune redundant heads in Transformer, such as Voita et al. (2019) and Michel et al. (2019). The third line of research focuses on knowledge distillation, including DistilBERT7 which distills BERT using a smaller BERT and Tang et al. (2019) which distills BERT with BiLSTM (Hochreiter & Schmidhuber, 1997).
7https://github.com/huggingface/pytorch-transformers/tree/master/examples/distillation

heading: 6 CONCLUSION

text: In this work, we study lightweight BERT model with the goal of achieving both efficiency and effectiveness. We profile and analyze the memory bottlenecks of BERT, and focus on optimize dotproduct self-attention, which consumes quadratic memory with respect to the sequence length. To reduce both training time and memory consumption, we present BlockBERT, which sparsifies the attention matrices to be sparse block matrices. The proposed model achieves time and memory saving without significant loss of performance. In the future, we would like to explore more applications of BlockBERT on NLP tasks involving long sequences such as coreference resolution (Joshi et al., 2019b) and document-level machine translation (Miculicich et al., 2018), and also non-NLP tasks such as protein sequence modeling (Rives et al., 2019).

heading: A APPENDIX

text:
A.1 NOTATIONS AND PRE-TRAINING HYPER-PARAMETERS
The notations and pre-training hyper-parameters are listed in Table 5 and Table 6.
Hyper-parameter Value
Vocabulary Size 30,522 Dropout 0.1 Attention dropout 0.1 Warmup steps 10K Weight decay 0.01 Max steps 2.4M Initial learning rate 0.00025 Learning rate decay Linear Adam 1e-8 Adam β1 0.9 Adam β2 0.999 Gradient Clipping 1.0
Table 6: Pre-training hyperparameters.
A.2 PROFILER IMPLEMENTATION
Among the three types of training memory, model memory and optimizer memory is relatively easy to profile (can be computed by enumerate each tenor and summing up tensor.numel() * tensor.element size()). To calculate activation memory, Sohoni et al. (2019) traverse PyTorch’s autograd graph and sum up necessary storage space. They find that the summation of model memory, optimizer memory and activation memory matches PyTorch memory profiling tool torch.cuda.max memory allocated. Based on their observation, we use
torch.cuda.max memory allocated−model memory− optimizer memory (4)
as an estimate to activation memory. When profiling BERT, we first pre-train it for 1000 steps, and then compute its model and optimizer memory. Finally, we esitmate its activation memory according to Equation 4.
A.3 SPARSEBERT
The sparse masking matrices we use for Sparse Transformer (Child et al., 2019) are shown in Figure 5. We adopt the implementation from Fairseq8.
A.4 FINE-TUNING SETTINGS
Our fine-tuning is implemented based on code base from HuggingFace9 and SpanBERT (Joshi et al., 2019a). We use max sequence length=N , i.e., we allow fine-tuning task to input sequences as long as the pre-training model. If the input sequence is too long to fit the max sequence length=N constraints, we use a sliding window of stride 128 to split it. We grid search learning rate from {5e-6, 1e-5, 2e-5, 3e-5, 5e-5} and batch size from {16, 32}. The fine-tuning is performed for 4 epoches.
A.5 TEST EFFICIENCY
We benchmark test efficiency of RoBERTa and our proposed BlockBERT. The benchmark code follows huggingface10. All experiments are run 30 times on a 32GB V100 GPU with half precision (FP16). We report the average running time at Table 7. As we can see, BlockBERT does achieve speedup and memory reduction during test time. Take 8×1024, i.e., batch size B = 8, sequence length N = 1024, as an example, we can see that BlockBERT with 2 blocks saves 27.8% of test time, and BlockBERT with 3 blocks saves more (30.4%). As for memory, we can observe that RoBERTa can not handle an input of 16×1024, while it is possible for BlockBERT to work on it.
8https://github.com/pytorch/fairseq/blob/master/fairseq/modules/sparse multihead attention.py. 9https://github.com/huggingface/pytorch-transformers
10https://github.com/huggingface/transformers/blob/master/examples/benchmarks.py


summaries: This paper mainly studies the relationship between the generalization error and mean/variance of the test accuracy. The authors first propose a new score for pruning called E[BN]. Then, the authors observe the generalization error and the test accuracy mean/variance for pruning large score weights and small score weights for VGG11, ResNet18, and Conv4 models. From these experiments, the authors observe that pruning large score weights generates instable but high test accuracy and smaller generalization gap compared to pruning small score weights. The authors additionally study some other aspects of pruning (e.g., pruning as a noise injection) and conclude the paper.

summaries: The paper is an empirical study that looks into the effect of neural network pruning on both the model accuracy as well as the generalization risk (defined as the difference between the training error and the test error). It concludes that while some pruning methods work, others fail. The authors argue that such discrepancy can be explained if we look into the impact of pruning on "stability".

abstractText: Pruning neural network parameters is often viewed as a means to compress models, but pruning has also been motivated by the desire to prevent overfitting. This motivation is particularly relevant given the perhaps surprising observation that a wide variety of pruning approaches increase test accuracy despite sometimes massive reductions in parameter counts. To better understand this phenomenon, we analyze the behavior of pruning over the course of training, finding that pruning’s effect on generalization relies more on the instability it generates (defined as the drops in test accuracy immediately following pruning) than on the final size of the pruned model. We demonstrate that even the pruning of unimportant parameters can lead to such instability, and show similarities between pruning and regularizing by injecting noise, suggesting a mechanism for pruning-based generalization improvements that is compatible with the strong generalization recently observed in over-parameterized networks.

authors:
Laurent Dinh
Razvan Pascanu
Samy Bengio
Yoshua Bengio

title: Sharp minima can generalize for deep nets

venue: arXiv preprint arXiv:1703.04933,

authors:
Gintare Karolina Dziugaite
Daniel M Roy

title: Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data

venue: arXiv preprint arXiv:1703.11008,

authors:
Jonathan Frankle
Michael Carbin

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

venue: arXiv preprint arXiv:1803.03635,

authors:
Trevor Gale
Erich Elsen
Sara Hooker

title: The state of sparsity in deep neural networks

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:
Song Han
Jeff Pool
John Tran
William Dally

title: Learning both weights and connections for efficient neural network

authors:
Babak Hassibi
David G Stork

title: Second order derivatives for network pruning: Optimal brain surgeon

authors:
Kaiming He
Xiangyu Zhang
Shaoqing Ren
Jian Sun

title: Deep residual learning for image recognition

venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,

authors:
Geoffrey E Hinton
Drew Van Camp

title: Keeping the neural networks simple by minimizing the description length of the weights

venue: In Proceedings of the sixth annual conference on Computational learning theory,

authors:
Sergey Ioffe
Christian Szegedy

title: Batch normalization: Accelerating deep network training by reducing internal covariate shift

venue: arXiv preprint arXiv:1502.03167,

authors:
Nitish Shirish Keskar
Dheevatsa Mudigere
Jorge Nocedal
Mikhail Smelyanskiy
Ping Tak Peter Tang

title: On large-batch training for deep learning: Generalization gap and sharp minima

venue: arXiv preprint arXiv:1609.04836,

authors:
Diederik P Kingma
Jimmy Ba

title: Adam: A method for stochastic optimization

authors:
Alex Krizhevsky
Geoffrey Hinton

title: Learning multiple layers of features from tiny images

venue: Technical report, Citeseer,

authors:
Guillaume Leclerc
Manasi Vartak
Raul Castro Fernandez
Tim Kraska
Samuel Madden

title: Smallify: Learning network size while training

venue: arXiv preprint arXiv:1806.03723,

authors:
Yann LeCun
John S Denker
Sara A Solla

title: Optimal brain damage

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

title: Pruning filters for efficient convnets

venue: arXiv preprint arXiv:1608.08710,

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

title: Learning efficient convolutional networks through network slimming

venue: In Computer Vision (ICCV),

authors:
Christos Louizos
Karen Ullrich
Max Welling

title: Bayesian compression for deep learning

venue: In Advances in Neural Information Processing Systems,

authors:
Dmitry Molchanov
Arsenii Ashukha
Dmitry Vetrov

title: Variational dropout sparsifies deep neural networks

venue: arXiv preprint arXiv:1701.05369,

authors:
Pavlo Molchanov
Stephen Tyree
Tero Karras
Timo Aila
Jan Kautz

title: Pruning convolutional neural networks for resource efficient transfer learning

venue: arXiv preprint arXiv:1611.06440,

authors:
Ari Morcos
David GT Barrett
Neil C Rabinowitz
Matthew Botvinick

title: On the importance of single directions for generalization

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

authors:
Vaishnavh Nagarajan
J Zico Kolter

title: Generalization in deep networks: The role of distance from initialization

venue: arXiv preprint arXiv:1901.01672,

authors:
Sharan Narang
Gregory Diamos
Shubho Sengupta
Erich Elsen

title: Exploring sparsity in recurrent neural networks

venue: arXiv preprint arXiv:1704.05119,

authors:
Behnam Neyshabur
Ryota Tomioka
Nathan Srebro

title: In search of the real inductive bias: On the role of implicit regularization in deep learning

venue: arXiv preprint arXiv:1412.6614,

authors:
Behnam Neyshabur
Zhiyuan Li
Srinadh Bhojanapalli
Yann LeCun
Nathan Srebro

title: Towards understanding the role of over-parametrization in generalization of neural networks

authors:
Ben Poole
Jascha Sohl-Dickstein
Surya Ganguli

title: Analyzing noise in autoencoders and deep networks

venue: arXiv preprint arXiv:1406.1831,

title: Modeling by shortest data

venue: description. Automatica,

authors:
Shai Shalev-Shwartz
Shai Ben-David

title: Understanding machine learning: From theory to algorithms

venue: Cambridge university press,

authors:
Karen Simonyan
Andrew Zisserman

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

venue: arXiv preprint arXiv:1409.1556,

authors:
Nitish Srivastava
Geoffrey Hinton
Alex Krizhevsky
Ilya Sutskever
Ruslan Salakhutdinov

title: Dropout: A simple way to prevent neural networks from overfitting

venue: The Journal of Machine Learning Research,

authors:
Zhewei Yao
Amir Gholami
Qi Lei
Kurt Keutzer
Michael W Mahoney

title: Hessian-based analysis of large batch training and robustness to adversaries

authors:
Jianbo Ye
Xin Lu
Zhe Lin
James Z Wang

title: Rethinking the smaller-norm-less-informative assumption in channel pruning of convolution layers

venue: arXiv preprint arXiv:1802.00124,

authors:
Zhonghui You
Jinmian Yan
Kun
Meng Ma Ye
Ping Wang

title: Gate decorator: Global filter pruning method for accelerating deep convolutional neural networks

authors:
Ruichi Yu
Ang Li
Chun-Fu Chen
Jui-Hsin Lai
Vlad I Morariu
Xintong Han
Mingfei Gao
Ching-Yung Lin
Larry S Davis

title: Nisp: Pruning networks using neuron importance score propagation

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

authors:
Chiyuan Zhang
Samy Bengio
Moritz Hardt
Benjamin Recht
Oriol Vinyals

title: Understanding deep learning requires rethinking generalization

venue: arXiv preprint arXiv:1611.03530,

authors:
Michael Zhu
Suyog Gupta

title: To prune, or not to prune: exploring the efficacy of pruning for model compression

venue: arXiv preprint arXiv:1710.01878,

heading: 1 INTRODUCTION

text:
Pruning weights and/or convolutional filters from deep neural networks (DNNs) can substantially shrink parameter counts with minimal loss in accuracy (LeCun et al., 1990; Hassibi & Stork, 1993; Han et al., 2015a; Li et al., 2016; Molchanov et al., 2017; Louizos et al., 2017; Liu et al., 2017; Ye et al., 2018), enabling broader application of DNNs via reductions in memory-footprint and inference-FLOPs requirements. Moreover, many pruning methods have been found to actually improve generalization (measured by model accuracy on previously unobserved inputs) (Narang et al., 2017; Frankle & Carbin, 2018; You et al., 2019). Consistent with this, pruning was originally motivated as a means to prevent over-parameterized networks from overfitting to comparatively small datasets (LeCun et al., 1990).
Concern about over-parameterizing models has weakened, however, as many recent studies have found that adding parameters can actually reduce a DNN’s generalization-gap (the drop in performance when moving from previously seen to previously unseen inputs), even though it has been shown that the same networks have enough parameters to fit large datasets of randomized data (Neyshabur et al., 2014; Zhang et al., 2016). Potential explanations for this unintuitive phenomenon have come via experiments (Keskar et al., 2016; Morcos et al., 2018; Yao et al., 2018; Belkin et al., 2018; Nagarajan & Kolter, 2019), and the derivation of bounds on DNN generalization-gaps that suggest less overfitting might occur as parameter counts increase (Neyshabur et al., 2018). This research has implications for neural network pruning, where a puzzling question has arisen: if larger parameter counts don’t increase overfitting, how does pruning parameters throughout training improve generalization?
To address this question we first introduce the notion of pruning instability, which we define to be the size of the drop in network accuracy caused by a pruning iteration (Section 3). We then empirically analyze the instability and generalization associated with various magnitude-pruning (Han et al., 2015b) algorithms in different settings, making the following contributions:
1. We find a tradeoff between the stability and potential generalization benefits of pruning, and show iterative pruning’s similarity to regularizing with noise—suggesting a mechanism unrelated to parameter counts through which pruning appears to affect generalization.
2. We characterize the properties of pruning algorithms which lead to instability and correspondingly higher generalization.
3. We derive a batch-normalized-parameter pruning algorithm to better control pruning stability.

heading: 2 RELATED WORK

text:
There are various approaches to pruning neural networks. Pruning may be performed post-hoc (LeCun et al., 1990; Hassibi & Stork, 1993; Han et al., 2015b; Liu et al., 2017), or iteratively throughout training, such that there are multiple pruning events as the model trains (Hochreiter & Schmidhuber, 1997; Narang et al., 2017; Zhu & Gupta, 2017). Most methods prune parameters that appear unimportant to the function computed by the neural network, though means of identifying importance vary. Magnitude pruning (Han et al., 2015b) uses small-magnitude to indicate unimportance and has been shown to perform competitively with more sophisticated approaches (Gale et al., 2019).
Many pruning studies have shown that the pruned model has heightened generalization (Narang et al., 2017; Frankle & Carbin, 2018; You et al., 2019), consistent with the fact that pruning may be framed as a regularization (rather than compression) approach. For example, variational Bayesian approaches to pruning via sparsity-inducing priors (Molchanov et al., 2017; Louizos et al., 2017) can describe weight removal as a process that reduces model description length, which in theory may help improve generalization (Rissanen, 1978). Similarly, the idea that models may be described more succinctly at flat minima has motivated pruning in service of flat minimum search (Hochreiter & Schmidhuber, 1997). Though Dinh et al. (2017) notes, however, that flatness can be arbitrarily modified by reparameterizing the function, and sharp minima can generalize well.
VC dimension (a measure of model capacity) has motivated the use of iterative pruning to improve generalization (LeCun et al., 1990; Hassibi & Stork, 1993). Overfitting can be bounded above by an increasing function of VC dimension, which itself often increases with parameter counts, so fewer parameters can lead to a guarantee of less overfitting (Shalev-Shwartz & Ben-David, 2014). Unfortunately, such bounds can be so loose in practice that tightening them by reducing parameter counts need not translate to better generalization (Dziugaite & Roy, 2017).
Rather than support parameter-count-based arguments for generalization in DNNs, our results suggest iterative DNN pruning may improve generalization by creating various noisy versions of the internal representation of the data, which unpruned parameters try to fit to, as in noise-injection regularization (Srivastava et al., 2014; Poole et al., 2014). Dropout creates particularly similar noise, as it temporarily sets random subsets of layer outputs to zero (likely changing an input’s internal representation every epoch). Indeed, applying dropout-like zeroing noise to a subset of features during training can encourage robustness to a post-hoc pruning of that subset (Leclerc et al., 2018; Gomez et al., 2018). Iterative DNN pruning noise ultimately differs, however, as it is: applied less frequently, not temporary (except in algorithms with weight re-entry), usually not random, and less well studied.

heading: 3 APPROACH

text:
Given a neural network and set of test data, let t be the top-1 test accuracy, the fraction of test data examples correctly classified multiplied by 100. We define a pruning algorithm’s instability on pruning iteration i in terms of t measured immediately before (tpre,i) and immediately after (tpost,i) pruning: instabilityi = tpre,i − tpost,i. In other words, the instability is the size of the accuracy drop caused by a particular pruning event.
Pruning Iterationsi− 1 i i+ 1
Train Test(tpre,i−1) Prune Test(tpost,i−1)
Train Test(tpre,i) Prune Test(tpost,i)
This measure is related to a weight’s importance (sometimes referred to as “saliency”; LeCun et al. (1990); Hassibi & Stork (1993)) to the test accuracy, in that less stable pruning algorithms target more important sets of weights (all else equal). The stability of a pruning algorithm may be affected by many factors. Our experiments (Section 4) explore the effects of the following: pruning target, pruning schedule, iterative pruning rate, and model. The remainder of this section provides an overview of these factors and demonstrates a need for a novel pruning target, which we derive.

heading: 3.1 PRUNING TARGET

text: In all of our experiments, we use iterative magnitude pruning (Han et al., 2015b), which removes weights according to some magnitude-based rule, retrains the resulting smaller network to recover from the pruning, and repeats until the desired size reduction is met. We denote pruning algorithms that target the smallest-magnitude parameters with an "S" subscript (e.g. pruneS), random parameters with an "R" subscript, and the largest-magnitude parameters with an "L" subscript. The usual approach to pruning involves removing parameters that have the smallest magnitudes (Li et al., 2016; Gale et al., 2019), or, similarly, those parameters least important to the loss function as determined by some other metric (LeCun et al., 1990; Hassibi & Stork, 1993; Molchanov et al., 2016; 2017; Louizos et al., 2017; Ye et al., 2018; Yu et al., 2018; You et al., 2019).

heading: 3.1.1 IDENTIFYING IMPORTANT BATCH-NORMALIZED PARAMETERS

text:
The correlation between parameter magnitude and importance weakens in the presence of batch normalization (BN) (Ioffe & Szegedy, 2015). Without batch normalization, a convolutional filter with weights W will produce feature map activations with half the magnitude of a filter with weights 2W : filter magnitude clearly scales the output. With batch normalization, however, the feature maps are normalized to have zero mean and unit variance, and their ultimate magnitudes depend on the BN affine-transformation parameters γ and β. As a result, in batch normalized networks, filter magnitude does not scale the output, and equating small magnitude and unimportance may therefore be particularly flawed. This has motivated approaches to use the scale parameter γ’s magnitude to find the convolutional filters that are important to the network’s output (Ye et al., 2018; You et al., 2019). Here, we derive a novel approach to determining filter importance/magnitude that incorporates both γ and β.
To approximate the expected value/magnitude of a batch-normalized, post-ReLU feature map activation, we start by defining the 2D feature map produced by convolution with BN:
M = γBN(W ∗ x) + β.
We approximate the activations within this feature map as Mij ∼ N (β, γ). This approximation is justified if central limit theorem assumptions are met by the dot products inW ∗x, and we empirically show in Figure A.1 that this approximation is highly accurate early in training, though it becomes less accurate as training progresses. Given this approximation, the post-ReLU feature map
R = max{0,M}
has elements Rij that are either 0 or samples from a truncated normal distribution with left truncation point l = 0, right truncation point r =∞, and mean µ where
µ = γ φ(λ)− φ(ρ)
Z + β,
λ = l − β γ , ρ = r − β γ , Z = Φ(ρ)− Φ(λ),
and φ(x) and Φ(x) are the standard normal distribution’s PDF and CDF (respectively) evaluated at x. Thus, an approximation to the expected value of Rij is given by
E[Rij ] ≈ Φ(λ)0 + (1− Φ(λ))µ.
We use the phrase "E[BN] pruning" to denote magnitude pruning that computes filter magnitude using this derived estimate of E[Rij ]. E[BN] pruning has two advantages. First, this approach avoids the problematic assumption that filter importance is tied to filter `2 norm in a batch-normalized network. Accordingly, we hypothesize that E[BN] pruning can grant better control of the stability of the neural network’s output than pruning based on filters’ `2 norms. Second, the complexity of the calculation is negligible as it requires (per filter) just a handful of arithmetic operations on scalars, and two PDF and CDF evaluations, which makes it cheaper than a data-driven approach (e.g. approximating the expected value via the sample mean of feature map activations for a batch of feature maps).

heading: 3.2 SUMMARY OF MODELS

text: We consider three basic model classes: a simple network with convolutions (2x32, pool, 2x64, pool) and fully connected layers (512, 10) that we denote Conv4, VGG11 (Simonyan & Zisserman, 2014) with its fully-connected layers replaced by a single fully-connected layer, and ResNet18 (He et al., 2016). All convolutions are 3x3. We trained these models using Adam (Kingma & Ba, 2014) with initial learning rate lr = 0.001, as we found Adam more helpful than SGD for recovering from unstable pruning (seemingly consistent with the observation in Zhu & Gupta (2017) that recovery from pruning is more difficult when learning rates are low).

heading: 3.3 ITERATIVE PRUNING RATE AND SCHEDULE

text:
For Conv4, we apply pruning to its first linear layer (which contains 94% of Conv4’s 1,250,858 parameters). For VGG/ResNet, pruning targets the final four convolutional layers (which contain 90% of VGG11’s 9,231,114 parameters, and 74% of ResNet18’s 11,173,962 parameters). Pruning focused on later layers partly because, as also found in Li et al. (2016); You et al. (2019), it allowed the network to recover more easily.
The pruning algorithms we consider are iterative: we define a pruning schedule that describes the epochs on which pruning events occur, and set a corresponding (constant) iterative pruning rate that will ensure the total pruning percentage is met by the end of training (please see Appendix A.5 for rate and schedule details). Thus, throughout training, pruning steadily removes DNN parameters, with the iterative pruning rate determining the number pruned per event. While our plots label each pruning configuration with its iterative pruning rate, the total pruning percentages were: 42% of VGG11, 46% of ResNet18, and 10% of Conv4 (except in Appendix A.4, wherein we prune 85% of Conv4).

heading: 4 EXPERIMENTS

text: Pruning studies often aim to compress pre-trained models that generalize well, and consequently, much work has focused on metrics to identify parameter importance: if you can find the parameters that matter the least to the function computed by the DNN, then you can prune more parameters without significantly harming accuracy. As a bonus, such pruning methods can sometimes even increase generalization (Narang et al., 2017; Frankle & Carbin, 2018; You et al., 2019). However, the mechanism by which pruning induces higher generalization remains unclear. Here, rather than investigate how to best maintain accuracy when pruning the network, we instead focus on understanding the mechanisms underlying these generalization improvements.

heading: 4.1 THE GENERALIZATION-STABILITY TRADEOFF

text:
Can improved generalization in pruned DNNs be explained by parameter-count reduction alone, or rather, do the properties of the pruning algorithm play an important role in generalization? As removing parameters from a DNN via pruning may make the DNN less capable of fitting to the noise in the training data, as originally suggested in LeCun et al. (1990); Hassibi & Stork (1993), we might expect that the generalization improvements observed in pruned DNNs are entirely explained by the number of parameters removed. In which case, methods that prune equal amounts of parameters would generalize similarly.
Alternatively, perhaps some aspect of the pruning algorithm itself is responsible for increased generalization. This seems plausible as the reported generalization benefits of pruning vary widely across studies. One possible explanation for this variability is differences in the pruning algorithms themselves. A key differentiator of these algorithms is their stability: more stable approaches may compute a very close approximation to the way the loss changes with respect to each parameter and prune a single parameter at a time (Hassibi & Stork, 1993), while less stable approaches may assume that parameter magnitude and importance are roughly similar and prune many weights all at once (Han et al., 2015b). Therefore, to the extent that differences in pruning algorithms explain differences in pruning-based generalization improvements, we might expect to observe a relationship between generalization and pruning stability.
To determine whether pruning algorithm stability affects generalization, we compared the instability and final top-1 test accuracy of several pruning algorithms with varying pruning targets and iterative pruning rates (Figure 1). Consistent with the nature of the pruning algorithm playing a role in generalization, we observed that more unstable pruning algorithms created higher final test accuracies than those which were stable (Figure 1, right; VGG11: Pearson’s correlation r = .84, p-value = 1.6e−11; ResNet18: r = .65, p-value = 5e−5). While many pruning approaches have aimed to induce as little instability as possible, these results suggest that pruning techniques may actually facilitate better generalization when they induce more instability. Furthermore, these results suggest that parameter-count based arguments may not be sufficient to explain generalization in pruned DNNs, and suggest that the precise pruning method plays a critical role in this process.
Figure 1 also demonstrates that pruning events for pruneL with a high iterative pruning rate (red curve, pruning as much as 13% of a given convolutional layer per pruning iteration) are substantially more destabilizing than other pruning events, but despite the dramatic pruning-induced drops in performance, the network recovers to higher performance within a few epochs. Several of these pruning events are highlighted with red arrows. Please see Appendix A.2 for visualization of the epoch-wise instabilities of each method in VGG11, and Appendix A.3 for an `2-norm pruning version of Figure 1, which has qualitatively similar results.
Interestingly, we initially observed that ResNet18 adapted to pruning events more quickly than VGG11 (accuracy rebounded after pruning then flattened soon after instead of climbing steadily). Thinking that shortcut connections were allowing the network to adapt to pruning events too easily, we tried pruning a larger amount of the penultimate block’s output layer: this reduced the number of shortcut connections to the final block’s output layer, lengthened the adaptation period, and improved generalization. This simple improvement of pruning hyperparameters suggests a potential for further optimization of the results shown. Please see Appendix A.5.1 for all hyperparameters/details of these experiments.

heading: 4.2 THE ROLE OF WEIGHT MAGNITUDE IN PRUNING REGULARIZATION

text:
We have demonstrated that, perhaps surprisingly, pruning larger magnitude weights via the E[BN] algorithm can result in larger test accuracy improvements (Figure 1). This suggests a positive correlation between pruning target magnitude and pruning’s regularization effect. However, it’s not clear whether this relationship holds more generally; i.e., perhaps it was caused by a feature of our
E[BN] algorithm or the networks examined. Alternatively, this effect may be dependent on whether nodes/filters (structured pruning) or individual parameters (unstructured pruning) are pruned.
As such, we tested whether target weight magnitude correlates with pruning’s regularizing effect when using both unstructured and structured magnitude pruning on the penultimate linear layer of a small network without batch normalization (Conv4). Specifically, we constructed a pruning target for each weight-magnitude decile (see Appendix A.5.2 for details), used each target to prune ten separate networks as they trained, and compared the generalization gaps (test-train accuracy) of the pruned networks to the target pruned (Figure 2).
For both unstructured and structured pruning (Figure 2 left and right, respectively), we found that pruning larger weights led to better generalization gaps, though, interestingly, this effect was much more dramatic in the context of unstructured pruning than structured pruning. One possible explanation for this is that, in structured pruning, the `2 norm of pruned neurons did not vary dramatically past the fifth decile, whereas the unstructured deciles were approximately distributed exponentially. As a result, the top 50% of filters for the structured case were not clearly distinguished, making magnitude pruning much more susceptible to small sources of noise. These results suggest that, when weight magnitudes vary considerably, pruning large magnitude weights may lead to improved generalization.
Interestingly, for ResNet18, we actually found that structured pruneL (red line) performed better than unstructured pruneL (green line) (Figure 3). The worse performance of unstructured pruneL may stem from its harming the helpful inductive bias provided by convolutional filters (i.e., perhaps removing the most important connections in all convolutional filters degrades performance more than pruning the same number of connections via removal of several entire filters) or its lower instability.

heading: 4.3 THE ROLE OF ITERATIVE PRUNING RATE IN PRUNING INSTABILITY

text:
While pruning large magnitude weights appears to play a role in pruning’s ability to improve generalization, more commonly used pruning algorithms often see generalization improvements when targeting the smallest magnitude or least important parameters, suggesting that target magnitude/importance is not the only characteristic of pruning algorithms relevant to generalization. One possibility is that, given a pruning target, pruning more parameters per pruning iteration (while holding constant the total pruning percentage) may lead to greater instability. If this is the case, the generalization-stability tradeoff suggests that the increase in instability from raising the iterative pruning rate would coincide with improved generalization performance. Alternatively, if the pruning target or total pruning percentage is all that matters, we may expect that changing the iterative
pruning rate (while keeping the pruning target and total pruning percentage fixed) would not affect generalization.
To test this, we plotted mean instability and test accuracy as a function of different iterative pruning rates for both `2-norm and E[BN] pruning (Figure 4). Consistent with iterative pruning rate playing a role in instability, we find that (given a pruning target) more instability is induced by using larger iterative pruning rates (Figure 4 left). Moreover, pruning random or small magnitude parameters performs best at the largest iterative rate (30%), supporting the idea that these methods require a source of instability to boost generalization. Note this suggests that, when targeting less important weights, higher iterative pruning rates during training can be an effective way to induce additional instability and generalization. (Algorithm and experiment details are available in Appendix A.5.4.)
Perhaps strangely, higher iterative pruning rates did not translate to improved generalization when targeting the largest magnitude weights (pruneL) with `2-norm pruning. The fact that pruneL does not generalize the best at the highest iterative pruning rate may be due to the reduction in pruning iterations required by the large iterative pruning rate (i.e., when the iterative rate is at 30%, the number of pruning events is capped at three, which removes 90% of a layer). Thus, while this rate grants more instability (Figure 4 left) per iteration, pruning affects the network less often. The idea that the regularizing effect of pruning is enhanced by pruning more often may also help explain the observation that methods that prune iteratively can generalize better (Han et al., 2015b).
Another possibility is that, since raising the iterative pruning rate (and consequently the duration between pruning events) tends to make the `2-norm worse for differentiating parameters by their importance to accuracy1, raising the iterative pruning rate causes pruneL with `2-norm pruning to target less important weights. Consequently, pruneL with `2-norm pruning may generalize worse at higher iterative rates by leaving unpruned more important weights, the presence of which can harm model generalization (Hinton & Van Camp, 1993; Morcos et al., 2018). Relatedly, this also means that pruneS with `2-norm pruning may increase (in networks with batch normalization at least) instability and generalization by failing to avoid the pruning of important parameters.

heading: 4.4 ITERATIVE PRUNING AS NOISE INJECTION

text:
Our results thus far suggest that pruning improves generalization when it creates instability throughout training. These prior results, though, involved damaging model capacity simply by the nature of pruning, which decreases the number of model parameters. It therefore remains possible that the generalization benefits we’ve seen rely upon the reduction in capacity conferred by pruning. Here, we examine this critical question.
We first note that iterative pruning can be viewed as noise injection (Srivastava et al., 2014; Poole et al., 2014), with the peculiarity that the noise permanently zeroes a subset of weights. Removing
1In Figure 4 left: with E[BN] pruning, pruneL tends to become more unstable relative to pruneS as the iterative rate increases, but with `2-norm pruning, pruneL and pruneS create similar instabilities for all rates.
the permanence of this zeroing can mitigate some of the capacity effect of pruning2 and, therefore, help us isolate and study how iterative pruning regularizes through noise injection.
As a baseline, we consider pruneL applied to VGG11, judging filter magnitude via the `2-norm (additional experimental details are in Appendix A.5.5). We then modify this algorithm such that, rather than permanently prune filters, it simply multiplies the filter weights by zero, then allows the zeroed weights to immediately resume training in the network ("Zeroing 0" in Figure 5 Left). However, by allowing pruned weights to immediately recover, this experiment also removes a key, potentially regularizing aspect of pruning noise: the requirement that the rest of the network adapts to fit the new representations generated by pruning. To encourage this potentially important facet of pruning noise, we also added variants that held weights to zero for 50 and 1500 consecutive batches3. As a related experiment, we also measured the impact of adding Gaussian noise to the weights in Figure 5, right. Noise was applied either once (Gaussian 0) or repeatedly over a series of training batches (Gaussian 50/1500).
If the capacity effects of weight removal are not necessary to explain pruning’s effect on generalization, then we would expect that the generalization behavior of these non-permanent noise injection algorithms could mimic the generalization behavior of pruneL. Alternatively, if weight removal is a necessary component of pruning-based generalization improvements, then we would not expect close similarities between the generalization phenomena of pruneL and non-permanent pruning noise injection.
Consistent with the capacity effects of weight removal not being necessary to explain generalization in pruned DNNs, applying zeroing noise for 50 batches to filters (rather than pruning them completely) generates strikingly similar accuracy to pruneL (Figure 5 Left). Specifically, the patterns in instability are qualitatively and quantitatively similar, as are the generalization levels throughout training.
Importantly, we found that applying zeroing noise once (Zeroing 0; brown line) was not sufficient to generate better performance, suggesting that the regularization induced by forcing weights to adapt to noised representations is critical to pruning’s ability to improve generalization. Moreover, we found that, while applying Gaussian noise could increase generalization if applied for long enough (Gaussian 1500; purple line), it still did not match the performance of pruneL, suggesting that multiplicative zeroing noise is substantially more effective than additive Gaussian noise4. Together, these results demonstrate that pruning induced generalization benefits are not merely explained by weight removal, but rather are dependent on the regularization conferred by forcing networks to adapt to noised representations over a sufficiently long period throughout training.
2This approach still effectively removes any weights that do not learn after reentering the model. However, we observed (see Figure 5) that pruning the reentered weights at convergence resulted in a marked drop in performance (for all noise schemes except “Zeroing 1500”), showing that the reentered weights had typically learned after reentry, and that temporary zeroing is therefore less harmful to capacity than permanent pruning.
3The pruning noise we apply affects a particular weight for no more than one segment of training time (where a segment is a series of N consecutive batches)
4At the scale we used.

heading: 5 DISCUSSION

text:
In this study, we defined the notion of pruning algorithm instability, and applied several pruning approaches5 to multiple neural networks, assessing the approaches’ effects on instability and generalization. Throughout these experiments, we observed that pruning algorithms that generated more instability led to better generalization (as measured by test accuracy). For a given pruning target and total pruning percentage, instability and generalization could be fueled by raising iterative pruning rates (Figure 4, Section 4.3). Additionally, targeting more important weights, again holding total parameters pruned constant, led to more instability and generalization than targeting less important weights (Figure 1, Section 4.1).
These results support the idea that the generalization benefits of pruning cannot be explained solely by pruning’s effect on parameter counts—the properties of the pruning algorithm must be taken into account. Our analysis also suggests that the capacity effects of weight-removal may not even be necessary to explain how pruning improves generalization. Indeed, we provide an interpretation of iterative pruning as noise injection, a popular approach to regularizing DNNs, and find that making pruning noise impermanent provides pruning-like generalization benefits while not removing as much capacity as permanent pruning (Figure 5, Section 4.4).

heading: 5.1 CAVEATS AND FUTURE WORK

text:
While not emphasized in our discussion, pruning algorithm stability can be a desirable property, as recovery from pruning damage is not guaranteed. Indeed, pruning too many large/important weights can lead to worse final generalization (Li et al., 2016). Recovery appears to be a function of several factors, including: learning rate (Zhu & Gupta, 2017)); presence of an ongoing regularization effect (Figure 3, Section 4.2); preservation of helpful inductive biases (Figure 3, Section 4.2); and damage to network capacity (e.g., removing too much of an important layer could cause underfitting).
A better understanding of the factors which aid recovery from pruning instability could aid the design of novel pruning algorithms. For example, pruning methods that allow weights to re-enter the network (Narang et al., 2017) could perhaps prune important weights occasionally to enhance generalization improvements, without risking permanent damage to the pruned networks (see Appendix A.4).
In describing how pruning regularizes a model, we touched on similarities between pruning and noise injection. Our results, however, may also be consistent with other parameter-count-independent approaches to understanding generalization in neural networks, as pruning may reduce the information stored in the network’s weights (Hinton & Van Camp, 1993), and make the network more distributed (Morcos et al., 2018; Dettmers & Zettlemoyer, 2019). This raises the possibility that pruning noise engenders helpful properties in DNNs, though it remains unclear whether such properties might be identical to those achieved with more common noise injection schemes (Srivastava et al., 2014; Poole et al., 2014). Further exploration will be necessary to better understand the relationship between these approaches.
One important caveat of our results is that they were generated with CIFAR-10, a relatively small dataset, so future work will be required to evaluate whether the presented phenomena hold in larger datasets. Relatedly, we only studied pruning’s regularizing effect in isolation and did not include commonly used regularizers (e.g., weight decay) in our setups. In future work, it would be interesting to examine whether pruning complements the generalization improvements of other commonly used regularization techniques.

heading: A APPENDIX

text:
A.1 QUALITY OF NORMALITY APPROXIMATION BY LAYER AND TRAINING TIME
5.0 2.5 0.0 2.5 5.0 Activation Magnitude
0.0
0.2
0.4
0.6
0.8
PD F
5.0 2.5 0.0 2.5 5.0 Activation Magnitude
0.0
0.2
0.4
0.6
0.8
PD F
5.0 2.5 0.0 2.5 5.0 Activation Magnitude
0.0
0.2
0.4
0.6
0.8
PD F
5.0 2.5 0.0 2.5 5.0 Activation Magnitude
0.0
0.2
0.4
0.6
0.8
PD F
5.0 2.5 0.0 2.5 5.0 Activation Magnitude
0.0
0.2
0.4
0.6
0.8
PD F
5.0 2.5 0.0 2.5 5.0 Activation Magnitude
0.0
0.2
0.4
0.6
0.8
PD F
5.0 2.5 0.0 2.5 5.0 Activation Magnitude
0.0
0.2
0.4
0.6
0.8
PD F
5.0 2.5 0.0 2.5 5.0 Activation Magnitude
0.0
0.2
0.4
0.6 0.8 PD F
5.0 2.5 0.0 2.5 5.0 Activation Magnitude
0.0
0.2
0.4
0.6
0.8
PD F
5.0 2.5 0.0 2.5 5.0 Activation Magnitude
0.0
0.2
0.4
0.6
0.8
PD F
5.0 2.5 0.0 2.5 5.0 Activation Magnitude
0.0
0.2
0.4
0.6
0.8
PD F
5.0 2.5 0.0 2.5 5.0 Activation Magnitude
0.0
0.2
0.4
0.6
0.8
PD F
5.0 2.5 0.0 2.5 5.0 Activation Magnitude
0.0
0.2
0.4
0.6
0.8
PD F
5.0 2.5 0.0 2.5 5.0 Activation Magnitude
0.0
0.2
0.4
0.6
0.8
PD F
5.0 2.5 0.0 2.5 5.0 Activation Magnitude
0.0
0.2
0.4
0.6
0.8
PD F
5.0 2.5 0.0 2.5 5.0 Activation Magnitude
0.0
0.2
0.4
0.6
0.8
PD F
VG G1
9 La
ye r D
ep th
Untrained VGG19 Trained VGG19
Figure A1: We examined the normalized activations (shown in blue histograms) of feature maps in the final eight convolutional layers of VGG19 before (left) and after (right) training to convergence. We found that the approximation to standard normality (shown in orange) of these activations is reasonable early on but degrades with training (particularly in layers near the output).
The main drawback to the E[BN] approach (Section 3.1.1) is the sometimes poor approximation Mij ∼ N(β, γ). In Figure A.1, the approximation’s quality depends on layer and training epoch. A less serious drawback is that this approach does not account for the strength of connections to the post-BN feature map, which could have activations with a large expected value but low importance if relatively small-magnitude weights connected the map to the following layer.
A.2 EPOCH-WISE INSTABILITIES
50 100 150 200 250 Epoch
10 1
10 2 0
10 2
10 1
100 101 M ea n In st ab ilit
y Pruning Style Prune_S 1% Prune_S E[BN] 1% Prune_S E[BN] 13% Prune_L E[BN] 13%
Figure A2: In VGG11, pruneS E[BN] is more stable than pruneS, which uses filter-`2-norm to compare parameter magnitudes. Methods with higher iterative pruning rates create more instability on a given iteration. Means reduce along the run dimension (10 runs per configuration).
Note that this graph uses a method (pruneS) that was included in Figure 1 right, but was not displayed in Figure 1 left due to its similarity to pruneS E[BN]. Additional experimental details are in Section A.5.1.
A.3 `2-NORM-PRUNING VERSION OF FIGURE 1
50 100 150 200 250 300 Epoch
83.0 83.0
83.5 83.5
84.0 84.0
84.5 84.5
85.0 85.0
85.5 85.5
86.0 86.0 86.5 86.5 To p1 Ac cu ra cy % 5% Pruning Style No Pruning Prune_S 1% Prune_S 13% Prune_L 13%
0.0 0.2 0.4 0.6 0.8 1.0 Mean Instability
85.4
85.6
85.8
86.0
86.2
86.4
To p-
1 Ac
cu ra
cy %
pearsonr = 0.67; p = 0.00023
50 100 150 200 250 300 Epoch
84 84
85 85
86 86
87 87
88 88
To p-
1 Ac
cu ra
cy %
28% 41%
Pruning Style No Pruning Prune_S 1% Prune_S 13% Prune_L 13%
0 1 2 3 4 Mean Instability
86.5
87.0
87.5
88.0
88.5
To p-
1 Ac
cu ra
cy %
pearsonr = 0.43; p = 0.017
Figure A3: Pruning instability improves generalization of (Top) VGG11 and (Bottom) ResNet18 when training on CIFAR-10 (10 runs per configuration). (Left) Test accuracy during training of several models illustrates how adaptation to more unstable pruning leads to better generalization. (Right) Means reduce along the epoch dimension (creating one point per run-configuration combination).
Here we use the same training and pruning configurations that were used for Figure 1, but we replace E[BN] pruning with `2-norm pruning. Qualitatively, the two figures’ results are similar. Interestingly, though, the correlation between instability and generalization is somewhat weaker with `2-norm pruning. This may be explained by the fact that `2-norm pruning generates a narrower spectrum of instabilities, which is perhaps due to `2-norm scoring’s inability to accurately assess parameter importance (illustrated in Figure 4).
A.4 THE GENERALIZATION-STABILITY TRADEOFF AT 85% SPARSITY IN CONV4
0 5 10 15 20 25 Epoch
20
40
60
80
100
To p-
1 Ac
cu ra
cy %
Pruning Style No Pruning Prune_S Prune_L Train/Test Train Test
Figure A4: When training Conv4 on CIFAR10, unstable pruning can significantly improve the baseline’s generalization. The training accuracies and test accuracies (the latter were calculated immediately after pruning) illustrate how much each pruning algorithm disturbs the neural network’s output during training.
While it’s unclear how much unstable pruning, which is particularly damaging to capacity, can be sustained at high sparsity levels, pruneL can lead to generalization several percentage points above the baseline/pruneS when pruning 85% of Conv4. Please see Section A.5.6 for experimental setup details.
A.5 EXPERIMENTAL DETAILS
Unstructured magnitude pruning entails removing individual weights (subsets of filters/neurons), which are selected for pruning based on their magnitude. Our unstructured pruning approach does not allow previously pruned weights to reenter the network (Narang et al., 2017; Zhu & Gupta, 2017; Gale et al., 2019). Structured magnitude pruning removes entire filters/neurons, which are selected based on their `2-norms or via the E[BN] calculation. Except where noted, we use structured pruning for VGG11 and ResNet18.
We denote the pruning of n layers of a network by specifying a series of epochs at which pruning starts s = (s1, ..., sn), a series of epochs at which pruning ends e = (e1, ..., en), a series of fractions of parameters to remove p = (p1, ..., pn), and an inter-pruning-iteration retrain period r ∈ N. For a given layer l, the retrain period r and fraction pl jointly determine the iterative pruning percentage il. Our experiments prune the same number of parameters il × size(layerl) per pruning iteration, ultimately removing pl × 100% of the parameters by the end of epoch el. Our approach is designed to study the effects of changing factors such as the iterative pruning rate and lacks some practically helpful features, e.g. hyperparameters indicating how many parameters can be safely pruned (Liu et al., 2017; Molchanov et al., 2017). When layerwise iterative pruning percentages differ (i.e., when there exists an a and b such that ia and ib are unequal), our figures state the largest iterative pruning rate that was used in any of the layers.
For ResNet, our pruning algorithms did not account for the magnitude of incoming shortcut connections when judging filter magnitude/importance. Though we did prune the incoming and outgoing shortcut connections associated with any pruned feature maps.
We used only the CIFAR-10 dataset (Krizhevsky & Hinton, 2009) in our experiments, a limitation of our study. We used batch size 128, and only used data augmentation in the decile experiment (Figure 2). For some experiments, we give multi-step learning rate schedules lrs = (x, y), which means we shrink the learning rate by a factor of 10 at epochs x and y.
A.5.1 FIGURE 1
We used E[BN] pruning in all models that were pruned, except for one model that used `2-norm magnitude pruning, which was included in Figure 1 right but not displayed in Figure 1 left due to its qualitative similarity to pruneS E[BN]. We leave out "E[BN]" in the legend of Figure 1 left, but all models nonetheless used E[BN] pruning.
The models were trained on CIFAR-10 with Adam for 325 epochs with lrs = (150, 300). The error bars are 95% confidence intervals for the mean, bootstrapped from 10 distinct runs of each experiment.
Since the layerwise pruning percentages varied, pruning required multiple iterative pruning percentages, the largest of which is denoted in the legend (rounded to the nearest integer).
VGG11 Pruning targeted the final four convolutional layers during training with (layerwise) starting epochs s = (3, 4, 5, 6), ending epochs e = (150, 150, 150, 275), and pruning fractions p = (0.3, 0.3, 0.3, 0.9). To allow for the same amount of pruning among models with differing iterative pruning percentages, we adjusted the number of inter-pruning retraining epochs. The models with the maximum iterative pruning percentage of 1% had r = 4, while the models with the maximum iterative pruning percentage of 13% had r = 40. The model pruned with `2-norm magnitude pruning, which only appeared in Figure 1 right, had r = 4 as well.
ResNet18 Pruning targeted the final four convolutional layers of ResNet18 during training with (layerwise) starting epochs s = (3, 4, 5, 6), ending epochs e = (150, 150, 170, 275), and pruning fractions p = (0.25, 0.4, 0.25, 0.95). As noted in Section 4.1, we increased the pruning rate of the output layer of the penultimate block to remove shortcut connections to the last layer, thinking that it should increase the duration of adaptation to pruning. The models with the maximum iterative pruning percentage of 1% had r = 4, while the models with the maximum iterative pruning percentage of 13% had r = 40.
A.5.2 FIGURE 2
Each experiment in Figure 2 targeted one of ten weight-magnitude deciles in the post-convolutional linear layer of the Conv4 network during training on CIFAR-10 with data augmentation.
While there are just ten deciles, the iterative nature of our pruning algorithms required the creation of eleven different pruning targets: ten methods pruned from the bottom of the decile upward (one experiment for each decile’s starting point: 0th percentile, 10th percentile, etc.), and one (D10) pruned from the last decile’s ending point downward (pruning the very largest collection of weights each iteration). In other words, D9 and D10 targeted the same decile (90th percentile to maximum value), but only D10 actually removed the largest weights on a given iteration (weights in the 100th-99th percentiles, for example). The D9 experiment would target weights starting from the 90th percentile (e.g. it may prune the 90th-91st percentiles on a particular iteration).
The training/pruning setup used the Adam optimizer, s = (4), e = (52), p = (0.1), r = 3, and lrs = (30, 60). We calculated the generalization gap on epoch 54 and sampled average pruned magnitudes on epoch 35. We obtained qualitatively similar results regardless of whether we used fewer training epochs or data augmentation. The error bars are 95% confidence intervals for the means, bootstrapped from 10 distinct runs of each configuration.
A.5.3 FIGURE 3
In Figure 3, pruneL was applied to the final four convolutional layers of ResNet18 during training with (layerwise) starting epochs s = (3, 4, 5, 6), ending epochs e = (150, 150, 170, 275), and pruning fractions p = (0.25, 0.4, 0.25, 0.95). Since the layerwise pruning percentages varied, pruning required multiple iterative pruning percentages, the largest of which is denoted in the legend (rounded to the nearest integer). The models with the maximum iterative pruning percentage of 1% had r = 4, the models with the maximum iterative pruning percentage of 13% had r = 40, and the "One Shot" model pruned all its targeted parameters at once on epoch 246.
When performing unstructured pruning, we pruned individual weights from filters based on their magnitude. The structured pruning experiments used E[BN] pruning.
The models were trained on CIFAR-10 with Adam for 325 epochs with lrs = (150, 300). The error bars are 95% confidence intervals for the means, bootstrapped from 10 distinct runs of each experiment.
A.5.4 FIGURE 4
In Figure 4, pruning targeted the final four convolutional layers of VGG11 during training with (layerwise) starting epochs s = (3, 4, 5, 6), ending epochs e = (150, 150, 150, 275), and pruning fractions p = (0.3, 0.3, 0.3, 0.9). To create the different iterative pruning rates, we used models with inter-pruning retrain periods r = 4, r = 20, r = 40, r = 60, and r = 100. Since the layerwise pruning percentages varied, pruning required multiple iterative pruning percentages, the largest of which is denoted on the horizontal axis. An unpruned baseline model average (10 runs) is plotted on the dotted line.
The models were trained on CIFAR-10 with Adam for 325 epochs with lrs = (150, 300). The error bars are 95% confidence intervals for the means, bootstrapped from 10 distinct runs of each experiment.
A.5.5 FIGURE 5
In Figure 5, pruning targeted the final four convolutional layers of VGG11 during training with (layerwise) starting epochs s = (3, 4, 5, 6), ending epochs e = (150, 150, 150, 275), pruning fractions p = (0.3, 0.3, 0.3, 0.9), and inter-pruning-iteration retrain period r = 40. When injecting pruning noise, we used the same pruning schedule and percentages, but applied noise to the parameters instead of removing them. The Gaussian noise had mean 0 and standard deviation equal to the empirical standard deviation of a noiseless filter from the same layer. PruneL used `2-norm pruning.
The models were trained on CIFAR-10 with Adam for 325 epochs with lrs = (150, 300). The error bars are 95% confidence intervals for the means, bootstrapped from 10 distinct runs of each experiment.
A.5.6 APPENDIX A.4
Each experiment in Appendix A.4 targeted the post-convolutional linear layer of the Conv4 network during training on CIFAR-10 with the Adam optimizer. The pruning algorithms start on epoch s = (3), end on epoch e = (18), prune the percentage p = (0.9), and prune every epoch via retrain period r = 1. These relatively simple experiments were conducted to show that, at higher sparsity (pruning 85% of the model’s parameters), unstable pruning can improve the generalization of the baseline. The error bars are 95% confidence intervals for the means, bootstrapped from 20 distinct runs of each configuration.

title: THE GENERALIZATION-STABILITY TRADEOFF IN NEURAL NETWORK PRUNING