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

Shashank Kapadia, Deep Narayan Mishra, Sujal Reddy Alugubelli, Haoan Wang, Saipraveen Vabbilisetty, Rishi Bhatia, Anupriya Sharma

Keywords: early exit, knowledge distillation, efficient inference, sentence embeddings, adaptive computation, transformer acceleration, representational convergence

TL;DR: LEAP adds an auxiliary training objective to knowledge distillation that preserves early-exit capability, achieving 1.61x wall-clock speedup where standard distilled models yield none

Abstract:

Layer-aligned distillation and convergence-based early exit represent two predominant computational efficiency paradigms for transformer inference; yet we establish that they exhibit fundamental 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.

Submission Type: Emerging

Submission Number: 162