HydraViT: Stacking Heads for a Scalable ViT
Authors: Janek Haberer
Ali Hojjat
Olaf Landsiedel
Venue: NeurIPS 2024 poster (2024)
Keywords: Deep Learning
Transformers
Vision Transformers
Scalable Transformers
Primary area: machine_vision
TL;DR: By sorting attention heads during training, we enable flexible inference that adapts to diverse hardware constraints by dropping the least important heads.

Abstract: The architecture of Vision Transformers (ViTs), particularly the Multi-head Attention (MHA) mechanism, imposes substantial hardware demands. Deploying ViTs on devices with varying constraints, such as mobile phones, requires multiple models of different sizes. However, this approach has limitations, such as training and storing each required model separately. This paper introduces HydraViT, a novel approach that addresses these limitations by stacking attention heads to achieve a scalable ViT. By repeatedly changing the size of the embedded dimensions throughout each layer and their corresponding number of attention heads in MHA during training, HydraViT induces multiple subnetworks. Thereby, HydraViT achieves adaptability across a wide spectrum of hardware environments while maintaining performance. Our experimental results demonstrate the efficacy of HydraViT in achieving a scalable ViT with up to 10 subnetworks, covering a wide range of resource constraints. HydraViT achieves up to 5 p.p. more accuracy with the same GMACs and up to 7 p.p. more accuracy with the same throughput on ImageNet-1K compared to the baselines, making it an effective solution for scenarios where hardware availability is diverse or varies over time. The source code is available at https://github.com/ds-kiel/HydraViT.

HydraViT: Stacking Heads for a Scalable ViT
Janek Haberer', Ali Hojjat", Olaf Landsiedel
Kiel University, Germany
Equual contribution
{janek. haberer, Kjanek.haberer, ali jojjaa, olaf landsiedel)@cs
Abstract
The architecture of Vision Transtormers (ViT') particularly the Multi-head Atten- [ion (MHA) mechanism, imposes substantial hardware demands. Deploying ViTs on devices with varying constraints, such as mobile phones, requires multiple mod- els of different sizes. However, this approach has limitations, such as training and storing each required model separately. This paper introduces HydraViT, a novel approach that addresses these limitations by stacking attention heads to achieve a scalable[Vi] By repeatedly changing the size of the embedded dimensions through- out each layer and their corresponding number of attention heads in MHA during training, HydraViT induces multiple subnetworks. Thereby, HydraViT achieves adaptability across a wide spectrum of hardware environments while maintaining performance. Our experimental results demonstrate the efficacy of HydraViT in achieving scalable ViT with up to subnetworks, covering a wide range resource constraints. HydraViT achieves up to 5 p.p. GMACs and up to p.p. more accuracy with the same throughput on ImageNet- IK compared to the basclines, making an effective solution for scenarios where hardware availability diverse or varies over time. The source code available at https: tig........... com ds kiel/ wwweeeeeeeeeaaeeee HyddraViT
accuracy with the same
82
Accuracy
83.0
82
Throughout vs. Accuracy
Deeaterr
ball
83.0
Zoomed
82.5
Henommasoo
82.0
Accuracy
81.5
Accurayy
somedher
76
81.0
80.5
80.5
74
80.0
72
Malforme
79.5
79.5
70
79.0
79.0
10
15
10
15
2000
4000
6000
8000
10000
3000
GMACs
GMACs
Throughppt (e/s)
Throughout [8S)
(a) GMACs vs. Accuracy
(b) Throughput vs. Accuracy
Figure 1: Performance comparison of HydraViT and basclines ImageNet-IK in terms of GMACs
(a) and throughput (b) evaluated on NVIDIA A 100 80GB PCle. HydraViT trained on 3-12 heads
demonstrates superior performance
DynaBERT (Hou et al. 2020) and SortedNet (Valipour
jet al. 2023). While MatFormer (Kudugunta et al. 2023) shows higher performance than HydraViT within its limited scalability range, but when we train on narrower scalability range (9-12 heads), HydraViT surpasses MatFormer. We also show that training HydraViT for more epochs can further improve accuracy. Note that each line corresponds one model, and changing the number of heads in the vanilla DeiT models significantly drops their accuracy to less than 30%.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Introduction
Following the breakthrough of Trransformers (Vaswani al. 2017). Dosovitskiy et al. (202T) established the Vision Transformer ( J as the base transformer architecture for computer vision tasks. As such. numerous studies build on top of ViTs as their base (Liu et 2021 Tolstikhin et al. 2021 Yu ct al.) 2022). In this architecture, MacrcttuurrMMuiieeeeeee
Motivation
Multi-head Attention (MHA) plays an important part, capturing global relations between different parts the input image, However, ViTS have much higher hardware demand due to the size of the attention matrices in MHA which makes it challenging find a configuration that fits heterogeneous devices.
To accommodate devices with various constraints, ViTs offer multiple independently trained models with different sizes and hardware requirements, such as the number of parameters, FLOPS, MACs, and hardware settings such latency and RAM, with sizes
Table 1: VIT Configurations
VII-Ti
VIE-S
VIEBB
# Layers
12
12
12
typically increasing nearly at logarithmic scale (Kudugunta et al.
Dim
768
2023), see Table ] Overall, in the configurations of ViTss the
1imm per Head
64
64
number of heads and their corresponding embedded dimension in
@ Params
5.7 M
22 M
86.M
MHA emerges as the key hyperparameter that distinguishes them.
While being a reasonable solution for hardware adaptability, this approach has two primary disadvantages: (1) Despite larger models, e.g., VIT-S and ViT-B, not having a significant accuracy difference, each of these models needs to be individually trained, tuned, and stored, which not suitable for downstream scenarios where the hardware availability changes over time. (2) Although the configuration range covers different hardware requirements, the granularity is usually limited to small selection of models and cannot cover all device constraints.
Observation By investigating the architecture of these configurations, we notice that VIT-Ti, ViT-S, and ViT-B share the same architecture, except they differ in the size of the embeddings and the corresponding number of attention heads they employ, having 3, 6, and 12 heads, respectively. In essence, this can be expressed as ViTr C ViT's C ViTR, see Table I
Research question In this paper, we address the following question: Can we train a universal ViT model with H attention heads and embedding dimension F,, such that by increasing the embedded dimension from e1 to e2, where e1 < <2 S E, and its corresponding number of heads from /1j to h2, where h1 h2 < H, the model's accuracy gracefully improves?
Approach In this paper, we propose HydraViT, stochastic training ap- proach that extracts subsets of embeddings and their corresponding heads within MHA] across a universal ViT architecture and jointly trains them. Specifically, during training, we utilize uniform distribution to pick a
MLP
value k. where S H. Subsequently, we extract the embedded dimension ([O k x H Dim)), where I Dim is the size of each head, and its corresponding first & heads (i0 : k)) and only include these in both the backpropagation and forward paths of the training process. To enable the extraction of such subnetworks, we reimplement all components of the ViT including MHA Multilayer Perceptron (MLP), and Normalization Layer (NORM) see Fig. 2 By using this stochastic approach, the heads will be stacked based on their importance, such that the first heads capture the most significant features and the last heads the least significant ones from the input image.
Norm
Silnee Selection
Multi Head
Attention
Norm
After the training phase completed, during inference, HydraViT can dy- namically select the number of heads based on the hardware demands. For example, if only p% of the hardware is available, HydraViT extracts subnet- work with the embedded size of [p X H X HeadDim and the first [p H heads and runs the inference. This flexibility particularly advantageous scenarios such as processing a sequence of input images, like a video stream, where latency is critical, especially on constrained devices such as mobile devices. In such environments, where various tasks running simultancously, and hardware availability dynamically fluctuates, or we need to meet
Figure 2: Architec- ture of HydraViT

---

H1
H2
H3
H4
Hs
H
important
The least important head
Embedding values ordered based the imporiance
subnctwork with hoads
In this figure, we illustrate an example how we extract subnetwork with heads in MHA with total number of heads. In HydraViT, with the stochastic dropout training, we order the attention heads MHA and consequently their corresponding embedding vectors based on their importance.
Figure 3:
deadline, the ability adapt the model's 's configuration without loading model offers significant benefits.
Contributions:
1. We introduce HydraViT, stochastic training method that extracts and jointly trains subnet- works inside the standard ViT architecture for scalable inference.
2. In a standard ViT architecture with H attention heads, HydraViT can induce submodels within universal model.
3. HydraViT outperforms its scalable baselines with up p.p. more accuracy at the same throughput and performance comparable to the respective standard models Dei'T-tiny, DeiT- small, and DeiT-hase, see Figure for details.
2 Related Work
The original Vision Transformer (ViT) Dosovitskiy al. 2021 has become the default architecture of Transformer- based vision models. While many works improve upon the original implementation
by changing the architecture training process (Liu al.) 2022 Touvron al. 2022 Wang et al.)
20216), none these works yield scalable architecture and need multiple separate sets of weights be able to deploy an efficient model on devices with various constraints.
For Convolutional Neural Networks (CNNs) Fang et al. (2018) create a scalable network by pruning unimportant filters and then repeatedly freezing the entire model, adding filters, and fine-tuning Thereby, they achieve a network that can be run with a flexible number of filters. Yu et al. (2018) achieve the same, but instead of freezing, they train network for different layer widths at once. For Transformers, Chavan et al. 2022) use sparsity to efficiently search for subnetwork but then require fine-tuning for every extracted subnetwork acquire good accuracy.
Beyer al. (2023) introduce small change in the training process by feeding differently sized patches to the network. Thereby, they can reduce or increase the number of patches, affecting the speed and accuracy during inference. Other works use importance each patch to prune the least
important patches during inference achieve dynamic VIT
Yin et al. 2022: Rao al. 2021 Tang
ctt al. 2022 Wang et al. a)
Matryoshka Representation Learning (Kusupati et al. 2022) and Ordered Representations with Nested Dropout (Kippel al. 2014 Hojjat al. (202)) are techniques make the embedding dimension of Transformers flexible, I.e., create Transtormer, which can also run partially. Kudugunta et al. (2023) use Matryoshka Learning to make the hidden layer of the MLP each Transformer block flexible.
Hou et al. (2020) change the hidden layer of the MLP
and and the MHA but still use the original dimension between Transformer blocks and also between MHA and MLP Salehi et al. (2023) make the entire embedding dimension in Transformer block flexible. However, they rely on a few non-flexible blocks followed by router that determines the embedding dimension for the flexible blocks, which adds complexity and hinders the ability to choose with which network width to run.

---

