SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
p _ {\Pi , m} \left(y \mid \boldsymbol {x}\right) := \frac {1}{Z _ {m}} \prod_ {i = 1} ^ {m} \left[ f _ {i} (\boldsymbol {x}) _ {y} > b \right] ^ {w _ {i}}, \quad Z _ {m} = \sum_ {y ^ {\prime} \in \{1, \dots , K \}} \prod_ {i = 1} ^ {m} \left[ f _ {i} (\boldsymbol {x}) _ {y ^ {\prime}} > b \right] ^ {w _ {i}}
$$
where  $f_{i}(\pmb{x}) \in \mathbb{R}^{K}$  denotes a vector of logits at the  $i$ -th early-entry,  $b$  represents a chosen threshold, and  $w_{i} \in \mathbb{R}^{+}$  represent ensemble weights. Since we assume that class  $y$  is in the support of the full NN ensemble, i.e.,  $p_{\Pi, M}(y \mid \pmb{x}) > 0$ , it follows that
$$
\prod_ {i = 1} ^ {m} \left[ f _ {i} (\boldsymbol {x}) _ {\tilde {y}} > b \right] ^ {w _ {i}} = \prod_ {i = 1} ^ {m + 1} \left[ f _ {i} (\boldsymbol {x}) _ {\tilde {y}} > b \right] ^ {w _ {i}} = 1, \forall m = 1, \dots , M - 1.
$$
It then remains to show that the normalization constant is non-increasing, i.e.,  $Z_{m}\geq Z_{m + 1},\forall m$  . To this end observe that  $\left[f_{m + 1}(\pmb {x})_{y^{\prime}} > b\right]^{w_{m + 1}}\in \{0,1\} ,\forall y^{\prime}\in \mathcal{V}$  using which we proceed as
$$
Z _ {m} = \sum_ {y ^ {\prime} \in \{1, \dots , K \}} \prod_ {i = 1} ^ {m} \left[ f _ {i} (\boldsymbol {x}) _ {y ^ {\prime}} > b \right] ^ {w _ {i}} \geq
$$
$$
\sum_ {y ^ {\prime} \in \{1, \dots , K \}} \prod_ {i = 1} ^ {m} \left[ f _ {i} (\boldsymbol {x}) _ {y ^ {\prime}} > b \right] ^ {w _ {i}} \cdot \left[ f _ {m + 1} (\boldsymbol {x}) _ {y ^ {\prime}} > b \right] ^ {w _ {m + 1}} = Z _ {m + 1},
$$
where  $Z_{m} = Z_{m + 1}$  if  $\left[f_{m + 1}(\pmb{x})_{y'} > b\right]^{w_{m + 1}} = 1 \forall y' \in \mathcal{Y}$  and  $Z_{m} > Z_{m + 1}$  otherwise.
B.1 Conditional Monotonicity
In the main paper, the ground-truth probability was our focus when studying conditional monotonicity in early-exit networks. Here, we supplement these findings with additional results, utilizing alternative estimators of prediction quality  $\gamma_{m}$ .
B.1.1 Full-Model Prediction Probability
Suppose only an unlabeled hold-out dataset is available. In that case, the anytime properties of EENNs can still be investigated by relying on the full model prediction  $\hat{y} := \arg \max_y p_M(y|x)$ , instead of  $y^*$ . Naturally, the merit of this alternative hinges on the underlying model being accurate in most cases. Furthermore, besides serving as a substitute for the ground-truth probability in the absence of a labeled dataset, the evolution of the full-model prediction probability can be used as a stand-alone measure of how well early-exits in an EENN approximate the full model.
In Figure 7, we present conditional monotonicity results using  $\hat{y}$  probability as the prediction quality measure. Similar to the results presented in Section 6.1 for  $y^{*}$  probability, we conclude that our PA outperforms both the baseline models and CA approach across all datasets considered.
Max Full-Model Prediction  $(\hat{y})$  Probability Decrease
Figure 7: The percentage of test examples with a full-model prediction probability drop exceeding a specified threshold. PA significantly improves the conditional monotonicity properties of ground-truth probabilities across various datasets and backbone models. To better illustrate the different behavior of various methods, a log scale is used for y-axis. We report average monotonicity together with one standard deviation based on  $n = 3$  independent runs. For DViT [Wang et al., 2021], we were unable to perform multiple independent runs and report the results for various model instantiations instead; see Appendix B.11.
B.1.2 Correctness of Prediction
We next shift our focus to the correctness of prediction as an estimator of performance quality. Recall from Section 2 its definition:  $\hat{\gamma}_m^c (\pmb {x})\coloneqq [\arg \max_yp_m(y|\pmb {x}) = y^* ]$  . It is a binary measure that can only take on the values of O or 1, and thus provides a less nuanced signal about the evolution of the performance quality compared to probability-based measures that evolve more continuously in the interval [0, 1].
To aggregate the results of the conditional monotonicity analysis, we consider the following two quantities:
-  $\%_{\mathrm{mono}}$ : the percentage of test data points that exhibit fully monotone trajectories in the correctness of the prediction, i.e.,  $\hat{\gamma}_{m}^{c}(\boldsymbol{x}) \leq \hat{\gamma}_{m+1}^{c}(\boldsymbol{x}), \forall m$ . This implies that if the model managed to arrive at the correct prediction at the  $m$ th exit, it will not forget it at later exits. An oracle model with perfect conditional monotonicity will result in  $\%_{\mathrm{mono}} = 100\%$ .
-  $\%_{\mathrm{zero}}$ : the number of points for which the EENN model is wrong at all exits. This is used to account for trivial solutions of the monotonicity desideratum:  $\hat{\gamma}_m^c (\pmb {x}) = 0$ ,  $\forall m$ . Naturally, the smaller the  $\%_{\mathrm{zero}}$ , the better.
Table 1: Aggregated results of conditional monotonicity analysis based on correctness of prediction as a performance quality measure. MSDNet [Huang et al., 2018] is used as a backbone model for PA and CA for the results reported here. See Section B.1.2 for definition of  $\%_{\mathrm{mono}}$  and  $\%_{\mathrm{zero}}$ .
CIFAR-10	CIFAR-100	ImageNet
%mono (↑)	%zero (↓)	%mono (↑)	%zero (↓)	%mono (↑)	%zero (↓)
MSDNet	91.8	3.6	70.8	12.6	84.9	20.4
PA (ours)	96.2	4.6	87.2	16.4	93.5	24.1
CA	97.8	5.5	90.8	18.6	94.5	25.6
Results are displayed in Table 1. Comparing our PA model against the MSDNet baseline, we observe a significant improvement in terms of the percentage of test data points that exhibit perfect conditional monotonicity, i.e.,  $\%_{\mathrm{mono}}$ . However, this improvement comes at a cost, as fewer points have a correct prediction for at least one exit. Despite this trade-off, the magnitude of improvement in  $\%_{\mathrm{mono}}$  is larger compared to the degradation of  $\%_{\mathrm{zero}}$  across all datasets considered, which demonstrates that our PA method improves the conditional monotonicity of  $\hat{\gamma}_{m}^{c}(\boldsymbol{x})$  in a non-trivial manner.
Although the Caching Anytime (CA) method outperforms PA in terms of  $\%_{\mathrm{mono}}$  across all datasets, it is worth noting that a considerable portion of CA's outperformance can be attributed to trivial solutions (i.e.,  $\hat{\gamma}_m^c (\pmb {x}) = 0$ ,  $\forall m$ ), which is reflected by CA's poor performance in terms of  $\%$  zero. Furthermore, the fact that neither CA nor PA fully achieves conditional monotonicity in prediction correctness (i.e.,  $\%_{\mathrm{mono}} = 100$ ) highlights that there is still potential for further improvement in endowing EENNs with conditional monotonicity.
B.2 Product Anytime (PA) Ablations
To better understand the role of different components in our proposed PA method, we conduct several ablations and present the results in Figure 8. In accordance with the theoretical results we presented in Section 4.2, we observe that the PoE ensemble combined with the Heaviside activation  $(-)$  yields a model that exhibits complete conditional monotonicity in the full model prediction (i.e.,  $\hat{y}$ ) probability. However, this combination also leads to a model with inferior performance in terms of overall test accuracy (see left plot). As mentioned, this occurs because the model assigns an equal probability to all classes in its support. The ReLU activation resolves this issue by performing on par with Heaviside in terms of monotonicity and on par with exponential, i.e., softmax, in terms of accuracy, thus combining the best of both worlds.
Figure 8: Test accuracy and conditional monotonicity properties of various PA model instantiations for the CIFAR-100 dataset using MSDNet as a backbone. Specifically, we vary (1) the type of activation function for mapping logits to probabilities and (2) the type of ensembling. Our proposed PA method uses ReLU activation function with product ensembling  $(-)$  and is the only method to perform well in terms of both accuracy and monotonicity.
Note that  $\hat{y} \coloneqq \arg \max_y p_M(y|\boldsymbol{x})$  is by definition part of the support of the final classifier  $p_{\Pi, M}$  (unless it degenerates to a zero distribution). However, this is not necessarily the case for  $y^*$ , which is the reason for a slight violation of conditional monotonicity in the PoE-Heaviside model, as seen in the right plot of Figure 8.
Delving deeper into the ReLU activation function and examining various ensembling styles, we also observe in Figure 8 that Mixture-of-Experts with ReLU (MoE-ReLU; ...) demonstrates performance comparable to our proposed PA (PoE-ReLU; —) in terms of both overall accuracy and conditional monotonicity. The MoE predictive distribution at the  $m$ th exit is defined as:
$$
p _ {\mathrm {M o E}, m} (y \mid \boldsymbol {x}) := \sum_ {i = 1} ^ {m} w _ {i} \cdot p _ {i} (y \mid \boldsymbol {x}), \quad p _ {i} (y \mid \boldsymbol {x}) = \frac {a \left(f _ {i} (\boldsymbol {x}) _ {y}\right)}{\sum_ {y ^ {\prime} \in \{1 , . . . , K \}} a \left(f _ {i} (\boldsymbol {x}) _ {y ^ {\prime}}\right)}
$$
where  $a: \mathbb{R}^K \to \mathbb{R}_{\geq 0}$  represents a chosen activation function, such as ReLU, and we utilize uniform weights  $w_i = 1/m$ . Given this observation, one might question why we prefer PoE over MoE in our final PA model. The rationale for this decision is two-fold. First, the proof presented in Appendix A, which serves as the inspiration for our method, is only valid when PoE ensembling is employed. Second, unlike the PoE-based model, the MoE-based model does not exhibit anytime uncertainty. Instead, it demonstrates constant underconfidence across all exits, as illustrated in Figure 9. Considering these factors, we find PoE to be superior to MoE for anytime prediction purposes and incorporate it into our final proposed solution.
Figure 9: Expected Calibration Error (ECE; left) and average maximum class probability over the test dataset (right) for various instantiations of our PA model using the MSDNet model as a backbone on the CIFAR-100. Only product ensembling (PoE) gives rise to a model that gets progressively more confident with more early exits, as exemplified by decreasing ECE and increasing magnitude of probability. On the contrary, mixture ensembles (MoE) result in constant underconfidence.
Early-Exit Ensembles Baselines MoE-exponential  $(\dots)$  and PoE-exponential  $(-)$  can also be interpreted as early-exit ensemble baselines for our PA approach. Instead of relying only on the current predictive distributions (as we do when reporting backbone EENNs results in Section 6), we ensemble distributions across all exits up to and including the current one. However, it is clear from Figure 8 that ensembling softmax predictions is insufficient for achieving conditional monotonicity, as both approaches significantly underperform compared to our PA. This underscores the importance of the activation function that maps logits to probabilities for improving monotonicity properties in EENNs. In Figure 10, we verify that these findings hold on the ImageNet dataset as well, where we again find that our PA is the only configuration that performs well both in terms of marginal accuracy and conditional monotonicity.
Sigmoid Activation Function So far, we have discussed why the ReLU function is preferred over the Heaviside or exponential (softmax) functions for transforming logits into probabilities in our setting. We also considered the sigmoid activation function; however, upon further examination, we identified two drawbacks that make it less suitable in our context. Firstly, the sigmoid function lacks the "nullifying" effect of ReLU. As a result, the support for classes diminishes more slowly, and we do not observe the same anytime uncertainty behavior as we do with ReLU (see Section 6.2). Secondly, logits with larger values all map to the same value, approximately 1, under the sigmoid function. This can negatively impact accuracy, as it does not preserve the ranking. Figure 11 illustrates the results using the sigmoid activation function for the CIFAR-100 dataset with MSDNet as a backbone: our choice of ReLU clearly outperforms the sigmoid activation.
Figure 10: Test accuracy and conditional monotonicity properties of various PA model instantiations for the ImageNet dataset using MSDNet as a backbone. Our PA  $(-)$  is the only method to perform well in terms of both accuracy and monotonicity. Note that the monotonicity curves for ReLU and Heaviside activation functions (almost perfectly) overlap at  $y = 0$ .
Figure 11: Test accuracy, conditional monotonicity and ECE results for MSDNet model  $(-)$  on CIFAR-100 dataset. Using sigmoid activation function  $(-)$  hurts both accuracy as well as calibration. Hence, we utilize ReLU activation in our proposed PA  $(-)$ .
B.3 Closing the Calibration Gap
Here we outline two modifications to our proposed PA method aimed at reducing the calibration gap reported in Section 6.3.
B.3.1 Adaptive Thresholds
We first introduce an extension of our Product Anytime (PA) method, which enables navigating the monotonicity-calibration trade-off in scenarios when a hold-out validation dataset is available.
Recall that ReLU is used as an activation function in PA to map logits to probabilities. This implies that 0 is used as a threshold, determining which classes 'survive' and get propagated further, and which do not. However, 0 is arguably an arbitrary threshold, and one could consider another positive constant  $b \in \mathbb{R}^{+}$  instead. It is important to note that the choice of  $b$  directly influences the size of the support of the predictive likelihood. As illustrated in Figure 13, a larger value of  $b$  results in a smaller number of classes with positive probability (as expected). This larger threshold does also lead to a reduction in the coverage of PA's predictive sets with respect to the ground-truth class. However, the
Figure 12: Calibration analysis for MSDNet on CIFAR-100. Left: expected calibration error (ECE) during anytime evaluation. Right: Oracle-Collaborative AUROC with  $0.5\%$  deferral [OC AUROC; Kivlichan et al., 2021], which simulates a realistic scenario of human-algorithm collaboration by measuring the performance of a classifier when it is able to defer some examples to an oracle, based on its maximum class probability.
Figure 13: Static threshold analysis for PA using MSDNet as a backbone on CIFAR-100 dataset. PA's predictive set is defined as:  $s_m(\boldsymbol{x}_n) \coloneqq \{y \in \mathcal{Y} \mid p_{\mathrm{PA},m}(y \mid \boldsymbol{x}_n) > 0\}$ . Right: Average size of PA's predictive set at each exit. As expected, increasing  $b$  leads to smaller sets. Middle: (Empirical) Coverage of ground-truth class at each exit:  $c_m \coloneqq \sum_n [y_n \in s_m(\boldsymbol{x}_n)] / N_{\mathrm{test}}$ . As  $b$  increases, it becomes more likely for the true class  $y^*$  to fall outside of the support giving rise to a smaller coverage. Left: The marginal test accuracy at each exit is shown as the threshold parameter  $b$  varies. Increasing  $b$  negatively affects accuracy, though the impact is minor for smaller values of  $b$ .
Early Exit
decrease in coverage does not significantly af
fect the overall accuracy of the model, as shown in the far right plot.
Expanding on this concept, one can also consider adaptive thresholding: rather than using the same threshold value for all data points, the model could dynamically adjust the threshold in the activation function based on the current example. Ideally, when the model is confident in its prediction, a higher value of  $b$  is employed, resulting in smaller predictive sets. Conversely, when faced with a more challenging example, a lower threshold is utilized to capture the model's uncertainty.
Switching to ReLU with adaptive thresholding, the predictive likelihood at  $i$ th exit (before ensembling) is given by:
$$
p _ {i} (y \mid \boldsymbol {x}) = \frac {\max  \left(f _ {i} (\boldsymbol {x}) _ {y} , \tau \left(f _ {i} (\boldsymbol {x})\right)\right)}{\sum_ {y ^ {\prime} \in \{1 , \dots , K \}} \max  \left(f _ {i} (\boldsymbol {x}) _ {y ^ {\prime}} , \tau \left(f _ {i} (\boldsymbol {x})\right)\right)}
$$
where  $\tau : \mathbb{R}^K \to \mathbb{R}^+$  maps the vector of logits  $f(\pmb{x}) \in \mathbb{R}^K$  to a non-negative threshold.10
Here, we consider threshold functions of the form  $\tau(f(\pmb{x})) = C \cdot p_{\psi}(f(\pmb{x}))$ , where  $p_{\psi} : \mathbb{R}^K \to [0,1]$  is a model that estimates the probability of the datapoint  $\pmb{x}$  being classified correctly, and  $C$  is a constant that ensures the outputs of  $\tau$  have a similar magnitude to that of logits. To estimate  $C$  and  $\psi$ , we rely on the hold-out validation dataset. Specifically, we fit a binary logistic regression model using labels  $y_n^\tau \coloneqq [\hat{y}_n = y_n^*]$ . Additionally, we only use the  $K = 3$  largest logits as input to the regression model, as we have empirically found that excluding the rest does not significantly impact the performance. To find  $C$ , we perform a simple grid search and select the value that yields the best monotonicity-calibration trade-off on the validation dataset. We fit a thresholding model only at the first exit and then reuse it at later exits, as this approach yielded satisfactory performance for the purposes of our analysis. It is likely that fitting a separate thresholding model at each exit would lead to further performance benefits.
In Figure 14, we observe that PA with adaptive thresholding still improves over baseline in terms of monotonicity (middle plot, note the log-scale of  $y$ -axis). Although the adaptive thresholding approach demonstrates worse monotonicity when compared to the PA algorithm with a static threshold (i.e.,  $\tau = 0$ ), it outperforms the PA method in terms of calibration, as measured by the Expected Calibration Error (ECE). To further elucidate the influence of adaptive thresholding, Figure 15 presents the ground-truth probability trajectories for a random selection of test data points. In general, it is apparent
In general,  $\tau$  could depend on the datapoint  $x$  itself; however, since we desire an efficient model (to avoid adding excessive overhead during inference time), we opt for a more lightweight option based on the processed input, i.e., logits.
Figure 14: Results for CIFAR-100 using MSDNet as the backbone model. We see that adaptive thresholding can be effective in reducing PA's calibration gap, though at some cost to both test accuracy and conditional monotonicity. Note that the baseline results in this figure differ slightly from those in Figures 3 and 4, as less data (90%) was used to fit the model, with a portion of the training dataset (10%) used to fit the thresholding model  $\tau$ .
Figure 15: Ground-truth probability trajectories (i.e.,  $\{\hat{\gamma}_m^p (\pmb {x})\}_{m = 1}^M$ ) for 10 random test CIFAR-100 data points using MSDNet as a backbone. Both instantiations of PA yield more monotone ground-truth probability trajectories.
that adaptive thresholding can serve as an effective strategy for improving the calibration of the PA method, albeit at some expense to both conditional monotonicity and overall accuracy.
B.3.2 PA Finetuning
The post-hoc nature of our proposed PA method is attractive due to its minimal implementation overhead. However, as mentioned in Section 6.3, it seems to adversely affect calibration, according to ECE-like metrics, when compared to baseline models. This discrepancy might be due to a mismatch between the objectives during training and testing in the PA method. To examine this hypothesis, we consider training our model using the PA predictive likelihood (as given by Equation 4) and assess its effect on calibration.
Given the product structure and the non-smooth activation function (ReLU) employed in PA's likelihood, the training can be quite challenging. To enhance the stability of the training process, we introduce two modifications. Firstly, instead of training the model from scratch, we begin by training with a standard EENN objective (see Equation 1) with a softmax likelihood and then finetune the model using the PA likelihood at the end.[11] Secondly, we substitute the ReLU activation function with a Softplus function[12], which we found aids in stabilizing the model's convergence during training. Consequently, we also utilize Softplus to map logits to (unnormized) probabilities during testing. Moreover, we found that using ensemble weights  $w_{m} = 1$  is more effective for models finetuned
<sup>11</sup>We train with EENN objective for the first 2/3 of epochs and with PA objective for the last 1/3 of epochs.
12Softplus is a smooth approximation of the ReLU function: Softplus(x) := log(1 + exp x)
Figure 16: Results for various image classification datasets using MSDNet as the backbone model. We see that finetuning with PA (as opposed to only applying it post-hoc) can be effective in reducing PA's calibration gap, though at some cost to both test accuracy and conditional monotonicity.
with the PA objective, as opposed to our default choice of  $w_{m} = m / M$  presented in the main paper for post-hoc PA.
The results are presented in Figure 16. Focusing initially on (marginal) test accuracy (right), we note that PA training  $(-)$  has a slight negative impact on performance, particularly when contrasted with the post-hoc PA  $(—)$ . Similarly, it underperforms post-hoc PA in terms of conditional monotonicity (middle), but it considerably surpasses the baseline MSDNet model  $(-)$ . When considering ECE, it becomes clear that fine-tuning with PA leads to significant improvements over post-hoc PA, thus substantiating our hypothesis that the calibration gap mainly arises from the discrepancy between train-time and test-time objectives. On the CIFAR-100 dataset, PA with finetuning even outperforms the baseline MSDNet in terms of ECE.
Based on these observations, we conclude that integrating PA during the training of EENNs is an effective strategy for bridging the calibration gap, though it incurs a minor cost in terms of overall accuracy and conditional monotonicity. We note that these costs could be mitigated (or ideally, completely eliminated) through further stabilization of product training. As such, we view this as a promising avenue for future research.
B.3.3 Relevance of ECE in Anytime Prediction Setting
In Section 6.3, we pointed out poor Expected Calibration Error (ECE) in the initial exits as one of the limitations of our proposed PA method. It is important to note, however, that no decision in the anytime setting is based on model (or prediction) confidence, i.e.,  $\max_y p_m(y|x)$ . This contrasts with the budgeted batch classification setting [Huang et al., 2018], where the decision on when to exit could be based on the model confidence. Thus, we believe that our PA potentially undermining
model confidence, as evident by poor ECE in the earlier exits, is not a limitation that would seriously impact the merit of our solution. Especially since we focus solely on the anytime setting, where, we would argue, non-monotone behavior in terms of performance quality is a much more serious issue.
However, it is true that model/prediction confidence, i.e.  $\max_y p_m(y|\boldsymbol {x})$ , could be useful in the anytime scenario as a proxy for model uncertainty. This allows the model—when prompted by its environment to make a prediction—to also offer an estimate of uncertainty alongside that prediction. However, given that our PA potentially adversely affects model confidence in earlier exits, we caution against its use. Instead, we recommend alternative measures of uncertainty that better align with our model, such as the conformal set size (c.f., Section 6.2).
B.4 NLP Experiments
In this section, we conduct NLP experiments with the aim of demonstrating that our proposed approaches for adapting EENNs for anytime inference are generalizable across different data modalities.