How to Train Your Multi-Exit Model? Analyzing the Impact of Training Strategies

We sincerely appreciate the reviewer’s thoughtful evaluation of our work and their recognition of its significance. We apologize for brief answers necessitated by the response length limit. Kindly if we resolved the concerns raised, we would be grateful if the reviewer would consider raising their score accordingly.

train time passes

We emphasize that the number of phases of training does not have to affect training time. As noted in the manuscript, in all of our experiments we used early-stopping. We report the resulting number of training epochs for the Tinyimagenet/ResNet50 experiment:

Disjoint: 523 +/- 156 Joint: 1610 +/- 395 Mixed: 1166 +/- 136

In practice the opposite is true, joint training is slower to converge.

benefits of the findings are quite incremental

In most cases the difference over joint regime is definitely significant, e.g. 2 percentage points for the ImageNet experiment. For low budgets the difference over disjoint is immense - and we emphasize that the disjoint regime is also in popular use.

LLMs

Please see the answer to reviewer cDMi.

mutual information should decrease monotonically across layers

The monotonic decrease shown in Shwartz-Ziv & Tishby (2017) and inferred through the Data Processing Inequality assumes exact calculations. In our work, we estimate mutual information using a Monte Carlo method (Kawaguchi et al., 2023) in high-dimensional spaces. This approximation can introduce estimation noise, which may result in non-monotonic fluctuations in the absolute MI values.

However, to further verify the MI impact, we perform additional experiment for another vision (CIFAR-100) and NLP (BERT-B/Newsgroups) datasets. In this case, the MI has a more decreasing tendency. However what we want to emphasize is the relative relationships between regimes present in all the figures that hint that higher Mutual Information may be an indicator for better performance of the mixed regime. results

Different Early Exit Adapters:...

As suggested, we run additional experiments where each head has an additional transformer block:

Due to the size of this head architecture, the model is not able to meet the 25% of the original model’s budget, so we omit the first column. We emphasize that the mixed regime has better performance as before.

Figure 3:...

When linearly interpolating the weights between the (models trained with) disjoint and mixed regime, we do not encounter a region of high loss. This means they lie in the same optimization basin. In contrast, when interpolating between the joint regime model and any other model, we do encounter a region of high loss (yellow color).

Table 3…? Table 6: Why the difference in accuracy…is 1.02%…while in Table 3 the difference is almost 30%?

Each dataset and model combination has a different budget-difficulty characteristic (the shape of FLOPs vs accuracy curves, see Figure 1). Tables 6 and 7 were computed for ViT-T architecture and the ImageNette dataset (a 10-class subset of ImageNet, with the same input image size as ImageNet), a combination which turns out to be a relatively easy problem for a model of this size. Table 3 was computed for CIFAR-100/ResNet-34, which turns out to be a relatively harder problem for this architecture.

For comparison, we repeat the experiment from Table 6 for Tinyimagenet below:

…"Mixed-gradual training" and "Alternating training"

"Mixed-gradual" trains the model with n ICs in n phases, each phase including a larger number of ICs, starting from the deepest ones. "Alternating" switches every training step between training with the last IC, and with all the ICs.

Line 666: …x and y?

x and y represent scalar coefficients used to perturb the model parameters θ∗ along randomly chosen directions (δ,η) to visualize the loss landscape around the trained model (Li et al., 2018).

What is…?

PBEE - Patience-based Early Exit - EE method where we exit for each sample after n consecutive ICs return the same class (Zhou et al., 2020), GPF - Global Past-Future - an EE method that incorporates hidden states from the preceding and deeper ICs (Liao et al., 2021) Entropy - means we base our exit decision on entropy of the prediction probabilities (Teerapittayanon et al., 2016) rather than maximum softmax probability

We sincerely appreciate the reviewer's detailed feedback. We incorporate all suggestions in the current version of our manuscript.

We sincerely appreciate and agree with the reviewer’s assessment that our work addresses a meaningful gap in the early-exiting literature.

numerical rank informativeness

Firstly, the key insight of the numerical rank metric is that placing multiple early exits increases the rank and also the expressive representation of the network. The difference in the ranks between the plain model (disjoint regime) and the mixed regime in Figure 4a is significant as it is substantially larger than standard deviation. In Fig. 4a we only show the mixed result for clarity but Fig. 4b shows both mixed and joint regimes obtain similar ranks (given that both regimes foster enriched representations), and both are respectively higher than the backbone.

