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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Kapadia, Shashank, Mishra, Deep Naryan, Alugubelli, Sujal Reddy, Wang, Haoan, Vabbilisetty, Saipraveen, Bhatia, Rishi, Sharma, Anupriya
Format: Preprint
Published: 2026
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866914526039900160
author Kapadia, Shashank
Mishra, Deep Naryan
Alugubelli, Sujal Reddy
Wang, Haoan
Vabbilisetty, Saipraveen
Bhatia, Rishi
Sharma, Anupriya
author_facet Kapadia, Shashank
Mishra, Deep Naryan
Alugubelli, Sujal Reddy
Wang, Haoan
Vabbilisetty, Saipraveen
Bhatia, Rishi
Sharma, Anupriya
contents 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 $θ$=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.
format Preprint
id arxiv_https___arxiv_org_abs_2605_01058
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle LEAP: Layer-wise Exit-Aware Pretraining for Efficient Transformer Inference
Kapadia, Shashank
Mishra, Deep Naryan
Alugubelli, Sujal Reddy
Wang, Haoan
Vabbilisetty, Saipraveen
Bhatia, Rishi
Sharma, Anupriya
Machine Learning
Artificial Intelligence
Computation and Language
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 $θ$=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.
title LEAP: Layer-wise Exit-Aware Pretraining for Efficient Transformer Inference
topic Machine Learning
Artificial Intelligence
Computation and Language
url https://arxiv.org/abs/2605.01058