PALBERT: Teaching ALBERT to Ponder
Authors: Nikita Balagansky
Daniil Gavrilov
Venue: NeurIPS 2022 Accept (2022)
Keywords: Early exit
ALBERT
GLUE

Abstract: Currently, pre-trained models can be considered the default choice for a wide range of NLP tasks. Despite their SoTA results, there is practical evidence that these models may require a different number of computing layers for different input sequences, since evaluating all layers leads to overconfidence in wrong predictions (namely overthinking). This problem can potentially be solved by implementing adaptive computation time approaches, which were first designed to improve inference speed. Recently proposed PonderNet may be a promising solution for performing an early exit by treating the exit layer's index as a latent variable. However, the originally proposed exit criterion, relying on sampling from trained posterior distribution on the probability of exiting from the $i$-th layer, introduces major variance in exit layer indices, significantly reducing the resulting model's performance. In this paper, we propose improving PonderNet with a novel deterministic Q-exit criterion and a revisited model architecture. We adapted the proposed mechanism to ALBERT and RoBERTa and compared it with recent methods for performing an early exit. We observed that the proposed changes can be considered significant improvements on the original PonderNet architecture and outperform PABEE on a wide range of GLUE tasks. In addition, we also performed an in-depth ablation study of the proposed architecture to further understand Lambda layers and their performance.

PALBERT: Teaching ALBERT to Ponder
Nikita Balagansky, Daniil Gavrilov Tinkoff
n. n. w...balagganskiyttttko... d. gevvriloovtiinnoff.
Abstract
Currently, pre-trained models can be considered the default choice for wide range of NLP tasks. Despite their SoTA results, there is practical evidence that these models may require a different number of computing layers for different input sequences, since evaluating all layers leads to overconfidence in wrong predictions (namely overthinking). This problem can potentially be solved by implementing adaptive computation time approaches, which were first designed improve inference speed. Recently proposed PonderNet may be a promising solution for performing an early exit by treating the exit layer's 's index as a latent variable. However, the originally proposed exit criterion, relying on sampling from trained posterior distribution on the probability of exiting from the i-th layer, introduces major variance in exit layer indices, significantly reducing the resulting model's performance. In this paper, we propose improving PonderNet with a novel deterministic Q-exit criterion and a revisited model architecture. We adapted the proposed mechanism to ALBERT and RoBERTa and compared it with recent methods for performing an early exit. We observed that the proposed changes can be considered significant improvements on the original PonderNet architecture and outperform PABEE on wide range of GLUE tasks. addition, also performed in-depth ablation study of the proposed architecture to further understand Lambda layers and their performance.
Introduction
These days, fine-tuning pre-trained models on downstream tasks became de facto standard technique for training NI.P models (Devlin et al.) 2019 Liu et al., 2019 Lan et al.) 2020 RRafford et al., (019). Although, model sizes increase, it becomes harder to them for real world applications due to computational footprint. Because of this fact, practitioners may desire to perform an early model exit (i.e., with adaptive computation time (Graves 2016 to reduce the computation needed for specific model input.
One model that is widely used in real-world applications is ALBERT Lan et al. 2020), which is based on the Transformer architecture (Vaswani al.) [017] with shared layers (i.e., the same layer is evaluated several times to provide an output). Sharing layer weights allows us to think of ALBERT as of recurrent model that simplify development of early exit methods. Most of them (Teerapittayanon et al. 2016: Liu et al.) 2020) Xin et al. (2000) rely on performing an early exit based on entropy of predictions during the model evaluation. Although, Zhou et al. (2020) showed that running the ALBERT-Bass block for fixed number of times (10) could increase the accuracy of the fine-tuned model on specific tasks (e.g., MRPC) despite the entropy of predictions monotonously decreasing for all 12 layers. A recent PABEE solution (Zhou et al. 2020) was designed to overcome this issue by performing an early exit based on the consensus between different classifier heads from different layers. The model stops evaluating when several classifiers in row produce the same result.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

