Research map / Architectures and efficiency / RNNs and normalization
Jointly-Learned Exit and Inference for a Dynamic Neural Network
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 JEI-DNN, which jointly learns lightweight gates and early-exit inference modules on a frozen pretrained ViT via a bi-level optimization, improving accuracy-cost tradeoff and uncertainty calibration/conformal intervals.
Problem
Early-exit dynamic neural networks train gating mechanisms and intermediate inference modules separately, causing train-test mismatch, and threshold gates produce poorly calibrated, biased probabilities that hinder uncertainty estimation.
Approach
Model exit probabilities sequentially via learnable gates g^l with a min-based parameterization so that P(G=l|x) is a valid distribution, then jointly train gates and IMs with a bi-level optimization in which gates are trained on surrogate binary tasks derived from the lowest-cost exit, while IMs are trained with gate-weighted cross-entropy.
Key results
- On CIFAR100 with T2T-ViT-14, JEI-DNN significantly outperforms baselines in the region above 80% of full-model accuracy, with tighter 95% confidence intervals.
- JEI-DNN yields tighter conformal intervals (lower inefficiency) at a given empirical coverage than BoostedNet and L2W-DEN on CIFAR100.
- Adding JEI-DNN to DynPerceiver (MobileNetV3 backbone) improves accuracy on CIFAR100 and ImageNet, with added compute below 0.01% of backbone Mul-Adds.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Jointly train the gates and inference modules instead of decoupling their training | JEI-DNN · CIFAR10, CIFAR100, CIFAR100-LT, and SVHN | accuracy-inference cost trade-off | Helped |
| Use learnable gates with fixed inference modules instead of joint training (ablation) | JEI-DNN · CIFAR10 (T2T-7) | accuracy-inference cost trade-off | Hurt |
| Replace learnable gates with threshold gates and train inference modules separately (ablation) | JEI-DNN · CIFAR10 (T2T-7) | accuracy-inference cost trade-off | Hurt |
| Use learnable gates with fixed inference modules rather than threshold gates (ablation) | Decoupled JEI-DNN · CIFAR10 (T2T-7) | accuracy-inference cost trade-off | Helped |
| Ignore the gate-dependence term in the inference-module gradient (ablation) | JEI-DNN · CIFAR10 (T2T-7) | test accuracy | No clear effect |
| Ignore the gate-dependence term in the inference-module gradient (ablation) | JEI-DNN · CIFAR10 (T2T-7) | convergence speed | Helped |
| Add JEI-DNN joint gate and inference-module training to Dynamic Perceiver | Dynamic Perceiver with MobileNetV3 · CIFAR100 and ImageNet | accuracy | Helped |
| Use learnable gates instead of threshold gates | JEI-DNN and early-exit baselines · CIFAR100 (T2T-14) | expected calibration error | Helped |
| Use JEI-DNN instead of threshold-gated baselines | Early-exit models · CIFAR100 (T2T-14) | conformal interval inefficiency | Helped |
| Use JEI-DNN instead of threshold-gated baselines | Early-exit models · CIFAR100 (T2T-14) | empirical coverage deviation from requested coverage | Helped |
Surprising findings
- Learnable gates with fixed IMs (decoupled) were less stable than expected, and ignoring the gate-dependent term of the IM gradient did not hurt performance while converging faster.
Limitations
- A single accuracy/cost trade-off point requires retraining with a different lambda, unlike threshold gates that can be adjusted at deployment.
- Cannot directly target a specific accuracy or inference cost constraint.
Experimental setup
Datasets: CIFAR10 CIFAR100 CIFAR100-LT SVHN ImageNet
Benchmarks: CIFAR10 CIFAR100 CIFAR100-LT SVHN ImageNet
Metrics: accuracy Mul-Add inference cost expected calibration error conformal inefficiency empirical coverage
Baselines: BoostedNet L2W-DEN MSDNet RANet DynPerceiver DiffBranch
Builds on: BoostedNet L2W-DEN MSDNet DynPerceiver temperature scaling
Closest papers by meaning
- Fast yet Safe: Early-Exiting with Risk Control (NeurIPS 2024 · similarity 0.88)
- How to Train Your Multi-Exit Model? Analyzing the Impact of Training Strategies (ICML 2025 · similarity 0.87)
- Beyond Greedy Exits: Improved Early Exit Decisions for Risk Control and Reliability (NeurIPS 2025 · similarity 0.86)
- Initial Guessing Bias: How Untrained Networks Favor Some Classes (ICML 2024 · similarity 0.84)
- BEEM: Boosting Performance of Early Exit DNNs using Multi-Exit Classifiers as Experts (ICLR 2025 · similarity 0.84)
- The Butterfly Effect: Neural Network Training Trajectories Are Highly Sensitive to Initial Conditions (ICML 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.