AUBER: Automated BERT Regularization
Authors: Hyun Dong Lee
Seongmin Lee
U Kang
Venue:  (2021)
Keywords: BERT Regularization
Reinforcement Learning
Automated Regularization

Abstract: How can we effectively regularize BERT? Although BERT proves its effectiveness in various downstream natural language processing tasks, it often overﬁts when there are only a small number of training instances. A promising direction to regularize BERT is based on pruning its attention heads based on a proxy score for head importance. However, heuristic-based methods are usually suboptimal since they predetermine the order by which attention heads are pruned. In order to overcome such a limitation, we propose AUBER, an effective regularization method that leverages reinforcement learning to automatically prune attention heads from BERT. Instead of depending on heuristics or rule-based policies, AUBER learns a pruning policy that determines which attention heads should or should not be pruned for regularization. Experimental results show that AUBER outperforms existing pruning methods by achieving up to 10% better accuracy. In addition, our ablation study empirically demonstrates the effectiveness of our design choices for AUBER.

Under review as a conference paper at ICLR 2021
AUBER: AUTOMATED BERT REGULARIZATION
Anonymous authors
Paper under double-blind review
ABSTRACT
How effectively regularize BERT? Although BERT proves its effective- ness in various downstream natural language processing tasks, it often overfits when there are only small number of training instances. A promising direc- tion to regularize BERT is based on pruning its attention heads based on a proxy score for head importance. However, beuristic-based methods are usually subop- timal since they predetermine the order by which attention heads are pruned. In order to overcome such limitation, we propose AUBER, an effective regulariza- tion method that leverages reinforcement learning to automatically prune atten- tion heads from BERT. Instead of depending on heuristics or rule-based policies, AUBER learns pruning policy that determines which attention heads should or should not be pruned for regularization. Experimental results show that AUBER outperforms existing pruning methods by achieving up to 9.39% better accuracy. In addition, our ablation study empirically demonstrates the effectiveness of our design choices for AUBER.
INTRODUCTION
How can we effectively regularize BERT (Devlin et al. (2018))? In natural language processing (NLP), fine-tuning a large-scale pre-trained language model has greatly enhanced generalization. In particular, BERT has demonstrated effectiveness through improvements many downstream NLP tasks such as sentence classification and question answering.
Despite its recent success and wide adoption, fine-tuning BERT on downstream task is prone to overfitting due to weerpprrameterizztion; BERT-base has 110M parameters and BERT-large has 340M parameters. The overfitting worsens when the target downstream task only has small number of training examples. Devlin et al. (2018) and Phang et al. (2018) show that datasets with 10,000 or less training examples sometimes fail to fine-tune BERT.
To mitigate this critical issue, multiple studies attempt to regularize BERT by pruning parameters or using dropout to decrease its model complexity (Michel et al. (2019); Voita et al. (2019); Lee et al. (2020)). Among these approaches, we regularize BERT by pruning attention heads since pruning yields simple and explainable results and it can be used along with other regularization methods. In order to avoid combinatorial search, whose computational complexity grows exponentially with the number of heads, the existing methods measure the importance of each attention head based on heuristics such as an approximation of sensitivity of BERT to pruning a specific attention head. However, these approaches are based on hand-crafted heuristics that are not guaranteed to be directly related to the model performance, and therefore, would result a suboptimal performance.
In this paper, we propose AUBER, an effective method for regularizing BERT. AUBER overcomes the limitation of past attempts to prune attention heads from BERT by leveraging reinforcement learning. When pruning attention heads from BERT, our method automates this process by learning policies rather than relying on a predetermined rule-based policy and heuristics. AUBER prunes BERT sequentially in layer- wise manner. For each layer, AUBER extracts features that are useful for the reinforcement learning agent to determine which attention head to be pruned from the current layer. The final pruning policy found by the reinforcement learning agent is used prune the corresponding layer. Before AUBER proceeds to process the next layer, BERT is fine-tuned to recapture the information lost due to pruning attention heads. An overview of AUBER transitioning from the second to the third layer of BERT is demonstrated in Figure 1.
Our contributions are summarized follows:

---

