Paper: Jointly-Learned Exit and Inference for a Dynamic Neural Network (ICLR-2024)

Reviewer comment (Presentation):
I find it hard to fully understand the method before reading through the entire paper. In particular: (a) the abstract has no mention of the paper's contributions; (b) the "Contributions" also does not explain how the method works; (c) there is no visual schema of the model; (d) some important details are described very late, for example "Gate design" is part of the experimental section but it is a crucial design decision. Overall it is understandable, but I think some minor reorganization and some additional visual descriptions (maybe a shortened pseudo-code?) can significantly help the reader.

Author response:
Thank you for the suggestions to improve the clarity of the paper. Unfortunately, it is forbidden to modify the abstract. We agree that an additional sentence would help to clarify the contribution. However, following your suggestion, we have added a sentence to the introduction in the ''Contributions'' section that provides a high-level description of how the proposed method works. Also in line with your suggestion, we have added Figure 1, which provides a visual schema of the model and describes the core steps of the methodology. In the original version of the paper, we elected to include ''Gate design'' in the experimental section to separate the core components of the methodology from what we considered to be implementation decisions. For example, it is possible to use the softmax output directly rather than construct uncertainty features (this achieves similar performance but at the expense of a larger gate). We have now moved the gate and IM design into the methodology, as recommended.

