# PALBERT: Teaching ALBERT to Ponder

arXiv: https://arxiv.org/abs/2204.03276 · HTML: https://mlsanity.com/paper/2204.03276/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Placed in Early exit and edge inference, one of 12 directions in Efficient LLM inference and serving.
- Builds on 12 of its 20 references that are in the map.
- First posted on arXiv 7 Apr 2022.

## Placed in

- Efficient LLM inference and serving / [Early exit and edge inference](https://mlsanity.com/inference/early-exit-and-edge-inference/)

## Builds on (most cited first, 12 of 12)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [ALBERT: A Lite BERT for Self-supervised Learning of Language Representations](https://mlsanity.com/paper/1909.11942/) (2019)
- [BranchyNet: Fast Inference via Early Exiting from Deep Neural Networks](https://mlsanity.com/paper/1709.01686/) (2017)
- [Adaptive Computation Time for Recurrent Neural Networks](https://mlsanity.com/paper/1603.08983/) (2016)
- [DeeBERT: Dynamic Early Exiting for Accelerating BERT Inference](https://mlsanity.com/paper/2004.12993/) (2020)
- [BERT Loses Patience: Fast and Robust Inference with Early Exit](https://mlsanity.com/paper/2006.04152/) (2020)
- [FastBERT: a Self-distilling BERT with Adaptive Inference Time](https://mlsanity.com/paper/2004.02178/) (2020)
- [PonderNet: Learning to Ponder](https://mlsanity.com/paper/2107.05407/) (2021)

## Built on by (most cited first, 6 of 6)

- [Adaptive Computation with Elastic Input Sequence](https://mlsanity.com/paper/2301.13195/) (2023)
- [BEEM: Boosting Performance of Early Exit DNNs using Multi-Exit Classifiers as Experts](https://mlsanity.com/paper/2502.00745/) (2025)
- [Diffusion Language Models Generation Can Be Halted Early](https://mlsanity.com/paper/2305.10818/) (2023)
- [Beyond Greedy Exits: Improved Early Exit Decisions for Risk Control and Reliability](https://mlsanity.com/paper/2509.23666/) (2025)
- [Improving Prediction Certainty Estimation for Reliable Early Exiting via Null Space Projection](https://mlsanity.com/paper/2506.17249/) (2025)
- [VARAN: Variational Inference for Self-Supervised Speech Models Fine-Tuning on Downstream Tasks](https://mlsanity.com/paper/2508.12061/) (2025)

## Abstract (arXiv)

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.