Exit Transformer Classifiers Early with Entropy Thresholds

A transformer classifier normally sends every input through every layer, even when an intermediate representation already supports a concentrated class prediction. Entropy-based early exit changes that fixed-depth behavior. Prediction heads attached to intermediate layers estimate class distributions, and inference can stop once a distribution passes a configured entropy threshold.

The mechanism makes model depth input-dependent. Some inputs may leave after relatively few layers, while uncertain inputs continue through more of the network. That flexibility also introduces a new source of error: an intermediate head can be confident and still be wrong.

Intermediate heads turn depth into a decision point

Consider an encoder with L transformer layers. Standard classification uses the final hidden representation h_L and a classifier head. A multi-exit model adds heads at selected intermediate layers:

h_i -> classifier_i -> logits_i -> softmax -> p_i

At an exit layer i, the probability vector p_i describes the current class prediction. Inference either returns that prediction or computes the next transformer layer.

This differs from truncating the model at a fixed depth. A truncated encoder applies the same compute budget to every input. Early exit evaluates a stopping rule separately for each input, so the realized depth depends on intermediate predictions.

The placement of exit heads matters. A head after every layer creates many possible stopping points but also adds classifier computations and state that must be trained. Sparse exits reduce that overhead but make depth control coarser. The useful arrangement depends on the cost of each encoder layer, the number of classes, and the latency characteristics of the serving system.

Entropy measures distribution concentration

For a classifier with C classes and probabilities p_1, ..., p_C, predictive entropy is:

H(p) = -sum_j p_j * log(p_j)

Entropy is low when probability mass is concentrated and high when it is spread across classes. For example, the distributions below have the same top class but very different concentration:

A = [0.96, 0.02, 0.02]
B = [0.40, 0.35, 0.25]

A has lower entropy than B. An entropy-based policy can therefore stop on A while sending B to another layer.

A simple rule is:

if entropy(p_i) < threshold:
    return argmax(p_i)
else:
    continue

The threshold controls how readily the model exits. A larger threshold admits less concentrated distributions and tends to permit earlier exits. A smaller threshold demands greater concentration and tends to send more inputs deeper.

Raw entropy also depends on the number of classes. Its maximum for a uniform distribution is log(C), assuming the same logarithm base is used throughout. Comparing a single raw threshold across tasks with different class counts can therefore be misleading. A normalized form such as H(p) / log(C) places the uniform distribution at 1, but normalization does not turn entropy into a calibrated error probability.

Confidence is not correctness

An early-exit policy assumes that distribution concentration contains useful information about whether more computation is needed. That assumption can fail.

Softmax probabilities are derived from logits, and a sharply peaked distribution does not prove that the predicted label is correct. An intermediate head may confidently select the wrong class. If the exit rule accepts that prediction, later layers never get an opportunity to change it.

This creates a distinction between two measurements that should not be collapsed into one. Exit rate describes how often inference stops at each depth. Task quality describes whether those exits preserve acceptable predictions. A threshold can produce many early exits while degrading accuracy on exactly the cases that matter to an application.

Calibration also matters. Two intermediate heads can have similar accuracy but different logit scales, producing different entropy distributions. A single threshold applied blindly to every exit can then represent different effective confidence requirements at different layers. Per-exit thresholds are one way to account for this behavior, provided they are selected on held-out data representative of the intended input distribution.

Training determines whether early heads are usable

Adding a classifier to an intermediate hidden state does not guarantee that the state supports the final task as cleanly as the last layer does. Intermediate heads need an explicit training strategy.

One approach trains the full-depth task model first, freezes its existing parameters, then trains the added exit heads. This isolates head training from the backbone. Another approach optimizes intermediate and final objectives jointly. Joint optimization lets intermediate losses influence shared representations, which changes the model being optimized rather than merely attaching probes to a fixed encoder.

These strategies are not interchangeable. Freezing preserves the previously fitted backbone but limits what intermediate heads can extract. Joint training can improve support for earlier predictions, yet the weighting of exit losses becomes part of the model objective and can affect later representations.

Distillation is another possible design: an intermediate head can be trained against a deeper head’s distribution rather than only against hard labels. That supplies information about relative class scores, but it also transfers properties of the deeper predictor. It does not remove the need to evaluate actual early-exit errors.

Average depth is more informative than the threshold alone

An entropy threshold has no useful performance meaning without the resulting exit distribution. Two models using the same numerical threshold can produce very different depths because their heads have different entropy scales.

For exits at layers 1 through L, let q_i be the fraction of inputs that stop at layer i. The average executed depth is:

average_depth = sum_i q_i * i

This quantity gives a simple compute-oriented view when layer costs are roughly comparable. If layer costs differ or the runtime includes substantial fixed overhead, measured latency is more informative than layer count.

Batching complicates the picture further. A single input that exits early can stop immediately in a request-at-a-time execution path. In a dense batch, samples may reach different stopping decisions at different layers. The runtime must compact active samples, mask completed items, or keep computing padded work. As a result, a reduction in average executed layers does not automatically translate into the same proportional reduction in wall-clock latency.

Threshold selection is an evaluation problem

The stopping threshold should be treated as part of the deployed decision policy, not as a cosmetic inference parameter. Its effect is visible only through the joint distribution of exit depth and prediction error.

A useful evaluation records the final prediction, exit layer, intermediate entropy, and whether the prediction is correct for each held-out input. From those records, developers can examine task quality against average depth or measured latency across candidate thresholds. Per-class slices can expose cases where an aggregate metric hides concentrated regressions.

Distribution shift deserves separate attention. A threshold fitted on one input distribution may encounter different entropy behavior after the data changes. Monitoring only the final class frequencies can miss this shift; changes in exit-layer frequencies and entropy distributions can reveal that the stopping policy is operating in a different regime.

Entropy-based early exit is therefore not merely a shortcut around transformer layers. It adds a confidence-sensitive control decision inside inference. Its value depends on intermediate heads that carry useful task signal, thresholds validated against prediction quality, and a serving path capable of converting reduced depth into actual compute or latency savings.