04. Early-Exit Neural Networks

• Primary References:

1. 1
   Teerapittayanon, S., McDanel, B., & Kung, H. T. (2016). BranchyNet: Fast Inference via Early Exiting from Deep Neural Networks. In 23rd International Conference on Pattern Recognition (ICPR), pp. 2464–2469.
2. 2
   Huang, G., Chen, D., Li, T., Wu, F., van der Maaten, L., & Weinberger, K. Q. (2018). Multi-Scale Dense Networks for Resource Efficient Image Classification (MSDNet). In ICLR 2018.
3. 3
   Kaya, Y., Hong, S., & Dumitras, T. (2019). Shallow-Deep Networks: Understanding and Mitigating Network Overthinking. In ICML 2019 (PMLR 97:3301–3310).
4. 4
   Schuster, T., Fisch, A., Gupta, J., Dehghani, M., Bahri, D., Tran, V., Tay, Y., & Metzler, D. (2022). Confident Adaptive Language Modeling (CALM). In NeurIPS 2022.

• Why I Wrote This Note: Adaptive computation is one of my core research directions. Specifically, I am investigating how calibrated predictive uncertainty (Note #01) can serve as a theoretically grounded halting controller for multi-branch vision and language backbones, achieving Pareto-optimal latency-accuracy curves on edge hardware.

• Conditional Depth per Input: Instead of fixing depth \(L\) for all \(x \in \mathcal{X}\), an early-exit architecture evaluates exits \(m = 1, \dots, M\) sequentially and halts at the earliest exit \(m^\star(x)\) satisfying a confidence or entropy criterion.
• Expected FLOPs Reduction: Because \(60\%\text{–}80\%\) of natural test samples are canonical ("easy"), routing them out at Stage 1 or Stage 2 cuts average inference FLOPs by \(2\times\text{–}3.5\times\) with \(< 0.3\%\) accuracy loss.
• Mitigating Overthinking: Kaya et al. (2019) proved that on a non-trivial subset of inputs, shallow exits predict the correct label \(y\) while deeper layers corrupt that prediction by attending to spurious fine-grained background textures ("overthinking").
• The Feature Readiness Conflict: Placing naive linear heads on early convolutional layers degrades the final exit's accuracy because early layers are forced to learn high-level semantic features prematurely, starving later layers of multi-scale spatial detail (solved by MSDNet dense multi-scale connectivity or self-distillation).

4.1 Multi-Exit Architecture and Joint Training Objective

Let a backbone consists of \(M\) sequential stages \(\{f_1, f_2, \dots, f_M\}\). Given input \(x\), the hidden representation at stage \(m\) is defined recursively as \(h_m(x) = f_m(h_{m-1}(x))\) with \(h_0(x) = x\). At each stage \(m \in \{1, \dots, M\}\), an auxiliary exit head \(g_m\) maps \(h_m(x)\) to class probabilities \(\hat{p}_m(y \mid x) = \text{softmax}(g_m(h_m(x)))\).

flowchart LR
    X["Input x"] --> S1["Stage 1<br/>h₁ = f₁(x)"]
    S1 --> S2["Stage 2<br/>h₂ = f₂(h₁)"]
    S2 --> SM["Stage M<br/>h_M = f_M(h_{M-1})"]
    S1 --> E1["Exit 1: g₁(h₁)<br/>H(p̂₁) < τ₁ ?"]
    S2 --> E2["Exit 2: g₂(h₂)<br/>H(p̂₂) < τ₂ ?"]
    SM --> EM["Final Exit M: g_M(h_M)<br/>Fallback / Hard Samples"]

During training, all \(M\) exits are optimized jointly via a weighted sum of cross-entropy losses, often augmented with self-distillation from the deepest exit \(M\) (acting as an online teacher) to the shallower exits \(m < M\):

The weighting schedule \(\{w_m\}_{m=1}^{M}\) controls gradient pressure on early layers:

• Linearly Increasing Weights (\(w_m = m / M\)): Prioritizes the representational capacity of the deeper backbone.
• Uniform Weights (\(w_m = 1\)): Maximizes early-exit accuracy when operating under strict latency budgets.

4.2 Entropy & Patience-Based Halting Criteria

At inference time, given a vector of per-exit thresholds \(\boldsymbol{\tau} = (\tau_1, \dots, \tau_{M-1})\), the halting exit index \(m^\star(x) \in \{1, \dots, M\}\) is determined by the Normalized Shannon Entropy Rule:

Alternatively, in Patience-Based Early Exiting (PABEE, Zhou et al., NeurIPS 2020), inference halts when \(k\) consecutive exits agree on the predicted class \(\arg\max_c \hat{p}_{m}(y=c\mid x) = \dots = \arg\max_c \hat{p}_{m-k+1}(y=c\mid x)\), which eliminates sensitivity to uncalibrated logit magnitudes.

The expected computational cost \(\mathbb{E}[\text{Cost}]\) over the test distribution \(p(x)\) is:

4.3 PyTorch Implementation of an Entropy-Gated Early-Exit Backbone

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

class EarlyExitClassifierHead(nn.Module):
    # Lightweight exit head with global pooling and temperature-calibrated logits
    def __init__(self, in_channels: int, num_classes: int, temperature: float = 1.0):
        super().__init__()
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(in_channels, num_classes)
        self.register_buffer("temperature", torch.tensor(float(temperature)))

    def forward(self, feat: torch.Tensor) -> torch.Tensor:
        pooled = self.pool(feat).flatten(1)
        return self.fc(pooled) / self.temperature

class AdaptiveEarlyExitNetwork(nn.Module):
    def __init__(self, stages: nn.ModuleList, exit_heads: nn.ModuleList, num_classes: int):
        super().__init__()
        assert len(stages) == len(exit_heads)
        self.stages = stages
        self.exit_heads = exit_heads
        self.num_classes = num_classes
        self.norm_factor = math.log(max(num_classes, 2))

    def forward_train(self, x: torch.Tensor) -> list[torch.Tensor]:
        # Execute all stages during training to compute joint multi-exit + KD loss
        logits_per_exit = []
        h = x
        for stage, head in zip(self.stages, self.exit_heads):
            h = stage(h)
            logits_per_exit.append(head(h))
        return logits_per_exit

    @torch.no_grad()
    def forward_adaptive(self, x: torch.Tensor, entropy_thresholds: list[float]) -> tuple[torch.Tensor, int, float]:
        # Sample-adaptive inference (batch size = 1 for true hardware early halting)
        h = x
        for idx, (stage, head) in enumerate(zip(self.stages, self.exit_heads)):
            h = stage(h)
            logits = head(h)
            probs = F.softmax(logits, dim=-1)
            norm_entropy = float(-(probs * probs.clamp_min(1e-12).log()).sum(dim=-1).item() / self.norm_factor)

            if idx < len(entropy_thresholds) and norm_entropy < entropy_thresholds[idx]:
                return probs, idx + 1, norm_entropy
        return probs, len(self.stages), norm_entropy

Early-Exit Networks sit right at the intersection of my research on uncertainty quantification and real-time computer vision:

• In real-time textile fabric defect detection (ICCA 2024), \(90\%+\) of loom frames are defect-free uniform fabric. An early exit after the second convolutional stage can clear defect-free frames in \(< 1.5\text{ ms}\), routing only anomalous texture patches to the full deep detector and ensemble verifier.
• In my ongoing research on uncertainty-guided adaptive inference, I replace naive softmax confidence halting with calibrated epistemic + aleatoric thresholds and conformal risk bounds (Risk-Constrained AI Decision-Making), ensuring that early halting provably controls the excess error rate \(\mathbb{P}(\hat{Y}_{\text{exit}} \neq \hat{Y}_M) \le \epsilon\).

• Never Set \(\tau_m\) Without Per-Exit Temperature Scaling: Shallow exits (\(m=1\)) and deep exits (\(m=M\)) have vastly different logit norms and calibration profiles. Applying a uniform raw entropy threshold \(\tau_1 = \tau_2 = \tau_3\) causes Exit 1 to either reject everything or accept uncalibrated mistakes. Always run L-BFGS Temperature Scaling independently per exit first.
• Geometric Ensemble of Visited Exits: When an input reaches Exit \(m\), we have already computed the logits for Exits \(1, \dots, m-1\) at zero extra cost! Computing a reliability-weighted geometric average \(\hat{p}_{\le m}(y\mid x) \propto \prod_{j=1}^{m} \hat{p}_j(y\mid x)^{\alpha_j}\) acts as a "free" depth-wise ensemble that consistently boosts accuracy by \(+0.5\%\text{–}1.2\%\).

• Implement AdaptiveEarlyExitNetwork with per-exit temperature buffers and normalized Shannon entropy halting.Verified
• Verify knowledge distillation loss from Exit \(M\) to shallow exits \(m \in \{1, \dots, M-1\}\).Verified
• Implement Cumulative Exit Ensembling (\(\sum_{j=1}^{m} \alpha_j \mathbf{z}_j\)) and benchmark its ECE against single-exit halting.Pending
• Apply Conformal Risk Control (Angelopoulos et al., 2022; Schuster et al., 2022) to select \(\boldsymbol{\tau}\) with finite-sample consistency guarantees.Pending

"Overthinking occurs when a deep neural network already knows the right answer at layer 12, but layers 13 through 50 talk it out of the right answer by fixating on background noise." — Adapted from Kaya et al. (ICML 2019)

Cumulative Depth-Ensemble Logit Aggregation:

• Connected Research Notes:
• Understanding Predictive Uncertainty — Provides the calibrated entropy \(\tilde{\mathcal{H}}(\hat{p}_m)\) used in the halting rule.
• Mixture-of-Experts Architecture — Complements depth-wise adaptive computation (Early-Exit) with width-wise conditional computation (MoE).
• Risk-Constrained AI Decision-Making — Formalizes the risk-coverage trade-off of halting early vs. deferring to deeper layers.