Reviewer comment (Novelty):
On the novelty, I think a few methods could be added and discussed more. For example, the zero-time waste model (cited in the paper) has a geometric ensemble of the early exits which is very similar to their gating probability, if I understand correctly. As another example, differentiable branching (https://ieeexplore.ieee.org/document/9054209) considers a trainable exit strategy that also combines the EEs and that can be trained end-to-end with the rest of the network. Some of these methods can be also added to the experimental comparison.

Author response:
Thank you for this suggestion. We have updated the discussion of the related work in the main paper to refer more prominently to the two indicated works. We have also added a slightly lengthier discussion in the appendix.
    
The ZTW paper by Wolczyk et al. (2021) focuses on learning weights to ensemble the earlier IMs (or ICs in their work). The weights are not actually used for early exiting. To early exit, ZTW uses the same fixed threshold mechanism as the baseline in our paper, as can be seen in their Equation 3. There is a fixed threshold, and the sample exits if the maximum probability of the prediction exceeds it. The setting that ZTW considers is to reuse the IM values in the best way possible. Our setting is to improve the interaction between the gating mechanism and the IM training. As such, we do not view it as a competing baseline, but addressing another problem of early exiting. It would be an interesting extension to incorporate ZTW into our joint learning framework to determine if reusing earlier IM results can improve performance further.

There is some similarity between the geometric ensembling employed by Wolczyk et al. and our proposed method, but there is a key difference in that ZTW does not model exit probabilities but instead ensemble weights. When combining the predictions of the IMs, the output probability at IM $l$ is raised to the power $w_l$, where $w_l$ is a learnable parameter. This allows the model to learn to place less emphasis on the predictions from the weaker IMs at lower layers during the fusion process. 
    
In our proposal, we do not form ensembles, but we do learn the exit probabilities. If we interpret each IM softmax output as a probability distribution across labels, then the overall label distribution for any given sample is a mixture of the individual IM label distributions, where the mixture weights are the learned exit probabilities. Thus, although our learning procedure does not perform ensembling for any specific classification decision, its probabilistic output can be viewed as being drawn from an ensemble of classifiers.
   
As for the differentiable branching paper (Scardapane et al., 2020), thank you for bringing it to our attention. It is true that the method is highly related, and we should have included it as a baseline. The critical differences with our method are (i) the way in which the exit probability is modelled; and (ii) the use of bi-level optimization. Scardapane et al. model the exit probability in the product form. We found that this made optimization very challenging and the change to a constrained sum (cumulative exit probabilities summing to 1) was critical. The bi-level optimization, although perhaps not essential, leads to a more robust learning behaviour. 
    
We have implemented a version of the differentiable branching approach and 
we include preliminary results in Appendix 14. The results will be finalized in the next few days. Since the code is not available and some key implementation details are missing (it is a 4-page paper), we are not absolutely sure that our implementation matches what Scardapane et al. did in practice. To give the approach the best chance, we incorporate the same warm-up that we use in our procedure. However, we can see that the performance is still not good, and training is very unstable. Indeed, the results in the paper are limited, but they also suggest that the behaviour may be unstable. For example, we see that for VGG-13, 68.8\% of samples exit at IM 6 and 20.3\% at IM 8, but close to zero at IM 7. This bi-modality is unexpected (considering that the IMs are close so the inference cost does not differ by a substantial amount), and if the learning were stable, we would anticipate a smoother exit behaviour as exhibited by our approach and all of the other baselines we have experimented with.

Reviewer comment (Experiments):
Apart from the topics discussed above, I have a general question about the bilevel task. Can you clarify why this is needed? There are techniques to train probabilistic blocks (e.g., reparameterization, score function estimators), which would allow to train all blocks simultaneously, or I am missing something? Also, what is the overall training time of your model compared to the alternatives?

Author response:
All blocks could be trained simultaneously, but we found that training via the bi-level optimization was more stable and yielded better results. We suspect that one characteristic of the problem that is challenging is the strong interconnections and dependencies between the different components. The model has to simultaneously learn the IM parameters and assign probabilities to sample exit points. Changes in exit probabilities can significantly impact the IM model parameters, but during the early stages of training, the determination of exit probabilities is based on inaccurate IMs. We suspect that the space of the loss has many irregularities since small changes can have a significant effect on the optimal solution. Although the judicious use of probabilistic block training strategies can assist with this, experimentally, we found that training the blocks simultaneously was generally unstable. In summary, we do not believe that bi-level optimization is essential, but it does lead to a reasonably stable and robust solution without too many hyperparameters to configure. We tried multiple techniques for the joint learning, but if you have had good experience with a particular method, we would be pleased to learn about it.


 In general, all the methods share similar training time, with BoostedNet begin the quickest, then L2W, the JEI-DNN. The training time, as well as the number of parameters, is provided in Table 1 of the Appendix.

Reviewer comment (Experiments):
**Lack of experiments on ImageNet**. To my understanding, the learning approach does not need to be applied in a downstream task on the toy small-scale datasets. Experiments on ImageNet would be more convincing.

Author response:
We have commenced experiments on ImageNet and obtained preliminary results. The experiments are ongoing, and we will finalize and report the results at least one day before the end of the response period. Although our introduced gates and IMs are lightweight and have few parameters, training with ImageNet is time-consuming, even with relatively powerful servers.

Reviewer comment (Presentation):
**Presentation**. It is recommended to use figures to clearly show the motivation and the method pipeline.

Author response:
Following your recommendation, we have added Figure 1 at the start of the methodology section to help the reader understand the approach.

Reviewer comment (Experiments):
A few technical aspects of the paper remain unclear. Specifically, the use of the min operator between two terms in Eq.8 is not fully justified in the manuscript, nor experimentally verified by a relevant ablation.

Author response:
Thank you for highlighting this issue. We require that the $P_{\phi} (G | \mathbf{x_i})$ sum to one. This is important because it means that our loss function in Eq. (6) is a correct approximation of the expectation in Eq. (5). Moreover, during inference, we use Eq.(4) to determine $ P(E^l=1 | \mathbf{x_i} ) $ at a gate, and this calculation can lead to probabilities greater than 1 if $P_{\phi}(G | \mathbf{x_i})$ does not represent a correctly normalized probability distribution. The question is then how to map from the unnormalized $\{g^l_{\phi}\}$ provided as output from the neural networks to normalized $P(G|\mathbf{x}_i)$. An immediate thought would be to divide by the sum. But unfortunately, during inference, we are evaluating $g^l$ successively and we do not have access to the sum when we decide to exit. Therefore, we employ the $\min$ operator to ensure that the sum of  $P(G^l|\mathbf{x}_i)$ never exceeds $1$. We have added some text around equation (8) to explain this design choice. We could consider other forms of normalization, provided they can be implemented sequentially, i.e., we need to be able to perform the proposed normalization at layer $l$ without the knowledge of any $g^{l'}(\mathbf{x}_i)$ values for $l'>l$. We would be happy to conduct and include an ablation study if there is an approach that seems a natural alternative to you, but we are not sure what the obvious alternative would be.

Reviewer comment (Experiments):
The proposed methodology is only evaluated on frozen-backbone CNNs. Although this a practical sub-category of EE-models, further evaluation on end-to-end EE models (jointly training the backbone and exits), which are shown to achieve superior speed-accuracy trade-off would increase the contribution of the paper.

Author response:
This is a valid point. We intentionally focused on fixed backbones because our goal was to develop a mechanism suitable even for very large models where even fine-tuning is a major burden. But we have conducted some additional experiments to explore this direction. In the Observation 5 in the revised paper, we include results where we combine the Dynamic Perceiver with JEI-DNN (our proposed method). We use the Dynamic Perceiver framework to jointly train the backbone and the IMs and then use JEI-DNN to jointly train the gates and the IMS. In our current experiment, this is not true end-to-end joint learning, because we do not jointly train IMs, gates, and backbone, but the combined process does train all three components. Figure 6 shows that using JEI-DNN achieves an accuracy improvement of approximately 0.5-1\% compared to Dynamic Perceiver on the CIFAR100 dataset when using MobileNetV3 as a backbone, when using 50-60 percent of the mul-add operations that are used in the original architecture. We are conducting similar experiments on ImageNet and hope to report the results before the end of the response period. We are also exploring performance for true end-to-end training when we allow the backbone model parameters to update as well as the IMs and gates. Here we retain our JEI-DNN bi-level optimization learning procedure, but update the backbone model parameters at the same time as the IM parameters.

Reviewer comment (Experiments):
Although the proposed approach is compared with a wide set of baseline, ImageNet experiments are omitted, rendering the presented results less convincing.

Author response:
We have commenced experiments on ImageNet and obtained preliminary results. The experiments are ongoing, and we will finalize and report the results at least one day before the end of the response period. Although our introduced gates and IMs are lightweight and have few parameters, training with ImageNet is time-consuming, even with relatively powerful servers.

Reviewer comment (Novelty):
It is unclear what the major contribution of the paper is. The use of learned early exiting gates and intermediate classifiers, trained jointly with the backbone is common practice for many early exiting architectures. After all the approximations described in 4.1, the loss term for the gating modules turns into the common practice of summing the losses of independent binary early-exiting classifiers. The main task loss is similarly aligned with prior work from Han et al. (2022b).

Author response:
### 1.1 Contributions: 
The major contributions of the paper are as follows. (1) We propose a novel, robust approach to jointly train the inference modules and decision gates in early-exit neural networks. As explained below, we disagree with the reviewer that this is common practice. (2) We introduce a novel modeling of the probability of exiting, representing it as a constrained sum, rather than a product. This leads to considerably more robust learning performance. (3) We provide a bi-level optimization formulation of the learning task. By decomposing the problem this way, we further encourage stable and efficient training. (4) We demonstrate that our approach results in much better calibration leading to significantly improved conformal intervals. 

### 1.2 Prior Work: 
We disagree that the use of learned early exiting gates and intermediate classifiers, trained jointly with the backbone is a common practice for early exiting architectures. If the reviewer could provide a reference, we would be very happy to discuss more concretely and compare with our approach. We compiled an extensive table (Table 4 in the Appendix) of the existing literature with this specific categorization (non-fixed backbone) to highlight this.
    
For those methods that train their proposed backbone architecture and the IMs jointly (non-fixed backbone), the architectures are specifically designed for early-exit. Even the more flexible models, such as the Dynamic Perceiver (Han et al., 2023) are restricted in their application; for example the Dynamic Perceiver, in its proposed form, can only be used with CNNs. In contrast, our approach can be easily applied to a transformer, a CNN, or a graph neural network. We demonstrate this better in supplementary experiments that we have done at the request of reviewers. But beyond this, to the best of our knowledge, all of the state-of-the-art methods that train the backbone use fixed-threshold gating mechanisms. For instance, the SOTA early-exit models (e.g. Dynamic Perceiver - Han et al. (2023)), do not train gates.

For the fixed backbone case, which is the setting of our work, the most common approach is to either focus (i) on training the IMs using fixed gates; or (ii) on training only the GMs with a fixed backbone and pre-trained IMs. For example, the method proposed by Han et al. (2022b), L2W, does use a trainable gating mechanism. It trains the inference modules. It is this specific difference in training approach that leads to a major difference in exit behaviour, which the reviewer correctly pointed out in Figure 2. 

Another reviewer kindly pointed us to a work - Scardapane et al. (2019) -that does perform joint training of the gates and the IMs (with a fixed backbone). We have added it in our Table 4, and provided results comparing to it in Section 8.13, showing that our JEI-DNN significantly outperforms.
Compared to the proposal in Scardapane et al. (2019), a key innovation of our work is (i) a novel modelling the exit probabilities; and (ii) the bi-level optimization strategy. These combine to make the learning procedure much more robust.

Reviewer comment (Novelty):
Constructing surrogate binary targets for the learned gates is also common, e.g. FrameExit by Ghodrati et al. CVPR 2021.

Author response:
### 1.4. Constructing surrogate problems:
We agree that constructing a surrogate problem is not novel and is not part of our claimed contribution (it does not appear in the paragraph at the end of the introduction). It is merely a detail of the optimization process.

Reviewer comment (Experiments):
Figure 2a) The fact that samples do not exit at all from the first few early exiting layers in JEI-DNN is puzzling and I am wondering if it is due to the choice of specific hyperparameters that prevent early exiting from these layers. As can be seen, the accuracy of the IM at layer 5 & 6 is far higher than the overall performance. It is conceivable that a good accuracy higher than the green dashed line is still achievable by the earlier classifiers at least for a proportion of samples. Is there any hyperparameter that could potentially give more control into the exiting patterns of JEI-DNN?

Author response:
### 1.5 Figure 2a exit behaviour:
This is a very interesting point and we appreciate that you have raised it. When examining the accuracies, we need to be careful because the accuracies at layers 5 \& 6 in Figure 2a) are only *for the samples that exit there*. These accuracies are very high, showing that the gates have effectively learned to exit the correct samples.  The accuracy at layers 5 \& 6 for all the samples of the dataset is much lower, as we can see in Figure 2b). This figure is for a specific $\lambda=1$ (the corresponding operating point is marked with a star in the left panel of Figure 1). For that particular parameter, it is not unexpected that the gates have learned to avoid the first 3 gates entirely and only start to exit a very few samples at gate 4. 
The decision being made is whether the drop in cross-entropy loss (in going to gate 5 from gate 4, for example) is worth the additional cost of inference. For the choice $\lambda=1$, it is not the case for many samples. While IM 4 may still be able to classify some subset of the points with accuracy of 90 percent, IM 5 would classify that same subgroup with accuracy 95 percent. In Appendix 8.15. of the revised paper we now provide the same figure for a higher $\lambda=3$,  corresponding to a greater penalty on computation, and $\lambda=0.5$, corresponding to a reduced penalty on computation, with more focus on accuracy. For the higher $\lambda=3$, we observe that now the gates decide to exit many more samples at earlier gates (2,3,4).  This behavior is a consequence of our novel early-exit framework. These results indicate that the $\lambda$ parameter does provide meaningful control over where samples exit.

