Skip to content

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

Conference: ACL 2026 (Industry Track · Emerging)
arXiv: 2605.01058
Code: TBD
Area: Model Compression / Early-Exit Inference / Knowledge Distillation / Sentence Embeddings
Keywords: Early-Exit Inference, Layer-wise Distillation, MiniLM, Sentence Embeddings, Inference Acceleration

TL;DR

This work demonstrates theoretically and empirically that "layer-wise alignment distillation" and "convergence-based early exit" are systemically incompatible in standard deployments: distilled models utilize every layer with no redundancy for early exiting. It proposes LEAP, a parameter-free auxiliary training objective that forces intermediate layers to approximate final layer representations early. This achieves a 1.61× actual wall-clock speedup on MiniLM-L12 (batch=1, with 91.9% of samples exiting at L7).

Background & Motivation

Background: Dense text embeddings are central to modern retrieval, semantic search, RAG, and recommendation systems. Two mainstream acceleration routes have been refined over years: (a) Knowledge Distillation: MiniLM, DistilBERT, and TinyBERT use layer-wise alignment to compress large teachers into small students; (b) Early-Exit Inference: Works like DeeBERT, FastBERT, PABEE, BERxiT, and CALM observe intermediate representation "convergence" to exit early. Intuitively, these should be combined for dual acceleration.

Limitations of Prior Work: The authors find that in industrial practice, attaching early-exit infrastructure to distilled models like MiniLM triggers convergence thresholds at intermediate layers, but actual wall-clock time increases instead of decreasing. The overhead of layer-wise similarity monitoring outweighs the gains from early termination. Essentially, "early exit appears to work but never actually exits significantly."

Key Challenge: Layer-wise alignment distillation (\(\mathcal{L}_{\text{distill}}=\sum_l \text{KL}(\mathbf{h}_s^{(l)} \| \mathbf{h}_t^{(\pi(l))})\)) uniformly distributes teacher capacity across all student layers, optimizing for "every layer is important." Early exit requires "subsequent layers to do progressively less" to enable termination. These goals are contradictory. Consequently, the inter-layer similarity \(\cos(\mathbf{e}_l, \mathbf{e}_L)\) in distilled models remains \(< 0.3\) until L12, leaving no natural exit points.

Goal: (1) Formalize this "distance-exit incompatibility" and provide measurable diagnostic metrics; (2) Design a training objective that requires no architectural changes or additional inference parameters to allow distilled models to retain both compression gains and early-exit capabilities; (3) Provide a deployment guide for practitioners (thresholds, wall-clock, fallback conditions).

Key Insight: For early exit to be effective, intermediate representations must approximate final representations. The authors propose adding an explicit approximation constraint alongside distillation losses, forcing intermediate layers to match both the teacher's final layer and the student's own final layer using a progressive soft margin via a sigmoid function.

Core Idea: In addition to standard distillation and final alignment, an "Exit Quality Loss" \(\mathcal{L}_{\text{exit}}\) is added. It uses a dual target (teacher final + student final with stop-gradient) and a sigmoid soft margin to push intermediate layers past the \(\tau=0.98\) similarity line, proactively creating exit points. At inference, a parameter-free patience-based convergence criterion (\(\cos(\mathbf{p}_l, \mathbf{p}_{l-k}) \geq \theta=0.95\)) is used.

Method

Overall Architecture

LEAP modifies only the training loss, requiring no architectural changes or additional inference parameters:

  • Training Phase: Teacher BERT-large (NLI fine-tuned) → Student MiniLM-L12. Objective: \(\mathcal{L}_{\text{LEAP}} = \mathcal{L}_{\text{final}} + \alpha \mathcal{L}_{\text{inter}} + \beta \mathcal{L}_{\text{exit}} + \delta \mathcal{L}_{\text{contrast}}\), with \(\alpha=0.3, \beta=0.4, \delta=0.3\).
  • Inference Phase: Starting from \(l_{\min}=6\), calculate \(s_l = \cos(\mathbf{p}_l, \mathbf{p}_{l-k})\) (patience \(k=1\)) at each layer. Exit at the first \(s_l \geq \theta=0.95\).
  • Training threshold \(\tau=0.98\) is stricter than inference threshold \(\theta=0.95\) to provide headroom for distribution shift.

