Slicing Vision Transformer for Flexibile Inference

Authors: yitian  zhang
huseyin  coskun
xu  ma
huan  wang
ke  ma
stephen ) chen
derek hao hu
yun  fu
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-12 18:12:02

Abstract: Vision Transformers (ViT) is known for its scalability. In this work, we target to scale down a ViT to fit in an environment with dynamic-changing resource constraints. We observe that smaller ViTs are intrinsically the sub-networks of a larger ViT with different widths. Thus, we propose a general framework, named Scala, to enable a single network to represent multiple smaller ViTs with flexible inference capability, which aligns with the inherent design of ViT to vary from widths. Concretely, Scala activates several subnets during training, introduces Isolated Activation to disentangle the smallest sub-network from other subnets, and leverages Scale Coordination to ensure each sub-network receives simplified, steady, and accurate learning objectives. Comprehensive empirical validations on different tasks demonstrate that with only one-shot training, Scala learns slimmable representation without modifying the original ViT structure and matches the performance of Separate Training. Compared with the prior art, Scala achieves an average improvement of 1.6% on ImageNet-1K with fewer parameters.

Main paper: 
 introduction: 
 Vision Transformers (ViTs) [9] are renowned for its scalability and various avenues [41, 5, 7] have been explored to scale up ViT models. In contrast, to tailor ViTs to run on devices with limited resources, some recent progress [36, 35] utilize knowledge distillation [13] to scale down ViT. Particularly, DeiT [29] introduces two smaller variants of DeiT-B: DeiT-Ti and DeiT-S which have been widely used in resource-limited applications. Although these small ViTs exhibit enhanced efficiency, they lack the flexibility to implement customized adjustments that accommodate dynamically changing resource constraints in real-world scenarios, e.g., the computation budget of mobile phones depends on the energy level (low-power mode) and number of running apps. Consequently, the standard Separate Training (ST) protocol trains models with different sizes separately to provide a spectrum of options with diversified performance and computation. ST requires repetitive training procedures to produce multiple model choices, and the challenge is amplified for foundation models [26, 24, 17] . From users' perspective, they are only offered limited model choices, that might not cater to all scenarios.Analyzing the architectures of ViT-Ti/S/B, we observe that these ViTs are the same architecture with the only difference in the number of embedding dimensions (we ignore the difference in the number of heads as it does not impact the overall model size), indicating smaller ViTs are intrinsically the sub-networks of larger model with different widths (see Fig. 1 ). This suggests that a large ViT can be transformed to represent small models by uniformly slicing the weight matrix at each layer. Given a width ratio r, we adjust the size of the network by this single hyperparameter, allowing a single ViT to represent multiple small variants with the weights of those sub-networks shared in a nested nature, e.g., ViT-B (r=0.25) equals ViT-Ti and ViT-B (r=0.5) corresponds to ViT-S. In this manner, we empower ViTs for flexible inference capability, and we aim to slice a ViT within a broad slicing Figure 1 : Illustration of different means to slice the ViT architecture. Irregular slicing [2, 38, 4] results in unconventional structures while uniform slicing [37] aligns with the inherent design of ViT to vary from widths. bound and fine-grained slicing granularity so that the diversity and number of sub-networks can be ensured for higher flexibility. This problem is non-trivial as fully training all the sub-networks within a constrained budget is nearly infeasible. Consequently, it is quite challenging for these subnets to match the performance of separate training.Although various approaches have delved into slicing deep networks for flexible inference, the problem we target to resolve, i.e., uniformly slicing ViTs within a large slicing bound and fine-grained slicing granularity, is intrinsically different from others in three perspectives: (1) slicing strategy: as shown in Fig. 1 (i), the supernet training techniques [2, 38, 4] in NAS usually slice through multiple dimensions with a small slicing bound, resulting in irregularities in model architectures and a minor computational adjustment space. (2) slicing granularity: recent width slicing approaches [14, 18] either slice specific portions of the network or utilize a considerably large slicing granularity, leading to a limited number of models produced. (3) network architecture: US-Net [37] shares a similar vision with us but it has only demonstrated success in the CNN architecture.It is crucial to note the fundamental differences between slicing CNN and ViT: (1) vanilla small ViTs such as ViT-Ti/S/B, are inherently designed to vary based on widths, aligning with our approach. Conversely, many CNNs are structured to vary from depths, like ResNet-18/34 [12] , and slicing them by width brings unconventional architectures. (2) slimmable CNN necessitates calibration [39, 37] for each sub-network pre-inference due to Batch Normalization [15] , unlike slimmable ViTs that can be directly utilized for evaluation. (3) the transformer architecture [33] has wider applications than CNN in this era, e.g., MAE [11] , CLIP [26] , DINOv2 [24] , LLMs [31, 32, 1] . Nevertheless, ViTs have much less image-specific inductive bias than CNN and their slimmable ability remains unclear. As shown in Fig. 2 , we empirically implement US-Net on ViT-S and observe substantial performance gaps at most width ratios compared to ST, indicating that the available solution of uniform slicing does not work well on the transformer architecture.To investigate the underlying causes of this phenomenon, we conduct analyses in Sec. 3 which are briefly summarized in two folds: (1) ViTs display minimal interpolation ability, indicating that the optimization of intermediate subnets falls notably short compared to Separate Training (ST); (2) sustained activation of the smallest sub-network poses a negative effect on other subnets, which affects the overall performance as their weights are shared in a nested nature. To resolve these issues, we propose a general framework, named Scala, to enforce ViTs to learn slimmable representation. Specifically, we propose Isolated Activation to disentangle the representation of the smallest subnetwork from other subnets while still preserving the lower bound performance. Besides, we present Scale Coordination to ensure each subnet receives simplified, steady, and accurate learning objectives. In this manner, the slimmable ViT can be transformed into multiple smaller variants during inference and match the performance of ST.Compared to ST which trains all the subnets individually, Scala reduces the storage and training costs remarkably since the weights of smaller ViTs are shared with the full model and we only need one-shot training without extending the training duration. Further, Scala has a very large slicing bound and fine-grained slicing granularity, enabling diverse sub-network choices during evaluation. In this way, the delivered system can make tailored adjustments that accommodate dynamically changing resource constraints in real-world scenarios, promising the application on edge devices. Compared with the prior art SN-Net [25] which supports flexible inference on ViTs, Scala clearly outperforms it under different computation with fewer parameters. Moreover, Scala matches the performance of ST on various tasks without modifying the network architecture, demonstrating its generalizability and potential to replace ST as a new training paradigm. The contributions are summarized as follows:• Although slicing ViTs exhibits multiple advantages, we provide detailed analysis and practical insights into the slimmable ability between different architectures (Sec.  