À = 0.05, sample
A 0.1, sample
To early random exit
A 0.05, CDF 0.05 A
=
A= 0.6, CDF 0.658
Deterministic exit
CDF 0.145 A
A
A
A
h
h2
y
h
h2 2
h 3
y
Figure 1: A comparison of the original sampling exit criterion of PonderNet (on the left) and the proposed Q-exit criterion (on the right). PonderNet performs sampling from the Bernoulli distribution obtained from the Lambda layer each step, possibly exiting model too early or too late. For Q-exit, we evaluate the Cumulative distribution function (CDF) of the probability of exiting at layer i. Once CDF becomes greater than the threshold value (0.5 this example), we perform an early exit. With such deterministic criterion, we can perform an early exit from model more robustly without introducing variance in the exit layer's index during inference.
variational approach A possible orthogonal solution for this task is PonderNet (Banino et al. 202 that treats the exit layer's index as a latent variable. By maximizing the lower bound of the likelihood predict whether it necessary exit from
of the training data, PonderNet trains model which specific layer during evaluation. However, [Banino et al. (2021) proposed to sample from the trained posterior distribution of exiting from each layer during inference, which leads major variance in exit layer indices.
This paper proposes improving PonderNet adapted for ALBERT (Lan et al. 2020) and RoBERTa (Liu et al.) 2019) fine-tuning (PALBERT and PRoBERTa). Instead of performing an early exit by sampling from the trained posterior distribution during evaluation, exit criterion, namely Q-exit, which evaluates the CDF of the exit layer's 's probability distribution and performs deterministic carly exit. We also revisited the architectural choices of Lambda layers used to predict the probability of exiting from the current layer to make them aware of dynamics hidden states across previous layers and the number of currently running layers.
used novel
We experimented with PALBERT and PRoBERTa on the GLUE Benchmark datasets (Wang et al.. 2018). The ablation study showed that PALBERT produced significantly better results than the original PonderNet architecture adapted for ALBERT fine-tuning. Furthermore, the proposed methods outperformed PABEE and are comparable to plain ALBERT fine-tuning, while also exceeding it in speeds. We also analyzed the trained model and provided insights on further improvement of the variational approach for early exiting]
Related Work
Most of the approaches used to perform an early exit from model are based on the probability distribution of predictions: BranchyNet (Teerapittayanon et al.) 2016). FastBERT (Liu et al.) 2020).
DeeBERT (Xin al. 2020). which can be seen an entropy criterion, Xin et (2021) proposed BERxiT with LTE layer used to estimate the correctness of predictions during the inference based on the hidden state from the specific layer.
However, there is strong practical evidence that classification models' overthinking causes a reduction in predictions' entropy, making these methods diffieult to (Zhou et al.] (020). Furthermore, is unclear how to adapt entropy methods for regression tasks (Zhou et al. 2020). Zhoou et al. (2020) proposed PABEE a method perform an early exit based on several classifiers from the different levels of a model. Once several classifiers in a row (the number of these classifiers is determined by the patience hyperparameter t) produce the same result, can perform an early exit. Several recent works also utilized the idea of a consensus-based exiting strategy (Zhu 202 Liao al. 2021). Although, most of them mainly augment the basic consensus scheme with orthogonal heuristics. Because of this, we did not include them in our experiments and only used PABEE as consensus- based method.
An alternative way to perform an early exit is the Ponder architecture (Banino et al. 2021 which uses auxiliary Lambda layers to predict whether model should exit from specific layer during
Source code: http:://ichab..om//iioffffaaaa/

---

runtime, Inputs to Lambda layers used in PonderNet are hidden states from the current layer of model. PonderNet can be seen as a model with the latent variable that corresponds the exit layer index, which is trained by maximizing the lower bound of the marginalized likelihood of the data.
During inference, [Banino et al. (22TT) proposed sample from the trained posterior distribution of exit layer probabilities. However, this exit criterion can lead to uncertainty in exit layers for the same input (see Section 4.5 Even if the Lambda layer produced a probability equal to 0.1 of exiting from the first layer, we could still exit a model too early in one of ten cases, even though the probability was small. We also hypothesize that predicting exiting from a layer based entirely on single hidden state could be sub-optimal since performing early exit could also depend on the dynamics in hidden states across layers (i.e., the Lambda layer should know how hidden states change during the evaluation).
3 Improving PonderNet
In this section, we describe the adaption of PonderNet to the ALBERT model (PALBERT). Although, it could be easily extended to other pre-trained models such as RoBERTa (which we also used in our experiments).
The usual ALBERT evaluation can be defined as computation of n hidden states h; S(h,1)
from the input embeddings fig of an input sequence x, where i [1;n]. Once h, is obtained, is passed to a classifier block (h,,,) to get the parameters of an output distribution p(y|zr. A common way to fine-tume tune this architecture on downstream tasks is initialize the embeddings and the S layer by using ALBERT (pre-trained on Masked Language Modelling) while initializing C randomly and then optimizing all parameters by maximizing the likelihood of the training data.
While plain ALBERT performs a fixed number of computational steps, is possible to perform an arbitrary number of evaluations of the layer S. Banino et al. (20TT) proposed extend each Trans- former layer with a so-called Lambda layer. More precisely, for each layer i, after S outputs new h;, it is then passed to the classifier and Lambda layers to get parameters C(h,) of output distributions p(y|æ, i) and the probability of exiting from the i-th layer A Ahha), which induces a generalized geometric distribution probability of exiting from layer equal to p(i]z) d I1 Xy).
Then, having the probability distribution from each layer p(y|æ, i), the parameters of the model are optimized to maximize
L(x v) = Ee-p(ela) log (p(vl,,)) :)) BK (pL log (p(vlæ:)
(I)
Here, p(-[A) is a prior geometric distribution of exiting from each layer, parametrized by the hyper- parameter X, and (x=p((l))) log (p(ulz, i)) is evaluated analytically by averaging likelihoods from different layers with posterior exit probabilities. If we treat the exit layer index as a latent variable, then optimizing L from the Equation I could be seen as maximizing the lower bound of marginalized log likelihood log (p(y(z)) (Kingma and Welling 2014).
Note that following Banino et al. (2021) the prior probability of exiting from the last layer n is normalized as p(n|A)- = En= p() in order to make p(i[A) sum into with a finite number of steps, while p(i]z) it truncated by the lowest layer index j, s.t. E p(ilz) > 1 0.05 (thus,
KL....) is evaluated only with first j layers). Also note that weights of Lambda layers are shared across layers of the model.
3.1 Exit Criterion
During inference, Banino et al. (2021) proposed to sample the exit layer index from p(ilz) (i.e., by sampling iteratively from Bernoulli distribution with parameter X²). While sampling-based exit criterion correlates with the variational view of PonderNet's training objective, such estimation introduces the randomness in the inference process of PonderNet.
3

---

Method
SST-2
RTE
CoLA
ALBERT
9.3
77.0
57.0 2.1
ALBERT PonderNet
73.5
50.8 2.2
ALBERT PonderNet (Closed-Form Expectation)
9..3
3.0
2.2
Q-exit
Lamhda I.R
3-Layerr Lambda
concat.
9..
77.3 1.4
0.9
9.47
77.3 1.4
56.5 * 1.2
92.6
77.0
56.3
+
+
93.0
0.3
76.5 1.6
56,9 1.9
92.9
77.8 1.2
$7.4 1.7
Table I: An ablation study of the proposed PALBERT architecture. "Lambda LR" corresponds to fine-tuning the Lambda layer with its learning rate, "3-layer Lambda" refers to making the Lambda layer have three MLP layers instead of one, and "h concat." stands for concatenation of two hidden states as input to the Lambda layer.
To overcome the issue of randomness, we propose Q-exi a novel deterministic criterion of performing early exit, which used for PALBERT. Instead of sampling from the distribution p(ilz) during inference, we evaluate its CDF by accumulating p(i]z) from each layer. Once the CDF is greater than the threshold hyperparameter schematic comparison the sampling criterion with Q-exit. Threshold & can be seen as trade-off' between underthinking and overthinking. Therefore, q should be selected during the validation of the trained model to choose the best-performing value.
perform early exi See Figure for
Based our experiments, on various tasks compared to the original sampling criterion (see Sections 4.1 4.2). while also being more practical than the original sampling criterion.
found that the proposed criterion produced significantly better accuracy
3.2 Lambda Layer Architecture
While the original PonderNet used single layer MLP to obtain logit of exiting probability, we hypothesize that making the Lambda layer understand the dynamics of changing ALBERT hidden states is crucial for achieving good performance. To do so, instead passing single hidden state h; from the i-th layer in A, we concatenate it with h Le., for PALBERT, we evaluate the probability of exiting from i-th layer as X; = A ((h,,, h1-1)).
We used a layer MLP with tanh activation for the Lambda layer to operate with more complex input. Based on the ablation study, we observed that increasing the capacity improves the accuracy of the trained model (See section 1) We also found it beneficial to fine-tune the Lambda layer with a different learning rate than all other parameters.
4 Experiments
4.1
Ablation Study
We performed an ablation study of the proposed changes in PonderNet architecture with PALBERT.
We experimented with adding the proposed Q-exit criterion, Lambda layer architecture, and fine- tuning strategies. These methods were benchmarked on SST-2, RTE, and CoLA tasks from the GLUE Benchmarking dataset (Wang et al.) 2018).
We compared proposed changes with plain ALBERT and PonderNet adapted for ALBERT, Also, in these experiments we compared with PonderNet evaluated in expectation of predictions (e.g..
'Qeexit stands for Quantile
'The proposed exit criterion is the simplest deterministic criterion which we came to, since there is no ability to estimate mean or argmax statistics during the inference without running all layers of the model, which is impractical if want perform early exit.

---

Method
SST-2
RTE
QNLI
CoLA
MRPC
MNLI
QQP
STS-B
Macro
Dev set
ALBERT
92.7
76.5
91.5
56.6
90.5
84.8
88.9
90.6
84.0
ALBERT (9 L.)
91.1
74.2
91.2
56.9
89.7
84.3
89.7
90.1
83.4
ALBERT PABEE
92.7
76.9
91.5
55.6
88.3
84.5
88.9
89.9
83.5
ALBERT PonderNet
91.3
74.0
88.3
51.3
87.1
81.7
87.7
88.2
81.2
PALBERT (ours)
93.1
78.3
91.0
58.1
89.3
84.7
88.9
89.9
84.2
Test set
ALBERT
93.4
70.0
92.1
50.5
85.6
79.0
84.7
87.4
80.3
ALBERT PABEE
92.7
71.1
91.3
46.0
84.3
79.2
83.7
86.5
ALBERT PonderNet
90.2
68.6
88.5
43.7
84.3
81.1
77.4
83.6
77.2
PALBERT (ours)
93.0
73.5
91.7
48.6
87.1
79.8
84.4
86.5
80.6
Table 2: A comparison PALBERT with recent approaches on the GLUE benchmark. In the Macro column, present the average results across tasks. We bolded the best results. Note that did not bold the ALBERT rows since there no early exit applied and we reported it for reference.
Method
SST-2
RTE
QNLI
CoLA
MRPC
MNLI
QQP
STS-B
Macro
Dev set
RoBERTa
94.5
79.4
92.4
63.1
91.3
86.7
89.8
90.7
86.0
RoBERTa PABEE
93.8
76.9
91.9
60.7
89.6
86.4
91.1
89.9
85.0
PRoBERTa (ours)
93.7
79.1
91.6
62.2
90.2
86.7
91.1
89.9
85.6
Test set
RoBERTa
94.5
74.1
93.1
58.9
89.1
86.5
80.9
88.4
83.2
RoBERTa PABEE
95.4
71.7
92.0
51.2
87.6
86.2
80.4
86.2
81.3
PRoBERTa (ours)
94.5
72.1
91.8
58.6
88.3
86.
79.6
85.7
82.1
Table 3: A comparison of PRoBERTa with recent approaches on the GLUE benchmark. In the Macro column, we present the average results across tasks. bolded the best results. Note that we not bold the RoBERTa rows since there is early exit applied and reported it for reference.
evaluated p(yl=) E()((())) p(y) a. ar, i)] in closed-form). While such a model is impractical for early exiting setup (it does not perform any early exit but evaluates all layers estimate the expectation), it could show a gap between closed-form expectation of model predictions and its single sample Monte Carlo estimation.
For evaluation, we performed a grid hyperparameter search on an appropriate metric score on the dev split for each dataset. Following the PABEE training setup, we trained all models with fixed learning rate until validation metries stopped increasing for epochs. We used Adam optimizer (Kingma and Ba 2015) for all experiments. fixed 0.5 on models with the Q-exit criterion, as well as a fixed classifier dropout value equal to 0.1 (Srivastava al., 2014), and 0.1.
We trained each model times with the best hyperparameters and reported the mean and std values. full list of the methods' hyperparameter ranges can be found in Table
See Table T for the full list the results our ablation study. Based on these experiments, PonderNet architecture is seen performing worse than vanilla ALBERT fine-tuning, Also, PonderNet with closed-fore foreed-fomm evaluation of predictions expectation performed better than the sampling criterion, indicating that randomness in exit layers leads to poor performance.
At the same time, deterministic Q-exit criterion improves PonderNet accuracy when compared to a random sampling of the exit layer and is comparable to the evaluation of PonderNet in expectation. more complex Lambda layer that can handle the dynamics of hidden state changes can further improve model accuracy when compared to the original PonderNet.
5

