LEAP: Layer-wise Exit-Aware Pretraining for Efficient Transformer Inference

TL;DR

LEAP introduces a new pretraining method that aligns intermediate transformer layers with final representations, enabling efficient early exit inference without architectural changes.

Contribution

It proposes LEAP, a training objective that reconciles distillation and early exit, improving inference speedup while maintaining performance.

Findings

LEAP achieves 1.61× wall-clock speedup on NVIDIA L4 hardware.

91.9% of samples exit by layer 7 with LEAP-MiniLM.

LEAP provides effective layer reduction and operational deployment guidance.

Abstract

Layer-aligned distillation and convergence-based early exit represent two predominant computational efficiency paradigms for transformer inference; yet we establish that they exhibit systematic incompatibility under standard deployment conditions for convergence-based early exit. Distillation objectives that align intermediate student layers to teacher representations suppress the representational convergence that early-exit mechanisms exploit, rendering such mechanisms ineffective on distilled models. We introduce LEAP (Layer-wise Exit-Aware Pretraining), an auxiliary training objective that reconciles this incompatibility. LEAP requires no architectural modifications; it augments standard distillation with a single constraint ensuring intermediate layers approximate final-layer representations. LEAP-MiniLM achieves 1.61$\times$ measured wall-clock speedup (batch=1, NVIDIA L4) at $\theta$=0.95, with 91.9% of samples exiting by layer 7 and 1.80$\times$ theoretical layer reduction, where standard distilled models achieve zero effective speedup. We validate across sentence similarity (STS-B: 0.760 $\pm$ 0.006) and retrieval benchmarks (BEIR), providing operational guidance including latency measurements, decision thresholds, and deployment criteria.