Under review as a conference paper at ICLR 2021
Accuracy (Reward)
Replay Memory
Accuracy (Reward)
Replay Memory
f
Layer 12
Layer
Attention Head
Embedding (state)
Layer
Attention Head
Fine-Tune
Layer 3
DQN
Embedding (state)
Prune Attention
Layer
DQN
Layer
Heads (action)
Prune Attention Heads (action)
Layer
Layer 1
BERT-base
BERT-base
Pruned
Not Yet Pruned
Figure 1: An overview of AUBER transitioning from Layer 2 to Layer of BERT-base.
Method. We propose AUBER for automatically learning to regularize BERT by reinforce-
ment learning. AUBER is designed to carefully represent the state of BERT and reduce the action search cost by dually-greedy search space pruning.
Analysis. We theoretically justify our design choice of using L1 norm of the value matrix of each attention head element of state (see Theorem 1).
Experiments. We perform extensive experiments, and show that AUBER successfully regularizes BERT improving the performance by up 9.39% and outperforms other head pruning methods. Through ablation study, we empirically show that our design choices for AUBER are effective.
2
PRELIMINARY
We describe preliminaries on multi-headed self-attention (Section 2.1), BERT (Section 2.2), and deep Q-learning (Section 2.3).
2.1 MUL.TI-HEADED DELF-ATTENTION
An attention function maps a query vector and set of key-value vector pairs to an output. We com- pute the query, key, and value vectors by multiplying the input embeddings Q, K, V E N with the parameterized matrices WºQ E Rdxi, W E Raxn, and W Raxm respectively, where N is the number of tokens in the sentence, and m, and d are query, value, and embedding dimension respectively. In multi-headed attention, independently parameterized attention heads are applied in parallel to project the input embeddings into multiple representation subspaces. Each attention head contains parameter matrices W.º E Rdxx, W.K RRdxn and W. Rdxmm Output matrices of H independent attention heads concatenated and projected by matrix Wo E R Hmxd to obtain the final result. This process can be represented as:
M MultiHeadAAtlQ. K,V) V) = Concat (AUt...AAll...11 (Q, K. V))WO,
(1)
where
(QW )(KW)T vWW.
Att, KV softmar
(2)
Vn
multi-headed self-attention follows the same mapping methods as general multe- attention function except that all the query, key, and value embeddings come from the same sequence.
2.2 BERT
BERT (Devlin et al. (2018)) multi-layer Transformer (Vaswani et al. (2017)) pre-trained on masked language model and next sentence prediction tasks. It is then fine-tuned specific tasks including language inference and question answering. BERT-base has 12 Transformer layers and each layer has 12 self-attention heads. Despite its success in various NLP tasks, BERT sometimes overfits when the training dataset is small due to over- paraaaatttttiiiioonnn. Thus, there has been a growing interest in BERT regularization through various methods such as dropout (Lee el. (2020)).
2

---

Under review as a conference paper at ICLR 2021
Regeat times
Reward
[4] Evaluate
1
the reward
Replay Memory
[5] Store the tuple (state, action, state, reward)
O
Update
Stage 1 Train DQN for Layer
DON
Layer
State vector
(2)
Action
DQN
[3] Prune an attention head and update state vector referring to the action
Until all the heads in the layer pruned or Action 33
State vector
Action
Layer
DQN
Stage 2
Prune Layer i
[3) Prune attention head and update state vector referring to the action
Untäl the heads in the layer pruned Action
Figure 2: Overall flow of AUBER while pruning layer i. first trains DQN and then uses the trained DQN to make the final decision which heads to prune.
2.3 DEEP Q-LEARNING
deep network (DQN) is multi-layer neural network that outputs vector action-value pairs for given state It is function that maps d,-dimensional state space ddddiennso action space. Two important features of the DQN algorithm are target network and experience replay (Mnih et al. (2013)). The target network has the same architecture as that of the policy network, and its parameters are copied every 7 steps from the policy network. For experience replay, transition tuples (i.e. (state, action, reward, next state)) are stored in a first-in-first-out memory buffer, and subset of those tuples, uniformly sampled from the buffer, update the parameters the policy network.
PROPOSED METHOD
We propose AUBER, our method for regularizing BERT by automatically learning to prune at- tention heads from BERT, After presenting the overview of the proposed method in Section 3.1, we describe how we frame the problem of pruning attention heads into a reinforcement learning problem in Section 3.2. Section 3.3 explains how states are represented in AUBER and provides justification for the process. Section 3.4 describes how AUBER reduces the extremely large search space, and describe the full algorithm.
3.1 OVERVIEW
We observe that BERT is prone to overfitting for tasks with few training data. However, the existing head pruning methods rely on hand-crafted heuristics and hyperparameters, which give sub-optimal results. The goal of AUBER is to automate the pruning process for successful regularization. De- signing such regularization method entails the following challenges:
1. Automation. How resorting to sub-optimal heuristics?
automate the head pruning process for regularization without
2. State representation. When automating the regularization process
reinforcement learning problem, how can we represent the state of BERT in a way useful for the pruning?
3. Action search space scalability. BERT has many parameters and attention heads in each layer. How can we handle prohibitively large action search space for pruning?
We propose the following main ideas address the challenges:

---