---

4.2 GLUE Experiments
We compared PALBERT with PonderNet architecture adapted for ALBERT fine-tuning and PABEE with fixed patience value 6 on all GLUE tasks. We also trained Ponder RoBERTa (PRoBERTa) and compared it with PABEE and adapted for RoBERTa. Note that we re-implemented all baselines used in our experiment for consisteney in comparison. See Appendix Section A for the details of reimplementation.
we trained 5 models with the best hyperparameters across the hyperparameter search and reported the median task score on the dev set. We evaluated the test scores on the best models, selected based on their dev scores. We report the two metrics' mean for the MRPC, QQP, and STS-B tasks. For the MNLI task, we report the mean accuracy across matched and mismatched datasets.
Following the experimental setup from Section[4.] 4.1
See Tables 2 for the full list of results. We observed that PALBERT outperformed PABEE wide range of tasks. Vanilla PonderNet with the sampling exit criterion performed the worst. Vanilla ALBERT outperformed PABEE most tasks and is comparable PALBERT, while the latter has the highest score averaged across all tasks (see Macro column in Tables
2 3
For experiments with RoBERTa, we observed that PRoBERTa outperformed PABEE on Dev split, while either performed better or marginally worse for the Test set. Thus, PRoBERTa showed a higher macro score compared to PABEE.
Note that PABEE is performing poorly for tasks with a small dataset (e.g., CoLA, RTE). We hypothesize that this is caused by several independent classifiers at each layer C, failing to train well enough, whereas PALBERT was capable of utilizing knowledge sharing between layers.
Parameter
Values range
Learning rale
[le-5, 2e-5, 3e-5, 5e-5]
Batch size
[16. 32, 128]
Lambda learning rate
[lc-5, 20-5, 3c-5] [0.5]
9
Table 4: Hyperparameter search ranges used in all of our experiments.
4.3 Understanding the Threshold of Q-exit
As noted previously Section 3.1 we treat the threshold value g of the Q-exit criterion as a trade-off
between underthinking and overthinking, where increasing q forces a model to evaluate more layers, and vice versa.
Therefore, it is necessary to find the besf-performing threshold for each task where a model has the highest accuracy. To do so, we evaluated trained PALBERT models from the ablation study (see Section ).1) on dev splits of tasks with different values of q. We then averaged obtained metrics and reported the mean and std values for various thresholds (See Figure for the results).
We observed that exiting models with q = 0.5 shows the best overall performance for different tasks. Making q greater than 0.5 leads to reduction in accuracy and can often force models to evaluate all 12 layers of ALBERT-Base.
It is also notable that PALBERT, having large threshold value that performs constant exit on the last layer, has better accuracy than vamilla ALBERT fine-tuned for the SST-2 task. This indicates that adding auxiliary tasks for each layer of ALBERT improves performance compared to plain fine-tuning with a single classifier on the last layer of the model.
4.4 Speed Analysis
While making < 0.5 improves inference speed, it can also lead to underthinking and lower accuracy (see Figure 3 We compared PALBERT using different threshold values q to PABEE with different patience values I, which stands for the number of layers necessary to output the same result in a row to perform an early exit. We evaluated task scores for the specified hyperparameters as well as the increase in speed when compared to vanilla ALBERT inference of full model with 12 layers.

---

0.80
0.92
0.75
0.90
Accuracy RTE
0.
Accuracy
0
65
0.86
0.60
0.84
ALBERT
0.55
ALBERT
0.8
0.5
0.7
Quantile
0.6
0.90
(Accuracy F112 MRAC
0.88
Matthew's come ation
0.86
0.4
0.2
PALBERT ALBERT
PALBERT
ALBERT
0.1
0.2
0.3
0.4
0.5
0.7
0.8
0.1
0.2
0.3
0.4
0.5
0.5
0.7
0.8
0.9
Quantile
Quantile
Figure 2: PALBERT score dependency on the Q-exit threshold. We report the mean and std values of task metrics across 5 trained models. See section 4.3 for more details
Overall, we observed that PALBERT mostly produced higher scores on different tasks while also being slightly faster than PABEE. For the CoLA and MRPC datasets, PALBERT performed better. The proposed method outperformed PABEE while achieving the same increase in speed.
We observed questionable results on the SST-2 dataset: the best score for the PABEE model is slightly higher than that of PALBERT. However, it was obtained with a negligible increase in speed compared to ALBERT, as the best-performing patience for this setup is 11 layers while the whole model has only 12 layers.
Furthermore, unlike PALBERT, PABEE performed significantly worse than plain ALBERT fine-tuned on the CoLA and RTE tasks.
4.5 On the Prior of PALBERT
While we used a fixed value of X 0.1 our main experiment following Banino et al. (2021) (see Section 4.2), also investigate the effect of changing the parameter of the prior distribution exit layer indices.
We followed the experimental setup from the previous sections and trained models with E 10.1, 0.1 0.25, 0.5]. For each X, we had its own best hyperparameter set, which was then used to train 5 models.
See Figure for the results. We observed that varying affects the performance of the trained model since it has huge impact on the distribution of exit layer indices. Although the best metric for the RTE task reached with À 0.08, such model mostly performs exit from the last layer. Meanwhile, A = 0.08 performs marginally worse than À = 0.1 for MRPC.
Note that while Figure Ex~D [[]][=] (see Figure which is the distribution of exit layer indices with vanilla sampling criterion proposed by [anino et al. (202I). Even À is low enough so that the mode of distribution the 12-th layer, in almost half of the cases, the sampling criterion forces the model to exit from one of the first 10 layers. We hypothesize that such randomness in the index the exit layer leads the poor performance observed during previous experiments (see Section
could also estimate obtained using the Q-exit criterion,
7

---

0.93
0.92
0.725
Acouracy RTE
0.90
0.700
Accurdiy
0.88
0.625
PALBERT
0.86
PABEE
0.600
0.85
ALBERT
ALBERT
Speed
Speed up
0.55
0.90
0.50
Accuracy MAPC
0.86
0.40
0.35
Matthew's
0.82
0.25
0.80
PABEE
PABEE
20
ALBERT
ALBERT
x1
x2
X3 4x
6x
Speed-up
Speed-up
Figure 3: A comparison between PALBERT and PABEE models on SST-2, RTE, CoLA, and MRPC tasks. We varied the threshold value of Q-exit for PALBERT and the patience hyperparameter for PABEE obtain the plots of task scores of inference increasing in speed. Ix stands for plain ALBERT inference without performing an early exit. The horizontal line corresponds to plain ALBERT fine-tuning. See Section 4.4 for the analysis of these plots.
From such perspective, dependency the explicit prior distribution could be the main limitation of the proposed method, since treating the exit layer index as latent variable justified its prior distribution. proper prior distribution was found, then the well-performing model could be trained, and vice versa.
5 Future Work and Limitations
In this paper, we proposed improving the PonderNer architecture to perform an early exit using fine-tuned ALBERT and RoBERTa models with the novel Q-cxit criterion and a revisited Lambda layer architecture. This approach is orthogonal to the consensus-based approaches (Zhou et al. 2020 Zhu 202T) highly represented the field.
While PALBERT and PRoBERTa outperformed some recent State-of- State-of-The-Aee methods used for an early exit, there a clear direction for further improvement this method, was not capable of outperforming plain ALBERT and RoBERTa on some GLUE tasks.
We believe that PALBERT could benefit from the development of new parameterization of the prior distribution on exiting from each layer since it directly affects the resulting posterior distribution used perform an early exit (see Section 4.5 However, the development of an appropriate prior distribution for early exiting models is orthogonal to the development of a proper exit criterion.
Furthermore, there is theoretical justification for the Q-exit threshold value. Although observed that y 0.5 performed best, it is without clear explanation as to why that so. We hypothesize that bringing more insights into developing deterministic exit criteria could further improve the proposed method.
In addition, adding more auxiliary tasks could also make it possible to improve PALBERT further. This way, PALBER training can be made more PABEF-like by making independent classifiers for each layer of the model or adding self-distillation across layers.

---

350
300
0.00
200
250
Etamptes
Exmmltes
150
100
50
MRPC
RTT
0.50
0.55
0.88
0.250
0.86
Accurcy
0.700
2.64
037
0.62
0.550
0.00
0.555
Am05 A-025
2-1-3
2=0000
A=055
A=025
A-1
2=03
2-0.00
Figure 4: Histogram of exit layer indices with Q-exit criterion and different prior distribution parameters X (top) alongside with task metrics for trained models (bottom) for MRPC and RTE tasks. See Section for more details.
Figure 5: An estimation of Eg>D p(ilz)] where p(i|z) is trained posterior probability of exiting from layer of PALBERT models across different tasks, 0.1, and D is the distribution of the training dataset. We took 5 models trained these tasks and sampled exit layer indices for the training dataset's inputs. We smoothed the obtained probabilities for visibility. Note that these probabilities could be seen as the distribution of exit layers across different tasks for plain PonderNet with sampling criterion, which indicates on high variance in indices of output layers.
While the variational view of the early exiting mechanism allows us to train models which could be exited early, we observed that the overall performance of the trained model highly depends on the prior distribution of exit layer indices (see Section 4.5), Because of this, the proposed method requires more computation for hyperparameter search compared PABEE (Zhou et al.) 2020). We believe that, in order to further develop training-level early exiting, it may be beneficial to move away from variational inference to some other form of trainable mechanism without relying on exit layers' explicit prior knowledge.

---

References
Andrea Banino, Jan Balaguer, and Charles Blundell. 2021. Pondernet: Learning to ponder In 8th ICML Workshop Automated Machine Learning (AutoML).
Jacob Devlin, Ming-Wer Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume 1 (Long and Short Papers), pages 4171-4186, Minneapolis, Minnesota. Association for Computational Linguistics.
Alex Graves. 2016. Adaptive computation time for recurrent neural networks Cite 103x0:1603.0888..
Diederik P. Kingma and Jimmy Ba. 2015. Adam: A method for stochastic optimization In 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings.
Diederik Kingma and Max Welling. 2014. Auto-Encoding Variational Bayes) In 2nd International Conference on Learming Representations, ICLR 2014, Banjf, AB, Canada, April 14-16, 2014, Conference Track Proceedings.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. 2020. Allert: A lite bert for self-supervised learning of language representations) In International Conference on Learning Representations.
Kaiyuan Liao, Yi Zhang, Xuancheng Ren, Qi Su, Xu Sun, and Bin He. 2021. A global past-future early exit method for accelerating inference of pre-trained language models Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguis- tics: Human Language Technologies, pages 2013-2023, Online. Association for Computational Linguisties.
Weijie Liu, Peng Zhou, Zhiruo Wang, Zhe Zhao, Haotang Deng, and Qi Ju. 2020. Fastbert: self distilling BERT with adaptive inference time In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, ACL 2020, Online, July 5-10, 2020, pages 6035-6044. Association for Computational Linguistics.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis,
Luke Zettlemoyer, and Veselin Stoyanov. 2019. Roberta: robustly optimized bert pretraining
approach] Cite
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. 2019. Language models are unsupervised multitask learners.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, llya Sutskever, and Ruslan Salakhutdinov. 2014. Dropout: A simple way to prevent neural networks from overfitting Journal of Machine Learning Research, 1555)):129-1155:
Surat Teerapittayanon, Bradley McDanel, and H. T. Kung. 2016. BranchyNet: Fast inference via carly exiting from deep neural networks In Proceedings of the 23rd International Conference IEEE.
Pattern Recognition, pages 2464-2469.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in Neural Information Processing Systems, pages 5998-6008.
Alex Wang, Amapreet Singh. Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. 2018. Glue: A multi-task benchmark and analysis platform for natural language understanding. Cite 184Y:1880400746 Commment: https://glueeencchartr
Xin, Raphael Tang, Jaejun Lee, Yaoliang Yu, and Jimmy Lin. 2020. DeeBERT: Dynamic early exiting for accelerating BERT inference, In Proceedings of the 58th Annual Meeting of the Associ- ation for Computational Linguistics, pages 2246-2251, Online, Association for Computational Linguisties.
10