apatchie
embed_dim
hidden dim
embed_cde dlim
embed dim
embed dlim
#ipatchee
spatchnes
A2
X
W
2
batch_sth
MLP
(b) NORM
Figure 4: An illustration of subnetwork extraction within MLP and NORM layers, introduced in HydraViT. Fig. Aa demonstrates how HydraViT slices activations, denoted as Ai and Ag. along with their respective weight matrices, denoted as W1 and W2, based on the number of utilized heads. Also, Fig. [4]]shows how HydraViT applies normalization on the activation corresponding to the used heads. For simplicity, only subnetworks with 3, 6, and 12 heads, corresponding to ViT-Ti, ViT-S, and ViT-B respectively, are presented.
Valipour et al. (0233) propose SortedNet that trains networks to be flexible in depth and width. However, they mainly focus on evaluating with CNNs on CIFARIO (Krizhevsky et al.) 2009) and Transformers on Natural Language Processing (NLP) tasks in contrast to us. Additionally, they keep the number of heads MH| at 12, whereas we show that reducing the number of heads coupled the embedding dimension, i.e., the first 64 values of the embedding always belong to the first head in MHA the second 64 values belong to the second head, and so on, removes inconsistencies in the scaling of the MHA and improves performance.
Motivated by these previous works, in HydraViT, we propose a flexible ViT in which we, unlike previous works, adjust every single layer, and except for one initial training run, there is no further fine-tuning required. Additionally, we show that reducing the number of heads coupled to the embedding dimension, weighted subnetwork sampling distribution, and adding separate classifier heads improves the performance of subnetworks.
HydraViT
In this section, we introduce HydraViT, which builds upon the archiitetttu We start by detailing how general ViTs function. Next, we explain how HydraViT can extract subnetworks within the MHA NORM and MLPllayers. Finally, we describe the stochastic training regime used in HydraViT to simullaneously train universal VTT architecture and all of its subnetworks.
Vision Transformer HydraViT is based on the ViT architecture (Dosovisskiy et al. 2021). We start by taking the input image a and breaking it down into P patches. Each patch is then embedded into a vector of size E' using patch embedding, denoted as E². Positional encoding is subsequently applied to the embeddings. Following these preprocessing steps, it passes the embeddings through blocks consisting of MHA with H heads denoted as A NORM layer denoted as NP, MMMP denoted as M xxxMxE to predict the class of the input image x, where M the dimension of the hidden layer of the MLP With the model parameters 0, we can formulate this architecture as follows:
x
Vo(x; MXMxE NP)
(1)
ViTraVIT HydraViT is able induce any subnetwork with & S H. heads within the standard architecture of ViT To do so, HydraViT extracts the first heads denoted as (00) and the embed- dings corresponding to these heads in MHA and NORM layers. Additionally, it extracts the initial E ki] neurons from the first and last layers of the MLP and the first !
A k) neurons from the
hidden layer of Therefore, we can formulate the subnetwrrk extracted from [q as follows:
Vo, (z; gl0:( xk)], (A:AA, (0:A), (0 ( xA)) (0:( # xA) (0: & xk)), 1/10: # x A [1,2,.... H} (2)
Figure 4 illustrates how HydraViT extracts subnetworks within NORM and ML.P layers. For simplicity, we demonstrate subnetworks with 3, 6, and 12 heads, representing configurations for

---

Algorithm 1: Stochastic dropout training
MHA
Data: HydraViT: Vo.:
MHA
Number of batches: Nbathh-
Subnetwwrkk Samplingg ling
MHA
Number of the heads of the universal model: H.
Uniform distribution: U. for 1 < E; S Nepoch do
MHA
MHA
for 1 S b; I Nbotch do
MHA
V sample a subnetwork #1
Vg kmea(())
Voa k E 2 3:;; calculate single-objective loss #1
MHA
1440
(Ve, (
After Trainin::
MHA
Back-propagation through subnetwork VaL:
Sorted baved ce the importance
end end
Figure Stochastic tail-drop training.
ViT-Ti, VIT-S, and VIT-B, respectively. Additionally, in Figure 3 we present an example of how HydraViT extracts subset of heads and their corresponding embeddings in [MH]] layers. By designing HydraViT this way, we can deploy only a subnetwork, e.g., HydraViT with 6 heads, and still have the option runtime with all weights, which is necessary for deployment on more constrained devices.
with even fewer heads. It is not necessary deploy HydraViT
Stochastic dropout training Ideally, to achieve a truly scalable model, we need to extract all the possible subnetworks, calculate their loss, sum them up, and minimize it. This yields the following multi-objective optimization problem:
N
min [01-.0m]
C(Vo, (2:). 9s)
where N is the number of samples, x; is the imput and y is the ground truth. However, optimizing this multi-objective loss function has complexity of O(N x H) and requires at least H times more RAM compared to a single-objective loss function to store the gradient graphs, a demand that exceeds the capacity of current GPU. To address this issue, we suggest employing stochastic training: On each iteration, instead of extracting all of the H possible subnetworks and optimizing multi-objective loss function, we sample value k E 21, HH based on a uniform discrete probability distribution (k)). Then we extract its respective subnetwork Ve,, and minimize only this loss function, Alg. [ This approach decreases the complexity of 3 O(N). In this training regime, the first parts of embeddings and their corresponding attention heads become more involved in the training process, while the later parts are less engaged. After training. due to this asymmetric training. which can also be seen as an order-aware biased version of dropout, the embedding values and their respective attention heads are ordered based on importance, see Fig. 5 Note that despite the similarity to dropout we do not need scaling during training as our training and inference phases are identical. Thereby, we can simplify the Eq. Fas follows:
N
k -~ u(k); min
(VVo
Separate classifiers We implement a mechanism to train separate classifier heads for each subnet- work. This adds few parameters to the model, but only during training or when running the model in dynamic mode, i.e., having the ability to freely choose for each input with how many heads to run the model. The advantage that we do not need to find a shared classifier that can deal with the different amounts of features each subnetwork provides. However, if we fix the number of epochs, each classifier gets fewer gradient updates than the shared one, which why we only use this when training HydraViT with 3 subnetworks.
5

---

Table The accuracy HydraViT with our different design choices. "3 Heads" corresponds to subnetwork that has the same architecture as DeiT-tiny, "6 Heads" corresponds to DeiT-smalle and "12 Heads" corresponds DeiT-base.
Weighted Sampling?
Separate Classifier?
Acc [ % 3 Heads
Acc Heads
Acc [%]
Epochs
12 Heads
X
300
72.56
79.35
80.63
400
73.16
79.63
80.90
X
500
73.54
80.09
81.30
300
72.02
79.35
80.98
V
400
72.45
79.85
81.49
500
72.50
79.89
81.63
X
300
72.78
79.44
80.52
X
400
73.24
79.88
81.13
X
500
73.42
80.12
81.13
300
72.13
79.45
81.18
400
72.46
79.93
81.58
500
72.65
80.08
81.77
DeiT-tiny/smal/hase
72.2
79.9
81.8
Subnetwork sampling function When trying to train single set of weights containing multiple subnetworks, we expect an accuracy drop compared each subnetwork had its own of weights. While we mention that we use uniform discrete probability distribution to sample subnetworks, we can also use weighted distribution function. With weighted subnetwork sampling, we can guide the model to focus certain submodels more than others. This is useful in deployment scenario in which we have many devices with similar resources and want maximize accuracy for them while maintaining good accuracy for other devices with different resources.
4
Evaluation
In this section, we evaluate the performance of HydraViT and compare it to the baselines introduced Sec. 2 We assess all experiments and baselines on ImageNet- 1K (Deng et al. 22009) at resolution of 224 x 224. We implement top of timm (Wightman) 2019) and train according to the procedure of Touvron et al. (202I) but without knowledge distillation. We use an NVIDIA A A100 80GB PCIe to measure throughput. For RAM, measure the model and forward pass usage with batch size of 1. We also calculate GMACs with a batch size of I, i.e. the GMACs needed to classify a single image.
For the experiments, we used an internal GPU cluster, and each epoch took around 15 minutes. During prototyping, we estimate that we performed an additional 50 runs with 300 epochs.
First, we show that we can attain one set of weights that achieves very similar results as the three separate DeiT models DeiT-tiny, DeiT-small, and DeiT-base (Touvron al.) (02T). Then, we look how our design choices, i.c., changing the number of heads coupled to the embedding dimension, weighted subnetwork sampling, and adding separate classifiers for each subnetwork, impact the accuracy. Afterward, we compare HydraViT to the following three baselines:
MatFormer Kudugunta et al. 2023 focus only on the hidden layer of the MLP achieve flexible Transformer and do not change the heads in MHA or the dimension intermediate embeddings.
DynaBERT Hou et (200)) adjust the heads MLP and, as a result, make both flexible. However, the intermediate embedding dimension
addition to the dimension of MHA
is the same the original one in between each Transformer block and between MHA and
MLP which results in more parameters and MACs.
SortedNet Valipour et al. (202)) change every single embedding, including the ones between
MHA and ML. and between Transformer blocks. However, they keep the number of heads in MHA fixed, resulting less information per head and introducing inconsistencies the scaling of the heads in MHA

---

RAM vs. Accuracy
82
30
3H
20
5H
80
I-SNE Component
10-
7H
BH
Accuracy
2H
76
1011
DeiT
DynABERT
74
HydrakiT MytraaMT1800))
-10
1211
MotFomer
-20
100
200
300
400
500
RAM [MB]
-20
-10
1.0
1:5E Companent
Performance comparison of HydraViT and baselines in terms of RAM us- age. Similar to Figure [I] HydraViT achieves the best accuracy per RAM usage.
Figure 6:
Figure 7: t-SNE representation of the last layer of HydraViT with different numbers of heads. As we can see, having more heads leads to more compact representations.
In contrast, instead of keeping the number of heads fixed, change it coupled to the embedding dimension, such that each head gets the same amount of information as in the original ViT] We also evaluate adding separate classifiers and employing weighted subnetwork sampling during the training. Finally, we perform an attention analysis on our model to showcase the effect of adding heads in MHA
4.1
One set of weights is as good as three: Tiny, Small, and Base at once
For this experiment, train HydraViT for 300, 400, and 500 epochs with a pre-trained DeiT- tiny checkpoint. We show how our design choices, i.e., changing the number of heads coupled the embedding dimension, weighted subnetwork sampling, and adding separate heads for each subnetwork. impact accuracy. Table shows each subnetwork's accuracy for all the combinations of our design choices. Note that subnetworks of HydraViT with heads result in the same architecture as DeiT-tiny. subnetworks with 6 heads result in the same smmal,, and subnetworks with 12 heads result in the same as DeiT-base.
To evaluate weighted subnetwork sampling, we show Table that with 25% weight for training the subnetwork with heads, 30% weight for heads, and 45% weight for 12 heads, we can achieve improvement of 0.3 to nearly 0.6 p.p. for the subnetwork with 12 heads depending on the number of epochs compared to uniform subnetwork sampling. Meanwhile, we get change of -0.2 +0.2 p.p. for the subnetwork with 6 heads and decrease of 0.5 to 1.0 p.p. for the subnetwork with heads compared uniform subnetwork sampling. Therefore, at the cost of an overall accuracy decrease. Keep in mind that removing only one head in the vanilla DeiT-base significantly drops accuracy less than 30%, whereas HydraViT achieves more than 72% at 3 heads and 79% at 6 heads and is therefore more versatile.
can increase accuracy at 12 heads
To evaluate separate classifiers for each subnetwork, show in Table 2 that it helps, in some cases, improve each subnetwork's 's accuracy by up to 0.2 percentage points. But it can also reduce the overall accuracy because each classifier gets fewer gradient updates than shared classifier.
Finally, we can combine weighted subnetwork sampling and separate classifiers to achieve a high 12-head accuracy, reaching up to 81.77% accuracy at 500 epochs while maintaining good accuracy at and 6 heads. We notice that compared to only weighted subnetwork sampling, all the accuracies up to 0.15 p.p. higher. Due to starting with pre-trained DeiT-tiny, the classifier for 3 heads needs fewer gradient updates, and the weighted subnetwork sampling shifts the gradient bias to the larger subnetworks, which leads overall better accuracy, Table [
To summarize, we show that with HydraViT, we can create one set of weights that achieves, on average, the same accuracy as the three separate models DeiT-tiny, Deit-small, and DeiT-base. To attain this one set of weights, we need least 300 fewer training epochs than are necessary train the three different DeiT models. The subnetworks have identical RAM usage, throughput, MACs, and

---

Table Comparison of HydraViT with the baselines MatFormer, DynaBERT, and SortedNet. The table shows for selected subnetworks the RAM usage. MACs, model parameters, throughput and their corresponding accuracy when trained from seratch (when applicable) and from the initial DeiT-tiny checkpoint. Note that DynaBERT relies on Knowledge Distillation in every block, which is why it reaches less than 1% accuracy when trained from scratch.
Mothod
Dim
RAM [MB]
MACs [G]
Params [M]
Throughput [9/s]
from scratch
Acc [%]
from Dell-tiny
Ace [%)
MatFormer
Kadugonta
192
294.9
17.56
9.2
86.5
44.1
2601 4
2251.6
1728.3
$1.89 79:40
80:48
82.04
DinaHEK
384
287.62
17.56
7.45
86.6
44.1
3014.6
1725.7
80.16
.00
Hou et al. 2020
192
177.2
3.43
22.8
4944.5
73,00
SuneddNet
768
508.45
17.56
86.6
1753.0
79.71
80.80
Valipour
169.6
63.87
5898.2
79:94
70.20
1754.1
$1.10
Bydrw/117
384
169.6
4.6
22.1
4603.
78.40
79.28
192
63.87
1.25
5.7
10047.6
67,34
70.56
800 Flpccs
HydraVis
384
768
508.45
169.6
17.56
4.6
86.5
22.1
4603.6
1754.1
79.84
$1.93
80.15
$1.50
63.87
71.67
82.25
HydraVl11
440.
14.82
72.9
1916 1
$1.55
82.22
9-12 Heads
640
376.63
12.31
60.3
2242.0
81.51
82.21
317.8
10.01
49.0
2385.2
81.21
81.92
model parameters compared to the DeiT models. While in this section, we investigated HydraViT with only subnetworks, we evaluate HydraViT with more subnetworks in the next section.
4.2 HydraViT vs. Baselines
For the next experiment, we train HydraViT and the baselines introduced at the beginning of this section for 300 epochs, once from scratch and once with DeiT--iny as initial checkpoint. While all of these baselines reduce the embedding dimension, the difference is they reduce of the model. We choose 10 subnetworks for each model, setting the embedding dimension from 192 to 768 with steps of 64 in between. These steps correspond to having from 3 to 12 attention heads, with steps of in between. While HydraViT supports up 12 subnetworks, choose exclude the two smallest ones, as their accuracy drops too much.
different parts
Table shows how each baseline compares to HydraViT relative to their RAM usage, GMACs, model parameters, and throughput when training from scratch and when starting with pre-trained DeiT-tiny checkpoint. Figure and o the results of all subnetworks when starting with a pre-trained DeiT-tiny checkpoint. Besides HydraViT, only SortedNet can run with less than 150 MB of RAM while achieving on average 0.3 to 0.7 p.p. worse accuracy than HydraViT, The other baselines, which have a more limited range of subnetworks, achieve a better accuracy when running higher embedding dimensions. The limited range, however, has the downside of not having smaller subnetworks for devices with fewer resources. And if we limit HydraViT to a similar range as MatFormer, training on 9 to 12 heads, we show that HydraViT reaches the overall highest accuracy at 82, 22.25% compared to MatFormer's ' 82.04%. We also notice that HydraViT cannot reach the exact same performance as the three DeiT models. This is because training for 10 subnetworks with shared classifier for only 300 epochs has its toll on the overall performance, One option is to train longer, which we demonstrated for HydraViT with 3 subnetworks in Section 4.1 We repeat the same here and train HydraViT for 800 epochs, showing that even with 10 subnetworks, we can still achieve near-similar performance as the three different DeiT models. This is while having another subnetworks with similar accuracy per resource trade-off points in between. One caveat, however, is that when training from scratch, HydraViT struggles to get good accuracy 3 heads. This most likely due to a sampling bias as the subnetworks with one and two heads are not included in training and due to training hyperparameters as they differ when training DeiT-tiny compared to DeiT-base. For detailed results, c.g., each subnetwork for every bascline, See Table 4 in Appendix A
8

---

Heads
4 Hea
Heads
Heads
7 Heads
Heads
Heads
10 Heads
Heads
Heads
Hound
Ralian grey Arorna
English Inahound
whippet
wh whippe
hound
Beagl English Inthound hound
baagie English sesiosnd Wather hound
Istagin Enginn fastound Walker hourdd
pargl Enplish hahound Wicher hound
besgle Engishh fodhoans ford Hainer hound
English fnshound Halie hound
beseit
English fostound hound
English fnshiuing auid Wast hound
Biogle
(a)
3 Heedss
Hieads
leeass
6 Hleads
7 Heads
Hieads
9 lieads
10 lleads
Heads
12 Heaas
gual
mongause
glace
qual
pastridge grauce
nutted grocee
qualt
price
goal
barotdee frosse
auas
Barliage roffed gr Pousc
auas
Partrisse ylauce
dusl
partiage
qual
parosdee de saace
qual
partideer
Hilted Gresce
qual
Heads
Heads
Heads
Heads
Heads
Heads
Heads
10 Heads
Heads
12 Heads
lench
100
anu dal
Leach
haraa aaracodda
poott an
tomm
gusition
lanacauaa
leam
kamac karacacc
goudin
tan aacaaa
posam famacaur
lench
dolatin
fatractiga
Leach
futra Siraccura
team
gusition fahacour oug
tomm
modition
banacauga
Team
Figure 8: Attention relevance maps Cheffer al. 2021 of 3 samples from ImageNet-I 1K for HydraViT with different number of heads. Increasing the number of heads leads to more confident classification and more condensed attention distribution.
In summary, HydraViT achieves, average, better accuracy than its baselines except for MatFormer within limited scalability range. However, we show that training HydraViT on similar scalability range outperforms MatFormer.
4.3 Analyzing HydraViT's inner workings
Fig. 8 displayy the attention relevance map (Chefer et al. 2021 ) of selected subnetworks of HydraViT, allowing us visually investigate how the model's attention shifts when increasing the number of heads. Fig. shows that fewer heads lead to more scattered attention, whereas increasing the number of heads makes the attention maps more compact and focused on the main object. Additionally, adding more heads enhances classification confidence. For instance, in Fig. Sal the model misclassifies the input with heads, but as we add more heads, the classification gradually shifts to the correct label and increases in confidence. We also illustrate the t-SNE visualization of the final layer for different subnetworks, see Fig. 7 The figure shows that subnetworks with more heads exhibit a denser representation, while having fewer heads results in more sparse representation. This indicates that increasing the number of heads enhances focus on the main object, which results in less entropy and, thereby, compact t-SNE representation. It worth noting that the outliers in this figure occur due to the high norm values of the embeddings (Darcet et al. 2024
4.4 Robustness of HydraViT
To show the robustness of HydraViT we also evaluate on different ImageNet variants: ImageNet- v2 (Recht et al. 2019), ImageNet-R (Hendrycks et al. 2020), ImageNet-A (Hendrycks et al. 2019

---

ImageNet-V2
ImageNet-R
47.5
45.0
Accuracy (%6
68
Accuracy [%6
40.0
DeiT
DynaBERT
62
HydrakiT
35.0
MydrohiT:5000
60
MotFomer
32.5
Sortechiet
MydroliT
MatFormer
30.0
10
15
GMACS
GMACs
Figure 9: Performance comparison of HydraViT and baselines in terms of GMACs on ImageNet-v2..
Figure 10: Performance comparison of HydraViT and baselines in terms GMACs ImageNet-R.
ImageNet-Sketer (Wang al., 2019), and ImageNet-Real (Beyer et al. 2020 Russakovsky et 2015). On four of these five ImageNet variants, HydraViT achieves the overall best results. Figure 2 shows this for ImageNet-v2. Figure shows the only ImageNet variant, i.e., ImageNet-R, where HydraViT trained with 9 to 12 heads is not able to achieve the best results. Nevertheless, HydraViT reaches a competitive accuracy on these difficult variants and has on average the best results. See Appendix for full results.
4.5 Limitations
Training complexity HydraViT optimizes 10 loss functions simultaneously, which increases the computational load on the optimization progress. As a result, we require more training iterations to achieve accuracy comparable to that of individually trained models such as DeiT-tiny, DeiT-small, and DeiT-base. However, by training multiple models within a unified framework, HydraViT ultimately requires much less total training time compared to training each of these 10 models for 300 epochs individually. See Appendix C for more details on training complexity.
Evaluation on different hardware Our main focus with HydraViT is on the efficiency and scalability on a single device, rather than the deployment on smaller hardware. However, metries such as GMACs and params, are consistent across different platforms. Additionally, the skeleton of HydraViT identical to DeiT,
others have evaluated the latency and performance metrics of DeiT on smaller devices. For instance, FistViT (Vasu et al. 2023) evaluates DeiT the iPhone 12 Pro. MobileVIT (Mehta and Rastegari, 2023) on the iPhone 12, SPVIT (Kong et al. 2022) on the ZCU102 FPGA and Galaxy S20, and GhostNetV3 (Liu al. 2024h the Huawei Mate 40 Pro. These studies provide insight into the expected performance and latency of HydraViT on different hardware, indirectly supporting claims about HydraViT's adaptability.
Evaluation on other models While HydraViT has been evaluated on DeiT-tiny, DeiT-small, and DeiT-base configurations, which have the same number of layers, we have not yet applied it to larger models like DeiT-large with more layers. We plan to explore this in future works.
5
Conclusion
We introduce HydraViT, a novel approach for achieving scalable ViT architecture. By dynamically stacking attention heads and adjusting embedded dimensions within the MHA]layer during training, HydraViT induces multiple subnetworks within single model. This enables HydraViT to adapt to diverse hardware environments with varying resource constraints while maintaining strong perfor- mance. Our experiments on ImageNet- IK demonstrate that HydraViT achieves significant accuracy improvements compared to baseline approaches, with up to percentage points higher accuracy at the same computational cost and up to 7 percentage points higher accuracy at the same throughput. This makes HydraViT practical solution for real-works deployments where hardware availability diverse changes over time.
10

---

Acknowledgments and Disclosure of Funding
This research has received funding from the Federal Ministry for Digital and Transport under the CAPTN-Förde 5G project grant no. 45FGU139 H and the Federal Ministry for Economic A Affairs and Climate Action under the Marispace-X X project grant no. 68GX21002E. It supported in part through high-performance computing resources available at the Kiel University Computing Centre.
References
Lucas Beyer, Olivier Henaff, Alexander Kolesnikov, Xiaohua Zhai, and Aaron van den Oord. 2020.
Are we done with ImageNet? arXiv preprint arXiv:2002.05709 (2020).
Lucas Beyer, Pavel Izmailov, Alexander Kolesnikov, Mathilde Caron, Simon Kornblith, Xiaohua Zhai, Matthias Minderer, Michael Tschannen, Ibrahim Alabdulmohsin, and Filip Pavetic. 2023.
Flexivit: One model for all patch sizes. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition. 14496-14506.
Arna Chavan, Zhiqiang Shen, Zhuang Liu, Zechun Liu, Kwang-Ting Cheng, and Eric P Xing. 2022.
Vision transformer slimming: Multi-dimennion searching in continuous optimization space. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 4931
4941.
Hila Chefer, Shir Gur, and Lior Wolf. 2021. Transformer interpretability beyond attention visual- ization. Proceedings the IEEE/CVF conference on computer vision and pattern recognition. 782-791.
Timothée Darcet, Maxime Oquab, Julien Mairal, and Piotr Bojanowski. 2024. Vision Transformers Need Registers. In The Twelfth International Conference on Learning Representations. https
//openreview.nette wttpt//foooomdddddnnn......
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. 2009. ImageNet: A large- scale hierarchical image database. In 2009 JEEE Conference on Computer Vision and Pattern Recognition. 248 255. https: org/ 10. 1109/ CVPR. 20..2000992202 5206848
Alexey Dosovitskiy. Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn. Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. 2021. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. In International Conference on Learning Representations. https://openrev//// net/ Forum?idd-YccbaNTTYy
Biyi Fang, Xiao Zeng, and Mi Zhang. 2018 Nestdnn: Resource-aware multi-tenant on-device deep learning for continuous mobile vision. In Proceedings of the 24th Annual International Conference on Mobile Computing and Networking. 115-127.
Dan Hendrycks, Steven Basart, Norman Mu, Saurav Kadavath, Frank Wang, Evan Dorundo, Rahul Desai, Tyler Zhu, Samyak Parajuli, Mike Guo, Dawn Song, Jacob Steinhardt, and Justin Gilmer. 2020. The Many Faces of Robustness: A Critical Analysis of Out-of- -iittrr-iiiittoionnn Generalization. arXiv preprint arXiv:2006. 16241 (2020).
Dan Hendrycks, Kevin Zhao, Steven Basart, Jacob Steinhardt, and Dawn Song. 2019. Natural Adversarial Examples. arXiv preprint arXiv:1907.07111 (2019).
Ali Hojjat, Janek Haberer, and Olaf Landsiedel. 2023. ProgDTD: Progressive Learned Image Compression With Double-Tail-Dra Training. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) Workshops. 1130-1139.
Lu Hou, Zhiqi Huang, Lifeng Shang, Xin Jiang, Xiao Chen, and Qun Liu. 2020. Dynabert: Dynamic bert with adaptive width and depth. Advances in Neural Information Processing Systems 33 (2020), 9782-9793.
Zhenglun Kong, Peiyan Dong, Xiaolong Ma, Xin Meng, Wei Niu, Mengshu Sun, Xuan Shen, Geng Yuan, Bin Ren, Hao Tang, et al. 2022. SPVIT: Enabling Faster Vision Transformers via Latency- Aware Soft Token Pruning. In Computer Vision-ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23-27, 2022, Proceedings, Part XI. Springer, 620-640.
11

---

Alex Krizhevsky, Geoffrey Hinton, et al. 2009, Learning multiple layers of features from tiny images. (2009).
Sneha Kudugunta, Aditya Kusupati, Tim Dettmers, Kaifeng Chen, Inderjit Dhillon, Yulia Tsvetkov, Hannaneh Hajishirzi, Sham Kakade, Ali Farhadi, Prateek Jain, et al. 2023. MatFormer: Nested Transformer for Elastic Inference. arXiv preprint arXiv:2310.07707 (2023).
Aditya Kusupati, Gantavya Bhatt, Aniket Rege, Matthew Wallingford, Aditya Sinha, Vivek Ramanu- jan, William Howard-Snyder Kaifeng Chen, Sham Kakade, Prateek Jain, et al. 2022. Matryoshka representation learning. Advances in Neural Information Processing Systems 35 (2022), 30233- 30249.
Zhenhua Liu, Zhiwei Hao, Kai Han, Yehui Tang, and Yunhe Wang. 2024. GhostNetV3: Exploring the Training Strategies for Compact Models. arXiv:2404. 1102 CV] https://arxiv, org/ abs/ 11202
Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, et al. 2022. Swin transformer v2: Scaling up capacity and resolution. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition. 12009- 12019.
Ze Liu, Yutong Lin, Yue Cao. Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. 2021. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEF/CVF international conference on computer vision. 10012-10022.
Sachin Mehta and Mohammad Rastegari. 2023. Separable Self-attention for Mobile Vision Trans- formers. Transactions on Machine Learning Research (2023). https://opennrevier nnet
Forum?id=tB14yBEJJ
Yongming Rao, Wenliang Zhao, Benlin Liu, Jiwen Lu, Jie Zhou, and Cho-Jui Hsieh. 2021. Dy- namicvit: Efficient vision transformers with dynamic token sparsification. Advances in neural information processing systems 34 (2021), 13937-13949.
Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. 2019. Do ImageNet Classifiers Generalize ImageNet?. In International Conference on Machine Learning. 5389- 5400.
Oren Rippel, Michael Gelbart, and Ryan Adams. 2014. Learning ordered representations with nested dropout. International Conference on Machine Learning. PMLR, 1746-1754.
Olga Russakovsky, Jia Deng. Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. 2015. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (ICV) 115, 3 (2015), 211-252. https:://oi. org/10 10....00//11111111111 015 01--10666--------------1------
Mohammadreza Salehi, Sachin Mehta, Aditya Kusupati, Ali Farhadi, and Hannaneh Hajishirzi. 2023. Shares: Efficient transformers through routing with dynamic width sub-networks. arXiv preprint arXiv:2310.122 12126 (2023).
Yehui Tang, Kai Han, Yunhe Wang, Chang Xu, Jianyuan Guo, Chao Xu, and Dacheng Tao. 2022. Patch slimming for efficient vision transformers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 12165-12177
Ilya O Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Un- terthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, et al. 2021. Mlp-mixer: An all-mlp architecture for vision. Advances in neural information processing systems 34 (2021), 24261-24272.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. 2021. Training data-efficient image transformers & distillation through attention. In International conference on machine learning. PMLR, 10347-10357.
Hugo Touvron, Matthieu Cord, and Hervé Jégou. 2022. Deit ii:: Revenge of the vit. In European conference on computer vision. Springer. 516-533.
12

---

Mojtaba Valipour, Mehdi Rezagholizadeh, Hossein Rajabzadeh, Marzich Tahaci, Boxing Chen, and Ali Ghodsi. 2023, Sortednet, place for every network and every network in its place: Towards generalized solution for training many-in-one neural networks. arXiv preprint arXiv:2309.00255 (2023).
Pavan Kumar Anasosalu Vasu, James Gabriel, Jeff Zhu, Oncel Tuzel, and Anurag Ranjan. 2023.
FastViT: A fast hybrid vision transformer using structural reparameterization. In Proceedings of the IEEE/CVF International Conference on Computer Vision. 5785-5795.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. Advances in neural information processing systems 30 (2017).
Haohan Wang. Songwei Ge, Zachary Lipton, and Eric P Xing. 2019. Learning Robust Global Repre- sentations by Penalizing Local Predictive Power. In Advances in Neural Information Processing Systems. 10506-10518.
Wenhai Wang, Enze Xie, Xiang Li, Deng-Ping Fan, Kaitao Song, Ding Liang, Tong Lu, Ping Luo, and Ling Shao. 2021b, lb. Pyramid vision transformer: versatile backbone for dense prediction without convolutions. In Proceedings of the IEEE/CVF international conference computer vision. 368-578.
Yulin Wang, Rui Huang, Shiji Song, Zeyi Huang, and Gao Huang. 2021a. Not all images are worth 16x16 words: Dynamic transformers for efficient image recognition. Advances in neural information processing systems 34 (2021), 11960-11973.
Ross Wightman. 2019.
https://gittuub.coo//iiiii pytorch- -mrorch-image-e iaa-g-modeeseeessse-e models https://aoi.org org/ 10. 528. 201.00
PyTorch Image Models.
Hongxu Yin, Arash Vahdat, Jose Alvarez, Arun Mallya, Jan Kautz, and Pavlo Molchanov. 2022. A- vit: Adaptive tokens for efficient vision transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 10809-10818.
Jiahui Yu, Linjie Yang, Ning Xu, Jianchao Yang. and Thomas Huang. 2018. Slimmable neural networks. arXiv preprint arXiv:1812.08923 (2018).
Weihao Yu, Mi Luo, Pan Zhou, Chenyang Si, Yichen Zhou, Xinchao Wang, Jiashi Feng, and Shuicheng Yan. 2022. Metaformer is actually what you need for vision. In Proceedings of the IEEE/CVF conference on computer vision and patlern recognition. 10819-10829.
13

---

A Detailed Results of Submodels
Table 4: Detailed results of all subnetworks for each Baseline and HydraViT. Note that DynaBERT relies on knowledge distillation in every block, which is why it reaches less than % accuracy when trained from scratch.
RAM [MB]
MACs
Params [M]
Throughput
Acc [ from scratch
Acc [%]
Method
Dim
[G]
[#/
from DeiTT-tiny
768
508.45
17.56
86.6
1728.3
81.89
82.04
704
484.73
16.63
81.8
1822.5
81.89
82.04
640
461.0
15.7
77.1
1860.3
81.89
82.04
576
437.27
14.78
72.4
1960.1
81.87
81.99
512
413.55
13.85
67.7
2020.7
81.79
81.95
Matformer
448
389.81
12.92
63.0
2128.4
81.66
81.94
384
366.08
11.99
58.2
2231.6
81.52
81.80
342.36
53.5
2356.4
81.09
81.71
256
318.63
10.13
48.8
2444.8
80.47
81.25
192
294.9
9.2
44.1
2601.4
79.40
80.48
768
508.45
17.56
86.6
1725.7
81.30
704
471.65
15.68
79.5
1876.8
81.29
640
434.85
13.88
72.4
1995.7
81.20
576
398.03
12.16
65.3
2213.7
81.22
512
361.23
10.51
58.2
2412.4
81.05
DynaBERT
448
324.43
8.94
$1.2
2709.3
80.71
384
287.62
7.45
44.0
3014.6
80.16
250.81
6.03
37.0
3522.5
78.78
256
214.01
4.69
30.0
4157.1
77.04
192
177.2
3.43
22.8
4944.5
73.00
768
508.45
17.56
86.6
1753.0
79.71
80.80
704
440.19
14.82
72.7
1629.5
79.79
80.70
640
376.63
12.31
60.3
1846.6
79.82
80.74
317.8
10.04
49.0
2318.3
79.69
80.60
512
263.68
7.99
38.8
2466.7
79.28
80.43
SortedNet
448
214.27
6.18
29.9
2612.1
78.88
79.93
384
169.6
4.6
22.1
3874.8
77.79
78.94
320
129.63
3.25
15.4
3886.0
75.85
77.55
256
94.4
2.14
4654.6
72.55
74.92
192
63.87
1,25
5.7
5898.2
66.64
70.20
768
508.45
17.56
86.6
1754.1
80.45
81.10
704
440.19
14.82
72.7
1916.1
80.47
81.08
640
376.63
12.31
60.3
2242.9
80.42
81.03
576
317.8
10.04
49.0
2503.2
80.36
80.99
512
263.68
7.99
38.8
3141.5
80.03
80.80
HydraViT
448
214.27
6.18
29.9
3616.4
79.45
80.35
384
169.6
4.6
22.1
4603.6
78.40
79.28
320
129.63
3.25
15.4
5652.0
76.66
77.80
256
94.4
2.14
7558,2
73.39
75.40
192
63.87
1.25
5.7
10047.6
67.34
70.56
508.45
17.56
1754.1
81.93
$1.60
704
440.19
14.82
72.7
1916.1
81,90
81.57
640
376.63
12.31
60.3
2242.9
81.84
81.63
576
317.8
10.04
49.0
2503.2
81.73
81.48
HydraViT
512
263.68
7.99
38.8
3141.5
81.54
81.36
800 Epochs
448
214.27
6.18
29.9
3616.4
80,98
80.92
384
169.6
4.6
22.1
4603.6
79.84
80.15
320
129.63
3.25
15.4
5652.0
78.07
78.43
256
94.4
2.14
10.0
7558.2
74.83
75.95
192
63.87
1.25
5.7
10047.6
68.78
71.67
14

---

Validation loss of all submodels of HydraViT (800 epochs) on INet-1K
3 heads
10
heads
5 heads
6 heads
8
heads
Loss
8 heads
6
9 heads
heads
4
11 heads
heads
2
100
200
300
400
500
600
700
800
Epoch
Figure 11: Training loss for all subnetworks when training HydraViT for 800 epochs.
Table 5: Average training throughput of HydraViT and DeiT models for one epoch.
Model
Throughput [#/s]
DenT-tiny
1610
DeiT-small
821
DeiT-base
340
HydraViT
1093
B Generalization of Submodels
Figure shows that we do not overfit on any subnetwork even when training for 800 epochs. Instead, HydraViT, stochastic dropout training minimizes the loss more less uniformly across all subnetworks. This is especially true when we have more subnetworks, only one subnetwork is optimized per batch and our loss objective gets more complicated. Therefore HydraViT needs more epochs to reach a specific accuracy with subnetwork than the individually trained subnetwork.
C Training Complexity
Table shows that we need to take model size into account when comparing training epochs. While DeiT-tiny by far the fastest to complete one epoch, HydraViT is actually on average faster train with its 10 subnetworks than DeiT-small and DeiT-base.
D Model Loading Time
Table 6 shows that the loading time, i.e., the time it takes to load model into RAM, is for many applications low enough to enable model switching during runtime.
Table 6: HydraViT ViT loading times, each model was loaded six times.
Model
Heads
Dim
Latency [ms]
ViT
192
74.1 + 1.4
ViT
384
133.3 + 6.1
HydraViT
12
768
138.6 + 5.8
15

---

Table 7: The accuracy of HydraViT when initialized with DeiT-tiny vs DeiT--bse. While the accuracy at the 12 heads is higher with DeiT-base initialization the average accuracy is lower than with DeiT-tiny initialization.
Acc [%]
Ace [%]
Ace [%]
Initialization
Epochs
Heads
Heads
12 Heads
DeiT-tiny
400
73.16
79.63
80.90
DesT-base
400
70.83
79.62
81.84
E Initializing HydraViT with DeiT-base instead of DeiT--iny
In HydraViT, the attention heads are treated as stacks, with each head built on top of the previous ones. The smallest submodel, with 3 heads, is equivalent DeiT-tiny. Initializing with DeiT-tiny positions this 3-head submodel at its optimal starting point and ensures that larger submodels, which also contain these 3 heads, also start from good local optimum. By employing the stochastical dropout training introduced in HydraViT, additional heads (4th, 5th, etc.) are trained iteratively on top of the initial three heads, creating layered stack of attention heads. DeiT-base initialization helps HydraViT at 12 heads but does not yield strong smaller submodels, as Table 7 shows.
F Results on ImageNet variants
Figure 12 shows all results of HydraViT and baselines in terms of GMACs and throughput on all ImageNet variants. Note that Figure 9 and Figure 10 contain the results for accuracy per GMACs on ImageNet-v2 and ImageNet-R.
16

---

ImageNet-ReaL
ImageNet-A
88
30
B
86
25
Accuracy [%)
Accuracy [%)
84
82
DaiT
15
DeiT
DynaBERT
DynaBERT
80
HydrsVIT
HydraviT
HycraWI18000)
Hydravim8800e
Hycre/i1/9-12 heads]
HydraVimT 12 hesos)
78
former
MatFormer
SottacNet
SoriedNiat
5
0
5
10
15
0
5
10
15
GMACs
GMACs
ImageNet-Skee
ImageNet-A
30
35.0
8*
DeiT
DynaBERT
32,5
25
HydraviT
Hydrovimt@0ee
Accuracy [%6
30.0
Accuracy [%)
HydraViTTO heaes)
Mathormer
20
SortedNie
25.0
DHIT
15
22.5
HydroWIT
HydraNT(8000)
20.0
Hydr heac))
Maihurme
SortedNet
17.5
5
15
2000
4000
6000
8000
10000
GMACS
Throughput [#/s]
ImageNet-R
88
ImageNet-Real
47.5
DeiT
DynABERT
DynaBERT
45.0
HydraviT
86
Hydravif
Hydrak11/900e)
MydraviT(B00e
Accuracy [%)
42.5
Nathormer SortedNet
HydhowiT(9111 heacl)
Accuracy [%)
84
HatFormer
HydraviTTg heads)
SortedNiet
40.0
82
35.0
80
32.5
30.0
2000
4000
6000
8000
10000
2000
4000
6000
8000
10000
Throughput [#/s]
Throughput [#/s]
ImageNet- Sketch
ImageNet-V2
35.0
DynoBERT
DefT
8
DeTT
DynaBERT
32.5
84
HydraMI
70
HydraviT
HydravMT/800e
Accuracy [%6]
30.0
Marformer
Hydhab twach)
Accuracy [%)]
Matformer
HydroviT 12 hosos)
27.5
SortodNex
66
Sortedhtee
25.0
64
22.5
20.0
60
17.5
58
2000
4000
6000
8000
10000
2000
4000
6000
8000
10000
Throughput [#/s]
Throughput [#/s]
Figure 12: Full results of HydraViT and baselines in terms of GMACs and throughput on ImageNet variants.
17

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We claim that we are introducing scalable VIT, inducing multiple subnetworks into one model, which is what we design and then verify in the evaluation.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We mention limitations throughout the evaluation when discussing results and also include a separate section, see Section 4.5
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). ). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
18