Reviewer comment (Experiments):
All the results in the paper are limited to the T2T-ViT backbone. The authors should show the efficacy of their method for a larger variety of backbones, preferably to models that are established in the early exiting literature such as MSDNet, DenseNet, etc.

Author response:
### 1.6. Limitation to the T2T-ViT backbone:
 Thank you for the suggestion to conduct experiments with different backbones. During the response period, we have conducted experiments using the following two backbones: (i) Dynamic Perceiver on MobileNet-v3; (ii) GRIT - a graph transformer. Dynamic Perceiver is an early-exit architectural extension for CNNs that is shown to outperform MSDNet and DenseNet, so we consider that using it as a backbone is sufficient to demonstrate that the proposed JEI-DNN can be combined with state-of-the-art early-exit architectures and CNNs. The successful application of our method to a graph transformer illustrates the flexibility of the proposed approach. The results for these two architectural settings are reported in Appendix Sections 8.12 and Observation 5 of the revised paper.

Reviewer comment (Evaluation):
Comparison of the accuracies among T2T-Vit and MSDNet architectures in Fig 5 seems unfair. Most of the gain in accuracy comes because of the more powerful transformer-based backbone and not because of the efficacy of the early exiting approaches. In fact, MSDNet and RANet show more robust early exiting results compared to the results reported for T2T. E.g. MSDNet retains the original accuracy of the model after almost 50% compression. The performance of the proposed JEI-DNN approach in comparison drops very rapidly even with 25% compression.