Key Designs

  1. Exit Quality Loss \(\mathcal{L}_{\text{exit}}\) (Dual-target soft margin):

    • Function: Actively pushes intermediate representations to be "nearly equal to final layers," creating viable exit points. Abolishing this (Ablation C.5) nullifies LEAP.
    • Mechanism: Calculates two losses for each intermediate layer \(l\). Teacher-side: \(\mathcal{L}_{\text{exit}}^{(t)} = \frac{1}{L_s}\sum_l w_l \cdot \sigma(10\cdot(\tau - \cos(\mathbf{e}_s^{(l)}, \mathbf{e}_t^{(L_t)})))\), forcing closeness to the teacher's final embedding. Student-side: \(\mathcal{L}_{\text{exit}}^{(s)} = \frac{1}{L_s-1}\sum_l w_l \cdot \sigma(10\cdot(\tau - \cos(\mathbf{e}_s^{(l)}, \text{sg}(\mathbf{e}_s^{(L_s)}))))\), using stop-gradient to align with the student's own final output. Total \(\mathcal{L}_{\text{exit}} = \mathcal{L}_{\text{exit}}^{(t)} + 0.7\mathcal{L}_{\text{exit}}^{(s)}\). The sigmoid with coefficient 10 creates a "saturated once crossed" soft margin.
    • Design Motivation: (a) Relying only on teacher final targets conflicts with intermediate distillation \(\mathcal{L}_{\text{inter}}\) as they point to different "teacher levels"; (b) Adding the student final target with stop-gradient ensures self-consistency at inference. A 0.7 weight ensures the teacher dominates quality while the student dominates self-consistency.

  2. Decoupled Thresholds: \(\tau=0.98\) vs \(\theta=0.95\):

    • Function: Decouples "training strictness" from "inference aggressiveness" for robustness to distribution shift.
    • Mechanism: Training enforces a high bar of \(\tau=0.98\); inference uses a lower patience threshold of \(\theta=0.95\), providing 0.03 margin for real-world perturbations. Pareto curves show STS-B remains stable (0.753–0.762) for \(\theta \in [0.93, 0.97]\) while the average exit layer moves from 4.6 to 8.9.
    • Design Motivation: Engineering risks include "exiting during training but never reaching thresholds in production." Decoupling provides a safety buffer and a single knob for production tuning.

  3. Zero-parameter Patience-based Exit:

    • Function: Avoids additional learnable modules (unlike DeeBERT's classifiers or PABEE's exit heads) by relying purely on geometric similarity.
    • Mechanism: From \(l_{\min}=6\), compare \(\cos(\mathbf{p}_l, \mathbf{p}_{l-k})\) with \(k=1\). Exit when it exceeds \(\theta\). This adds only one mean-pool and one cosine calculation, much lighter than a classifier head.
    • Design Motivation: Learned exit heads require task-specific fine-tuning, which is unsuitable for "label-free" sentence embedding scenarios (DeeBERT STS-B Spearman 0.26 vs LEAP 0.76). Being task-agnostic is critical for industrial embedding services.

Loss & Training

Final objective: \(\mathcal{L}_{\text{LEAP}} = \mathcal{L}_{\text{final}} + 0.3\mathcal{L}_{\text{inter}} + 0.4\mathcal{L}_{\text{exit}} + 0.3\mathcal{L}_{\text{contrast}}\), where \(\mathcal{L}_{\text{final}}=1-\cos(\mathbf{e}_s^{(L_s)},\mathbf{e}_t^{(L_t)})\); \(\mathcal{L}_{\text{inter}}\) is layer-wise cosine alignment; \(\mathcal{L}_{\text{contrast}}\) uses KL alignment of the batch similarity matrix. Trained on AllNLI 1.5M pairs for 10 epochs (batch 64, lr \(5\times 10^{-5}\)) taking \(\sim\)14h on 4×L4 GPUs.

Key Experimental Results

Main Results

Comparison of MiniLM-L12 vs. LEAP-MiniLM-L12 on STS-B:

Model STS-B \(\rho\) Layer Ratio Wall-clock Speedup \(\mathbb{E}[\text{layer}]\) Exit@L7
Published MiniLM-L12-v2 0.831 1.00× 1.00× 12.0 0%
MiniLM-L12 (baseline, same pipeline) 0.777 1.00× 1.00× 12.0 0%
LEAP-MiniLM-L12 0.760 ±0.006 1.80× 1.61× 6.7 91.9%