---

Answer:
[NA]
Justification:
We include no proofs and only provided a formal definition for the problem.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We disclose all the information needed to reproduce the main experimental results. We submit the code to verify the basic results as supplementary material and also open-source at wottttt:://// gigggithh..... com ds www.eeeeeeeelallee
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
19

---

reproducible, they Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: We open-source the code at reproducing the results.
for
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/
Cooee//dddeeebmimiia ssi onPolicy) for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The main experimental settings in the paper, while provide the code for detailed experimental settings and details.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: Error bars reported because it would be computationally expensive.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
20

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We list the compute resources in Section
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make it into the paper).
9. Code Of Ethics
Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Question:
Answer: [Yes]
Justification: Our research does not involve human subjects or participants, and we adhere to licenses to prevent data-related concerns.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No. they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: HydraViT improves scalability of transformers, which makes deploying them devices with various constraints easier. We do not believe there is societal impact of that.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
21

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied partieular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: We used the ImageNet-I IK dataset in the computer vision domain and believe there no risk for misuse the trained models.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We build on top of the pytorch-image-monn repository, properly credit them, and follow their license. We also include a citation for the repository and the datasets used.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
22

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/dataseta
latasets.papersstt
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Question:
Answer: [Yes]
Justification: We introduce new model and training scripts to facilitate the training, as described in Section 3 We include everything necessary the code release, with respective documentation and licensing.
Guidelines:
The answer NA means that the paper does release assets.
Researchers should communicate the details the dataset/code/mooell as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsoureing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
23

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
24


