Jointly-Learned Exit and Inference for a Dynamic Neural Network

Thank you for the detailed and perceptive review, which has brought a number of issues to our attention, leading to modifications that we believe have improved the paper.

4.1. Reorganization and visual descriptions:

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.

We would like to express our appreciation for the thorough review, and in particular, the suggestions to conduct experiments with Dynamic Perceiver. We believe that the experimental results have strengthened the paper.

3.1. Lack of experiments on more advanced architectures; Dynamic Perceiver:

Thank you for the suggestion to experiment on more advanced architectures and, in particular, those tailored for early-exit. It prompted some interesting and fruitful experiments with Dynamic Perceiver. We initially selected T2T-ViT as an example transformer model that performs well on the tasks under study. But we agree that combination with a state-of-the-art early-exit model would be more compelling.

In 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. 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.

3.2. Overclaiming contributions:

Thank you for pointing this out. Indeed, we agree that we do not propose a new architecture, and we have removed the word ''architecture'' from the claimed contributions, both in the introduction and the conclusion, in the revised paper.

3.3. Lack of experiments on ImageNet:

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.

3.4. Presentation:

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

Thank you for your careful and insightful review of the paper and for your positive comments concerning its strengths. Your constructive criticisms have prompted us to make some improvements to the paper.

2.1. Use of the min operator between two terms in Eq.8:

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.

2.2 Evaluation only on frozen-backbones:

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.

2.3. ImageNet 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. Although our introduced gates and IMs are lightweight and have few parameters, training with ImageNet is time-consuming, even with relatively powerful servers.

Minor comments:

That was a mistake, thank you for pointing it out.

1.3. Similarity of Loss Terms to Prior Work:

The critical aspect of our loss formulation and bi-level optimization is the interaction of the learning for the gates and IMs. While there is a superficial similarity between the forms of the losses and the losses in previous work, the model parameterizations and learning process are considerably different. Specifically, it is the combination of the losses that is important. The two loss terms interact via the bi-level optimization; updates of the IM parameterss continually influences the target of the binary tasks of the GMs, and updates of the GM parameters change the weights in our weighted cross entropy loss.

Starting with the gate loss, while it is true that we sum the losses of independent binary variables $g^l_{\phi}$, our binary variables $g^l_{\phi}$ do not directly correspond to exit probability in the way that is common in the literature. In particular, we do not set $P(G=l)$ $= g^l_{\phi} \prod_{i=1} ^{l-1} (1-g^i_{\phi})$, as most learnable GMs approaches do. The parameterization of our exit probabilities given by the equations (7) and (8) (and therefore our whole gating mechanism) is different and novel. We acknowledge that the way we described our gate mechanism might have been too compact, so we have provided an additional figure to clarify our architecture.

Next, we turn to the IMs and our main task loss. By dividing the parameters in our bi-level approach, we obtain two losses, one for the gate parameters $\phi$ and one for the IMs parameters $\theta$. The loss for the IM parameters includes a term with a weighted cross entropy loss, but, critically, the weights are driven by the probability of a sample exiting at that given gate ($P_{\phi}(G=l|\mathbf{x}_i)$). In contrast, the weights in the formulation of Han et al. (2022b) are not connected to the exit probability.

We arrive at the IM loss by starting from a principled combined loss that encompasses both the GMs and the IMs, balancing accuracy and computation. In contrast, Han et al. directly start from the weighted IM loss formulation, with the weights being modeled by a weight predictor network (WPN). There is no interaction with the IM parameters during training. The WPN takes as input the losses at each IM to predict the weight, so it cannot be used at inference. As a result, whereas our weights are tied to the probability of exiting, as they should be, the weights in the loss of Han et al. (2022b) are heuristic, selected somewhat arbitrarily to reduce the train-test mismatch. Our loss incorporates the gate parameters that are actually used at inference, while Han et al. (2022b)'s loss does not.

All of these differences are major. We can confirm empirically that these differences are major because they lead to significantly different behavior, as shown in Figure 2.

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.

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.

Thank you for the detailed and thoughtful review. You raised some excellent points that made us think carefully, leading to modifications and experiments that have provided insights and improved the paper.

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.

Our most recent revised version contains the following:

• As promised, we have added results on ImageNet in Observation 5 of our main paper and in Appendix 8.13. The ImageNet experiments show JEI-DNN on two types of backbone architectures: a general architecture (T2T-ViT-14) and a state-of-the-art early-exit dedicated architecture (Dynamic Perceiver).
• For the general architecture (T2T-ViT-14), we compare to the best-performing training baseline, L2W-DEN. For equivalent accuracy, we can reduce the inference cost by 1-2 percent over an important operating regime. The results on ImageNet are thus in line with the results of our other experiments on the other datasets.
• For the early-exit specialized architecture (Dynamic Perceiver), which already achieves an impressive trade-off between inference cost and accuracy, we show that employing JEI-DNN can offer further improvement, reducing the inference cost by an additional 1-2 percent.
• Moreover, we additionally presented an experiment for end-to-end training where we combined the training of our JEI-DNN method with the Dynamic Perceiver in Appendix 8.17. Fully end-to-end training (JEI-DNN+backbone) offers the best performance for the CIFAR100 dataset.