Under review as a conference paper at ICLR 2021
1. Regularization with reinforcement learning (Section 3.2). We exploit reinforcement learning, specifically DQN, with accuracy enhancement as reward. DQN has shown a superior performance for many tasks, and is natural choice for model-free and off-po learning (Sutton & Barto (2018)), which is exactly our setting. Experience replay also allows efficient usage of previous experiences and stable convergence (Mnih et al. (2013) 3)). 2. L1 norm of value matrix as state representation (Section 3.3). We use LI norm of value matrix of each attention head to represent the state of a layer. When a head is pruned, the corresponding value is set to O.
3. Dually-greete search space pruning (Section 3.4). To reduce the search space, we use two greedy methods: 1) we prune layer-by-layee and an action is performed only within layer, and 2) in each layer, we prune one attention head at time to reduce the search space.
3.2 AUTOMATED REGULARIZATION WITH REINFORCEMENT LEARNING
AUBER leverages reinforcement learning for efficient search of regularization strategy without relying on heuristics. We exploit DQN among various reinforcement learning frameworks which have shown superior performance in model-free and off-policy environment. The overall flow is described in Figure 2. this section, we introduce the detailed setting of our reinforcement learning framework.
Initial state. Each layer of BERT has multiple attention heads, each of which has its own query, key, and value matrices. For layer of BERT, we derive the initial state st using LI norm of the value matrix of each attention head. Further details for this computation is elaborated in Section 3.3.
The action space of AUBER is discrete. For a BERT model with H attention heads per layer. the number of possible actions is И + 1 (i.c. a E {1.2...., H, H + 1}). When the action i E 21, ... H 1, H. is chosen, the corresponding 2th attention head is pruned. The action a = H + 1 signals the DQN agent to quit pruning. To facilitate exploration via off-policy learning, actions chosen in serayyng--psilon-geeee probability e that decays over time or otherwise selects an action based on the current policy, where
Action.
manner: the agent chooses random action with
Cánitial Cfinal
ffinal
(3)
ex/Caonny
Here, s is the total number of actions taken by the agent up to the current episode, Emnttial is the starting value of (i.e. when = 0), Efinal is the value that converges to - co, and Edecay is a hyperparameter that adjusts the rate of decay c.
Next State. After the tth head is pruned, the value of 5th index of st is set to 0. This modified state is provided as the next state to the agent. This mechanism allows the agent to recognize which attention heads have been pruned and decide the next best pruning policy based on past decisions. When the action a = H + 1, the next state is set to the terminating state which ends an episode.
Reward. To evaluate the reward, the training data split into sets: mini-training set and mini- dev set. We use the data in mini-training set for fine-tuning and data in mini-dev set for reward evaluation. The reward of AUBER is the change in accuracy
Aace current waw.untccccracyyyooocoo
(4) the accuracy of the current BERT model evaluated on the mini-dev set,
arre_accuracy
where eurrent-accuratt and prev accuracy is the accuracy obtained from the previous state or the accuracy of the original BERT model if no attention heads are pruned.
If we set the reward simply as current_accuron reward values the changes in accuracy are relatively small. Setting the reward as the change in accuracy has the normalization effect, thus stabilizing the training process of the DQN agent. The reward for action = + a hyper-parameter that can be adjusted to encourage or discourage active pruning. In AUBER, it is set to 0 to encourage the DQN agent to prune only when the expected change in accuracy is positive.
DQN cannot capture the differences among
Fine-tuning. After the best pruning policy for layer / of BERT is found, the BERT model pruned according to the best pruning policy is fine-tuned with a smaller learning rate. This fine-tuning step is crucial since it adjusts the weights of remaining attention heads to compensate for the information lost due to pruning. Then, the initial state of the layer / + 1 is calculated and provided to the agent. Since frequent fine-tuning may lead overfitting, separate the training dataset into two: mini-
4

---

Under review as a conference paper at ICLR 2021
validation dataset and mini-training dataset. The mini- aini-validation dataset is the dataset which the pruned BERT model is evaluated on to return a reward. After the optimal pruming policy is determined by using the mini-validation dataset, the mini-training dataset is used to fine-tune the pruned model. When all layers are pruned by AUBER, the final model is fine-tuned with the entire training dataset with early stopping.
3.3 STATE REPRESENTATION
The initial state of layer 1 of BERT is computed through the following procedure. We first caleu- late the LI norm of the value matrix of each attention head. Then, standardize the norm values to have mean = 0 and a standard deviation = 1. Finally, the so roftmar function is applied the norm values to yield 8i. The justification of using L1 norm is given by Theorem (proof in Appendix A. that L1 norm of the value matrix of head bounds the LI norm of its output matrix, which implies the importance of the head in the layer.
Theorem 1. For layer with H heads, let be the number of tokens in the sentence and m, n, and d be the value, query, and embedding dimension respectively. Let Q,K,V N be the input query, key, and value matrices, and W.O. W.K. and WV be the weight parameters of the tth head such that W W.K I and W.V E <dd Let O₂ be the output of the ihh head. Then, ||0,11 S C 11 where the constant C N
Algorithm 1: AUBER: Automatic BERT Regularization
Input BERT model B, fine-tuned on task t, # L of layers in BERT model, # of attention heads per layer of BERT model, and # E of episodes.
Output: Regularized Br.
for <- 1 to L do
Initialize policy network P and replay memory M for < to E do
B < copy(B.) B3.state_ prev_accuracy & eval(B;)
while action H + do
if Bi wooeeeeeae www.hhaddleee......e then +
action
else
action +- P.choose action(si)
end
13
action = H + then
t Terminal State
15
reward +
16
else
B (r-prueeheaallccion)
18
sf s --cop(sss) [actioo +-
20
current accuracy + eval(B))
reward current.accuracy prev_accuracy
prev_accuracy + current accuraacy end
23
M. M.push(s action, s7, reward)
25
81 + 87 end
26
P.optimiize(M)
28
end
Br.prune( P.final pI) finelune))
29
end
5

---