Author response:
### 1.7. Figure 5 - Unfair accuracy comparison between T2T-Vit and MSDNet architectures:
 We recognize that the comparison between the T2T-Vit and MSDNet architectures is unfair. The main point of the comparison was to demonstrate that the proposed JEI-DNN learning procedure can achieve performance improvement because it is sufficiently flexible that it can be incorporated with state-of-the-art architectures. The text in the original version of the paper was  insufficiently clear on this point. In contrast, MSDNet  and RANet are dedicated architectures that incorporate design decisions like skip connections and multi-scale density that are specifically tailored for early-exiting. We thus cannot use MSDNet in conjunction with a state-of-the-art model to improve performance. The two SotA baselines used in the original version of the paper (BoostNet and L2W-DEN) provide architecture-agnostic training for early-exiting that aims to improve the accuracy-cost trade-off. As such, a fairer and better comparison is to compare these two training procedures with ours on a given architecture as we did throughout most of the paper using T2T-ViT.  Another reviewer raised the point that the Dynamic Perceiver (Han et al., 2023) can probably be considered the state-of-the-art for dynamic early-exit architectures. The results presented by Han et al. show clear outperformance of MSDNet and RANet. We have now conducted experimental comparisons with Dynamic Perceiver, which can augment almost any CNN model to perform efficient early-exiting. The results show that our learning procedure improves upon the performance of Dynamic Perceiver, so we consider this to be a much fairer and more meaningful comparison with an existing, state-of-the-art dynamic early-exit architecture.