Reviews:
[{"review_id": "mwGBEOWQAw", "reviewer": "Reviewer_2Apv", "summary": "The paper proposes a new training scheme for ViTs to enable creating models with varying sizes (in the terms of the number of attention heads in multi-head attention layers). Full version of the model can be then used when more computational resources are available, and smaller portions of the model (with less attention heads operating sorted via importance) can be used when less computational power is available. The authors perform experiments on ImageNet-1k and compare the results with DeiT and some techniques to obtain more ‘dynamic’ models.", "questions": "- “We also show that training HydraViT for more epochs can further improve accuracy.” - isn’t it caused among others by the fact that not all images are seen by all attention maps during each epoch (as a result of subnetwork sampling)?\n- What possible impacts on generalization of the minimum and larger subnetworks can have the proposed sampling technique? The first part is always used in the subnetwork sampling procedure and the last attention heads are trained the least often.These aspects should be discussed more thoroughly.\n- Tab. 3: What is the reason for the large difference between a network trained from scratch and initialized with DeiT-T for the smallest model? They should be equivalents, as the first 3 attention maps are used in each training step.\n- Why are the results for HydraViT trained from scratch not given in Tab. 3?\n- How much time does it take to rebuild the model? Is it suitable for the example application mentioned in the paper - video streams?", "limitations": "The limitations could be better discussed. While the authors state in the checklist that \"while we did not include a separate section, we do discuss limitations for every result.\", it would be nice to squeeze it somewhere in the last section to make it more visible and transparent. I mention some example limitations below:\n- E.g. not having the error bars and comparing the performance of the models (Especially those trained from scratch) based on single runs can be perceived as a limitation. Although the authors state in the checklist (which is a plus) that “Error bars are not reported because it would be too computationally expensive”, it would be nice to state it as a limitation.  \n- Also, the authors do not discuss what possible impacts on learning/generalization/overfitting of the minimum subnetworks/maximum subnetworks can have increasing the number of epochs significantly. \n- The authors do not discuss the limitations of applying the dynamic model rebuilding in practice.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The idea is interesting, simple, straight-forward in interpretation and I believe that it could be useful in some practical applications. \n- It is a nice direction to treat deep neural networks in a more ‘modular’ manner.\n- The methodology is clearly described, the authors provide many helpful visualizations.", "weaknesses": "- The paper puts a lot of stress on suitability of different subnetworks in their approach and its adaptability for different hardware platforms and mobile devices (e.g. “HydraViT achieves adaptability across a wide spectrum of hardware environments while maintaining performance.”), however they do not test their approach on different hardware platforms (they test their approach on NVIDIA A100 80GB, for which the hard constraints are not valid at all).\n- The authors do not focus on the memory aspect in the part “during inference, HydraViT can dynamically select the number of heads based on the hardware demands” - while in terms of computation, dynamic reduction of the size can help, I assume that the parameters of the ‘dropped’ heads are still stored on a given hardware platform (therefore, take the disk space). Also, what about the delays caused by dynamic changes of the model architecture in the given example (video streams)? It would be nice to provide some numbers on how much time it takes to rebuild the model. \n- Although the authors state in the checklist that “Error bars are not reported because it would be too computationally expensive”, it should be stated as a limitation of the experiments, because it decreases the chance to fairly compare the algorithms. At least some smaller portions of the results could be given with error bars to support the findings and increase the soundness.\n- The authors do not focus on the fact that although DeiT-T is an equivalent of 5.7M parameter HydraViT (which they state), when trained from scratch, HydraViT obtains much lower accuracy (more than 3 p.p. less) compared to the case, in which DeiT’s weights are used as a base. Also, it is difficult to fairly compare some of the results as the setups are not equivalents.\n- The authors motivate their work by stating that their approach is the answer to building models with different sizes for the same model family “Despite not having a significant accuracy difference, each of these models needs to be individually trained”. In fact, smaller and larger models do exhibit large accuracy differences, which can be even seen in Tab. 2 in this paper (DeiT-T achieves 72.2 accuracy compared to 81.8 for DeiT base).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer 2Apv for their remarks and their thoughtful review. We answer the questions below:\n\n> Q1. The paper emphasizes the adaptability of different subnetworks for various hardwares but only tests on an A100, where constraints are not applicable.\n\nA1. We acknowledge this oversight and will revise the text to manage expectations more appropriately. Our main focus with HydraViT is on the efficiency and scalability on the same device, rather than explicitly fitting on smaller hardware. However, the metrics such as GMACs and params, are consistent across different platforms. Additionally, the skeleton of HydraViT is identical to DeiT, and other researchers have evaluated the latency and performance metrics of DeiT on various devices. For instance, FastViT [1] evaluates DeiT on the iPhone 12 Pro, MobileViT [2] on the iPhone 12, SPViT [3] on the ZCU102 FPGA and Galaxy S20, and GhostNetV3 [4] on the Huawei Mate 40 Pro. These studies provide insight into the expected performance and latency of HydraViT on different hardware, indirectly supporting our claims about HydraViT's adaptability. Please see the general rebuttal (G.4).\n\n> Q2: [....] I assume that the parameters of the ‘dropped’ heads are still stored on a given hardware platform.\n\nA2. Based on hardware demands, only a subset of the network is loaded into RAM for inference.  Alternatively, the entire model can be loaded into RAM with only a portion of the weights utilized according to the specific demands. We will clarify this in the paper. Additionally, as discussed in response to Q3, the loading latency should be sufficiently fast for most use cases. \n\n>Q3: The paper does not address delays caused by dynamic model changes.\n\nA3. We will add a table in the appendix of the paper reporting loading latency on the A100. However, loading latency also depends on PCI latency and other hardware factors, not just the GPU.\n\n| Model| latency |\n|:-:|:-:|\n| HydraViT-12H| 138.6ms ± 5.8 |\n| HydraViT-6H | 133.3ms ± 6.1|\n| HydraViT-3H | 74.1ms ± 1.4|\n\nAs you can see, the loading latency is very low and should be sufficiently fast for most use cases.\n\n>Q4: Lack of error bars\n\nA4: Error bars are often omitted due to computational costs. We provide code, initialization details, and weights to ensure reproducibility. Also, most papers in the field follow this practice. However, for some settings, we conducted multiple training runs and observed no noticeable differences. Please see the general rebuttal (G.4).\n\n> Q5: The authors do not focus on the fact that HydraViT with 5.7M is 3 p.p. worse when trained from scratch compared to using DeiT’s weights as initialization\n\nA5. Thank you for pointing this out. In HydraViT, the attention heads are treated as “stacks”, with each layer of heads built on top of the previous ones. By initializing HydraViT with DeiT-T, we essentially position the submodel with three heads to its “optimum position” in the loss space, i.e. 72.2% acc. Therefore, during training, “HydraViT from scratch” needs to move the accuracy of the submodel with three heads from 0%, while the “HydraViT initialized” starts at 72.2% and can leverage this to improve accuracy further. We will adjust our evaluation and focus on this difference more.\n\n> Q6: The paper claims minimal accuracy differences across model sizes, but there is a gap between smaller and larger models (DeiT-T and DeiT-B).\n\nA6. Thank you for pointing this out. We acknowledge this for DeiT-T and DeiT-B. However, DeiT-S and DeiT-B show only a 2% accuracy difference despite DeiT-B having 4x more parameters. We will clarify this in the manuscript to ensure a clear understanding .\n\n> Q7. The paper states that training HydraViT for more epochs improves accuracy. Could this be due to subnetwork sampling causing not all images to be seen by all attention maps \n\nA7. Yes, this is a valid observation. However, it is important to note that due to subnetwork sampling, the **training computational complexity** (number of parameters to be optimized) of HydraViT with 3-12 heads for one epoch is significantly lower than that of DeiT-B with 12 heads, see the general rebuttal. In addition to the mentioned reason, we believe the main factor is that HydraViT optimizes 10 different loss functions simultaneously, which adds complexity to the optimization progress and as a result needs more training. We will add these observation to the appendix of the paper.\n\n> Q8. The effect of the proposed sampling training on generalization?\n\nA8. We appreciate you bringing this to our attention. This aspect was missing from our evaluation, and we will address it in our revised manuscript. Please see the general rebuttal for more details (G2: Generalization of Submodels).\n\n> Q9. Accuracy difference between HydraViT scratch and HydraViT initialized with DeiT-T in Tab.3?\n\nA9. Please see A5.\n\n> Q10. Why are the results for HydraViT trained from scratch not given in Tab. 3?\n\nA10. Thank you for pointing this out. We initially wanted to focus on showcasing HydraViT's maximum potential by training it to 800 epochs, assuming that from-scratch results would not significantly impact evaluation. Nonetheless,  after receiving the reviews, we initiated the from-scratch training, which takes approximately 10 days (Thursday). Once this process is complete, we will update Table 3 and provide a comment with the new results.\n\n> Q11: Model rebuilding latency?\n\nA11.  Please see A3.\n\n> Limitations\n\nThe comments raised by the reviewer here were already pointed out by them in the weaknesses section, so we will not repeat them. Also please refer to the general rebuttal for more details.\n\n---\n[1]: FastViT: A Fast Hybrid Vision Transformer using Structural Reparameterization, CVPR2023\n\n[2]: Separable Self-attention for Mobile Vision Transformers, TMLR 2023\n\n[3]: SPViT: Enabling Faster Vision Transformers via Latency-aware Soft Token Pruning, ECCV2022\n\n[4]: GhostNetV3: Exploring the Training Strategies for Compact Model", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zunyqZcEIf", "reviewer": "Reviewer_Yk3V", "summary": "ViTs that allow users to dynamically select the amount of RAM consumed or latency at deployment are valuable. This paper proposes HydraViT, a training procedure that enables plain ViTs to perform relatively well when attention heads and embedding dimensions are dropped on demand. During training, HydraViT samples the number of attention heads (H), selects the corresponding embedding dimension (H times 64), extracts the subnetwork, and updates the subnetwork's parameters. HydraViT achieves a better trade-off between compute and performance than prior work.", "questions": "1. Some experiments initialize models from DeiT-tiny, which has dim=128. How are the weights of the full network, i.e., ViT-B, initialized in this case? Why not initialize from DeiT-B?\n2. If one wishes to train 3 models — say ViT-Tiny, ViT-Small, and ViT-Base — do you advise training these three models separately or train a HydraViT with 3 sub-networks of the same size as the independent models?", "limitations": "I don't believe limitations are adequately addressed. Potential limitations include: lack of scaling to larger models and datasets beyond ImageNet-1k.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "1. The paper is well motivated, well presented, and clear. Resource flexibility is often a missing feature of pretrained models and is an important topic of study.\n2. The experimental conditions are strong, i.e., the DeiT procedure is used, which is a commonly used recipe for training ViTs from scratch on ImageNet. DeiT III is a superior recipe, however I don’t expect an improved training recipe will change the conclusions of this paper.\n3. HydraViT significantly outperforms SortedNet, especially at higher throughputs. And it offers far more flexibility than MatFormer.\n4. More heads leading to more compact representations is an interesting finding.", "weaknesses": "1. The only reported results are on the ImageNet-1k validation set. There are no segmentation or object detection experiments, or other ImageNet-1k test sets — which are easy to implement and would be a valuable addition to the paper.\n2. To me, HydraViT is a straightforward combination of SortedNet (which adjusts embedding dimension) and DynaBERT (which adjusts the number of attention heads and MLP size).\n3. Missing comparison to DynaBERT trained from scratch (table 3). Considering DynaBERT is 1 of 3 baseline methods, I expect a direct comparison to it when trained from scratch.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer Yk3V for their remarks and their thoughtful review. We answer the questions below:\n\n> Evaluation on Additional Datasets?\n\nThank you for your suggestion. We have evaluated our model and baselines on 5 ImageNet variants. Please see the general rebuttal (G.1).\n> Comparison with SortedNet and DynaBERT?\n\nWe apologize for any confusion caused by the original explanation. To clarify these distinctions, we will revise our explanation and highlight the following differences. While HydraViT, DynaBERT, and SortedNet all address scalability, our approach differ significantly:\n\n1- DynaBERT is mainly a gradient-based method that determines the importance score of each attention head by evaluating variations in the loss function when a head is removed. 2- Due to its design, DynaBERT requires Knowledge Distillation (KD) in all layers, resulting in significant training overhead. 3- Since DynaBERT uses KD for all blocks, it must maintain the same dimensionality as the original (teacher) model for distillation. This limits its scalability to the number of heads and specific parts of the MLP layers. In contrast, HydraViT does not rely on gradients to find the important attention heads. Instead, it trains the model such that heads are inherently ordered by importance. Additionally, HydraViT does not use KD, allowing scalability across all blocks, including MLPs, CNNs, normalization, MHA and matrix multiplications. Furthermore, DynaBERT focuses on Transformers for NLP tasks, whereas our approach is for ViTs.\n\nSortedNet adjusts the embedding dimension but keeps the number of heads in MHA fixed at 12. This can result in inconsistent behavior, as noted by reviewer “aEuL.” In transformer blocks, each attention head corresponds to a subset of the embedding. For example, in the ViT, each head has a dimension of 64, so embeddings with dimensions [0:64] always correspond to the first head, [64:128] to the second head, and so on. In HydraViT, when we change the embedding dimension, we also distribute the embedding to the heads in a way that ensures [0:64] always assigns to the first head, [64:128] to the second head, and so on. This coupling of heads with the embedding dimension maintains consistency. In contrast, SortedNet, by always having 12 heads, results in varying embedding dimensions for each head when the embedding dimension is changed, leading to inconsistencies such as different scaling and therefore lower accuracy. Additionally, SortedNet primarily evaluates with CNNs on CIFAR-10 and Transformers on NLP tasks.\n\n> Resuls of DynaBERT from scratch is missing?\n\nDynaBERT uses KD for all its blocks. However, for fairness, we trained HydraViT and all baselines without KD. We attempted to train DynaBERT from scratch without KD, but the accuracy was nearly zero. We attribute this to its design, which necessitates guidance like KD or a proper initialization for effective training. We will add this to the table for clarity.\n\n>  Initialization of Weights from DeiT-Tiny?\n\nTo initialize the model with DeiT-T (3 heads), we load its weights into the corresponding positions, i.e. first 3 heads within the HydraViT (12 heads) skeleton and the remaining 9 heads are initialized with random weights. During stochastic training, when a subnetwork with 3 heads is extracted, it benefits from pre-trained weights, maintaining a starting performance equivalent to the pre-trained DeiT-T (72.2% accuracy). \n\n>  Why not initialize from DeiT-B?\n\nThis is a very interesting question. In HydraViT, the attention heads are treated as “stacks”, with each layer of heads built on top of the previous ones. The smallest submodel, with 3 heads, is equivalent to DeiT-T. Initializing with DeiT-T positions this 3-head submodel at its optimal starting point and ensures that larger submodels, which also contain these 3 heads, also start from a good local optimum. By employing the stochastical dropout training introduced in HydraViT, additional heads (4th, 5th, etc.) are trained iteratively on top of the initial three, creating a layered \"stack\" of attention heads. DeiT-B initialization benefits HydraViT with 12 heads but doesn't effectively leverage smaller submodels with fewer heads, as these 12 heads are not always used with smaller submodels. Nevertheless, in the initial steps of our research, we trained HydraViT with 3, 6, and 12 heads using both DeiT-T and DeiT-B for initialization. These were the results after 400 epochs of training (refer to Table 2):\n\n|Initialization| 3H|6H | 12H |\n|:-:|:-:|:-:|:-:|\n| DeiT-B |70.83|79.62 |81.84|\n| DeiT-T |73.01|79.82|80.95|\n\nAs the table shows, although the model initialized with DeiT-B has higher accuracy with 12 heads, the average accuracy of submodels initialized with DeiT-T is higher, which supports our explanation.\n\n> Should one train ViT-T, ViT-S, and ViT-B separately or use HydraViT with three sub-networks?\n\nThis very much depends on your use case. We argue that HydraViT is a competitive alternative that can be used instead of separately trained models. But ultimately with a unified model you will lose out on a bit of accuracy as you cannot specialize for each head configuration as you can with separately trained models. Since our evaluation primarily focuses on accuracy, we recommend HydraViT if accuracy is your main concern. However, there might be side effects that we have not yet investigated. For example, if your use case involves attention maps or embeddings for specific purposes, such as in VLMs, where two models are coupled together, HydraViT's behavior could differ from that of separately trained models. Additionally, HydraViT simplifies deployment by requiring only one model, which reduces memory usage and loading delays compared to using multiple separately trained models.\n\n> Lack of Limitation?\n\nPlease see the general rebuttal (G.4).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nioIi1tBdl", "reviewer": "Reviewer_1s2T", "summary": "Building on the concept of Matryoshka Representation Learning, the authors propose a novel variant of Vision Transformers (ViT) that improves quality-speed trade-offs. Specifically, the proposed ViT model, trains on subsets of embedding dimensions and associated attention heads, which are ordered by importance. In each training iteration a subset of dimensions is trained. The authors demonstrate in a series of qualitative and quantitative experiments the improved performance.", "questions": "As mentioned in the weaknesses section, I am interested how the stochastic training impacts training time and hardware utilization.", "limitations": "The authors do not address limitations of the paper. I do not see key missing limitations.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper addresses the efficiency of Vision Transformers. With the increasing size of current models and the ongoing trend that larger models perform better, it is of great importance to the research community to improve the quality / size trade-off. \n\nI want to particularly highlight the excellent presentation of the paper. The contributions are precisely defined and very clearly communicated and explained. \n\nRelated to the point above, the paper sports a good discussion of related work and clearly defines the relative positioning of the paper at hand to related work.\n\nAnother key strength of the paper is the experimental section. The paper includes both good quantitative and qualitative evaluations.", "weaknesses": "This paper does not have many weaknesses. If I had to pick anything, the qualitative gain over the baselines is relatively limited. \n\nI can imagine that making stochastic training perform efficiently on modern hardware. It would be great to see an extensive discussion about how to efficiently implement such stochastic training.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer 1s2T for their remarks and their thoughtful review. We answer the questions below:\n\n> It would be great to see an extensive discussion about how to efficiently implement such stochastic training\n\n**Rebuttal**: This is a very interesting question. In HydraViT, we have modified all components of the ViT architecture to accept an additional **cropping** input variable. This modification enables us to selectively crop out the weights and gradients for the forward and backward passes. Consequently, if during training the random number generator produces a value of $k$, the forward/backward pass for a batch in HydraViT with $k$ heads is equivalent to forward/backward pass of a batch in a separate model with $k$ heads. Therefore, the average training time for HydraViT with $i$ to $j$ heads, is roughly as follows:\n\\begin{equation}\n\\text{Average Training Time} = \\frac{\\sum_{k=i}^{j} T_k}{j - i + 1}\n\\end{equation}\nwhere $T_k$ ​represents the of training time for a ViT model with $k$ heads. For HydraViT with 3-12 heads, this average training time would be comparable to that of a separate model with 8 heads.\nWe will add this to the appendix of our paper for more clarity.\n\n\n> Limitations: The authors do not address limitations of the paper. I do not see key missing limitations.\n\nWe appreciate your feedback. As mentioned in the checklist, we have discussed limitations throughout the paper in relation to each result. We will incorporate a separate section in the final part of the paper to summarize all of the limitations for more clarity and transparency, please see the general rebuttal.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BGZSWjlNp5", "reviewer": "Reviewer_aEuL", "summary": "This study proposes a training scheme called stochastic dropout training, thereby presenting its resulting models, called HydraViT. From the observation that ViT with a smaller version can be understood as a subset of a larger version, the authors propose to sample a subnetwork of a larger ViT. Although this scheme trains the subnetwork of ViT, during inference, the size of ViT can be flexibly chosen considering hardware specs. Experiments demonstrate the validity of the proposed method.", "questions": "See weaknesses. Although I would like to admit the motivation and novelty of the proposed method, changing demensions of ViT may cause inconsistent behavior of ViT in self-attention or may require scale calibration for an intermediate feature map similar to Dropout. The proposed method may seem plausible at first glance, but technically, we should be careful when changing the dimension of ViT. I would like to know whether the authors addressed these issues.", "limitations": "The authors discussed several limitations as occasion arises.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "I would like to admit good motivation for choosing the subnetwork of ViT for training by slicing its dimension. Further, incorporating order importance can be useful for other purposes during training and inference. The flexible choice of dimension during inference can be useful for several practical scenarios, as described in the manuscript.", "weaknesses": "When we change the size of ViT, such as hidden dim, head size, and embedding dim, can we guarantee consistent behavior of ViT? In other words, if several dimensions of ViT change, ViT might exhibit inconsistent behavior. For example, the self-attention operation is defined as $softmax(\\frac{Q K}{\\sqrt{d_k}})V$ and calibrates its scale using $d_k$, the dimension of the quary and key. I would like to know how the authors incorporated different dimensions into the self-attention layers.\n\nThe concept of training subnetwork might be similar to Dropout. I would say that though Dropout applies stochastic dropping without order-awareness, the schochastic dropout training proposed in this study incorporates the ordering, starting from the first elements. In this regard, the proposed method can be understood as a special case of Dropout. I am not saying that the contribution is weak due to similarity; however, I would like to know whether the authors adopted a $1/p$ scaling similar to Dropout. Specifically, the Dropout scheme drops elements during training and uses all elements during the test phase. To have consistent output from Dropout, we apply $1/p$ scaling during training. For $p=0.5$, we amplify the output of Dropout by twice. By adopting this practice, we obtain a consistent output scale. In consideration of this, the training subnetwork of ViT may require scale calibration similar to Dropout. Similarly, though the size of ViT can be flexibly chosen during inference, this may require scale calibration with respect to the size of ViT to have consistent behavior of ViT.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer aEuL for their remarks and their thoughtful review. We answer the questions below:\n\n\n> When we change the size of ViT, such as hidden dim, head size, and embedding dim, can we guarantee consistent behavior of ViT? In other words, if several dimensions of ViT change, ViT might exhibit inconsistent behavior. For example, the self-attention operation is defined as  𝑠𝑜𝑓𝑡𝑚𝑎𝑥(𝑄𝐾𝑑𝑘)𝑉  and calibrates its scale using  𝑑𝑘 , the dimension of the quary and key. I would like to know how the authors incorporated different dimensions into the self-attention layers.\n\n1- There appears to be a misunderstanding. In Vision Transformers, each attention head corresponds to a **subset** of the embedding, and the scaling factor $d_k$ relates to the embedding dimension assigned to **each head**, which is $ \\frac{\\text{embedding dim}}{\\text{number of heads}}$, and for ViT’s variants this values is always 64. Unlike SortedNet, in the design of HydraViT, we adjust both the number of heads and the corresponding embedding dimensions (the numerator and denominator of the fraction) so we do not have this inconsistency.\n\nIn the following we have attached a table which details the  $d_k$ for each configuration of HydraViT:\n\n|Model|Embedding Dim|#Heads|Dim per Head ( $d_k$) for K,Q and V|\n|:---:|:---:|:---:|:---:|\n|HydraViT-12H|768|12|64|\n|HydraViT-11H|704|11|64|\n|HydraViT-10H|640|10|64|\n|HydraViT-9H|576|9|64|\n|HydraViT-8H|512|8|64|\n|HydraViT-7H|448|7|64|\n|HydraViT-6H|384|6|64|\n|HydraViT-5H|320|5|64|\n|HydraViT-4H|256|4|64|\n|HydraViT-3H|192|3|64|\n\n2- In Figure 7, we illustrated the t-SNE representation of the last layer of HydraViT with different numbers of heads. As we can see, by adding more heads, the sparsity of the output of the last layer becomes more compact and forms a subset of the previous layer's output, demonstrating consistent behavior. This behavior ensures that the model's behavior remains stable when varying the dimensions.\n\n3- Additionally, by closely examining the logits for each class reported in Figure 8, we observe that, for example, in Figure 8.b, which represents a relatively challenging image, the differences in the logits—i.e., the confidence of the classification—gracefully increase. This further indicates embeddings do not contradict each other and have consistent behavior across varying dimensions.\n\n\n> The concept of training subnetwork might be similar to Dropout. I would say that though Dropout applies stochastic dropping without order-awareness, the schochastic dropout training proposed in this study incorporates the ordering, starting from the first elements. In this regard, the proposed method can be understood as a special case of Dropout. I am not saying that the contribution is weak due to similarity\n\nThank you for your comment. Most current dynamic models typically use either pruning [1] methods to remove redundant weights from the model or routing [2] methods to identify the best possible subnetwork for an input. However, these approaches introduce their own complexities, such as additional parameters or GMACs, to find the best parts. Despite these methods performing really well, we were thinking about the following question: Instead of identifying these important parts, can we train the model in a way that **orders** the information based on its importance so we don’t have to look for it for each image? If dropout is a well-established technique for enhancing generalization, can a **biased** version of dropout be used to organize information in a structured manner?\n\nOur results show that, despite the simplicity of the dropout idea, it can achieve a superior performance. To enhance understanding and provide clearer insight into our motivation, we will clarify this design intuition in the revised paper.\n\n\n> I would like to know whether the authors adopted a 1/𝑝 scaling similar to Dropout. Specifically, the Dropout scheme drops elements during training and uses all elements during the test phase. To have consistent output from Dropout, we apply  1/𝑝 scaling during training. For  𝑝=0.5 , we amplify the output of Dropout by twice. By adopting this practice, we obtain a consistent output scale. In consideration of this, the training subnetwork of ViT may require scale calibration similar to Dropout. Similarly, though the size of ViT can be flexibly chosen during inference, this may require scale calibration with respect to the size of ViT to have consistent behavior of ViT.\n\nThis does not apply to HydraViT. In normal Dropout, the dropped neurons during training are used during inference, necessitating 1/p scaling to maintain consistent output scales. In HydraViT, however, the training and inference phases are **identical**, which means *neurons are being dropped during both training and inference*. For example, if $k$ heads are dropped during training, the model is evaluated with the same $k$ dropped heads. Therefore, there is no need for such scaling in HydraViT. We will add this clarification to the design discussion for better clarity.\n\n> Questions:\n\nWe have addressed this in the preceding sections, and we hope our clarifications provide a clearer understanding of HydraViT.\n\n------------\nReferences:\n\n[1]: DynamicViT: Efficient Vision Transformers and CNNs with Dynamic Spatial Sparsification, NeurIPS 2021\n\n[2]: Flextron: many-in-one flexible large language model, ICML2024", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We would like to extend our sincere thanks to all the reviewers for their time and valuable feedback. Your insights help us to improving the clarity and quality of our work. In the following sections, we address the key points raised by the reviewers and provide clarifications and additional details as requested. 

**G1.** HydraViT and Baselines on New Datasets: As suggested we have run new experiments to conduct an extended, comprehensive evaluation of HydraViT and its baselines across five new datasets: ImageNet-Real, ImageNet-V2, ImageNet-R, ImageNet-A, and ImageNet-Sketch. The comparisons are based on GMACs vs. Accuracy and Throughput vs. Accuracy across these datasets. Our findings indicate that, with the exception of HydraViT (9-12 heads) on ImageNet-R, HydraViT consistently outperforms its baselines on these strongly augmented datasets. The detailed results are visualized in the attached PDF to this comment and will be added to the paper for further reference.

**G2.** Generalization of Submodels: Our training logs indicate that the stochastic training method used in HydraViT does not lead to overfitting across subnetworks. Instead, it effectively reduces the loss for all submodels simultaneously. We have included a plot of validation loss for all subnetworks of HydraViT-800e in the attached PDF, which demonstrates that stochastic training minimizes loss uniformly across all networks. This plot and discussion will be added to the paper.

**G3.** Training Overhead of HydraViT:  The computational load for training HydraViT with 12 heads for one epoch is significantly less compared to training DeiT-B (12 heads). Specifically, HydraViT uses a uniform random number generator to select a subnetwork with 3 to 12 heads for each batch, updating “only” those weights. Consequently, the overhead of training HydraViT with a variable number of heads (3-12) is roughly comparable to training DeiT with 8 heads, which is much lower than the overhead for DeiT-B with 12 heads. For a more detailed comparison, please refer to our response to Reviewer "1s2T".

**G4.** Limitations: In response to the suggestions regarding addressing limitations, we will include a dedicated section in the revised paper to discuss the limitations of HydraViT:

- Training and Accuracy: HydraViT optimizes 10 loss functions simultaneously, which increases the computational load on the optimization progress. As a result, it requires more training iterations to achieve accuracy comparable to that of individually trained models such as DeiT-T, DeiT-S, and DeiT-B. However, by training multiple models within a unified framework, HydraViT ultimately requires much less total training time compared to training each of these 10 models for 300 epochs individually (a total of $10 \times 300 = 3000$ epochs).

- Evaluation on Different Models: While HydraViT has been evaluated on DeiT-T, DeiT-S, and DeiT-B configurations, which have the same number of layers, it has not yet been applied to larger models like DeiT-L with more layers. We plan to explore this in the future works.

- Efficiency Across Different Hardware: The main focus of HydraViT is on efficiency on the same device, rather than evaluating HydraViT across different hardware configurations. Although we assessed our model from various perspectives, such as GMACs, RAM usage, number of parameters, and throughput on an A100, we did not explicitly measure latency on heterogeneous devices. However, since HydraViT's skeleton is exactly like the general DeiTs, other existing literature [1,2,3,4] has already evaluated the latency and performance metrics of DeiT on various devices, which indirectly supports our claims about HydraViT's adaptability.

- Training Overhead and Reproducibility: Due to training overhead and as it is not common in the related community, we did not report error bars for our experiments. However, to ensure reproducibility, we have provided the code, weights, and initialization details.

We appreciate the reviewers’ feedback and will update the manuscript of the paper to address these points thoroughly.

[1]: FastViT: A Fast Hybrid Vision Transformer using Structural Reparameterization, CVPR2023

[2]: Separable Self-attention for Mobile Vision Transformers, TMLR 2023

[3]: SPViT: Enabling Faster Vision Transformers via Latency-aware Soft Token Pruning, ECCV2022

[4]: GhostNetV3: Exploring the Training Strategies for Compact Model

Meta-review: The paper received very positive scores. The reviewers mentioned motivation and importance of the task. They particularly highlighted the presentation and discussion. Performance is another aspect noted by the reviewers. There were some concerns initially, which were addressed during the discussion period, leading to increased scores. All in all, this is a strong piece of work and the decision is to accept the paper. Congrats!
Decision: Accept (poster)