Under review as a conference paper at ICLR 2021
3.4 DUALLY-GREEDY SEARCH SPACE PRUNING
The total number of attention heads in 12-layer BERT 144, when 12 attention heads are used. Naively designing actions would lead to 2 244 possible actions which are prohibitively large. Our idea to reduce the search space is dually-greedy pruning: we prune layer-by-laye in a greedy manner (from lower to upper layers), and in each layer we greedily prune 1 attention head at time.
For each layer with attention heads, the agent receives an initial layer embedding s which encodes useful characteristics (1 norm of the attention heads) of this layer. Then, the agent outputs the index of an attention head that expected to increase the training accuracy when removed. After an attention head is pruned, the value of the jth index of 8; is set to 0, and it is provided as the next state to the agent. This process is repeated until the action a becomes H + 1. The model pruned up to layer / is fine-tuned on the training dataset, and new initial layer embedding 8141 is calculated from the fine-tuned model.
Algorithm 1 illustrates the process of AUBER. Algorithm 1 receives BERT model B,, which is fine-tuned on a specific task t, and the parameters for BERT and reinforcement learning. The algorithm aims output the regularized B,. In line memory M for layer 1. Lines 3-28 train the policy network P for F episodes. For each episode, lines 8-25 choose an action based on the current state vector st, prune an attention head based on the action, compute the resulting reward and the next state sf and store the transition tuple into M. This process is repeated until action H + 1, which indicates the termination of pruning. In line 27, we optimize P with transition tuples sampled from M. Finally, in lines 29-30, we use the trained policy network find the optimal pruning policy for layer I, prune B, according to the policy, and finally fine-tune B, After pruning a layer I, we proceed to prune the next layer + up to layer
initialize policy network P and a replay
4 EXPERIMENTS
We conduct experiments to answer the following questions of AUBER.
Q1 Accuracy (Section 4.2). Given BERT model fine-tuned on specific NLP task, how well does AUBER improve the accuracy of the model?
State Representation (Section 4.3). How useful is the LI norm of the value matrices of attention heads in representing the state of BERT?
Q3 Order of Processing Layers (Section 4.4). How does the order which the layers are processed by AUBER affect regularization?
4.1 EXPERIMENTAL SETUP
Datasets. We test AUBER on four GLUE datasets (Wang et al. (2019)) MRPC, CoLA, RTE, and WNLI which contain less than 10,000 training instances; has been observed that datasets with 10,000 or less training examples often fail in fine-tuning BERT (Devlin et al. (2018); Phang et al. (2018)). Detailed information of these datasets described in Table 1.
Table 1: Datasets.
dataset of classes
# of train # of dev
MRPC
2
3668
408
CoLA²
2
8551
1043
RTE3
2
2490
277
W
635
https://cc....eaaeo/.//////lllee htm1

---

Under review as a conference paper at ICLR 2021
Table 2:
Performance of AUBER and competitors on 4 GLUE tasks. The performance after fine-
tuning is measured. AUBER gives the best performance for the same number of pruned heads. Bold font indicates the best accuracy among competing pruning methods.
MRPC
CoLA
RTE
WNLI
Average
Original
84.07
57.01
63.54
46.48
62.77
AUBER
8..5
0..59
65.34 1.30
56.06-100.63
6..90:1..79
Random
84.02=1.112
57.89+#0.90
63.47=1.29
54.08:22.14
14
1.3
Confidence
83.70=0.477
27.69-2.19
64.26=1.664
55.77=:0.77
65.36=1.27
Michel et al. (2019)
84.32+00.3
58.86+0.64
63.90+00.00
55.21+
65.55+00.00
Voita et al. (2019)
83.92+00.71
55.3440.81
64.12++1.65
52.96+
64.08++..17
We use the pre-trained bert-base-cased model with 12 layers and 12 attention heads per layer provided by huggingface*. We fine-tune this model each dataset mentioned in Table I to obtain the initial model. Initial models for MRPC, CoLA, and WNLI are fine-tuned on the corresponding dataset for epochs, and the initial model for RTE fine-tuned for 4 epochs. The maximum sequence length is set to 128, and the batch size per GPU is set to 32. The learning rate for fine-tuning initial models for MRPC, CoLA, and WNLI set to 0.00002, and the learning rate for fine-tuning the initial model for RTE is set to 0.00001.
BERT Model.
Reinforcement Learning. We use a 4-layer feedforward neural network for the DQN agent. The dimensions of input, output, and all hidden layers are set to 12, 13, and 512, respectively. LeakyReLU is applied after all layers except for the last one. We train the DQN agent for 150 episodes. For the epsilon greedy strategy for choosing actions, Cinstial and c final are set to 1 and 0.05 respectively, and the epsilon decreases exponentially with the decay rate Cdecoy of 256. The replay memory size is set 5000, and the batch size for training the DQN agent is set 128. The discount value for the DQN agent is set to 1. The learning rate is set to 0.000002 when fine-tuning BERT after processing layer. Before processing each layer, the training dataset randomly split into 1 : 2 to yield mini- mini-training dataset and mini-validation dataset. When fine-tuning the final model, the patience value of carly stopping is set to 20.
Competitors. We compare AUBER with other methods that prune BERT's attention heads. If AUBER prunes P number of attention heads from BERT, we prune P heads in all the competitors.
Random. Prune attention heads randomly.
Confidence. Prune P heads with the smallest confidence score, which the average of the maximum attention weight after a series of forward passes. A high confidence score indicates that the weight concentrated on a single token.
Michel et al. (2019). Perform forward and backward pass to calculate gradients and use them assign an importance score to each attention head.
Voita et al. (2019). Construct new loss function that minimizes both the classification er- ror and the number of used heads so that unproductive heads are pruned while maintaining the model performance.
We fine-tune the competitors after pruning in layer-wise manner for fair comparison. In other words, we prune the attention heads in layer-wise manner and fine-tune the model every layer.
Implementation. We construct all models using PyTorch framework. All the models are trained and tested on a GeForce GTX 1080 Ti GPU.
4.2 ACCURACY
We evaluate the performance of AUBER against competitors. Table 2 shows the results on four GLUE datasets listed in Table 1. Note that AUBER outperforms all of its competitors on regu- larizing BERT, providing the best accuracy for all the datasets. While most of the competitors fail improve performance of BERT on the dev dataset of MRPC and CoLA, AUBER improves the performance of BERT by to 9.39% on those datasets.
thub. wwt.bb.ccc///ggggggcc/g
7

