PALBERT: Adaptive ALBERT Transformer

• PALBERT is an adaptive extension of ALBERT that dynamically adjusts inference depth using a deterministic Q-exit criterion.
• The architecture innovates by replacing fixed-depth inference with per-input exit decisions, mitigating overthinking and improving computational efficiency.
• Evaluated on GLUE benchmarks, PALBERT achieves competitive accuracy with up to a 1.5× speed-up, demonstrating robust performance improvements.

PALBERT is an adaptive computation extension of the ALBERT-Base Transformer model, designed to address inefficiencies and inferential instability arising when standard pre-trained models are applied to a wide variety of NLP tasks. By introducing a mechanism that adaptively determines the number of transformer layers to execute for each input, PALBERT mitigates overthinking (unnecessarily deep computation on easy instances) and yields early-exit behaviors while preserving, and in some cases exceeding, the accuracy of deterministic full-depth baselines. The architecture draws on and refines the PonderNet framework for adaptive computation, introducing determinism and improved modeling fidelity through methodological and architectural innovations (Balagansky et al., 2022).

1. Model Architecture and Adaptive Computation Mechanism

PALBERT replaces ALBERT's fixed-depth inference paradigm with a variational, per-input exit decision system inspired by PonderNet. In vanilla ALBERT, the shared transformer block $S(\cdot)$ is applied $n$ times,

$h_i = S(h_{i-1}), \quad i = 1, \dots, n$

with a single classifier $C(h_n)$. PALBERT, in contrast, augments each iteration $i$ with:

• An auxiliary classifier head $C(h_i)$ producing $p(y\mid x, i)$,
• A Lambda module $\Lambda(h_i, h_{i-1})$ outputting stop probability $\lambda_i$.

This design elevates the exit index $z \in \{1, \dots, n\}$ to a latent variable, allowing the network to select, per instance, the optimal computational depth. PALBERT thereby dynamically trades off computation and predictive confidence, targeting efficiency and the avoidance of overthinking.

2. Exit Criteria: From Stochastic Sampling to Q-exit

PonderNet's original method chooses an exit layer by sequentially sampling Bernoulli random variables with probabilities $\lambda_i$, resulting in high-variance and occasionally suboptimal exit decisions:

$$p(z=i \mid x) = \lambda_i \prod_{j=1}^{i-1}(1-\lambda_j)$$

This stochasticity can lead, for example, to exiting at layer 1 with probability $\lambda_1=0.1$, imposing performance degradation and instability across runs.

PALBERT adopts a deterministic "Quantile-exit" (Q-exit) criterion. The cumulative posterior probability up to layer $i$,

$$F(i\mid x) = \sum_{k=1}^{i} \lambda_k \prod_{j=1}^{k-1}(1-\lambda_j),$$

is computed at each step. Inference exits at the first $i$ where

$F(i \mid x) \geq q$

for a user-specified quantile threshold $q \in (0,1)$; $q=0.5$ empirically achieves a balance between underthinking and overthinking across GLUE benchmarks. This rule eliminates sampling-induced variance and yields deterministically repeatable decisions.

3. Lambda Layer Design and Optimization

PALBERT introduces critical architectural modifications to the Lambda (stop-probability) layer. Instead of the single hidden state $h_i$ as input, the Lambda module receives the concatenation $[h_i; h_{i-1}]$, providing richer context for the stop decision. The module itself is implemented as a three-layer feedforward network with hidden widths matching ALBERT's hidden size and $\tanh$ activations:

$\lambda_i = \Lambda([h_i; h_{i-1}])$

Layer weights, including the Lambda layers, are shared across depth, maintaining parameter efficiency. To enhance training stability, a distinct (larger) learning rate is assigned to the Lambda parameters during fine-tuning.

4. Training Objective and Variational Inference

Training follows PonderNet's evidence lower bound (ELBO), treating $z$ as a latent variable:

$$\mathcal{L}(x,y) = -\left\{ \mathbb{E}_{i \sim p(i|x)}[\log p(y|x,i)] - \beta\, \mathrm{KL}(p(\cdot|x)\|p(\cdot|\lambda)) \right\} \leq -\log p(y|x)$$

where:

• $p(i|x)$ is the posterior (defined by the Lambda layers),
• $p(i|\lambda)$ is a geometric prior over layers (clipped to sum to unity),
• $\beta > 0$ controls the likelihood-regularization tradeoff ($\beta=0.5$ in practice).

The expectation over $i$ is computed in closed-form, weighting each layer’s cross-entropy loss by $p(i|x)$—no sampling is required. Hyperparameters include learning rate $$\in\{1\mathrm{e}{-5}, 2\mathrm{e}{-5}, 3\mathrm{e}{-5}, 5\mathrm{e}{-5}\}$$, batch size $\in\{16,32,128\}$, Lambda learning rate in $$\{1\mathrm{e}{-5}, 2\mathrm{e}{-5}, 3\mathrm{e}{-5}\}$$, $\lambda=0.1$, and $q=0.5$.

5. Empirical Evaluation: Performance and Efficiency

PALBERT was evaluated across all eight GLUE tasks (SST-2, RTE, QNLI, CoLA, MRPC, MNLI, QQP, STS-B), with baselines including ALBERT-Base, ALBERT+PonderNet, ALBERT+PABEE, and analogous RoBERTa-based models. Key metrics were dev/test macro scores, accuracy, Matthews correlation, F1/accuracy, Pearson/Spearman correlation, average exit layer, and inference speed-up.

Model	Dev Macro	Test Macro	Average Layers Used	Speed-up (vs ALBERT-Base)
ALBERT-Base	84.0	80.3	12	1× (full)
ALBERT + PonderNet (sampling)	81.2	77.2	varied	(stochastic, unstable)
ALBERT + PonderNet (closed-form)	83.7	79.9	varied	modest
ALBERT + PABEE ($t=6$)	83.5	79.3	6–9	1.2–1.5×
PALBERT (Q-exit, $q=0.5$)	84.2	80.6	6–9	1.2–1.5×

For the RoBERTa-based “PRoBERTa”, dev macro 85.6 (vs PABEE 85.0), and test macro 82.1 (vs 81.3) were recorded. PALBERT consistently outperformed stochastic PonderNet and PABEE across average macro metrics, with speed-up factors up to $1.5\times$. Lower $q$ further reduces average active layers, allowing control of the efficiency-accuracy trade-off (Balagansky et al., 2022).

6. Ablation Studies and Architectural Insights

Ablations demonstrated the impact of PALBERT’s key innovations:

• Replacing stochastic sampling with Q-exit improves accuracy by 1–2 points.
• Assigning a larger Lambda learning rate adds +0.3–0.5 points.
• Deepening $\Lambda$ (1→3 layers) yields +0.2–0.4 points.
• Concatenating $h_i, h_{i-1}$ into $\Lambda$ adds up to +0.5 points.

These interventions collectively close the performance gap with vanilla ALBERT (and, in some tasks, yield mild improvements) while enabling robust, deterministic early-exit. An auxiliary effect of per-layer classifiers is observed; PALBERT marginally outperforms vanilla ALBERT on some tasks even when always exiting at layer 12, suggesting a regularization benefit from multi-layer supervision.

7. Limitations and Future Directions

PALBERT's reliance on a geometric prior $\lambda$ introduces sensitivity: varying $\lambda$ shapes exit probabilities and can significantly influence final accuracy distributions. This suggests that data-driven or learned priors might yield more robust calibration. The chosen value of $q=0.5$ for Q-exit, while empirically balanced, lacks theoretical justification, indicating an opportunity for future principled or task-specific threshold selection.

Further, PALBERT's design could integrate self-distillation or cross-layer consistency objectives, as in PABEE, for possible additional gains. Moving beyond variational inference toward fully differentiable and deterministic exit procedures, potentially eliminating hand-tuned priors, is identified as a direction for reducing complexity and hyperparameter search demands (Balagansky et al., 2022).