Secondly, the backbone remains unchanged for the disjoint regime and so the numerical rank experiment seen in Fig. 4a is also valuable because it may explain the inferior performance of early ICs in the disjoint regime. Joint training results with layers that enable more expressive intermediate representations (higher numerical rank), thus allowing for coexistence of features relevant for the nearest IC, and those relevant for the deepest ICs. On the other hand, a consistently lower rank is obtained for the disjoint regime. As mentioned above when even an already trained backbone is allowed to be affected by the ICs’ gradients, the numerical rank of its representation rises (see Figure 4a), and so does its performance on lower budgets (see accuracy results for disjoint vs mixed in any setting).

In the current version of the manuscript we rewrite the explanation of the numerical rank results to make these points more clear to the reader.

While I understand it is not the focus on this work, it would still be valuable to say something about whether the findings on training dynamics presented in this paper translate to early-exit LLMs [1, 2, 3]

We appreciate this insightful suggestion. We agree that exploring early exit regimes in LLMs is a promising direction. However, we intentionally refrained from discussing generative LLMs as translating our observations to early-exit LLMs is challenging due to several fundamental differences (we also note that recent early-exiting studies [1, 2] refrain from studies on LLMs.). Methods like proposed CALM and FREE incorporate intricate confidence mechanisms, state-copying techniques, and synchronization strategies tailored specifically for token-level decisions, which are not directly analogous to the layer-wise early-exiting approach we analyzed. Moreover, EE-LLM emphasizes 3D parallelism and large-scale model training, aspects that differ substantially from our focus on training regimes and layer-wise representation dynamics.

However, as an early effort we add an experiment on the NLP classification problem using BERT-B trained on Newsgroups dataset. results For the backbone network (BERT-B), lower MI in earlier layers and higher in the rear layers used by the disjoint regime may explain its positive performance only at the deeper layers. A slightly increased MI for mixed regime is further indicative of better performance.

[1] Jazbec, Metod, et al. "Towards anytime classification in early-exit architectures by enforcing conditional monotonicity." [2] Meronen, Lassi, et al. "Fixing overconfidence in dynamic neural networks." [3] Yoshitomo, Levorato, Restuccia. "Split computing and early exiting for deep learning applications: Survey”.

The fact that joint strategy leads to suboptimal performance for larger computational budgets reminds me a bit of the "interference" …

We agree that interference between internal classifiers is a central issue - indeed, it’s what inspired our gradient dominance metric. Our experiments reveal that even MSDNet’s architectural adjustments, including its dense connectivity, does not fully resolve this issue. The same is true for gradient equilibrium from [1], as we show in Section 4.4. In fact, our mixed training regime outperforms these approaches, demonstrating more effective handling of interference and yielding superior performance, particularly under larger computational budgets.

[1] Li, Hao, et al. "Improved techniques for training adaptive deep networks."

We thank the reviewer for the effort spent on reviewing our paper and the valuable insights. If we further addressed the remaining concerns, we would kindly ask for possible score reconsideration.

The authors do not do experiment on SOTA settings. The results presented in Table 1~6 are problematic. They provide the x-axis as a ratio, make the reviewer hard to know they achieve this performance in which FLOPs

In this work, we followed the performance reporting scheme initiated by [1] due to its simplicity and readability (though without the absolute FLOPs values). However, as per the reviewer’s suggestion, we provide the same results as FLOPs vs accuracy plots (similar to the one from Figure 1) in this link. We also include them in the current version of our manuscript.

[1] Kaya at al. "Shallow-deep networks: Understanding and mitigating network overthinking."

We are also not sure if the reviewer might have meant it but we want to emphasize that our paper does not aim to achieve state-of-the-art performance on any setup. Its purpose is to analyze the training dynamics and compare the multi-exit model training strategies. We do this by evaluating all the considered regimes on multiple commonly-used architectures, on multiple modalities and datasets, and for multiple early-exiting approaches. Moreover, most of the results we present were computed for networks trained from scratch, for three different PRNG seeds. Our limited computational resources prevent us from running large SOTA models. Nevertheless, to demonstrate that our findings scale to larger models, we present the results for the ImageNet-1k experiment for the ViT-S variant of the vision transformer architecture in the link provided (Figure 11):

We can see that - as before - the disjoint regime still is significantly inferior on lower budgets, while the joint regime exhibits a noticeable performance gap on higher budgets.

the proposed method, though well-explained, lacks technical innovation - essentially combining two existing approaches.

We thank the reviewer for their feedback. Our mixed training strategy is intentionally technically simple, and we consider that as an advantage rather than a weakness of our work. We believe that clarity and effectiveness should take priority over unnecessary complexity. Rather than adding complexity for its own sake, we focus on a practical, well-reasoned method that consistently outperforms existing approaches. Moreover, technically simple methods are easy to implement, and - as a consequence - are more likely to be widely adopted by the community. We emphasize that the primary goal of our work is to systematically analyze early-exit training strategies. Since existing strategies are also simple, it is important to first analyze their limitations before considering more complex alternatives. Our paper demonstrates the weaknesses of the two commonly used training strategies, and highlights the large impact of this previously overlooked aspect. The proposed “mixed” approach is the most straightforward approach to alleviate the issues of the joint and disjoint regimes. While being technically simple, it is not necessarily obvious, as it was not used in any prior work.

