PALBERT: Teaching ALBERT to Ponder
Anonymous Author(s)
AffiliationAddress email
Abstract
Currently, pre-trained models can be considered the default choice for a widerange of NLP tasks.
Despite their SoTA results, there is practical evidence thatthese models may require a different number of computing layers for differentinput sequences, since evaluating all layers leads to overconfidence in wrongpredictions (namely overthinking).
This problem can potentially be solved byimplementing adaptive computation time approaches, which were first designed toimprove inference speed.
Recently proposed PonderNet may be a promising solution for performing an earlyexit by treating the exit layer’s index as a latent variable.
However, the originallyproposed exit criterion, relying on sampling from trained posterior distribution onthe probability of exiting from the i -th layer, introduces major variance in modeloutputs, significantly reducing the resulting model’s performance.
In this paper, we propose improving PonderNet with a novel deterministic Q-exitcriterion and a revisited model architecture.
We adapted the proposed mechanismto ALBERT and RoBERTa and compared it with recent methods for performing anearly exit.
We observed that the proposed changes can be considered significantimprovements on the original PonderNet architecture and outperform PABEE on awide range of GLUE tasks.
In addition, we also performed an in-depth ablationstudy of the proposed architecture to further understand Lambda layers and theirperformance.
Introduction
These days, fine-tuning pre-trained models on downstream tasks became a de facto standard techniquefor training NLP models (Devlin et al., 2019; Liu et al., 2019; Lan et al., 2020; Radford et al., 2019).
Although, as model sizes increase, it becomes harder to use them for real world applications due tocomputational 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 a specificmodel input.
One model that is widely used in real-world applications is ALBERT (Lan et al., 2020), which isbased on the Transformer architecture (Vaswani et al., 2017) with shared layers (i.e., the same layer isevaluated several times to provide an output).
Sharing layer weights allows us to think of ALBERT asof recurrent model that simplify development of early exit methods.
Most of them (Teerapittayanonet al., 2016; Liu et al., 2020; Xin et al., 2020) rely on performing an early exit based on entropyof predictions during the model evaluation.
Although, Zhou et al.
(2020) showed that running the
ALBERT-Base block for a fixed number of times (10) could increase the accuracy of the fine-tunedmodel on specific tasks (e.g., MRPC) despite the entropy of predictions monotonously decreasingfor all 12 layers.
A recent PABEE solution (Zhou et al., 2020) was designed to overcome this issue hs λ λ c hy λ = 0.05,sampleλ = 0.1,sample
To earlyrandom exit hs λ λ c hy λ = 0.05,CDF = 0.
λ =
0.1,CDF = 0.
s λ hλ =
0.6,CDF
= 0.
Deterministic exit
Once CDF becomes greater than the threshold value ( 0 . 5 in this example), we perform an early exit.
With such a deterministic criterion, we can perform an early exit from a model more robustly without introducing variance in the exit layer’s index during inference.
by performing an early exit based on the consensus between different classifier heads from differentlayers.
The model stops evaluating when several classifiers in a row produce the same result.
A possible orthogonal solution for this task is PonderNet (Banino et al., 2021) – a variational approachthat treats the exit layer’s index as a latent variable.
By maximizing the lower bound of the likelihoodof the training data, PonderNet trains a model which can predict whether it is necessary to exit from aspecific layer during evaluation.
However, Banino et al.
(2021) proposed to sample from the trainedposterior distribution of exiting from each layer during inference, which leads to major variance inmodel outputs.
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 bysampling from the trained posterior distribution during evaluation, we used a novel zero-varianceexit criterion, namely Q-exit , which evaluates the CDF of the exit layer’s probability distribution andperforms a deterministic early exit.
We also revisited the architectural choices of Lambda layers usedto predict the probability of exiting from the current layer to make them aware of dynamics in hiddenstates across previous layers and the number of currently running layers.
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 theoriginal PonderNet architecture adapted for ALBERT fine-tuning.
Furthermore, the proposed methodsoutperformed PABEE and are comparable to plain ALBERT fine-tuning, while also exceeding it inspeeds.
We also analyzed the trained model and provided insights on further improvement of thevariational approach for early exiting.
Related Work
Most of the approaches used to perform an early exit from a model are based on the probabilitydistribution of predictions: BranchyNet (Teerapittayanon et al., 2016), FastBERT (Liu et al., 2020),
DeeBERT (Xin et al., 2020), which can be seen as an entropy criterion.
Xin et al.
(2021) proposed
BERxiT with LTE layer used to estimate the correctness of predictions during the inference based onthe hidden state from the specific layer.
However, there is strong practical evidence that classification models’ overthinking causes a reductionin predictions’ entropy, making these methods difficult to use (Zhou et al., 2020).
Furthermore,it is unclear how to adapt entropy methods for regression tasks (Zhou et al., 2020).
Zhou et al.(2020)
proposed PABEE – a method to perform an early exit based on several classifiers from thedifferent levels of a model.
Once several classifiers in a row (the number of these classifiers isdetermined by the patience hyperparameter t ) produce the same result, we can perform an early exit.
Several recent works also utilized the idea of a consensus-based exiting strategy (Zhu, 2021; Liaoet al., 2021).
Although, most of them mainly augment the basic consensus scheme with orthogonalheuristics.
Because of this, we did not include them in our experiments and only used PABEE as aconsensus-based method.
An alternative way to perform an early exit is the Ponder architecture (Banino et al., 2021), whichuses auxiliary Lambda layers to predict whether a model should exit from a specific layer duringruntime.
Inputs to Lambda layers used in PonderNet are hidden states from the current layer of amodel.
PonderNet can be seen as a model with the latent variable that corresponds to the exit layerindex, which is trained by maximizing the lower bound of the marginalized likelihood of the data.
During inference, Banino et al.
(2021) proposed to sample from the trained posterior distribution ofexit layer probabilities.
However, this exit criterion can lead to uncertainty in exit layers for the sameinput (see Section 4.5.
Even if the Lambda layer produced a probability equal to 0 .
1 of exiting fromthe first layer, we could still exit a model too early in one of ten cases, even though the probability wassmall.
We also hypothesize that predicting exiting from a layer based entirely on a single hidden statecould be sub-optimal since performing early exit could also depend on the dynamics in hidden statesacross layers (i.e., the Lambda layer should know how hidden states change during the evaluation).
Improving PonderNet
The usual ALBERT evaluation can be defined as a computation of n hidden states h i
=
S ( h i − 1 )from the input embeddings h 0 of an input sequence x , where i ∈ [1; n ] .
Once h n is obtained, it ispassed to a classifier block C ( h n ) to get the parameters of an output distribution p ( y | x ) .
A commonway to fine-tune this architecture on downstream tasks is to initialize the embeddings and the S layerby using ALBERT (pre-trained on Masked Language Modelling) while initializing C randomly andthen optimizing all parameters by maximizing the likelihood of the training data.
While plain ALBERT performs a fixed number of computational steps, it is possible to perform anarbitrary number of evaluations of the layer S .
Banino et al.
(2021) proposed to extend each Transformer layer with a so-called Lambda layer.
More precisely, for each layer i , after S outputs a new h i ,it is then passed to the classifier and Lambda layers to get parameters C ( h i ) of output distributionsp ( y | x, i ) and the probability of exiting from the i -th layer
λ i =
Λ( h i ) , which induces a generalizedgeometric distribution on probability of exiting from layer i equal to p ( i | x )
=
λ i (cid:81)
i − 1 j =1 (1 − λ j ) .
Then, having the probability distribution from each layer p
( y | x, i ) , the parameters of the model areoptimized to maximize
[Equation]
Here, p ( ·| λ ) is a prior geometric distribution of exiting from each layer, parametrized by the hyperparameter λ , and E i ∼ p ( i | x ) (cid:104) p ( y | x, i ) (cid:105) is evaluated analytically by averaging likelihoods from differentlayers with posterior exit probabilities.
If we treat the exit layer index as a latent variable, thenoptimizing L from the Equation 1 could be seen as maximizing the lower bound of marginalizedlikelihood p ( y | x ) (Kingma and Welling, 2014).
Note that following Banino et al.
the prior probability of exiting from the last layer n isnormalized as p ( n | λ ) = 1 − (cid:80) n − 1 i =1 p ( i | λ ) in order to make p ( i | λ ) sum into 1 with a finite numberof steps, while p ( i | x )
it truncated by the lowest layer index j , s.t.
(cid:80) ji =1 p
( i | x ) ≥ 1 − 0 .
(thus,
KL ( . . . ) is evaluated only with first j layers).
Also note that weights of Lambda layers are sharedacross layers of the model.
While we implied using ALBERT above, all these derivations could be extended to any other BERTarchitecture.
E.g. we also adopted it for RoBERTa (namely PRoBERTa) for our experiments.
3.1 Exit Criterion
During inference, Banino et al.
(2021) proposed to sample the exit layer index from p ( i | x ) (i.e.,by sampling iteratively from a Bernoulli distribution with parameter λ i ).
While a sampling-basedexit criterion correlates with the variational view of PonderNet’s training objective, such estimationintroduces the randomness in the inference process of PonderNet.
[Table]
+ PonderNet [Table] -exit [Table]
LR
[Table]
+ - - -
[Table]
±
[Table] + + - - [Table] + + + -
[Table]
+ + - +
[Table]
+ + + +
[Table] 9 ± [Table] To overcome the issue of randomness, we propose Q-exit 1 : a novel deterministic criterion ofperforming an early exit, which we used for PALBERT.
Instead of sampling from the distributionp ( i | x ) during inference, we evaluate its CDF by accumulating p ( i | x ) from each layer.
Once the
CDF is greater than the threshold hyperparameter q , we perform an early exit 2 .
See Figure 1 for aschematic comparison of the sampling criterion with Q-exit.
Threshold q can be seen as a trade-offbetween underthinking and overthinking.
Therefore, q should be selected during the validation of thetrained model to choose the best-performing value.
Based on our experiments, we found that the proposed criterion produced significantly better accuracyon various tasks compared to the original sampling criterion (see Sections 4.1, 4.2), while also beingmore practical than the original sampling criterion.
3.2 Lambda Layer Architecture
While the original PonderNet used a single layer MLP to obtain logit of exiting probability, wehypothesize that making the Lambda layer understand the dynamics of changing ALBERT hiddenstates is crucial for achieving good performance.
To do so, instead of passing a single hidden state
h ifrom the i -th layer in Λ , we concatenate it with h
i − 1 .
I.e., for PALBERT, we evaluate the probabilityof exiting from i -th layer as λ i = Λ([ h i , h
i − 1 ]) .
We used a 3 layer MLP with tanh activation for the Lambda layer to operate with more complexinput.
Based on the ablation study, we observed that increasing the capacity improves the accuracy ofthe trained model (See section 4.1).
We also found it beneficial to fine-tune the Lambda layer with adifferent learning rate than all other parameters.
Experiments
4.1 Ablation Study
We performed an ablation study of the proposed changes in PonderNet architecture.
We experimentedwith adding the proposed Q-exit criterion, Lambda layer architecture, and fine-tuning strategies.
Wealso compared the proposed changes with fine-tuning vanilla ALBERT and PonderNet adapted for
ALBERT with the evaluation of predictions in expectation (i.e., averaging predictions based on valuesof p ( i | x ) ).
These methods were benchmarked on SST-2, RTE, and CoLA tasks from the GLUE
Benchmarking dataset (Wang et al., 2018).
For evaluation, we performed a grid hyperparameter search on an appropriate metric score on thedev split for each dataset.
Following the PABEE training setup, we trained all models with a fixed which is impractical if we want to perform an early exit. learning rate until validation metrics stopped increasing for 5 epochs.
We used Adam optimizer(Kingma and Ba, 2015) for all experiments, a fixed q = 0 .
5 on models with the Q-exit criterion, aswell as a fixed classifier dropout value equal to 0 .
1 (Srivastava et al., 2014), and λ = 0 .
We trained each model 5 times with the best hyperparameters and reported the mean and std values.
A full list of the methods’ hyperparameter ranges can be found in Table 3. [Table] ALBERT [Table]
ALBERT PABEE
[Table]
ALBERT [Table] PALBERT [Table]
RoBERTa
[Table]
RoBERTa PABEE
[Table]
PRoBERTa (ours)
[Table] ALBERT PABEE [Table] ALBERT PonderNet
[Table]
RoBERTa [Table] See Table 1 for the full list of the results of our ablation study. Based on these experiments, PonderNetarchitecture is seen as performing significantly worse than vanilla ALBERT fine-tuning.
At the sametime, the deterministic Q-exit criterion improves PonderNet accuracy when compared to a randomsampling of the exit layer and is comparable to the evaluation of PonderNet in expectation.
A morecomplex Lambda layer that can handle the dynamics of hidden state changes can further improve model accuracy when compared to the original PonderNet.
4.2 GLUE Experiments
We compared PALBERT with PonderNet architecture adapted for ALBERT fine-tuning, and PABEEwith a fixed patience value t = 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 baselinesused in our experiment for consistency in comparison.
See Ablation Section A for the details ofreimplementation.
Following the experimental setup from Section 4.1, we trained 5 models with the best hyperparametersacross the hyperparameter search and reported the median task score on the dev set.
We evaluated thetest scores on the best models, selected based on their dev scores.
See Table 2 for the full list of results.
We observed that PALBERT significantly outperformed PABEEon a wide range of tasks.
Vanilla PonderNet with the sampling exit criterion performed the worst.
Vanilla ALBERT outperformed PABEE on most tasks and is comparable to PALBERT, while thelatter has the highest score averaged across all tasks (see Macro column in Table 2).
For experiments with RoBERTa, we observed that PRoBERTa outperformed PABEE on Dev split,while it either performed better or marginally worse for the Test set.
Thus, PRoBERTa showed ahigher macro score compared to PABEE.
Note that PABEE is performing poorly for tasks with a small dataset (e.g., CoLA, RTE).
Wehypothesize that this is caused by several independent classifiers at each layer
C i failing to train wellenough, whereas PALBERT was capable of utilizing knowledge sharing between layers.
[Table]
4.3 Understanding the Threshold of Q-exit
As noted previously in Section 3.1, we treat the threshold value q of the Q-exit criterion as a trade-offbetween underthinking and overthinking, where increasing q forces a model to evaluate more layers,and vice versa.
Therefore, it is necessary to find the best-performing threshold for each task where a model has thehighest accuracy.
To do so, we evaluated trained PALBERT models from the ablation study (see
Section 4.1) on dev splits of tasks with different values of q .
We then averaged obtained metrics andreported the mean and std values for various thresholds (See Figure 2 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 a reduction in accuracy and can often force models to evaluate all12 layers of ALBERT-Base.
It is also notable that PALBERT, having a large threshold value q that performs constant exit on thelast layer, has better accuracy than vanilla ALBERT fine-tuned for the SST-2 task.
This indicatesthat adding auxiliary tasks for each layer of ALBERT improves performance compared to plainfine-tuning with a single classifier on the last layer of the model.
4.4 Speed Analysis
While making q < 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 differentpatience values t , which stands for the number of layers necessary to output the same result in a rowto perform an early exit.
We evaluated task scores for the specified hyperparameters as well as theincrease in speed when compared to vanilla ALBERT inference of a full model with 12 layers.
Overall, we observed that PALBERT mostly produced higher scores on different tasks while also beingslightly faster than PABEE.
For the CoLA and MRPC datasets, PALBERT performed significantlybetter.
The proposed method outperformed PABEE by a large margin while achieving the sameincrease in speed.
We observed questionable results on the SST-2 dataset: the best score for the PABEE model is slightlyhigher than that of PALBERT.
However, it was obtained with a negligible increase in speed comparedto ALBERT, as the best-performing patience for this setup is 11 layers while the whole model hasonly 12 layers.
Furthermore, unlike PALBERT, PABEE performed significantly worse than plain ALBERT fine-tunedon the CoLA and RTE tasks.
4.5 On the Prior of PALBERT
While we used a fixed value of λ = 0 . 1 in our main experiment following Banino et al.
(2021) (see
Section 4.2), we also investigate the effect of changing the parameter of the prior distribution on exitlayer indices.
We followed the experimental setup from the previous sections and trained models with λ ∈[0 .
For each λ , we had its own best hyperparameter set, which was then usedto train
See Figure 4 for the results.
We observed that varying λ affects the performance of the trained modelsince it has a huge impact on the distribution of exit layer indices.
Although the best metric for the
RTE task was reached with λ = 0 .
08 , such a model mostly performs an exit from the last layer.
Meanwhile, λ = 0 .
08 performs marginally worse than λ = 0 .
1 for MRPC.
Note that while Figure 4 was obtained using the Q-exit criterion, we could also estimate
E x ∼ D (cid:104) p ( i | x ) (cid:105) (see Figure 5), which is the distribution of exit layer indices with vanilla sampling criterion proposed by Banino et al.
Even if λ is low enough so that the mode of distribution isthe 12 -th layer, in almost half of the cases, the sampling criterion forces the model to exit from one ofthe first 10 layers.
We hypothesize that such randomness in the index of the exit layer leads to thepoor performance observed during previous experiments (see Section 4.2).
From such perspective, dependency on the explicit prior distribution could be seen as the mainlimitation of the proposed method, since treating the exit layer index as a latent variable justified itsprior distribution.
If proper prior distribution was found, then the well-performing model could betrained, and vice versa.
Future Work and Limitations
In this paper, we proposed improving the PonderNet architecture to perform an early exit usingfine-tuned ALBERT and RoBERTa models with the novel Q-exit criterion and a revisited Lambdalayer architecture.
This approach is orthogonal to the consensus-based approaches (Zhou et al., 2020;
Zhu, 2021) highly represented in the field.
While PALBERT and PRoBERTa outperformed some recent State-of-The-Art methods used for anearly exit, there is a clear direction for further improvement of this method, as it was not capable ofoutperforming plain ALBERT and RoBERTa on some GLUE tasks.
We believe that PALBERT could benefit from the development of a new parameterization of the priordistribution on exiting from each layer since it directly affects the resulting posterior distributionused to perform an early exit (see Section 4.5).
However, the development of an appropriate priordistribution for early exiting models is orthogonal to the development of a proper exit criterion.
Furthermore, there is no theoretical justification for the Q-exit threshold value.
Although we observedthat q = 0 .
5 performed best, it is without a clear explanation as to why that is so.
We hypothesize thatbringing more insights into developing deterministic exit criteria could further improve the proposedmethod.
In addition, adding more auxiliary tasks could also make it possible to improve PALBERT further.
This way, PALBERT training can be made more PABEE-like by making independent classifiers foreach layer of the model or adding self-distillation across layers.
While the variational view of the early exiting mechanism allows us to train models which couldbe exited early, we observed that the overall performance of the trained model highly depends onthe prior distribution of exit layer indices (see Section 4.5).
Because of this, the proposed methodrequires more computation for hyperparameter search compared to PABEE (Zhou et al., 2020).
Webelieve that, in order to further develop training-level early exiting, it may be beneficial to move awayfrom variational inference to some other form of trainable mechanism without relying on exit layers’explicit prior knowledge.
A Re-implementation of Baselines
While PABEE (Zhou et al., 2020) provided official implementation of the proposed method, we haveseveral concerns regarding its reproducibility and experimental setup.
The original implementationimplies running evaluation iteration-wise (i.e., the best performing model is selected once in niterations).
In this case, several iterations n between validation epochs could be seen as an additional hyperparameter, which has a large impact on the resulting performance.
Although, Zhou et al.
(2020) did notinclude the best performing value so that the results could be easily reproduced.
Because of this, we decided to re-implement PABEE strictly following the original paper and trainingdetails, while evaluating the model after every epoch, instead of evaluating after n training steps,which we believe is a fairer way to compare methods.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] See Section 4.
(b) Did you describe the limitations of your work?
[Yes] See Section
(c) Did you discuss any potential negative societal impacts of your work?
[N/A]
We
proposed a method that allows us to perform an early exit from an ALBERT model.
We believe that there are no any potential negative societal impacts of our work.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A]
(b) Did you include complete proofs of all theoretical results?
[N/A]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main exper-
imental results (either in the supplemental material or as a URL)?
[Yes] We ran all
experiments on openly available GLUE benchmark datasets and provide source code
for our paper.
We also provide WandB configs for hyperparameter search and best
model evaluation on dev split
of the data
so that the results could be easily reproduced.(b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes] We conducted all experiments on default splits of the GLUEbenchmarking datasets.
Hyperparameter ranges are available in Table 3.(c)
Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] We reported mean and std values across different seedsfor the ablation study (see Table 1), while for the main experiment, we decided toreport only median values across runs for consistency with recent works.
During theexploration of the proposed method in Sections 4.3, [-4.
 ,-]4.5, we also reported std valuesacross runs as well as mean values.(d)
Did you include the total amount of computing and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[No] We did not include this informationdue to NDA concerns.4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] See Section
(b) Did you mention the license of the assets
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL
[N/A
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiableinformation
or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
text instructions given participants and screenshots
[Table]
participant
Review
[Table]
[Table]
[Table]