related work: 
 Scaling Up ViTs. Like Transfromer [33] in NLP, scalability and performance improvements in ViTs [9] have been a central focus of recent research. Specifically, strategies have been explored to scale the depth of ViT [45, 30] and it is scaled to even larger sizes with almost 2 billion parameters and reaches new state-of-the-art results [41] . Afterward, ViTs have been scaled up to 4 billion [5] and 22 billion [7] parameters with extraordinary performance and enormous costs.Scaling Down ViTs. The advent of ViTs has also sparked interest in scaling down these models.Techniques such as knowledge distillation [13] have been explored to reduce the ViT model size [36, 35] . For example, DeiT [29] presents smaller ViTs with 5M parameters. Additionally, researchers have explored quantization methods [22, 21] to further compress ViTs for deployment on edge devices. Unfortunately, these static models cannot make customized adjustments for resourcechanging environments in real scenarios.Slimmable Neural Network. The derivation of multiple smaller models from a single network has been previously explored but most works focus on the CNN structure. Slimmable Networks [39, 37] and its variants [34, 3, 10] train a shared network which adapts the width to accommodate the resource constraints during inference. Later, this idea is adapted into two-stage NAS methods [2, 38, 4] for supernet training. The supernet is scaled at multiple dimensions with a small computation change, in contrast to our work where we only scale the width dimension with a large scaling bound. SN-Net [25] is a recently proposed method that constructs a supernet with several pre-trained models and inserts linear layers to build dynamic routes for flexible inference. Recently, several of these techniques have been extended to Transformer architecture [14, 18] , while they either scale part of the network or the scaling granularity is large which means they could only deliver very few models in the end. Differing from the previous works, our method is the first work to scale the ViT structure with large scaling bound and small scaling granularity which is intrinsically a more challenging problem. 
revisiting slicing in vision transformer: 
 Due to the excessive costs of constantly activating all the sub-networks during training, the sandwich rule is proposed in US-Net [37] to train the slimmable network at the smallest width, largest width, and 2 random intermediate widths in each iteration so that the performance of the lower bound and upper bound are guaranteed. Although the intermediate sub-networks are optimized less frequently compared to Separate Training (ST), US-Net manages to achieve comparable performance with ST on the CNN architecture. To have a better understanding of the distinction between CNN and ViT, we apply US-Net to MobileNetV2 [28] , a CNN, and DeiT-S [29] , a ViT, but constantly activate four sub-networks with the width ratio of {0.25, 0.5, 0.75, 1.0} at each iteration. Subsequently, we evaluate the pre-trained models at both inbound [0.25, 1.0] and outbound (0, 0.25) unseen width ratios to evaluate their interpolation and extrapolation abilities, respectively. Shown in Fig. 3 , CNN exhibits moderate interpolation and extrapolation capabilities by achieving acceptable performance at previously unobserved widths during training. In stark contrast, ViT fails entirely at unseen widths, suggesting that optimizing larger sub-networks does not directly benefit the performance of smaller ViTs, even though their weights are shared in a nested nature. We analyze the results from the expected training epochs for each sub-network. Let X represent the total number of networks to be delivered, wherein X -2 intermediate sub-networks are included, and according to the sandwich rule, two of these are randomly sampled during each iteration. Formally, the expected training epochs for the intermediate networks ξ can be expressed as:EQUATIONwhere η is the number of training epochs for the full model and it suggests that the optimization of most sub-networks falls notably short compared to ST. As ViT has demonstrated minimal interpolation ability at unseen widths compared to CNN, each sub-network within the slimmable ViT requires optimal utilization of every training iteration to achieve satisfactory performance. Nevertheless, the smallest subnet is constantly activated during training according to the sandwich rule and we hypothesize that the over-emphasis of the smallest sub-network, often exhibits the worse performance, may increase the training difficulty of other subnets as their weights are shared in a nested nature.To validate it, we implement US-Net [37] on DeiT-S without constantly activating the smallest sub-network. Fig. 4 verifies our hypothesis showing an accuracy drop at the smallest subnet but a significant performance improvement at other width ratios.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [NA]• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is 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 a 96% CI, if the hypothesis of Normality of 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 are 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. 