The authors provide some experimental analysis in section 3, but I can hardly find the relationship between them and the proposed method.

We appreciate the reviewer’s feedback and would like to clarify how the analyses in Section 3 are directly connected to the proposed mixed training strategy. Specifically, the experimental metrics - gradient dominance, mode connectivity, numerical rank, and mutual information - were chosen to reveal distinct aspects of the training dynamics under different regimes. For example, gradient dominance analysis shows that the mixed training regime shifts the optimization focus toward later exits. This shift directly correlates with improved performance under higher computational budgets, as later exits become more robust. Please note the relation to the gradient equilibrium experiments in Sec. 4.4. Gradient dominance results may explain the need to rescale the rear exists when trained in joint regime while mixed obviates the need for application of gradient equilibrium. Mutual Information and Numerical Rank show further aspects why the performance of the presented regimes may differ (e.g. for early vs rear exits). Finally mode connectivity reveals that mixed and disjoint regimes produce very similar solutions, despite having significantly different performance on low computational budgets.

We thank the reviewer for assessing our work and recognizing the importance of this previously overlooked aspect. We hope that the answers below adequately answered the reviewer's questions and concerns. If that is the case, we kindly ask for a reconsideration of the score.

number and placement of early exits, which have not be adequately considered in the submission … but not adequate variations in the configuration of the early-exit models … consider conduction an ablation on the number and placement of ICs

In Appendix B we provide this kind of an analysis for our ViT-T model trained on the ImageNette dataset. These results show that our findings generalize to different placement densities. However, to make our case even stronger, we rerun those experiments on Tinyimagenet/ResNet-50 and include additional placement schemes (Dense-Sparse places ICs at blocks: [1, 2, 3, 4, 5, 6, 7, 11], Sparse-Dense places ICs at blocks: [1, 4, 8, 9, 10, 11, 12, 13, 14]). We present the results below:

The results are consistent with the main findings of the paper, that is:

1. The mixed regime still presents generally better performance over the joint regime.
2. The disjoint regime is still inadequate for low budgets.

The proposed "mixed" approach lacks novelty, … starting with an ImageNet-pretrained backbone can be considered common practice… …essentially comprises a backbone pre-training followed by joint training, and cannot be considered novel as it is common practice in ML deployment…

We emphasize that the proposed “mixed” approach is not in common use. The common practice of using pre-trained models for transfer learning is independent of the choice of the training regime. Starting with a model pre-trained on dataset A does not preclude us from fine-tuning on dataset B with different training regimes.

In particular, let’s assume that a ML practitioner starts with a backbone model (e.g. ViT-B) pretrained on dataset A (e.g. ImageNet-1k). His aim is a multi-exit model that performs well on dataset B (e.g. CIFAR-100). If he takes the backbone, attaches classification heads, and trains (finetunes) everything jointly on dataset B, then that is still joint training (finetuning) according to our terminology.

To perform mixed regime training in such a setting, we first finetune (on dataset B) the backbone only, and only then proceed to finetune everything together. In the paper we compare all regimes in the transfer learning setting and present the results in Table 4. In this setting the behaviour of each regime and the findings are the same as in other experiments - both joint and disjoint training display a significant performance gap on some budgets.

If A == B, then indeed taking a pre-trained model, attaching ICs and training jointly indeed would be equivalent to our “mixed” approach. However, we argue that such cases are extremely uncommon.

Again, we apologize for the misunderstanding, which resulted from our insufficiently thorough description of the transfer learning experiment. In our current version of the manuscript we modify this section to present the pre-trained setup more clearly.

…alternative presentation…

We thank the reviewer for this valuable suggestion regarding the positioning of our contribution. Although it may not have been emphasized enough, our goal in this work is to provide a comprehensive comparative analysis of different early-exit training regimes, and investigate their training dynamics and implications. In the paper, we point out scenarios where the joint or even disjoint training regimes can be more suitable or advantageous. For instance, we observe that the joint regime may be preferable at very low computational budgets (where early classifiers dominate), and the disjoint regime can show good performance when the backbone model is already well-trained or fixed and further training resources are limited. Nevertheless, our empirical results consistently suggest that, in most practical scenarios and computational budgets, the mixed regime tends to outperform the others, hence our emphasis on highlighting its benefits.