---

Under review as a conference paper at ICLR 2021
Table 3: We compare AUBER with four variants: AUBER-QueTs Query, AUBER-Key AUBER-L2, and AUBER-Reverse on four GLUE datasets. AUBER-Query and AUBER-Key use the query and key matrices respectively, and AUBER-L2 uses the L2 norm of the value matrix to obtain the initial state. JUBER-Reverse processes BERT starting from the final layer (e.g. 12th layer for BERT-base). Bold font indicates the best accuracy among competing pruning methods.
(a) MRPC (Accuracy)
(b) CoLA (Matthew's correlation)
Policy
of heads pruned
dev
Policy
# of heads pruned
dev
AUBER
86.03
AUBER
53
61.28
AUBER-QUEE Queery
27
83,33
AUBER Query
72
55,52
AUBER-Key
42
84.56
AUBER-Key
55
55.78
AUBER-L2
83.33
AUBER-L2
63
54.85
AUBER-Reveres
58
86.52
AUBER-Revene
57
59.48
(c) RTE (Accuracy)
(d) WNLI (Accuracy)
Policy
of heads pruned
dev
Policy
# of heads pruned
dev
AUBER
87
66.43
AUBER
86
56.34
AUBER-QUe FUBBRR-ueryy
83
65.34
AUBER -uRRr
96
56.34
AUBER Key
63.18
AUBER-Key
101
56.34
AUBER-L2
61
66.43
AUBER-L-
94
53.52
AUBER-Reverge
99
64.62
Revee-e
101
54.93
4.3 EFFECT OF STATE REPRESENTATION
We validate that the LI norm of value matrix each attention head effectively guides AUBER to predict the best action. Table 3 shows the performances of the variants of AUBER.
4.3.1 AUBER WITH THE KEY/QUERY MATRICES AS THE STATE VECTOR
Among the query, key, and value matrices of each attention head, we show that the value matrix best represents the current state BERT. Here we evaluate the performance of AUBER against AUBER LUBER-Query and AUBER-Key AUBER- Query and AUBER-Ke. use the query and key matri- ces respectively to obtain the initial state. Note that AUBER, which uses the value matrix to obtain state vectors, outperforms AUBER-Query and AUBER-KO Key all four tasks.
4.3.2 AUBER WITH L.2 NORM OF THE VALUE MATRICES AS THE STATE VECTOR
LI norm of the value matrices is used to compute the state vector based on the theoretical derivation. In this ablation study, we experimentally show that the LI norm of the value matrices is appropriate for state vector. We set new variant AUBER-L2 which leverages L.2 norm of the value matrices compute the initial state vector instead of LI norm. The performance of AUBER is far more superior than AUBER-L2 in most cases bolstering that LI norm of the value matrices effectively represents the state of BERT.
4.4 EFFECT OF ORDER PROCESSING LAYERS
We empirically demonstrate how the order in which the layers are processed affects the final perfor- mance. We evaluate the performance AUBER against REBERReeverse which processes BERT layers in the opposite direction (i.e. starting from the 12th layer) to what AUBER does. As shown in Table 3, AUBER provides a better performance than AUBER- Reverse in 3 out of cases. This shows that pruning lower layers first and then move to upper layers by AUBER is effective. A possible explanation is that is easier to train task-specific parameters (those upper layers) after learning general parameters (those in lower layers), rather than that in the reverse order.
8

---

Under review as a conference paper at ICLR 2021
5 RELATED WORK
BERT Regularization. To prevent overfitting of BERT on downstream NLP tasks, various regu- larization techniques have been proposed. Variants of dropout improves the stability of fine-tuning large pre-trained language models even when presented with a small number of training examples (Lee et al. (2020); Fan et al. (2020)). Using slanted triangular learning rate schedule and discrimina- tive fine-tuning has been proven to effectively prevent overfitting (Howard & Ruder (2018)) 8)). Jiang al. (2019) propose SMART, which regularizes the model by smoothing it and preventing aggres- sive update while Liu et al. (2019) train BERT model with multi-task learning algorithm to remedy overfitting. Introducing adversarial training to enhance the generalization of BERT has been tackled in Zhu et al. (2020). Pre-traine Pre-training tasks have also been changed for a better regularization. Xu et al. (2020) modify the pre-training task, next sentence prediction, to additionally predict previous sen- tence that it can capture more correlation. Our method has definite advantages for it can be used along with any of the aforementioned methods.
number of studies have analyzed the effectiveness of pruning parameters in BERT. Gordon et al. (2020) experimentally prove overparameterization of BERT by showing that pruning 30 40% of parameters does not affect model performance. Chen et al. (2020) find out the best parameter pruning strategy from the viewpoint of the lottery ticket hypothesis while Guo et al. (2019) deploy reweighted LI regularization with proximal algorithm. However, these methods primarily aim to compress BERT model not regularize it, and, therefore, no specific method to enhance the model performance has been proposed.
BERT Pruning.
Thanks to the unique structure of BERT that consists of multi-headed attention, not only parameter pruning but also structured attention head pruning has been addressed (Voita et al. (2019); Michel et al. (2019); Kovaleva et al.; McCarley (2019)). Michel et al. (2019); Kovaleva et al.; Prasanna et al. (2020) evaluate the role and importance of each attention head by measuring heuristics: the average of its maximum attention weight, where average is taken over tokens in set of sentences used for evaluation, or the expected sensitivity of the model to attention head pruning. Their results show that large percentage of attention heads with low importance scores be pruned without significantly impacting performance. The approach to set LO regularization term to minimize both the training loss and the number of used attention heads has been also presented (Voita et al. (2019); McCarley (2019)). However, they usually yield suboptimal results since they predetermine the order in which the attention heads are pruned by using heuristies.
Automation of Neural Network Pruning To automate the process of Convolutional Neural Net- work pruning, Lin et al.; He & Han (2018) have leveraged reinforcement learning to determine the best pruning strategy for each layer. Important features that characterize a layer are provided to reinforcement learning agent to determine how much the current layer should be pruned. To the best of our knowledge, AUBER is the first attempt heads from Transformer-based models such as BERT.
reinforcement learning to prune attention
6 CONCLUSION
We propose AUBER, an effective method to regularize BERT by automatically pruning attention heads. Instead of depending heuristics rule-based policies, AUBER leverages reinforcement learning to learn a pruning policy that determines which attention heads should be pruned for better regularization. Experimental results demonstrate that AUBER effectively regularizes BERT, in- creasing the performance of the original model on the dev dataset by up to 9.39%. In addition, we experimentally demonstrate the effectiveness of our design choices for AUBER.