scala: 
 We first introduce the training and inference paradigms of Scala. Then, we describe Isolated Activation which disentangles the smallest subnet from other sub-networks while maintaining the lower bound performance. Further, we present Scale Coordination to ensure each subnet receives simplified, accurate, and steady learning objectives. Without any modification to the architecture, we deliver a general framework Scala which could be easily built on existing methods.Our goal is to build a general framework that makes a ViT F (•) slimmable, i.e., the delivered network can be transformed into different small variants for flexible inference. First, we introduce a hyperparameter r to denote the width ratio of the sub-network F r (•). Based on our analysis in Fig. 3 , ViTs have minimal interpolation ability, which suggests that all subnets have to be individually optimized to achieve decent performance. Following the sandwich rule [37] , we sample the smallest s, largest l (l = 1), and 2 random intermediate width ratios m 1 , m 2 at each iteration during training. The corresponding sub-networks are: F s (•), F l (•), F m1 (•) and F m2 (•). and we accumulate the gradients of those subnets at each iteration. At the inference stage, the network F (•) is evaluated at an arbitrary width ratio that has been optimized during training by adjusting r.Illustrated in Sec. 3, constant activation of the smallest sub-network F s (•) ensures its own accuracy at the cost of other subnets' performance. This is a dilemma as there is a significant accuracy drop of F s (•) if we do not constantly activate it (see Fig. 4 ), otherwise, the performance of other subnetworks are severely limited. To alleviate this issue, we propose Isolated Activation to disentangle the representation of F s (•) from other sub-networks while still constantly activating it. It not only ensures the performance of the lower bound but facilitates the optimization of other subnets as well.Formally, given the learnable weight θ ∈ R Co×Ci×H×W of a random layer in ViT (C o , C i stands for the output, input channel number, H, W represents the height and width of the convolution kernel, H = W = 1 for fully connected layers), the weight of F r (•) where r ̸ = s is selected as:EQUATIONIn contrast, the smallest sub-network F s (•) is activated as:EQUATIONwhere we slice the weights in a reverse direction so that we disentangle the representation of F s (•) from other sub-networks. With this simple but critical design, we not only ensure the performance of F s (•) with constant activation, but also alleviate the negative effects it brings.We present the training strategy of Scala in this section. We follow the setting of DeiT [29] to train the full model F l (•) with knowledge distillation and introduce a distillation token for knowledge transfer between sub-networks. As our goal is to scale down a given network F (•) to multiple smaller variants, we simply choose the pre-trained model itself F • (•) as the external teacher for the full network F l (•) to facilitate training. To optimize the sub-networks at different scales, we present the Scale Coordination training strategy, which is composed of three techniques: Progressive Knowledge Transfer, Stable Sampling, and Noisy Calibration, to ensure that each subnet receives simplified, accurate, and steady learning objectives.Progressive Knowledge Transfer. Given an input image v, the activated sub-network F r (•) produces two predictions:EQUATIONwhere p r cls and p r dis denote the prediction generated by the classification and distillation head, respectively. As we activate multiple sub-networks: F s (•), F l (•), F m1 (•) and F m2 (•) at each iteration during training, our idea is to utilize the predictions of the larger network to facilitate the optimization of smaller subnets.Given the sorted width ratio list R = [s, m 1 , m 2 , l], we utilize the KL divergence [19] loss to progressively distill the knowledge of the larger network into the smaller one:EQUATIONwhere K represents the number of classes, r ′ = R [ς (r) + 1] and ς (r) denotes the index of r in R. Instead of using p l dis as the optimization target for all smaller networks, we ensure each subnet receives simplified learning objective as small subnet have large capacity gap compared to F l (•) (e.g., F l (•) is almost 16 times larger than F s (•) if s = 0.25) and minimizing their KL loss complicates the optimization process and leads to inferior performance. With Progressive Knowledge Transfer, we simplify the optimization objective for small sub-networks by utilizing F m1 (•) and F m2 (•) as the teacher assistants to fill the gap between F l (•) and F s (•) and train them in a one-shot manner.Stable Sampling. As the knowledge is gradually transferred from the larger network to the smaller one, the two intermediate networks serve as the bridge to connect F l (•) and F s (•), as F m2 (•) is the student of F l (•) and F m1 (•) is the teacher of F s (•). Therefore, we need to carefully control the width ratios m 1 and m 2 to prevent the obvious model capacity variation.Concretely, we introduce the scaling granularity ϵ and the number of networks X we can deliver (including the full model) with a single ViT is denoted as:EQUATIONThen, we divide the scaling bound B = (s,l) ϵ into two smaller ones:EQUATIONwhere m = (s+l) 2 and m1 , m2 will be the random integer sampled from the uniform distribution B 1 , B 2 , respectively. Thus, m 1 and m 2 are defined as:EQUATIONand we ensure the model capacity gap between the four networks is stable and secure the learning objective for each subnet is steady.Noise Calibration. Although all the subnets receive guidance from larger networks, a notable issue is that the predictions from the teacher are not always accurate, sometimes even noisy, especially at the early training stage. To avoid the noisy signal dominating the optimization direction, we first calculate the Cross-Entropy loss by:EQUATIONwhere ŷk represents the one-hot label for class k. Then, we calibrate the noise by combining the KL divergence loss and Cross-Entropy loss:EQUATION)where λ is a hyperparameter used to balance the two losses and we empirically let λ = 1 in our implementations. By doing so, we mitigate the negative effects brought by the noisy predictions of the teacher model and ensure each subnet is guided by the accurate learning objectives. 