Ours achieves a 1.61× wall-clock speedup at the cost of 2.2% STS-B quality. Standard MiniLM exhibits 0% exit rates even with the LEAP inference protocol (L7 similarity is only 0.29 vs LEAP's 0.96), confirming incompatibility is intrinsic to the distillation objective.

Cross-Distillation Compatibility (Max Exit Rate):

Model Distillation Type Max Exit Rate
TinyBERT-6 Layer-wise (MSE on hidden) 0.0%
MiniLM-L6-v2 Layer-wise (KL on attention) 0.7%
DistilBERT-6 Output-only distillation 71.5%

Only DistilBERT, which lacks layer-wise alignment, retains natural early-exit ability, confirming layer-wise alignment is the root cause.

Key Findings

Layer-wise Similarity Comparison:

Layer MiniLM (baseline) Sim MiniLM Exit% LEAP Sim LEAP Exit%
6 0.162 0.0% 0.945 38.9%
7 0.215 0.0% 0.963 91.9%
8 0.285 0.0% 0.968 97.6%
10 0.547 0.0% 0.975 99.5%
12 1.000 100% 1.000 100%

LEAP maintains similarity \(>0.9\) starting from L6, while baseline remains at 0.86 even at L11.

Pareto Curve (\(\theta\) vs Quality/Speed):

\(\theta\) STS-B Layer Ratio \(\mathbb{E}[\text{layer}]\)
0.90 0.756 2.58× 4.6
0.95 (recommended) 0.763 1.80× 6.7
0.99 0.762 1.08× 11.1

Wall-clock vs Batch Size (NVIDIA L4):

Batch Full (ms) EE (ms) Speedup
1 8.46 5.25 1.61×
8 11.51 8.75 1.32×
32 13.14 10.61 1.24×

BEIR Retrieval (NDCG@10): LEAP full inference outperforms the baseline on 3/5 tasks (+3.3% avg), showing \(\mathcal{L}_{\text{exit}}\) has almost zero cost for sentence embedding quality. Early exit drops 24.7% on ArguAna (which requires deep semantic composition) but remains flat on NFCorpus/FiQA, indicating exit costs are task-dependent.

Highlights & Insights

  • Layer-wise alignment, not distillation itself, is the culprit: DistilBERT supports early exit naturally by focusing only on final output KD. TinyBERT/MiniLM "kill" exit paths by locking every layer with KL/MSE.
  • Gains depend on Batch=1: As batch size increases, GPU parallelism amortizes per-layer costs (1.61× → 1.24×). LEAP is a victory for real-time low-latency scenarios, not throughput.
  • Decoupled thresholds provide robustness: Training with \(\tau=0.98\) and inferring with \(\theta=0.95\) leaves a 0.03 buffer, keeping the Pareto curve flat for STS-B.
  • Scientific "Falsifiable Prediction": The authors state that any distilled model maintaining monotonic convergence towards the final layer would not need LEAP. This level of scientific rigor is rare in industry-track papers.

Limitations & Future Work

  • Backbone Scope: Validated only on 12-layer MiniLM; not yet tested on deeper SOTA like E5-large or multilingual variants.
  • Task Scope: Limited to sentence embeddings; does not touch token-level early exit (MT, generation), which requires more complex per-token decisions.
  • Training Overhead: Adds ~20% training cost, which may be significant for large-scale pretraining.
  • Fixed \(l_{\min}=6\): Every sample runs at least 6 layers; could explore dynamic lower bounds based on input difficulty.
  • Task Sensitivity: Large drop on ArguAna (24.7%) suggests a need for inference-time "difficulty predictors" to prevent premature exits on complex queries.
  • vs DeeBERT / PABEE / BERxiT: These adapt at inference time with learned heads; LEAP eliminates the root cause at training time and remains parameter-free at inference.
  • vs MiniLM / TinyBERT / DistilBERT: Rather than treating redundancy as a byproduct, LEAP makes "intermediate redundancy" an explicit optimization objective. It is an orthogonal enhancement to standard distillation.
  • vs Matryoshka Representations: Matryoshka provides "width adaptive" (dimension) embeddings; LEAP provides "depth adaptive" (layer) embeddings. Both are orthogonal and can be combined for multiplicative gains.

Rating

  • Novelty: ⭐⭐⭐⭐ Combination of "identifying incompatibility + parameter-free training intervention" is solid.
  • Experimental Thoroughness: ⭐⭐⭐⭐ Covers STS-B, BEIR, wall-clock, and cross-framework validation, though limited to one backbone.
  • Writing Quality: ⭐⭐⭐⭐⭐ Exemplary Industry Track style—clear problem statement and diagnostic framework.
  • Value: ⭐⭐⭐⭐⭐ High ROI for teams using MiniLM/DistilBERT in production for RAG/search.