---

Under review as a conference paper at ICLR 2021
REFERENCES
Tianlong Chen, Jonathan Frankle, Shiyu Chang, Sijia Liu, Yang Zhang. Zhangyang Wang, and Michael Carbin. The lottery ticket hypothesis for pre-trained BERT networks. CoRR, abs/2007.12223, 2020.
Jacob Devlin, Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. CoRR, abs/1810.04805, 2018.
Angela Fan, Edouard Grave, and Armand Joulin. Reducing transformer depth on demand with structured dropout. In ICLR, 2020.
Mitchell A. Gordon, Kevin Duh, and Nicholas Andrews. Compressing BERT: studying the effects of weight pruning on transfer learning. In RepLANLP@ACL, 2020.
Fu-Ming Guo, Sijia Liu, Finlay S. Mungall, Xue Lin, and Yanzhi Wang. Reweighted proximal pruning for large-scale language representation. CoRR, abs/1909, 12486, 2019.
Yihui He and Song Han. ADC: automated deep compression and acceleration with reinforcement learning. CoRR, 802.0892..3444 2018.
Jeremy Howard and Sebastian Ruder. Universal language model fine-tuning for text classification. In ACL, 2018.
Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Tuo Zhao. SMART: robust and efficient fine-tuning for pre-trained natural language models through principled regu- larized optimization. CoRR, abs/1911 1.04437,7 2019.
Olga Kovaleva, Alexey Romanov, Anna Rogers, and Anna Rumshisky. Revealing the dark secrets of BERT. In EMNLP-IJCNLP.
Cheolhyoung Lee, Kyunghyun Cho, and Wanmo Kang. Mixout: Effective regularization to finetune large-scale pretrained language models. In ICLR, 2020.
Lin, Yongming Rao, Jiwen Lu, and Jie Zhou. Runtime neural pruning. In NeurIPS.
Xiaodong Liu, Pengcheng He, Weizhu Chen, and Jianfeng Gao. Multi-task deep neural networks for natural language understanding. In ACL, 2019.
J. S. McCarley. Pruning bert-based question answering model. CoRR, abs/1910.06360, 2019.
Paul Michel, Omer Levy, and Graham Neubig, Are sixteen heads really better than one? In NuerIPS, 2019.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin A. Riedmiller. Playing atari with deep reinforcement learning. CoRR, abs/1312.5602, 2013.
Jason Phang, Thibault Févry, and Samuel R. Bowman. Sentence encoders on stilts: Supplementary training on intermediate labeled-data tasks. CoRR,
2018.
Sai Prasanna, Anna Rogers, and Anna Rumshisky. When BERT plays the lottery, all tickets winning. EMNLP. 2020.
Richard S. Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction. Bradford Book, Cambridge, MA, USA, 2018. ISBN 0262039249.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. NeurlPS, 2017.
Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. In ACL., 2019.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. GLUE: multi-task benchmark and analysis platform for natural language understanding. In ICLR, 2019.
10