experiments: 
 In this part, we mainly validate Scala with the plain ViT structure DeiT [29] . First, we analyze of the main property of Scala and the slimmable representation in detail. Further, we compare our method with the state-of-the-art method SN-Net [25] which supports flexible inference on ViT and Separate Training (ST) on a larger network. Moreover, we examine the transferability of Scala on video recognition and its application on dense prediction task, i.e., Semantic Segmentation. Finally, we provide ablations to validate the efficacy of our designs.All the object recognition experiments are carried out on ImageNet-1K [8] . We follow the training recipe of DeiT [29] and conduct the experiments on 4 V100 GPUs. For Scala, we set s = 0.25, l = 1.0, and ϵ = 0.0625 so that we could enable a single ViT to represent 13 different networks (X = 13) with a large slicing bound (i.e., F l (•) is almost 16 times larger than F s (•)). In this part, we conduct experiments over DeiT-S [29] for 100-epoch training to prove the concept.Comparison with scaling baselines. We compare Scala with multiple scaling baselines, including:(1) AutoFormer [4] : we apply this ViT-based supernet training method into our setting to scale Slicing Granularity and Bound. We empirically change the slicing granularity ϵ and the results are shown in Fig. 5 . First, Scala outperforms ST with different ϵ and the advantage at small width ratios is more obvious, which promises its application on edge devices. Moreover, it is shown that less fine-grained granularity ϵ results in better overall performance with the same slicing bound as the expected training epochs ξ for intermediate subnets increase correspondingly. Notably, X = 4 is the theoretical upper bound of Scala as the four networks are constantly activated during training and exhibit clearly better performance than others. We further conduct experiments with different s while fixing l and ϵ to study the effect of the slicing bound. Fig. 6 shows that smaller bounds lead to markedly better performance and it further verifies that slicing through a large bound is intrinsically more difficult, which distinguishes Scala from the supernet training methods [2, 38, 4] in NAS.Application on Hybrid Structures. We experiment Scala on the CNN-ViT hybrid architecture Uniformer-S [20] . As Uniformer contains Batch Normalization (BN) [15] which cannot be directly evaluated after slicing because of normalization shifting [39] , we calibrate the statistics of BN before inference following [37] . Shown in Fig. 7 , Uniformer-S is scaled down to 13 different variants with better performance compared to ST, demonstrating the generalization ability of Scala. However, performing BN calibration at each width ratio requires considerable extra effort. This highlights the benefit of ViT, as Layer Normalization (LN) allows direct evaluation without additional operations.Application on Lightweight Structures. We further validate Scala on lightweight structure Uniformer-XS [20] which integrates the design of token pruning and train these methods for 150 epochs. Shown in Fig. 8 , Scala still matches the performance of ST and exhibits a significant advantage at small width ratios, which promises its application on edge devices with a limited budget. Fast Interpolation of Slimmable Representation. Shown in Fig. 5 , slicing granularity ϵ has an essential impact on the overall performance of Scala and it is time-consuming to repetitively train models with different ϵ from scratch. Here we show that the slimmable representation of certain granularity can be scaled to others with a small amount of training epochs. Specifically, we train the model with the original ϵ for 70 epochs and decrease the value of ϵ in the last 30 epochs to deliver more sub-networks for higher inference flexibility. Fig. 9 shows the results of fast interpolation are similar to those trained from scratch and the newly appeared sub-networks are quickly interpolated to achieve decent performance. We further increase ϵ for sub-networks with higher performance and the phenomenon shown in Fig. 10 is similar to down interpolation. Besides, we observe that the accuracy of abandoned sub-networks gradually decreases but they maintain the performance to a great extent. Slimmable Ability across Architectures. We examine the slimmable ability of different architectures in Tab. 2 by applying Scala on different architectures and evaluating these networks at unseen width ratios to explore the interpolation ability. CNN exhibits very strong interpolation ability as the performance at unseen widths lies in the range of trained width ratios. In contrast, CNN-ViT and ViT suffer from remarkable performance decreases to different extents and ViT achieves almost zero accuracy which further validates that the problem we target to solve, i.e., slicing ViT, is the most challenging one.In this section, we perform training over DeiT-B [29] for 300-epoch training following the standard protocol on ImageNet-1K [8] to compare with the state-of-the-art.Comparisons with state-of-the-art. SN-Net [25] is state-of-the-art work that supports flexible inference on ViT. Specifically, it utilizes several pre-trained models (e.g., DeiT-Ti/S/B) to construct a supernet and inserts additional layers to build dynamic routes for flexible inference. Shown in Tab. 3, we empirically compare Scala with SN-Net over DeiT-B following the standard 300-epoch training protocol [29] . Scala obtains similar performance with SN-Net at large width ratios and clearly outperforms it at small computational budgets. Besides, SN-Net has to preserve the parameters of multiple models and additional layers, while Scala only needs to keep the weights of the full network. When adopting the stronger teacher network [27] as SN-Net does, Scala outperforms SN-Net with an average improvement of 1.6% across all width ratios. width ratios. This substantiates the effectiveness of Scala and the slimmable representation.To assess the transferability of Scala, we initially employ DeiT-B [29] as the backbone for a 300-epoch pre-training on ImageNet-1K. To acquire superior performance, we leverage the foundation model DINOv2-B [24] trained on extensive data to inherit the good behaviors. As depicted in Fig. 11a , Scala consistently outperforms Separate Training (ST) across all width ratios, despite the intermediate sub-networks being trained for approximately 55 epochs. After that, we conduct linear probing on video recognition dataset UCF101 with 8 evenly sampled frames and average their features for the final prediction. For the classification head added on Scala, we make it slimmable to fit the features with various dimensions and follow the same training protocol as in object recognition. In Fig. 11b , two notable observations emerge: (1) Scala consistently outperforms ST across different width ratios on the UCF101 dataset, implying the great transferability of the slimmable representation; (2) Scala retains its slimmable ability when applied to a new task and exhibits promising performance across a wide slicing range (10∼141 GFLOPs), promising its application on other downstream tasks. In previous sections, we have validated the effectiveness of Scala on classification tasks, we further examine whether the slimmable representation could be transferred for dense prediction task like semantic segmentation. We utilize the pre-trained model Uniformer-S [20] drawn from Fig. 7 , which has a hierarchical design and is obtained by 100-epoch training (our results lag behind official results where the backbone is trained for 300 epochs), and equip it with Semantic FPN [16] . To compare with Separate Training (ST), we extract four subnets from Scala (Uniformer-S) and train them separately. Shown in Tab. 5, Scala outperforms ST at all widths which verifies the slimmable representation Table 5 : Evaluation of slimmable representation on dense prediction task Semantic Segmentation over ADE20K [44] dataset. We equipped the pre-trained Uniformer-S [20] from Fig. 7 with Semantic FPN [16] and compare the sub-networks extracted from Scala with Separate Training (ST). benefits the downstream tasks. Note that we do not scale the decoder as it involves extra designs and is out of the scope of this work. However, we show that the slimmable representation can be generalized to semantic segmentation as feature extractors because the feature maps are spatially intact, promising its application as an end-to-end slimmable framework on dense prediction tasks.We conduct ablation to examine the effectiveness of our designs in Tab. 9. First, we build Scala without Isolated Activation so that the smallest sub-network will entangle with others and it shows an obvious performance drop at all width ratios. Then, we remove Progressive Knowledge Transfer (PKT) and pass the knowledge from F l (•) to smaller subnets through classification token following US-Net [37] . It shows much worse performance, especially at small ratios, which proves the strength of PKT as it implicitly introduces some teacher assistants to simplify the optimization objective for small sub-networks. Further, we random sample the width ratios of m 1 and m 2 between (s, l) and compare it with Stable Sampling (SS). The results are slightly inferior to SS which suggests SS is helpful in securing the steady learning objective for each sub-network. Finally, we remove Noise Calibration (NC) from Scala and only use the predictions from larger networks to guide the small subnets. It shows remarkable performance drops at small width ratios, where the noise from the teacher network is most obvious, demonstrating the effectiveness of NC in calibrating the noise and providing accurate signals for sub-networks. 
