Recursive Early-Exit Framework

• Recursive Early-Exit is a neural network inference framework that enables dynamic, goal-oriented decision-making via recursive multi-exit architectures.
• It integrates reinforcement learning-based policies to adaptively decide early exits, offloading, or continuation to optimize accuracy, latency, and wireless resource usage.
• Empirical results on CIFAR-10 with ResNet20 show 30% FLOP reduction or 1–2% accuracy gains, demonstrating its effectiveness for edge inference under time constraints.

Recursive early-exit is a framework for neural network inference designed to enable dynamic, goal-oriented decision-making in edge-device and server-collaborative scenarios. The approach combines a novel recursive mechanism for multi-exit neural architectures with reinforcement learning-based policies to adaptively determine when to halt computation, where to partition workloads, and whether to offload embeddings for remote inference. This framework jointly optimizes inference accuracy, computational latency, and wireless resource consumption, and is particularly suited to edge inference settings with fluctuating resource availability and strict real-time constraints [2412.19587].

1. Recursive Multi-Exit Neural Network Structure

The foundational component is a base model consisting of $b$ sequential blocks (convolutional or residual), noted as $\ell_1, \ldots, \ell_b$, followed by a terminal classifier $c_b$. In conventional form, inference is computed as:
$$
f_b(x) = c_b \circ \ell_b \circ \ell_{b-1} \circ \cdots \circ \ell_1(x)
$$
To enable early inference, a selected subset of indices $\mathcal{I} \subset {1, \ldots, b-1}$ is chosen for auxiliary "early-exit" branches. Each branch at index $i \in \mathcal{I}$ attaches an exit-specific classifier $c_i$, itself decomposed into a feature-processing head $e_i(x)$ (e.g., convolution + pooling) and a linear projection $p_i(\cdot)$ to output class scores.

Recursive prediction is implemented as follows. Let $f_i(x) \in [0,1]^{|\mathcal{C}|}$ denote the class probability vector at exit $i$. The combination mechanism is:
- For $i = 1$,
$f_1(x) = p_1(e_1(x))$
- For $1 < i < b$,
$f_i(x) = f_{i-1}(x) + M_i(x)$
- For $i = b$,
$f_b(x) = f_{b-1}(x) + (1 - f_{b-1}(x)) \odot p_b(e_b(x))$

Here, $M_i(x)$ is the "moving mass" term that reallocates class probabilities:
$$
M_i(x) = (1 - f_{i-1}(x)) \odot m_i^+(e_i(x)) - f_{i-1}(x) \odot m_i^-(e_i(x))
$$
with $m_i^+, m_i^-$ defined as linear plus sigmoid heads per exit. This recursive update allows intermediate exits to refine and compensate for class assignments based on new representations, yielding a path-dependent and efficiently updatable probability vector at each exit.

2. Confidence Evolution and Early-Exit Criterion

At each exit $i$, the model outputs class probabilities $f_i(x)$. The instantaneous confidence is defined as
$$
C^{(i)}(x) = \max_{c \in \mathcal{C}} f_i^c(x)
$$
The halting criterion employs a margin-based rule: denoting $t_1$ and $t_2$ as indices of the largest and second-largest entries of $f_i(x)$, the system halts at $i$ if
$$
f_i^{t_1}(x) - f_i^{t_2}(x) > m
$$
where $m > 0$ is a user-specified safety margin. This ensures halting only when the model is sufficiently more confident in the top prediction compared to subsequent candidates.

Alternative schemes used in classical early-exit models include simple thresholds $C^{(i)}(x) > \tau$ and difference-based strategies $C^{(i)}(x) - C^{(i-1)}(x) < \Delta$ in combination with $C^{(i)}(x) > \tau$. The margin-based halting rule, however, generalizes to the multiclass scenario and directly encodes the desired separation between leading predictions.

3. Reinforcement Learning-Based Online Exit Policy

The joint optimization of exit decisions, computation partitioning, and offloading is formulated as a Markov Decision Process (MDP), with each inference request constituting an episode of up to $K$ decision steps (number of exits). The state at step $t$ is
$$
s_t = (k_t, \mathrm{mcs}_t)
$$
where $k_t$ indexes the current exit, and $\mathrm{mcs}_t \in \mathcal{M}$ is the instantaneous wireless Modulation & Coding Scheme selected per the observed channel SNR.

At each state, agent actions $a_t \in {0, 1, 2}$:
- $0$ (exit now): output prediction at exit $k_t$ locally,
- $1$ (continue): compute $\ell_{k_{t+1}}$ and evaluate next exit,
- $2$ (offload): transmit current embedding to server and terminate.

