Research map / Label-efficient and robust learning / Test-time adaptation
COSEE: Consistency-Oriented Signal-Based Early Exiting via Calibrated Sample Weighting Mechanism
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes COSEE, a signal-based early exiting framework for BERT that uses a calibrated sample weighting mechanism over randomly sampled thresholds plus an online signal calibration loss, aligning training with test-time exiting across speed-up ratios.
Problem
Signal-based early exiting for pre-trained language models trains every internal classifier on all samples equally, creating a mismatch with test-time exiting behavior; router-based alternatives fix the speed-up ratio, so flexible acceleration is lost.
Approach
Simulate test-time early exiting at several randomly sampled thresholds during training to find each sample's likely exit layer, then weight each sample's cross-entropy across classifiers by a softmax decay around that layer, and add an online signal calibration loss that widens the exiting-signal gap between easy and hard samples.
Key results
- On GLUE with BERT-base, COSEE (energy) reaches 93.0 on SST-2 at 2.14x speed-up and 90.2 on QNLI at 2.56x, outperforming listed baselines on those tasks.
- COSEE outperforms the original BERT-base on RTE (68.7 vs 66.4) and QQP (71.4/89.4 vs 71.2/-) per the paper's table.
- With ALBERT-base averaged over 4 tasks, COSEE-energy reaches 87.3 average at 2.12x speed-up versus 87.5 for the full ALBERT-base model.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use sample-wise loss weights based on simulated exit layers (SWM) (ablation) | COSEE · GLUE tasks, performance-efficiency trade-off | task performance | Helped |
| Add the online signal calibration objective (OSC) (ablation) | COSEE · GLUE tasks, performance-efficiency trade-off | task performance | Helped |
| Add the OSC objective (ablation) | COSEE · SST-2 and QNLI training | convergence speed | Hurt |
| Use SWM (ablation) | COSEE · SST-2 and QNLI training | convergence speed | Helped |
| Use COSEE with energy scores instead of entropy or softmax scores (ablation) | COSEE · SST-2, QNLI, and QQP | task performance | Helped |
| Use excessively large or small β₀ values (ablation) | COSEE · SST-2 and QNLI | acceleration performance | Hurt |
| Increase α excessively (ablation) | COSEE · SST-2 and QNLI | task performance | Hurt |
| Vary α within the tested range (ablation) | COSEE · SST-2 and QNLI | acceleration performance | No clear effect |
| Set ε between 0.1 and 0.4 (ablation) | COSEE · SST-2 and QNLI | task performance | No clear effect |
| Increase the number of sampled thresholds K (ablation) | COSEE · SST-2 and QNLI | task performance | Helped |
| Add online signal calibration (ablation) | COSEE · SST-2 and QNLI | Difficulty Inversion Score | Helped |
Surprising findings
- COSEE slightly exceeds the full BERT-base on RTE and QQP, which the authors attribute to alleviating overthinking.
- Adding the OSC objective slightly slows convergence, though COSEE still converges faster than vanilla training.
Limitations
- Results reported mainly on GLUE with BERT-base; ALBERT-base experiments are limited to 4 tasks.
- Hyperparameter grids (alpha, beta0, epsilon, K) were tuned and sensitivity is shown per task, with optimal beta0 varying across tasks.
- Speed-up measured as saved layers rather than wall-clock runtime.
Experimental setup
Datasets: GLUE SST-2 GLUE MRPC GLUE QQP GLUE MNLI GLUE QNLI GLUE RTE
Benchmarks: GLUE SST-2 GLUE MRPC GLUE QQP GLUE MNLI GLUE QNLI GLUE RTE
Metrics: accuracy F1 speed-up ratio (saved layers) difficulty inversion score
Baselines: BERT-6L DeeBERT PABEE BERxiT LeeBERT GPFEE GAML-BERT PALBERT DisentangledEE ConsistentEE
Builds on: DeeBERT PABEE E-LANG ConsistentEE BERT
Closest papers by meaning
- ConsistentEE: A Consistent and Hardness-Guided Early Exiting Method for Accelerating Language Models Inference (AAAI 2024 · similarity 0.89)
- PRESTO: Preimage-Informed Instruction Optimization for Prompting Black-Box LLMs (NeurIPS 2025 · similarity 0.84)
- Signal and Noise: A Framework for Reducing Uncertainty in Language Model Evaluation (NeurIPS 2025 · similarity 0.84)
- EE-LLM: Large-Scale Training and Inference of Early-Exit Large Language Models with 3D Parallelism (ICML 2024 · similarity 0.84)
- Correctness-Optimized Residual Activation Lens (CORAL): Transferrable and Calibration-Aware Inference-Time Steering (ICML 2026 · similarity 0.84)
- Through the River: Understanding the Benefit of Schedule-Free Methods for Language Model Training (NeurIPS 2025 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.