---

Ji Xin, Raphael Tang. Yaoliang Yu, and Jimmy Lin. 2021. BERxiT: Early exiting for BERT with better tine-tuning and extension regression In Proceedings of the 16th Conference of the European Chapter of the Association for Computational Linguisties: Main Volume, pages 91-104, Online. Association for Computational Linguisties.
Wangchunshu Zhou, Canwen Xu, Tao Ge. Julian McAuley, Ke Xu, and Furu Wei. 2020. Bert loses patience: Fast and robust inference with early exit In Advances in Neural Information Processing Systems, volume 33, pages 18341.1834,, Curran Associates, Inc.
Wei Zhu. 2021. LeeBERT: Learned early exit for BERT with cross-level optimization In Proceed- ings of the 59th Annual Meeting of the Association for Compulational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume I: Long Papers), pages 1968-2980 Online, Association for Computational Linguistics.


Reviews:
[{"review_id": "uqmMABKOg2e", "reviewer": "Reviewer_yzdP", "summary": "This paper focuses on early exit mechanism for encoder-only Transformer models. Specifically, it uses the ALBERT architecture that is similar to BERT but shares parameters across all layers (12 layers for base size). \n\nThe authors adopt the PonderNet framework of Banino et al. that learns early-exit \"lambda\" classifiers for each layer, and suggest two modifications to it. First, instead of sampling from the geometric-like distribution that PonderNet emulates to decide if to exit, they compute the empirical CDF of the layers so far, and exit if the value is greater than some defined threshold q. Second, they change the module that computes lambda to be slightly larger and to use the last two hidden states instead of just the current one, and modify the learning rate.\n\nExperiments are completed on the GLUE tasks, where PonderNet and PABEE (exiting if several consecutive layers predict the same) are used as baselines (authors reimplemented these baeslines to use the same ALBERT backbone).  Also, a grid-search over threshold values for q and ablation studies on hyper-parameters are reported.", "questions": "See points above, also:\n* Why PonderNet in table 2 doesn't have test results?\n* How were the hyper-parameters and threshold decided for each model in table 2?", "limitations": "I don't see potential negative societal impact of this work.\n\nLimitations are described in the last section.", "rating": 4, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1. The paper deals with an important and impactful task of improving the efficiency of pretrained Transformer-based NLP classifiers\n2. Empirical results suggests that the proposed modifications to PonderNet could be effective, though I have concerns on the evaluation and comparisons (see below)\n3. The writing is overall clear and easy to follow, though there are several unsupported claims (see below)\n4. The ablation studies could help future studies\n\nWeaknesses:\n\n\n1. The authors claim that other methods are heuristics, but the suggested method is also a heuristic with no theoretical grounding (as the authors admit in section 5). In fact, sampling from the geometric-like distribution seems to me more grounded as it aims to simulate the probability of a \"good exit\" exactly at the current layer, compared to the CDF that computes the probability that one of the layers so far has been a \"good exit\" layer, but the prediction is taken from the current layer, so the chance of an earlier exit being good is not relevant anymore. Perhaps the authors can clarify this?\n2. The authors main motivation for using CDF instead of sampling from the geometric-like distribution is having a deterministic exiting rule (e.g. see bolded span lines 39, and lines 118, 126-127). However, I wasn't convinced (1) why determinism is better, and (2) why the CDF is needed for achieving determinism. Assuming I had a calibrated predictor [1], I don't see a particular issue with sampling as the training and evaluation are anyway in expectation. Also, why not make a deterministic threshold-based decision directly on the lambda values like in [2]. Though, to obtain theoretical-backed guarantees on the marginal performance of the model, there should be some MHT correction, or other calibration process for the threshold value like in [3].\n3. The main evaluation in table 2 presents a different tradeoff point between the models. The accuracy is bolded as being better, but the speedup for PALBERT is lower (also, how is the speedup computed is missing). It is impossible to compare the models this way as it is unclear how does the performance compare when tuned for the same speedup or vice versa. For example, figure 3 shows at least for SST-2 several points where PABEE is better than PALBERT, but in table 2 PALBERT is reported with higher accuracy.\n4. No details on how the reported speedup value is computed. Does it take into account the computation of the lambda function (that is now larger)?\n\n[1] On Calibration of Modern Neural Networks\n\n[2] Depth-Adaptive Transformer\n\n[3] Consistent Accelerated Inference via Confident Adaptive Transformers\n\n**update post author response**: I appreciate the author's detailed response. I've decided to slightly increase my score. I'm still not entirely convinced about the motivation behind the \"deterministic\" exiting rule. But more importantly, I think that a table with only accuracies is insufficient for evaluating such early exiting models since there is some tradeoff with the actual speedup of each model. According to figure 3, the gains over PABEE are inconsistent across tasks.\n", "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": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "hY2El8hJMOD", "reviewer": "Reviewer_Ndy3", "summary": "This paper aims to improve the inference time efficiency of ALBERT, a pretrained large transformer, through the early exit. In particular, a dedicated binary classifier is trained to decide, at each layer, whether or not to use the current layer’s output as the final representation and skip further forward propagation. If the model exits at an early layer, significant efficiency gains can be achieved. This work heavily builds on a previous work along this line—the PonderNet—and improves its test-time exit criterion: instead of drawing from a Bernoulli, this work proposes to accumulate the early-exit classifier’s output probability throughout the layers and determine the early exit by comparing it to a prespecified threshold. The motivation is to reduce the variance. This method is combined with other techniques aiming to improve the capacity of the early-exit classifier. The model is evaluated on the GLUE benchmark, with extensive ablation study, tradeoff curves, and other analyses. The proposed model outperforms previous early-exit methods designed for ALBERT, and improves the efficiency with little or no accuracy drop.\n\nI find the idea very intuitive and clearly presented. The detailed analysis and honest results are a big plus. However, I do have concerns about the technical contributions. It seems that most of the proposals heavily build on previous works and are rather engineering. Maybe the authors could help me and clarify their technical contributions in the response. Further, the method is presented in a way that can only be used in ALBERT. Exploring its applications in other pretrained models could improve its impact.\n", "questions": "- It may happen that different instances in a mini-batch exit at different layers. How is this resolved? Would it affect the efficiency?\n- The paper argues that a major drawback of PonderNet is that it has a higher variance. But the std is barely larger than ALBERT in Table 1. Could the authors comment on this? Also, it would be nice to include std too in Table 2.\n- It would be interesting to see how the exit layers might correlate with the inputs’ lengths, labels, or some other metrics measuring their difficulty.\n- Have the authors compared to a handcrafted exit strategy? One can dedicate that the model must exit after X layers to achieve the 1.29x speedup, and compare its accuracy with learned exit strategies.", "limitations": "The authors have discussed the limitations in detail. I don't have any further concerns about the potential negative societal impact.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "\nStrengths:\n- Very clear writing\n- Very detailed and informative analysis of the results.\n- The method is intuitive.\n- Honest and detailed discussion of limitations\n\nWeaknesses:\n- Early exit methods have been explored in other large pretrained models (e.g., [1]). Applying Ponder ALBERT in other pretrained models could improve its impact.\n- The 1.29/1.26x speedup without sacrificing accuracy is marginal.\n\n[1] https://arxiv.org/abs/2004.07453", "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": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "S3SB_T8wGo1", "reviewer": "Reviewer_id4o", "summary": "This work proposes an improvement over PonderNet, an early-exit method that can help overcome long inference times and \"overthinking\" in a neural network. The specific improvement is the change of exit layer selection during inference-time, using a deterministic criterion instead of the probabilistic method proposed originally in PonderNet. The paper compares the improvements to both PonderNet (partially) and PABEE, a heuristic based early-exit method. The results are presented using ALBERT and a variety of GLUE tasks. Overall, the work shows that there are small but consistent gains over earlier methods, with a 1.2-1.3x speedup in inference compared to vanilla ALBERT.\n\n", "questions": "Table 1 shows the overall variance of accuracies across several tasks. While the variance of PonderNet is slightly higher than PALBERT, it does not justify the claim that there is \"major variance\" (L11, L45) in PonderNet's outputs. Do you have some quantitative evidence for this claim?\n\nLine 159 mentions that the baselines were re-implemented, but PonderNet results are missing from the test set section in Table 2. What is the reason for these missing results?\nFigure 3's caption says that the results are for two tasks, while the results for 4 tasks are presented, please fix this.\n\nThe paper repeatedly mentions that PALBERT \"significantly outperforms\" or \"outperforms by a large margin\", while the actual overall improvements are not qualified by this. Consider simplifying the language and/or presenting more evidence to support this claim.\nL160 \"Ablation\" -> \"Appendix\"", "limitations": "Authors have addressed limitations adequately.\n\n", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\nConsistent gains over previous methods \nLogical improvements over previous work (using two layers to study change across them) \nEvaluation of the method itself is done well (ablation studies for the sub-component across a few tasks, speed comparison with one other method)\n\nWeaknesses\n\nPaper is not very clearly written, and required several passes to fully grasp. \nIt seems like the PonderNet results were added after the fact (e.g. they are missing from the speed vs performance graph, missing test set results)\nThe work does not shed much light on why PALBERT works better, and seems like a more empirically motivated and studied method. While there is nothing wrong in this in an itself, the improvements themselves are fairly incremental over PonderNet, and may not surpass the base quality of this particular venue.\n\n", "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": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "7Jwe45xXGo", "reviewer": "Reviewer_RdPQ", "summary": "Previous work found that the computation of a deep neural network can exit at earlier layers to prevent \"overthinking\". This improves upon previous method (PonderNet) of performing an early exit and proposes a new deterministic Q-exit criterion. This new method determines the exit layer index in a deterministic way by using the CDF of the layer exit distribution. Experiments of language understanding tasks (GLUE) and demonstrate improved performance over PonderNet.", "questions": "The questions are asked in the weaknesses section.", "limitations": "N/A.", "rating": 3, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n1. The paper is relatively well written and proposes a simple and effective method to alleviate the high-bias issue of sampling based layer exit method.\n2. Comprehensive analysis have been conducted to show the effects of the hyperparameter of threshold, the speedup and the distribution of examples that exit at different layers.\n\n**Weaknesses:**\n1. The original PonderNet paper which this paper is based on is evaluated on more complex tasks, e.g. QA and other multi-step reasoning tasks, which I find a more appropriate setup for evaluating this layer-exit approach, which is a better test bed for these approaches. With that being said, experiments on the GLUE task alone is not sufficient to demonstrate the effectiveness of the proposed method. Moreover, as discussed in the limitation section, the proposed method still lags behind the Albert model on some tasks, I think this paper still needs further improvements for a better contribution.\n2.  This paper uses Albert model as the base model for executing the proposed method which is a limitation of the proposed method. I was wondering if there is a specific reason for the authors not to use other pre-trained models, especially deeper networks, to implement the idea. If this is because Albert uses shared weights for layers, then is the proposed method a model-specific trick but not a generic method?\n3. The threshold p that decides when to exit seems to be a very important hyperparameter for the success of the proposed method and it is selected with the validation set. This directly contrasts with selecting the best exiting layer from 1 to 12 layers as a baseline. This simple baseline is not included in the paper. Could you also compare with that?\n4. A suggestion on the presentation of the paper: it's better to describe Albert and PonderNet more clearly in a background section to make the paper more clear to the general audience.", "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": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}]

Author rebuttal:

Decision: Accept