In the current version of the paper, we have moved the comparison with MSDNet, RANet, BoostedNet and L2W-Den to an Appendix and included a much clearer statement that the results should be interpreted with care because the performance discrepancy is largely due to the backbone. We would also be happy to remove this comparison from the paper entirely, since we consider the results to be somewhat redundant in light of the new dynamic perceiver results.

The new experiments demonstrate that JEI-DNN can be used to improve the performance of Dynamic Perceiver by adding small trainable gates to the architecture, and then jointly learning the IM and gate weights. Currently, we include the results for CIFAR100 in ``Observation 5'' in the results section, but we are compiling ImageNet results and hope they will be finalized before the end of the response period.

Reviewer comment (Experiments):
The method is only evaluated on three small-scale datasets: CIFAR10 & 100, SVHN. The authors should consider expanding the evaluation to ImageNet.

Author response:
### 1.8. Larger scale experiments: 
We have commenced experiments on ImageNet and obtained preliminary results. The experiments are ongoing, and we will finalize and report the results at least one day before the end of the response period.

Reviewer comment (Theory):
The assumption that a sample exited at a gate at layer $l$ should also exit from any late stage gate seems against the prevalent view of $overthinking$: "Overthinking is computationally wasteful, and it can also be destructive when, by the final layer, a correct prediction changes into a misclassification." by Kaya et al. ICML 2019. How did you make this assumption?

Author response:
### 1.9 Assumption that a sample exited at a gate at layer should also exit from any late stage gate vs. Overthinking:
 Overall, the entirety of our proposed approach is very much aligned with the quoted sentence by Kaya et al. (2019). We aim to train to exit a sample at the earliest possible layer where the probability of correct classification is sufficiently high. Having said that, we understand that you are referring to our training assumption for optimizing the gates in Section 5, where we state that ``we assume that it should also exit at any later gate''. Kaya et al. 2019 study shallow-deep networks, where intermediate backbone layers are trained to provide outputs that can be used effectively for classification. In contrast, the backbones we employ are not shallow-deep. For a shallow-deep architecture, it is much more likely that the IM based on layer 5, for example, might produce a correct answer, whereas the IM at layer 9 might not. In our backbones, where the intermediate features are not intended to be used for immediate classification without processing by subsequent layers, this phenomenon is much less common. To quantify this, we calculated a conditional accuracy for each IM, where the conditioning is on the previous IM correctly classifying the sample. The results are included in Figure 20 in Appendix 8.16 in the revised version of the paper. The figure shows that the conditional accuracy exceeds 90 percent for all IMs. This compares to unconditional accuracies of 40 and 65 for IMs 2 and 3. This observation supports the intuition behind the assumption we made, namely that once one IM manages to make a correct classification, all subsequent IMs are likely to do so and we should exit as soon as possible (set all the targets to 1). 

Having conducted this test, your insightful observation prompted us to consider the necessity of the assumption. Indeed, there is no reason that JEI-DNN cannot be used in conjunction with a shallow-deep network, and there these assumptions may be flawed. For this reason, we identified an alternative approach of setting a later gate's binary target to 1 only if its combined accuracy and inference cost (for the current IM parameters) is less than all subsequent gates. With this strategy, we avoid the assumption that all subsequent IMs will be correct. We have conducted preliminary experiments using this approach; for the t2t-vit-7 backbone on the SVHN dataset, the performance is almost the same as the original scheme.