PALBERT: Teaching ALBERT to Ponder

Why run all twelve layers every time

ALBERT reuses one transformer block: evaluation is $h_i = S(h_{i-1})$, the same layer applied twelve times, with a classifier $C$ reading the final state. That weight sharing makes the model a recurrence in disguise, and for a recurrence the number of iterations is a choice, not a constant. Stopping early is architecturally free; the question is when.

Running to the end is not even always the best choice. Zhou and colleagues showed that stopping ALBERT-Base at ten iterations can beat all twelve on MRPC while prediction entropy keeps falling, a failure they call overthinking: more computation makes the model more confident, including in its mistakes. Their PABEE exits when several per-layer classifiers agree in a row. We wanted a mechanism that learns when to stop instead of voting.

Treat the exit layer as a latent variable

PonderNet, introduced by Banino and colleagues, treats the exit layer as a latent variable. After each iteration a small Lambda layer reads the hidden state and emits a halting probability $\lambda_i = \Lambda(h_i)$, which induces a generalized geometric distribution over exit depths, $p(i\mid x)=\lambda_i\prod_{j<i}(1-\lambda_j)$. Training maximizes a variational lower bound on the likelihood:

The KL term pulls the learned exit distribution toward a geometric prior with parameter $\lambda$, the knob that says how early exits should happen a priori; we use $\lambda=0.1$ and $\beta=0.5$. We keep this training objective intact in PALBERT. What we change is everything around it: how the model decides to leave at inference time, and what the halting layer gets to see.

Sampling wastes what the model knows

At inference, PonderNet samples: at each layer it flips a coin with probability $\lambda_i$. Even a well-trained model pays for that randomness, because $\lambda_1=0.1$ still aborts at the first layer once in ten runs. Our cleanest evidence is a closed-form control that evaluates predictions in expectation over $p(i\mid x)$, removing sampling while changing nothing else: CoLA climbs from 50.8 to 55.9, and SST-2 from 91.1 to 92.3.

Q-exit makes the decision deterministic and keeps the early exit. We accumulate the exit distribution into its CDF, $\sum_{i\le k} p(i\mid x)$, and leave at the first layer where it crosses a threshold $q$. Zero variance, and $q$ becomes an interpretable dial between underthinking and overthinking; $q=0.5$ performed best across tasks in our sweeps.

We also rebuilt the halting layer itself. Whether to stop should depend on how the computation is moving, not only on where it is, so our Lambda layer reads the concatenation $[h_i, h_{i-1}]$ instead of $h_i$ alone, is a three-layer MLP with $\tanh$ activations rather than a single projection, and trains with its own learning rate.

Which change buys what

The ablation stacks the same way on the other tasks: Q-exit alone recovers most of the sampling gap, the separate Lambda learning rate adds a little, and the dynamics-aware input pushes the full model to 92.9 on SST-2, 77.8 on RTE, and 57.4 on CoLA, ahead of PonderNet's 91.1, 73.5, and 50.8 everywhere, and ahead of plain ALBERT on RTE and CoLA.

Across GLUE, determinism pays

On dev macro across the eight tasks PALBERT reaches 84.2 against 83.5 for PABEE and 84.0 for plain ALBERT; on the test split it holds 80.6 against 79.3 and 80.3. The pattern transfers to RoBERTa: PRoBERTa's dev macro of 85.6 beats PABEE's 85.0, and on test 82.1 beats 81.3, with plain RoBERTa at 86.0 and 83.2.

PABEE struggles exactly where data is scarce. Its per-layer classifiers are independent, and on CoLA and RTE they simply do not get enough signal to train well; PABEE lands below even plain ALBERT there. Our shared classifier and shared halting layer reuse everything across depths, which is where we believe the small-data advantage comes from.

The prior is the price

The honest cost of the variational frame is the prior. The exit behavior a model learns tracks $\lambda$ closely: sweeping it across 0.1, 0.15, 0.25, and 0.5 reshapes the exit histograms outright, and the metric-best RTE setting we found, $\lambda=0.08$, mostly exits at the last layer, which defeats the purpose. Choosing the prior is really choosing the behavior.

The same histograms indict sampling one more time: even when the learned distribution's mode is layer twelve, almost half of sampled runs would exit within the first ten layers. A deterministic read of the distribution is not a convenience; it is the difference between using what the model learned and gambling on it.

Two loose ends we state plainly. First, $q=0.5$ works well without a theory of why. Second, a PALBERT forced to always exit at layer twelve still beats vanilla ALBERT on SST-2, so part of our gain is the auxiliary per-layer supervision rather than early exiting itself. We suspect the next step is a trainable stopping rule that does not lean on an explicit prior at all.