English

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

Machine Learning 2026-05-05 v1 Artificial Intelligence Computation and Language

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.

Keywords

Cite

@article{arxiv.2605.01058,
  title  = {LEAP: Layer-wise Exit-Aware Pretraining for Efficient Transformer Inference},
  author = {Shashank Kapadia and Deep Naryan Mishra and Sujal Reddy Alugubelli and Haoan Wang and Saipraveen Vabbilisetty and Rishi Bhatia and Anupriya Sharma},
  journal= {arXiv preprint arXiv:2605.01058},
  year   = {2026}
}

Comments

Accepted at ACL 2026 (Industry Track). 14 pages, 5 figures

R2 v1 2026-07-01T12:45:54.465Z