Transitions for $a_t = 1$ yield a new state with incremented $k$ and a fresh channel sample; $a_t = 0, 2$ are terminal.

Terminal rewards aggregate three key performance indicators:
- Computation saving: $\Gamma_{\mathrm{comp,} k} = (F_K - F_k)/F_K$
- Communication saving: $\Gamma_{\mathrm{comm,} k} = (\max_j N_j - N_k)/\max_j N_j$ (embedding size in bits at $k$; set to 1 if $a_t = 0$)
- Goal-effectiveness: If total delay $D_\mathrm{tot} \leq D_\mathrm{max}$ and margin $\geq m_\mathrm{th}$, then indicator $\mathbf{1}\mathrm{succ} = 1$.
$$
r =
\begin{cases}
\gamma{\mathrm{comm}} \cdot \Gamma_{\mathrm{comm,} k}
+ \gamma_{\mathrm{comp}} \cdot \Gamma_{\mathrm{comp,} k} & \text{if } \mathbf{1}_{\mathrm{succ}}=1 \
-1 & \text{otherwise}
\end{cases}
$$
Tabular $Q$-learning is used to optimize policy $Q(s,a)$, enabling adaptive selection of exit, continuation, or offloading actions under varying device, channel, and task constraints.

4. Computation Partitioning and Offloading Strategy

Computation splitting is operationalized by model partitioning: after evaluating up to exit $k$, the device has performed $F_k$ FLOPs and holds the corresponding $N_k$-bit embedding. At this point, it can either:
- Exit locally, transmitting zero data and finalizing inference, or
- Offload the $N_k$-bit embedding to the server, which then computes the remaining $F_K - F_k$ FLOPs.

The detailed delay components are:
- Local compute delay: $D_{l, \mathrm{comp}} = F_k / f_l$
- Transmission delay: $D_{tx} = N_k / R$, with $R = B \cdot \mathrm{mcs}$, $B=20$ MHz, $mcs$ as instantaneous rate
- Remote compute delay: $D_{r,\mathrm{comp}} = (F_K - F_k) / f_r$
- Total delay: $D_{tot} = D_{l, \mathrm{comp}} + \text{(offload? } D_{tx} + D_{r, \mathrm{comp}} : 0)$

The wireless environment is modeled with path-loss plus Rayleigh fading. At each step, a fresh SNR determines the maximum achievable MCS and data rate, directly affecting $D_{tx}$ and thus $D_{tot}$.

5. Empirical Evaluation and Resource-Utility Trade-offs

The recursive early-exit method was benchmarked using CIFAR-10 and a ResNet20 backbone with nine early exits. Performance was compared against two canonical baselines:
- Highest-probability halting (softmax threshold $\tau$ on top-class probability)
- Patience-based halting (exit after $p$ consecutive exits agree on predicted class)

Results indicate that the recursive-mass approach dominated both baselines, delivering the same accuracy with roughly 30% fewer FLOPs, or yielding 1–2% higher accuracy for a fixed FLOPs budget.

In an edge inference scenario with randomized device locations (10–100m from access point), 20 MHz @3.5 GHz, $P_{tx} = 0.1$ W, device/server compute times of 50/10 ms, and delay constraint $D_{max} = 40$ ms, margin thresholds $m_{th}$ and communication weight $\gamma_{comm}$ were varied. Trade-off curves illustrate that as $\gamma_{comm}$ (the emphasis on communication savings) increases, the RL policy shifts from early offloading (maximal goal-effectiveness, low resource savings) to deeper local exits (higher computation saving, lower communication saving), eventually trading off goal-effectiveness if pushed to extremes.

Exit index distributions as a function of resource requirements showed that under strict communication-saving targets, the agent favored deep local exits; under abundant resources, early offloading was preferred for optimal accuracy. The learned policy effectively adapts exit, split, and offload behavior in real time, balancing accuracy, latency, and resource consumption.

6. Key Components and Hyperparameter Considerations

The recursive early-exit framework encompasses:
- A recursive, multi-exit network structure updating class probabilities per layer via moving mass
- A margin-based early-exit rule generalizing conventional confidence-based halting
- Integration of early-exit, computation partitioning, and offloading as a joint MDP with RL optimization
- Practical wireless channel and latency models for realistic deployment settings
- Tabular Q-learning policy for real-time adaptive exit, split, and offload decision-making

Hyperparameters including safety margins ($m$), reward coefficients ($\gamma_{\mathrm{comm}}, \gamma_{\mathrm{comp}}$), maximum delay ($D_{max}$), margin threshold ($m_{th}$), available MCS values, and device/server processing rates are all explicitly defined and tunable. These enable reproducibility and facilitate adaptation to alternative architectures, datasets, and wireless environments [2412.19587].