# Jointly-Learned Exit and Inference for a Dynamic Neural Network : JEI-DNN

arXiv: https://arxiv.org/abs/2310.09163 · HTML: https://mlsanity.com/paper/2310.09163/

## Key facts

- Cited by 15 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 58 references that are in the map.
- First posted on arXiv 13 Oct 2023.

## Builds on (most cited first, 30 of 37)

- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Searching for MobileNetV3](https://mlsanity.com/paper/1905.02244/) (2019)
- [BLOOM: A 176B-Parameter Open-Access Multilingual Language Model](https://mlsanity.com/paper/2211.05100/) (2022)
- [Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNet](https://mlsanity.com/paper/2101.11986/) (2021)
- [GLaM: Efficient Scaling of Language Models with Mixture-of-Experts](https://mlsanity.com/paper/2112.06905/) (2021)
- [Benchmarking Graph Neural Networks](https://mlsanity.com/paper/2003.00982/) (2020)
- [BranchyNet: Fast Inference via Early Exiting from Deep Neural Networks](https://mlsanity.com/paper/1709.01686/) (2017)
- [Adaptive Computation Time for Recurrent Neural Networks](https://mlsanity.com/paper/1603.08983/) (2016)
- [Dynamic Neural Networks: A Survey](https://mlsanity.com/paper/2102.04906/) (2021)
- [Revisiting the Calibration of Modern Neural Networks](https://mlsanity.com/paper/2106.07998/) (2021)
- [Confident Adaptive Language Modeling](https://mlsanity.com/paper/2207.07061/) (2022)
- [BERT Loses Patience: Fast and Robust Inference with Early Exit](https://mlsanity.com/paper/2006.04152/) (2020)
- [BlockDrop: Dynamic Inference Paths in Residual Networks](https://mlsanity.com/paper/1711.08393/) (2017)
- [Adaptive Neural Networks for Efficient Inference](https://mlsanity.com/paper/1702.07811/) (2017)
- [Graph Inductive Biases in Transformers without Message Passing](https://mlsanity.com/paper/2305.17589/) (2023)
- [Resolution Adaptive Networks for Efficient Inference](https://mlsanity.com/paper/2003.07326/) (2020)
- [Not All Images are Worth 16x16 Words: Dynamic Transformers for Efficient Image Recognition](https://mlsanity.com/paper/2105.15075/) (2021)
- [Glance and Focus: a Dynamic Approach to Reducing Spatial Redundancy in Image Classification](https://mlsanity.com/paper/2010.05300/) (2020)
- [Online Deep Learning: Learning Deep Neural Networks on the Fly](https://mlsanity.com/paper/1711.03705/) (2017)
- [Adaptive Inference through Early-Exit Networks: Design, Challenges and Directions](https://mlsanity.com/paper/2106.05022/) (2021)
- [Improved Techniques for Training Adaptive Deep Networks](https://mlsanity.com/paper/1908.06294/) (2019)
- [CF-ViT: A General Coarse-to-Fine Method for Vision Transformer](https://mlsanity.com/paper/2203.03821/) (2022)
- [Uni-Perceiver-MoE: Learning Sparse Generalist Models with Conditional MoEs](https://mlsanity.com/paper/2206.04674/) (2022)
- [Sparsely Activated Mixture-of-Experts are Robust Multi-Task Learners](https://mlsanity.com/paper/2204.07689/) (2022)
- [You Look Twice: GaterNet for Dynamic Filter Selection in CNNs](https://mlsanity.com/paper/1811.11205/) (2018)
- [Deep Ensembling with No Overhead for either Training or Testing: The All-Round Blessings of Dynamic Sparsity](https://mlsanity.com/paper/2106.14568/) (2021)
- [Learning to Weight Samples for Dynamic Early-exiting Networks](https://mlsanity.com/paper/2209.08310/) (2022)
- [Zero Time Waste: Recycling Predictions in Early Exit Neural Networks](https://mlsanity.com/paper/2106.05409/) (2021)

## Built on by (most cited first, 15 of 15)

- [Fast yet Safe: Early-Exiting with Risk Control](https://mlsanity.com/paper/2405.20915/) (2024)
- [Cascade-Aware Training of Language Models](https://mlsanity.com/paper/2406.00060/) (2024)
- [EE-Tuning: An Economical yet Scalable Solution for Tuning Early-Exit Large Language Models](https://mlsanity.com/paper/2402.00518/) (2024)
- [ThinkingViT: Matryoshka Thinking Vision Transformer for Elastic Inference](https://mlsanity.com/paper/2507.10800/) (2025)
- [Early-Exit Neural Networks with Nested Prediction Sets](https://mlsanity.com/paper/2311.05931/) (2023)
- [BEEM: Boosting Performance of Early Exit DNNs using Multi-Exit Classifiers as Experts](https://mlsanity.com/paper/2502.00745/) (2025)
- [RAEE: A Robust Retrieval-Augmented Early Exit Framework for Efficient Inference](https://mlsanity.com/paper/2405.15198/) (2024)
- [Beyond Greedy Exits: Improved Early Exit Decisions for Risk Control and Reliability](https://mlsanity.com/paper/2509.23666/) (2025)
- [Predicting Probabilities of Error to Combine Quantization and Early Exiting: QuEE](https://mlsanity.com/paper/2406.14404/) (2024)
- [S2D: Sorted Speculative Decoding For More Efficient Deployment of Nested Large Language Models](https://mlsanity.com/paper/2407.01955/) (2024)
- [Lightweight Remote Sensing Scene Classification on Edge Devices via Knowledge Distillation and Early-exit](https://mlsanity.com/paper/2507.20623/) (2025)
- [Rethinking Calibration for Early-Exit Neural Networks](https://mlsanity.com/paper/2508.21495/) (2025)
- [CalexNet: Soft Cascade-Aligned Training and Calibration for Lightweight Early-Exit Branches](https://mlsanity.com/paper/2509.08318/) (2025)
- [On-Demand Multi-Task Sparsity for Efficient Large-Model Deployment on Edge Devices](https://mlsanity.com/paper/2511.19986/) (2025)
- [Federated Personalization of Early-Exit Networks](https://mlsanity.com/paper/2601.10015/) (2026)

## Abstract (arXiv)

Large pretrained models, coupled with fine-tuning, are slowly becoming established as the dominant architecture in machine learning. Even though these models offer impressive performance, their practical application is often limited by the prohibitive amount of resources required for every inference. Early-exiting dynamic neural networks (EDNN) circumvent this issue by allowing a model to make some of its predictions from intermediate layers (i.e., early-exit). Training an EDNN architecture is challenging as it consists of two intertwined components: the gating mechanism (GM) that controls early-exiting decisions and the intermediate inference modules (IMs) that perform inference from intermediate representations. As a result, most existing approaches rely on thresholding confidence metrics for the gating mechanism and strive to improve the underlying backbone network and the inference modules. Although successful, this approach has two fundamental shortcomings: 1) the GMs and the IMs are decoupled during training, leading to a train-test mismatch; and 2) the thresholding gating mechanism introduces a positive bias into the predictive probabilities, making it difficult to readily extract uncertainty information. We propose a novel architecture that connects these two modules. This leads to significant performance improvements on classification datasets and enables better uncertainty characterization capabilities.