---

Under review as a conference paper at ICLR 2021
Weidi Xu, Xingyi Cheng, Kunlong Chen, and Taifeng Wang. Symmetric regularization based BERT for wise semantic reasoning. In SIGIR, 2020.
Chen Zhu, Yu Cheng, Zhe Gan, Siqi Sun, Tom Goldstein, and Jingjing Liu. Freelb: Enhanced adversarial training for natural language understanding. In ICLR, 2020.

---

Under review as a conference paper at ICLR 2021
A APPENDIX
A.I PROOF FOR THEOREM 1
Proof. For 7th head in the layer, let
(QW? (KWW
so fomtmar softmax
(5)
vn
and
0i = VWY.
(6)
The output of the head, O₁ is evaluated as O, = softmar,01. Then,
N
10:11
Σ
(O:)aa
j=1 k=
N
((softmarz))))) (wx)-xl
(8)
N
(9)
(soottma))))) Σ ((vx))+ll2
(10)
Since the LI norm vector is always greater than or equal to the L2 norm of the vector,
n
100111
Σ
llso (t
11
1(xx)-+ll
(11)
==1
N
((xx).oll
(12)
N
(())gal
(13)
j=1
m
N
(14)
j=
N
IVs.lall(WW) will2
(15)
N
Σ
112-112
(WW)).ll
(16)
J=1
m I(W ).ll
N
IV j=1
(17)
V || l|
(18)
where the norm of the matrices is entrywise norm, ||A||] E, Ee Ask- All heads in the same layer take the same V as input and N is constant. Thus,
||0'|| S C|
(19)
for the constant = N|II
12