conclusion and limitations: 
 In this paper, we observed that smaller ViTs are intrinsically the sub-networks of a large ViT with different width ratios. However, scaling ViT is very challenging due to its poor interpolation ability.  
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes]Justification: We have provided details of our computing resources in the appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or 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 or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: Our work has no such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 to 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 we encourage authors to take this into account and make a 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 have cited the papers that created the datasets.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 or dataset.• The authors should state which version of the asset is used and, if possible, include a 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. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: For Separate Training (ST), we follow the exact training strategy of the official DeiT [29] and Uniformer [20] setting. We use random horizontal flipping, random erasing [43] , Mixup [42] , CutMix [40] , and RandAugment [6] for data augmentation. AdamW [23] is utilized as the optimizer with a momentum of 0.9 and a weight decay of 0.05. We set the learning rate to 1e-3 and decay with a cosine shape. The models are trained on 4 V100 and 8 A100 GPUs with a total batch size of 1024. We adopt Exponential Moving Average (EMA) following the official setting.While we utilize the pre-trained ST model (r = 1.00) as the teacher for F l (•) to facilitate training as mentioned in the Scale Coordination section in the main paper, we adopt a larger learning rate (2e-3) and mild data augmentation (reduce the magnitude for RandAugment [6] to 1 and turn off repeated augmentation) because Scale Coordination already regularizes the network training strongly. At every training iteration, we activate four sub-networks separately based on Stable Sampling and accumulate their gradients for backpropagation. The rest hyperparameters are set as the same as those in Separate Training.Tab. 1 shows that ViT is not slimmable if we directly evaluate the vanilla pre-trained model at other widths. Here, we further explore the slimmable ability of vanilla representation and fine-tune the vanilla pre-trained model with Scala. Tab. 11 shows that fine-tuning obtains obviously worse performance at small width ratios compared to training from scratch, which denotes that the vanilla representation is not slimmable and is essentially different from the slimmable representation. As discussed in the main text, the slicing bound has a huge impact on the performance of Scala. Here we further expand the slicing bound from [0.25, 1.00] to [0.125, 1.000]. As shown in Fig. 12 , Scala suffers from an obvious performance drop at r = 0.25 as it is not constantly activated in the new setting. Nevertheless, our method still manages to outperform ST at all width ratios and shows a significant advantage at the smallest ratio r = 0.125. Previous experiments validate that Scala achieves achieves performance comparable to that of Separate Training (ST) on DeiT-B [29] in the 300-epoch training setting, even though the intermediate sub- We validated the effectiveness of Isolated Activation by removing this component and we further conduct more ablation studies on the activation methods where the designs are illustrated in Fig. 13 . As shown in Tab. 9, choice (b) leads to slightly better performance at F m2 (•), but the performance drops at F m1 (•) significantly as it is entangled with F s (•) and the over-emphasize of F s (•) adversely affect its performance. On the other hand, choice (c) results in a similar performance at F m1 (•), but the accuracy decreases significantly at F m2 (•) which further verifies our hypothesis that F s (•) should be isolated to reduce its negative impact on other sub-networks.𝑭 𝒍 (#) 𝑭 𝒔 (#) 𝑭 𝒎 𝟐 (#) 𝑭 𝒎 𝟏 (#) (𝒂) 𝑭 𝒍 (#) 𝑭 𝒔 (#) 𝑭 𝒎 𝟐 (#) 𝑭 𝒎 𝟏 (#) (𝒃) 𝑭 𝒍 (#) 𝑭 𝒔 (#) 𝑭 𝒎 𝟐 (#) 𝑭 𝒎 𝟏 (#) (𝒄)In the main text, we found that ViT has the minimal interpolation ability compared to the CNN structure. This suggests that optimizing larger sub-networks does not directly contribute to the performance improvement of smaller variants, even though their weights are shared in a nested nature. A further question is, whether ViT can still maintain the slimmable ability for the unseen width ratios during training.To verify this point, we respectively fix the width ratios of m 2 , m 1 to 0.8125 and 0.4375 during training, so that only one sub-network is optimized at each range. Fig. 14 shows that the accuracy of unseen sub-networks is very low due to the lack of interpolation ability. Nevertheless, the performance Figure 14 : Verification of Slimmable Ability over DeiT-S [29] on ImageNet-1K. We respectively fix the width ratio of m 2 , m 1 to 0.8125 and 0.4375, and observe the performance at other ratios is not affected.at other width ratios remains similar to the default setting even though their weights are shared with each other. This indicates that the correlation between sub-networks in ViT is weak and further highlights how challenging this problem is. While we have shown that Scala outperforms baseline methods US-Net [37] and Separate Training (ST), we further compare Scala with much stronger baselines, by adding an external teacher to their full network during training. Specifically, we adopt the pre-trained full model from ST as the teacher and conduct knowledge distillation for models with different widths separately. Tab. 10 shows that ST+KD exhibits similar performance at larger width ratios with Scala, despite that Scala clearly outperforms ST+KD at smaller widths, promising its application on edge devices. Although obtaining a better full model, US-Net+KD exhibits worse performance at smaller width ratios because it utilizes the full network as the teacher for all subnets and this phenomenon verifies our motivation of proposing Progressive Knowledge Transfer.Assuming to deliver 13 models in the end, we compare the training time (100 Epoch) of Scala with US-Net [37] , Separate Training on 8 A100 GPUs. The difference between US-Net and Scala is not large as the transformer architecture has been well-optimized on GPU and we do observe a significant time gap between Scala and Separate Training as they have to train 13 models iteratively. Moreover, Scala can be configured to deliver 25 models without an increase in training time as we sample 4 networks at each iteration in all scenarios which further highlights our strengths. A.9 Comparisons with MatFormerMatFormer [18] only slices the FFN block in the transformer architecture so it offers a minor computational adjustment space and we adapt their method on DeiT-S to compare with Scala. Fig. 15 shows that Scala achieves comparable performance with it (better in most cases) when s = 0.5 with a larger adjustment scope. NeurIPS Paper Checklist Guidelines:• The answer NA means that the paper has no 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 (e.g., independence assumptions, noiseless settings, model well-specification, 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 a few datasets or with a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to 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 fairness. • 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 acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.Justification: We have read and follow the rules.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 a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).Justification: There is no societal impact.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 or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness 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 to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to 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 as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 


Decision: After rebuttal and discussions, reviewers generally lean towards acceptance. I do agree with the originally more conservative reviewers that the related work is not discussed enough. For two examples: while MatFormer is discussed, there are several more works from Aditya Kusupati that are highly relevant. Furthermore, FlexiViT tackles all the same issues but along sequence length.

Thus, for the general case of "flexible amount of computation at inference time", too few baselines are compared to. That being said, the presented observations and introduced ideas are interesting and novel enough that, compared to their obvious baselines/ablations, they are worth publishing. Hence, we still recommend acceptance.

Finally, the AC agrees with reviewer rLHb that the title is both much too broad, and a wrong use of the word "scalable" compared to its typical meaning within the community. Not only that, also "representation learning" is not a good fit for this as it usually has a different meaning/focus. Please change the title to something more fitting and commong for this scenario. Some examples could be "Elastic training of Vision Transformers" or "Vision Transformers with Flexible Width" or something to that effect.

Reviews: review: paper_summary: paper_summary: This paper advances an approach for training Vision Transfomers (ViTs) such that at inference time they can be dynamically adjusted to fit different budget constraints with reduced drops of performance. To this end, the authors introduce Scala, a framework that allows a single network to encapsulate and train simultaneously multiple sub-networks of different capacities and widths. The methodological backbone of this work are the Universally slimmable networks (US-Net) [37], originally devised for CNNs. The authors identify and analyze a few flaws of US-Nets: difficulty to generalize to ViTs, small interpolation and extrapolation ability to sub-network size unseen during training, impact of sustained activation of the smallest sub-network that coupled with the sandwich rule for selecting sub-networks during training leads to an over-emphasis on it at the expense of the other sub-networks.
The authors propose two simple strategies towards such a method for ViTs: (i) Isolated activation that separates the smallest sub-network from the other sub-networks; (ii) scale coordination consisting of a set of heuristics to ensure that each sub-network gets simple, accurate and stable learning objectives: (a) progressive knowledge transfer from larger networks to smaller ones in gradual decrease of capacity, (b) stable sampling of intermediate width ratios to avoid large variations in capacities in the sandwich, (c) noise calibration, essentially a composite loss of supervised cross-entropy and distillation from the bigger sub-network.
Scala is evaluated on several settings on the ImageNet-1k dataset with ViT-Ti/S/B, hybrid CNN-ViT architectures, lightweight networks, but also for dense prediction on semantic segmentation and self-supervised pre-training with interesting results. The baselines used here were Separate Training,  Autoformer and US-Net.
strength_weakness: strength_weakness: ### "Scalable" naming
- I think that the framing of the method as _"scalable representation learning"_ is quite confusing as it's not representative for this task, it's not a name used by other related works. Importantly, it can be easily mistaken with most works that use "scalable" for depicting the ability/property of a system (method, architecture) to handle a growing amount of data, parameters, and the potential to accommodate this growth. In other words "scalable" is rather used for depicting scaling up, whereas this work depicts the property of the proposed approach to accommodate sub-networks of different lower sizes/scales from the original.

- maybe other names us in related works would be more appropriate here: slimmable, elastic, modular, flexibile inference, etc.


### Limited baselines and related work
- some relevant related works dealing with tranformer networks are either just briefly mentioned, e.g., Matformer  [18], or not mentioned at all, e.g., SortedNet [a], Early exit [b]

- One of the main baselines, US-Net is originally designed for CNNs and, as the authors mentioned, moving to ViTs is not straightforward. Matformer is criticized for the limited number of models produced, but can be considered in the several experiments with X=4 sub-networks. Matformer and SortedNet could be included in the experimental evaluation


### Scope of experiments
- While the authors considered several settings for computer vision tasks (image classification, segmentation, light architectures), transformer architectures are also encountered in NLP (as mentioned by the authors in L56). In such cases the original models can have much more parameters and elastic inference for lower computational budgets would be of high interest.

- It would be useful to include an experiment from NLP in the style of those from Matformer or SortedNet.

- The biggest architectures used here is a ViT-B (~86M params). Extending experiments to larger modern architectures would be definitely useful and interesting.

### Clarity
- it's not always clear in the text and cost estimations that Scala needs a pre-trained full network as teacher for the distillation. This add some cost in compute and time in the end. Besides it's not clear whether US-Net also needs and uses a pre-trained teacher in the reported results.

- in the intro, the authors mention that they address the issue of minimal interpolation ability of ViTs. Results from Table 2 show that the interpolation abilities of ViTs with Scala are still very low. However the fast interpolation strategy from $\S$A.2 is actually interesting for practical settings even though not fully solving this issue. It might be worth moving up in the main paper.

- the idea of the transferability experiment ($\S$5.4) with DINOv2 is nice. From the description it is not clear whether DINOv2 was used as teacher for the distillation or also as supervised pre-training on ImageNet-1k? Or the pre-training on ImageNet-1K was done in a supervised manner as in previous experiments?

- the ablation experiment from Table 6 is nice. However the presentation with removing one component at once offers only a partial understanding of the contributions of each module. Different configurations with different modules in on/off mode should give a better global understanding.



**References:**

[a] Valipour et al., SortedNet: A Scalable and Generalized Framework for Training Modular Deep Neural Networks, arXiv 2023

[b] Xin et al., Deebert: Dynamic early exiting for accelerating bert inference, ACL 2020
questions: questions: This paper takes an interesting direction of study: how to train ViTs such that a fine-grained elasticity in terms of sub-network sizes is possible at runtime.
The proposed Scala approach is well described, makes sense and achieves good results in several computer vision settings.

I do have a few concerns related to the phrasing of this type of works ("scalable representation learning") which can be confusing, the absence of larger architectures and of recent relevant baselines.
My current rating is mildly positive (rather on the fence though) and I'm looking forward for the rebuttal.

Here are a few questions and suggestions that could be potentially addressed in the rebuttal or in future versions of this work (please note that suggested experiments are not necessarily expected to be conducted for the short rebuttal period):

1. Please clarify the points raised in the clarity section: use of teacher model for Scala and US-Net, implementation of transferability experiment. 

2. Comparison of training cost between Scala (including teacher training), US-Net and Separate Training baselines.

3. Add a discussion of differences and when possible experimental comparison with Matformer and SortedNet baselines on image classification or semantic segmentation.

4. Extension of experiments to NLP architectures and tasks in the style of SortedNed, Matformer
limitations: limitations: The authors addressed some of the limitations in the conclusion section.
paper_summary: paper_summary: The paper introduces Scala, a novel framework designed to effectively scale down Vision Transformers (ViTs) for use in environments with fluctuating resource constraints. The key insight is that smaller ViTs can function as sub-networks within a larger ViT, differing mainly in width. Scala enables a singular network architecture that can emulate multiple smaller ViTs, thereby offering versatile inference capabilities while maintaining the structural principles of ViTs. The framework uniquely incorporates multiple sub-networks during its training phase, utilizes Isolated Activation to differentiate the smallest sub-network, and implements Scale Coordination to streamline the learning objectives for each sub-network, aiming for simplicity, stability, and accuracy. The empirical results across various tasks confirm that Scala can learn scalable representations efficiently with a single training iteration, maintaining the integrity of the original ViT architecture and achieving performance on par with networks trained separately.
strength_weakness: strength_weakness: 1. Recent papers[1,2,3] with "Scalable" usually scale ViT to billion size with large scale datasets like DFN, JFT, and Datacomp. Therefore, I suggest authors should reconsider if the experiments can support "Scalable".


[1] Zhai, X., Kolesnikov, A., Houlsby, N., & Beyer, L. (2022). Scaling vision transformers. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition (pp. 12104-12113).

[2] El-Nouby, A., Klein, M., Zhai, S., Bautista, M. A., Toshev, A., Shankar, V., ... & Joulin, A. (2024). Scalable pre-training of large autoregressive image models. arXiv preprint arXiv:2401.08541.

[3] Dehghani, M., Djolonga, J., Mustafa, B., Padlewski, P., Heek, J., Gilmer, J., ... & Houlsby, N. (2023, July). Scaling vision transformers to 22 billion parameters. In International Conference on Machine Learning (pp. 7480-7512). PMLR.
questions: questions: please refer the weakness.
limitations: limitations: This paper addresses the limitations in conclusion
paper_summary: paper_summary: The paper presents Scala, a novel framework for scalable representation learning developed from US-Net. It identifies the issues of directly applying US-Net to ViTs and proposes solutions including Isolated Activation, Scale Coordination, and Stable Sampling. These innovations enable Scala to output several sub-networks in one-shot learning. Extensive experiments on various network architectures and datasets demonstrate that the sub-networks produced by Scala consistently outperform those generated by separate training, with significantly reduced training time.
strength_weakness: strength_weakness: Originality: The novelty of Scala is somewhat constrained. For instance, Noise Calibration does not show a distinct difference from standard knowledge distillation. Essentially, Scala integrates US-Net with an alternative activation for the smallest subnet and fixed scaling ratios.

Quality: The authors might consider emphasizing results from a more standard 300-epoch ViT training schedule to align with common practices in the field.

Clarity: No further issues.

Significance: The challenge of scaling ViTs with arbitrary ratios remains unresolved.
questions: questions: 1. Regarding the issue shown in Fig. 4, is it always the smallest subnet that causes the issue, or does it occur with subnets having a scaling ratio near 0.25?

2. Can the authors clarify any differences between Noise Calibration and standard knowledge distillation?

3. What would be the impact if the distillation part were discarded and only Cross-Entropy (CE) loss were used for the initial epochs?
limitations: limitations: The novelty and significance of Scala are somewhat limited, as discussed in the weaknesses section. However, the extensive experimental results provide a robust foundation for the claims made in the paper. Overall, the work is well-executed and makes a valuable contribution to the field, justifying a recommendation for weak acceptance.
paper_summary: paper_summary: The paper targets scaling down Vision Transformers (ViT) to fit environments with dynamically changing resource constraints. The authors propose Scala, a framework enabling a single network to represent multiple smaller ViTs with flexible inference capability by activating various subnets during training. Scala introduces Isolated Activation to disentangle the smallest sub-network and uses Scale Coordination to provide stable and accurate learning objectives. Empirical validations on different tasks show that Scala achieves scalable representation with one-shot training, matching the performance of Separate Training without modifying the original ViT structure. Scala demonstrates an average improvement of 1.6% on ImageNet-1K compared to previous methods, using fewer parameters.
strength_weakness: strength_weakness: 1. My major concern is that, the same aim of adapting ViTs to dynamically changing resource constraints, can also be achieved by multi-exit networks, e.g., [*1, *2, *3]. However, the paper does not discuss these highly relevant works or compare with them. Hence, I vote for rejection.
2. The method seems to lack novelty. 'smaller ViTs are intrinsically the sub-networks of a larger ViT with different widths' is not a surprising observation. The key techniques (e.g., Isolated Activation and Knowledge Distillation) are not new (naive or have been widely adopted).


[*1] Huang, Gao, et al. "Multi-Scale Dense Networks for Resource Efficient Image Classification." International Conference on Learning Representations. 2018.

[*2] Wang, Yulin, et al. "Not all images are worth 16x16 words: Dynamic transformers for efficient image recognition." Advances in neural information processing systems 34 (2021): 11960-11973.

[*3] Han, Yizeng, et al. "Dynamic perceiver for efficient visual recognition." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.
questions: questions: Please refer to Weaknesses.
limitations: limitations: The authors have addressed the limitations and potential negative societal impacts of their work.

Comments: title: Thanks for the support
Kind Reminder of Deadline
Thanks for the support
Post-rebuttal comments
Deadline Approaching
comment: I highly appreciate the authors' response. However, I still think the lack of discussions and comparisons with dynamic networks in the original submission is an important limitation. Moreover, although the authors have provided some clarification, Isolated Activation seems like simply an engineering trick, and the proposed knowledge distillation method seems too incremental compared to existing works. These techniques may be insufficient as significant scientific contributions.

I increase my rating to 4. However, I still vote to reject this paper. The major reasons are (1) lack of discussions on important related works, and (2) weak technical novelty.
Dear reviewer, please have a look at the rebuttal and indicate whether it addressed some of your concerns.
Dear authors, thank you for your rebuttal. FYI: after the author discussion phase closed, there is still a bit over a week of discussion between ACs and Reviewers.
Dear Reviewer,

Thank you so much for the constructive comments for us to improve the work and we sincerely appreciate your support.

Thanks again for being with us so far.

Sincerely,  
Authors
Thank you for the detailed response. I have updated my vote to accept.
Dear Reviewer,

Thanks for your efforts so far in our paper review.

Currently, most of the concerns from other Reviewers have been well-addressed and we are eager to know whether our response has resolved your concerns. Concretely, we have added the discussion between Scala and dynamic networks which shows that our work is agnostic with this line of research. Moreover, we have clarified the contribution of our work and further expanded the application scenarios of Scala to foundation models like DINOv2.

Due to the coming discussion deadline, we would like to kindly remind the Reviewer if our response through the whole rebuttal has addressed any of your concerns and helped the Reviewer to reevaluate this work. 

We really appreciate it if you could give us any feedback and your opinions are rather important to us. Thank you very much for your time!

Sincerely,  
Authors
Dear Reviewer,

Thank you very much for your recent feedback. We truly appreciate your support and the time you have invested in reviewing our paper. Yes, we will definitely include the new findings in the main paper and we promise to release the code if this work is accepted. 

We are glad to hear that the reviewer has no concerns and is leaning toward a positive recommendation. If the reviewer finds that the work is improved during the rebuttal, we would be grateful if you could reflect the updated stance in the rating on OpenReview. This would greatly assist us in the final decision-making process.

Thank you once again for your thoughtful consideration and valuable insights.

Sincerely,  
Authors
I would like to thank the authors for the detailed and informative rebuttal. I imagine they invested significant effort into clarifying the concerns from the 4 reviewers and I'm confident they will improve this work.

I've read the other reviews and the responses of the authors to them and skimmed through the paper again.

The strong points I see in this rebuttal are:
- several clarifications around the implementation, in particular on the use of pre-trained teacher across settings: both on the regular setting but also on the transferability task where DINOv2 was used as teacher and not as a self-supervised objective
- the addition of a comparison with a more recent and related work, Matformer. Pity that the code of SortedNet is not available.
- more detailed information about training costs w.r.t. to the main baseline US-Net, that originally does not benefit from a pretrained teacher. The authors report results for US-Net + KD in the rebuttal and it still under-performs compared to Scala.
- new results using DINOv2 as teacher without cross-entropy loss showing the effectiveness of the option of using foundation models for this setting.

On the downside the evaluation is still limited to mid- to small-sized architectures (the biggest architecture used here is a ViT-B  with 86M parameters) and to computer vision settings. I can relate to the argument of the authors regarding computational cost. However given their results with DINOv2, one may devise a similar strategy for distilling a NLP foundation model to a different task than next-token prediction, e.g., text-classification. Hopefully the authors will find a solution in this direction for the paper update.

**Wrap-up**. I think the submission has improved over the rebuttal with the new information and results and I encourage the authors to include the new findings in the main paper, and to release implementation code (as they stated in the NeurIPS paper checklist, to help other works to build upon and compare against Scala). At this point I'm rather leaning for a positive recommendation for this work. I do not have other questions.
Dear Reviewers:

Thank you so much for delivering your valuable comments to us. 

We have tried our best to provide a detailed response based on your concerns. We would like to kindly remind you that Aug 13th is the deadline for the discussion phase and your opinions regarding our response matter a lot to us.

Thank you so much for helping us improve our paper.

Sincerely,  
Authors