Reviews:
[{"review_id": "t5bx0EOFKe4", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "========================\n\nPaper Summary:\n\nThis paper proposed to prune heads of multi-head attention in BERT to achieve regularization for tasks with a smaller dataset. Existing pruning papers on transformers focus on compression (model size) with a smaller loss in accuracy but this work claims that pruning can actually improve accuracy in some cases. The authors use DQN to learn a policy to prune attention head layer by layer. They demonstrate improvements on 4 small tasks from GLUE.\n\n========================\n\nReview\n\nThe central idea of this paper is interesting, but the experiment is not convincing enough. Whether we actually need to prune large pretrained transformer for regularization itself is debatable because researchers have found that over-parameterized deep neural language models, for example GPT-3, is good at few shot learning. This paper needs much stronger experiments to support their claim.\n\nPros:\n- The proposed method is interesting. Use pruning to improve performance is a new thing, at least with transformers and BERT.\n\nCons:\n- This work is limited to small tasks. No results are demonstrated on larger datasets.\n- Experiment is weak. The authors may improve this by running all tasks in GLUE, perhaps with a smaller training split to fit the authors’ setting. Also, it might make more sense to experiment on BERT-large or other larger, more powerful pretrained transformers to demonstrate the proposed idea can regularize over-parameterized networks.\n- Lack of comparison with other regularization techniques, for example those mentioned in related work.\n- The training overhead of this method is not discussed. It involves repeated finetuning after each layer is pruned. It will benefit the readers to measure the time needed for training AUBER.\n\n==========================\n\nOther Questions\n\n- Is Table 2 a fair comparison with other baselines? In AUBER, model is finetuned after each layer is pruned. It is not clear if this finetuning is done for other heuristic-based pruning. This might need another ablation experiment.\n- Why not directly use the whole value embedding V as the input to the DQN? Theorem 1 is also somewhat heuristic.\n- It is not clear how the training data is sampled in Algorithm 1 to train the DQN model. Does s contain all training examples? Also, in line 29, what data does the model see to decide pruning?\n\n=================================\n\nMinor Issues\n\n- Please fix the citation format (citep/citet in latex).\n- Perhaps Figure 1 can be improved. In my opinion, the current figure does not help the readers understand the method.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "gOq9tAX8pP", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\n\nThe paper focuses on reducing over-fitting for the BERT model by pruning the attention heads. For determining the order of pruning, the paper uses reinforcement learning along with greedy algorithm to get rid of unimportant attention heads. The experimental section shows the analysis on various design choices like using value matrix to represent the state instead of key or query matrices, using L1 norm of value matrix rather than L2 norm and lastly, the order in which the layers should  be handled I.e. from top to  bottom or the reverse.\n\n+ve \n\n- The paper is easy to follow and the authors have presented the details in smaller sub-sections which makes it easy to understand.  \n- Using RL to determine the head pruning sequence makes if better than the previous greedy methods like Michel et al. and Voita et al.\n\nConcerns\n\n- Training routine seems very time consuming as the network is fine-tuned after removing every single attention head. The authors can try to remove the heads in a batch in order to reduce the training time. \n- For reducing over-fitting simple methods like increasing the dropout probability or reducing number of transformer blocks could also be tried. It would be good if the authors can provide comprehensive comparison with simple techniques in order to justify the multi-step training for AUBER.", "extra_scores": {}, "extra_text": {}}, {"review_id": "XQiycnUDBPL", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: This work proposes learning to prune attention heads in BERT in order to achieve better accuracies on downstream tasks, especially when there are a small number of training examples. The authors employ reinforcement learning, or more specifically deep Q-learning, to learn the pruning policy in a layer-wise manner. The model is finetuned iteratively once pruning is done for each layer. Experiments show that after pruning through the proposed method, the performance improves. Ablations shows the designs of the state representation and order of pruning are effective.\n\nPros:\nThe paper is clearly presented and easy to follow. The use of RL for pruning policy is reasonable, and to my knowledge this is the first work on BERT models. The experiments shows the pruned models can achieve better performance than the original models on several benchmark datasets.\n\nConcerns and questions:\n1. The idea of using reinforcement learning for neural network pruning has been tackled in previous works and is thus not quite new. \n2. The reason of choosing L1-norm of the value matrix as the state vector does not fully convinces me. The comparison in Table 3 may not be fair since different numbers of heads are pruned under different strategies, therefore the models have different capacities. A related question is that since the value matrix varies across training examples, will the state be the average of L1-norm of value matrix across all the training data? Please the author clarify.\n3. Only four of the GLUE tasks are included in the experiments, it will be more interesting to know how the method works for other GLUE tasks with relatively large training data.\n4. The performance of finetuning GLUE tasks could be unstable. Are the results reported in the paper from a single run or average of multiple runs? Is it possible the baseline model also have chance to achieve better performance with multiple runs?\n5. AUBER finetunes models when heads are pruned in every layer, could this be the main reason that it works better than baselines? What if we conduct layer-wise finetuning for the baseline models?\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "ILSc7hfCj0", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes AUBER which applies reinforcement learning algorithms (DQN) to progressively prune attention heads from the lower layer to the higher layer in pre-trained transformer models (BERT) in order to improve the model fine-tuning. The state of DQN is the L1 norm of the value matrix of each attention head, the action space for each layer is the total number of attention heads with an additional quit action (H + 1).  To reduce the search space, AUBER prunes one head each time until it reaches the quit action in the lower layer and then repeats the process in the next layer. Experiments on 4 tasks from GLUE shows the effectiveness of AUBER against baselines. Overall, it is quite interesting and inspiring.\n\nHere are a few concerns.\n--There are many works such as [1, 2, 3, 4, 5, 6] and adapter-based approaches to improve the BERT fine-tuning. It is better to have a discussion compared with those algorithms.\n\n--The 4 tasks evaluated are small. What’s the performance on the larger tasks, such as MNLI? Does AUBER also help in these settings? I’d like to mention that there are some issues with WNLI and all the machine learning algorithms without specific additional data or hack cannot outperform a simple majority voting. Thus it is better to select other tasks for evaluation. Lastly, these tasks have large variance and I’d like to suggest reporting mean/var instead of a single number. \n \n--AUBER is complicated and requires L (# of layers) rounds to get the final model. Thus, it is better to report the comparison of training time between these baseline systems. Does AUBER use the dev set to compute rewards during training? If so, it is not fair to report the dev accuracy on these models. \n\n--Pruning attention heads in BERT is interesting. However, the attention block only takes  around 25% of total parameters of BERT models. Is there any way to apply the current approach to other components in BERT?\n  \n--In section 4.1, it claims that fine-tuning on small tasks e.g., RTE often fails. It is better to provide more evidences, since a simple fine-tuning approach on BERT leads SOTA on these tasks compared with the pre-BERT era models/algorithms. \n \n\n\n[1] Zhang et al. Revising Few-sample BERT fine-turning, https://arxiv.org/abs/2006.05987\n\n[2] Zhu et al. FreeLB: Enhanced adversarial training for natural language understanding, ICRL 2020.\n\n[3] Jiang et al. Smart: Robust and efficient fine-tuning for pre-trained natural language models through principled regularized optimization, ACL 2020.\n\n[4] Liu et al. Multi-task deep neural networks for natural language understanding, ACL 2019.\n\n[5] Prasanna et al. When BERT Plays the Lottery, All Tickets Are Winning, https://arxiv.org/abs/2005.00561\n\n[6] Chen et al. The Lottery Ticket Hypothesis for Pre-trained BERT Networks, https://arxiv.org/abs/2007.12223", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes to use RL to learn how to prune attention heads in BERT to achieve regularization for tasks with small dataset size. Specifically, the authors use DQN to learn a policy to prune heads layer by layer.  

This paper receives 4 reject recommendations with an average score of 4.5. Though the idea in this paper is interesting, the experiments in the current draft are far from convincing. The reviewers have raised many concerns regarding the paper. (i) Experiments are weak. Only 4 GLUE tasks are considered; it is necessary to also test the proposed methods on other GLUE tasks. (ii) The comparison with other regularization techniques is lacking. (iii) The training overhead of this method needs more careful discussion, as it involves repeated finetuning after each layer is pruned, therefore could be very time-consuming. (iv) More comprehensive related work discussion is needed.

The rebuttal unfortunately did not address the reviewers' main concerns. Therefore, the AC regrets that the paper cannot be recommended for acceptance at this time. The authors are encouraged to consider the reviewers' comments when revising the paper for submission elsewhere.
Decision: Reject