TACO Vision Models Can Be Efficiently Specialized via Few-Shot Task-Aware Compression
Authors: Denis Kuznedelev
Soroush Tabesh
Kimia Noorbakhsh
Elias Frantar
Sara Beery
Eldar Kurtic
Dan Alistarh
Venue: Accepted by TMLR ()

Abstract: Recent vision architectures and self-supervised training methods have enabled training  computer vision models that are extremely accurate, but come with massive computational costs. In settings such as identifying species in camera traps in the field, users have limited resources, and may fine-tune a pretrained model on (often limited) data from a small set of specific categories of interest. Such users may still wish to make use of highly-accurate large models, but are often constrained by the computational cost. To address this, we ask: can we quickly compress generalist models into accurate and efficient specialists given a small amount of data? 

Towards this goal, we propose a simple and versatile technique, which we call Few-Shot Task-Aware COmpression (TACO). Given a general-purpose model pretrained on a broad task, such as classification on ImageNet or iNaturalist datasets with thousands of categories, TACO produces a much smaller model that is accurate on specialized tasks, such as classifying across vehicle types or animal species, based only on a few examples from each target class. The method is based on two key insights - 1) a powerful specialization effect for data-aware compression, which we showcase for the first time; 2) a dedicated finetuning procedure with knowledge distillation, which prevents overfitting even in scenarios where data is very scarce. Specifically, TACO is applied in few-shot fashion, i.e. only a few task-specific samples are used for compression, and the procedure has low computational overhead. We validate this approach experimentally using highly-accurate ResNet, ViT/DeiT, and ConvNeXt models, originally trained on ImageNet and iNaturalist datasets, which we specialize and compress to a diverse set of ``downstream'' subtasks, with notable computational speedups on both CPU and GPU.

Under review as submission to TMLR
TACO:
Vision Models Can Be Efficiently Specialized
via Few-Shot Task-Aware Compression
Anonymous authors Paper under double-blind review
Abstract
Recent vision architectures and self-supervised training methods have enabled training com- puter vision models that are extremely accurate, but come with massive computational costs. In settings such as identifying species in camera traps in the field, users have limited re- sources, and may fine-tune pretrained model on (often limited) data from a small set of specific categories of interest. Such users may still wish to make use highly-accurate large models, but are often constrained by the computational cost. To address this, we ask: quickly compress generalist. models into accurate and efficient. specialists given small amount of data?
Towards this goal, propose simple and versatile technique, which we call Few-Shot Task-Award ware COmpression (TACO). Given a gencral-purpose model pretrained on broad task, such as classification on ImageNet or iNaturalist datasets with thousands of cate- gories, TACO produces a much smaller model that accurate on specialized tasks, such as classifying across vehicle types or animal species, based only on few examples from each target class. The method is based on two key insights 1) a powerful specialization effect for data-aware compression, which showcase for the first time; 2) dedicated fine- tuning procedure with knowledge distillation, which prevents overfitting where data is very scarce. Specifically, TACO is applied in few-shol fashion, i.e. only few task-specific samples are used for compression, and the procedure has low computa- tional overhead. We validate this approach experimentally using highly-accurate ResNot, ViT/DeiT, and ConvNeXt models, originally trained on ImageNet and iNaturalist datasets, which we specialize and compress a diverse set of "downstream" subtasks, with notable computational speedups on both CPU and GPU.
in sccnarios
1 Introduction
TACO
MMAGNEET
Sharks
Large Generalist Model
Colibration Set
Sparse Specialist Model
Figure 1: Few-Shot Bask-Awart COmpression (TACO) is done using calibration set, which is simply few-shot subset of the desired task. It consists of single-step compression procedure, followed by few-shot finetuning step.
The recent introduction of new training techniques that can leverage massive amounts of unlabelled data, such self-supervised (contrastive) learning (Chen et al. 2020: Radford et al. 2021), complemented by new

---

Under review as submission to TMLR
architectures, such Vision Transformers (ViTs) Dosovitskiy et al generation convolutional networks (e.g., ConvNeXt) (Liu et al. 2022) has led to unprecedented accuracy on standard benchmarks, such as the ImageNet image classification task, even in zero- or few-shot settings.
2020 Touvron et al. 2021 and next-
Yet, this impressive increase in accuracy comes at the price of large parameter counts and computational costs. For example, the highly accurate "large" variants of the Vision Transformer (ViT) (ViT/DeiT) model family, commonly used for pretraining, can easily count more than 200 Million parameters, requiring more than 50 GFLOPs for inference at standard input sizes. Thus, although these models reach extremely high accuracies for image classification e.g. close to 90% Top-1 accuracy ImageNet- they are impractical to deploy, especially in resourco-constrained settings, such as edge inference.
One popular way of bridging this computational divide is posl-training compression, e.g. Nagel et al. (2020) Hubara et al. (2021): ; Frantar et al.] (2022): : given a small but representative calibration dataset, the model is compressed to given sparsity or low-procision requirement in one compression step. However, highly- accurale models, such as the ones from the Vision Transformer family, are notoriously hard to compress Chen
et al.
202laa): intuitively, these models may be inherently parameter-heays across massive and often noisy datasets, and classifying images across thousands of output classes.
they are generalists, trained
In this paper, we seek to extend the reach of post-training compression to consider the requirement of model specialization. We start from the idea that, in practical deployment scenarios, such as an ecologist seeking to identify animal species in camera trap images (Beery et al. 2018 or a city planner seeking to identify vehicles using a traffic camera. (Ozkurt & Camei, 2009), users only need their models to categorize a smaller set of classes of interest. Thus, we shift the objective to producing efficient specialist models, which maintain high accuracy on narrower subset of the data distribution. Specialization should allow for higher compression, 8s only a fraction of the features should be required over narrower Lasks.
We propose a simple and versatile approach which we call Task COmpression (TACO), illustrated in Figure I Given a pre-trained generalist model and a "calibration set" containing a few samples for calegories of interest, TACO can produce an ellicient specialist model in two steps. The first a lask-aware compression step, which reduces the model size layer-wise, preserving the model's features on the tpsk-specific calibration data, but removing unnccessary paramcters. Here, our key insight is that existing highly-accurato compression solvers can also be leveraged for model specialization over task-specific calibration set. The second step is a few-shol luning procedure, where we finetune the compressed model on the same small calibration set, leveraging information from the uncompressed model to prevent overfitting.
We validate TACO experimentally on a wide range of pretraining settings, model architectures, and tar- get. tasks. Specifically, we start from models trained on diverse and general tasks such as ImageNet ILSVRCI2 (Ruseakovsky et al., 2015) and iNaturalist 2021 Van Horn et al. (2018). We consider highly- accurate Basc-size models on these datasets from the Vision Transformer (ViT), Swin (Liu et al.) 2021 | and ConvNeXt families, which specialize across tasks obtained by restricting the taxonomies to natural sub- sets such as animal species or vehicle classes. We also extend TACO to a standard transfer learning sctting, showing that it can produce accurate data-aware "lottery tickets" (Chen et al.) 2021b). We compress mod- els with both unstructured and structured pruning, and examine model specialization, as well as practical speedups due to compression. In addition, we also showcase TACO on the popular diffusers text-to-image gencrative task (Hugging Face, Inc. 2023).
Our results show that TACO can reduce the number of model parameters by 10-20x, with low to moderate accuracy loss on the target task. We show that this compression is achieved via task specialization effect: TACO models have significantly better accuracy at the same compression level than models compressed via standard post-training compression. As a rule of thumb, TACO can double the degree of sparsity, at the same accuracy level on the target task, relative to generic post-training compression.
The TACO procedure is computationally efficient: for instance, we can specialize an 86M parameter ViT- single GPU. For instance, this allows us to specialize over large fraction of the WordNet taxonomy (166 tasks) in reasonable time. Morcover, TACO supports different solvers for imposing sparsity: we instantiate it to structured and unstructured pruning, and propose "hybrid" solver, which is both accurate and eighly-effciient. For practical gains, TACO models provide end-to-end
Base model in 10 minutes
2

---

Under review as submission to TMLR
computational speedups 1.5-3x relative to the uncompressed models, at similar accuracy levels to dense, and 3-6x at moderate accuracy loss, on both CPUs and GPUs.
Related work. Few-shot learning (Song al. 2022 Parnami & Lee 2022) is task of adapting pre- trained generalist model to limited number of labeled examples. In many cases, large amount of data is not available due to the cost of data collection or privacy concerns. Therefore, the ability of the model to adapt to a task of interest given a small dataset is of significant importance. Similarly to work on few-shot learning, we adopt a few samples for each class. However, the problem addressed in this work different. Few-shot learning tackles the problem of data-cfficient model adaptation, i.c., learning a new ability. On the contrary, we address the problem of preserving specific knowledge already present in the model.
2
Method
TACO involves two stages; a task-aware compression step, and a few-shot tuning step. Both steps utilize the same small representative set of l samples per class corresponding to the task of interest (say 1 10).
2.1
Few-Shot Task-Aware COmpression
The first step of TACO performs one-shot compression on the given model and lask, reducing the model size in a layer-wise fashion by leveraging a set of task-specific calibration data D. At a high level, given an uncompressod model f, a sparsity solver S, task-spocific calibration data D, and compression constraint C, the TACO procedure produces a compressed version
fe = =ACC.
To implement this, we follow the post-training compression approach: we split the network layer-wise, and run task-specific calibration data D through the uncompressed model. For each network layer & e (1.2....,4) with weights We, we extract the imput X on the calibration data D. Compression is then applied to each layer by solving laser-wise optimization problem, ie. identifying the layer weights W which minimize the difference between the dense and compressed layer output on the provided calibration data, measured in L2-norm. For example, for a linear layer, our goal is to determine a set of weights W; satisfying a compression predicate C such that:
Wi = arg min W. WxXP112. W.ec
(1)
The key difference relative to prior work on post-training compression is specialization over task-aware calibration dota: While prior work uses generic calibration data, our hypothesis is that, when applied over task-spocific calibration data, a good solver for this constrained optimization problem can also isolate the features or paramelers that are specific to the specialized task. This should be done by considering the spec-fp calibration data during the solving, as the weights alone are initially generalists.
2.1.1 Solving the Task-Aware Compression Problem
Sparsity Solvers. To test the specialization hypothesis, we investigated the three existing solvers for unstructured sparsity in the post-training setting:
Magnitude Pruning (Han al. 2015). which sorts the weights by their absolute value, is data- agnostic, so use it reasonable baseline.
2. The Optimal Brain Compression (OBC) solver (Frantar et al. 2022) approximates optimal so- lution for the problem in Equation by greedily pruning the weight with the lowest impact on the layer's L2 loss, and updating all remaining weights to best compensate for the incurred compression error. The solver's asymptotic complexity is O(drow & where drow is the output dimension of the layer and deat cither the number of input features (for linear layer) or the number of chan- nels multiplied by kernel size for convolutional layer. This is the most accurate known solver for layer-wise sparsity but can take several hours to apply to model with more than 100M parameters.
3

---

Under review as submission to TMLR
3. The FastOBC (SparseGPT) (Frantar & Alistarh 2023 solver approximates the OBC approach by pruning in a computationalll-eficci same layer size, at the cost of lower accuracy in the provided solution at a given sparsity.
block pattern. The solver is >> 100 faster than OBC for the
A Hybrid Solver. We also propose HybridOBC, solver which can be seen best-of-both-wool composition of OBC and FastOBC, Most of the runtime cost of OBC comes from the layers with the largest imput dimensions, as its complexity is Oddow dea)). Therefore, we prune sensitive layers, with smaller input dimensions, via the accurate OBC pruner, and the ones with large values of deal (for example, fc2 layer in ViTs) with FastOBC. We will show that HybridOBC significantly reduces the runtime cost relative to OBC, with negligible loss of accuracy.
Alternative Solvers. We have also considered sparsity solvers that leverage the Fisher approximation, such as WoodFisher (Singh & Alistarh 2020 Kurtic et al. 2022) or CAP (uzzedelev et al. 2024). In our experiments, these performed worse in terms of accuracy, relative to the le layer-wise pruners above in the low-data regime we consider in the paper, i.e. 10 samples per class. This is natural since the Fisher approximation requires a large number of gradient samples to be accurate. However, as we show in Section 3.5 if the data is more abundant (i.e., 100 samples/ /class, one can obtain even higher accuracy, especially high sparsities. In addition, also consider the ZipLM (Kurtic et al. 2023) solver for structured (row-wise) sparsity, We observe task specialization effect in this case as well.
2.2 Few-Shot Task-Aware Tuning
At high compression rates, e.g. >90% sparsity, the above one-shot layer-wise approach can lead to significant accuracy degradation. Therefore, we employ a limited few-shot form of finetuning to re-calibrate the model and recover additional accuracy. Yet, naive finetuning on a small amount of data results in "overfitting:" the calibration samples are memorized, but validation performance does not improve. In addition, sparse models tend be more difficult optimize (Evci et al. | 2019). We address both challenges by introducing celf-distillation-basee based approach (Zhang et al.) 2019 2021) recover the accuracy of the sparsified model relative the dense one, which does not overfit when applied over small calibration set.
Recall that, in standard distillation, the student network is trained to reproduce the distribution of the teacher model via minimization of the KL-divergence:
1 Σ po, (x:) log
po, (x:) po, (x:)
Llogit DKL
(2)
B
However, standard distillation may not provide sufficient information in the case of extremely little data. Instead, we found it beneficial to transfer intermediate activations between the dense model and the sparse one. Specifically, we apply normalized mean squared error loss (NMSE) between teacher and student features at the output of each residual block:
MSE f!
Cleat
(3)
MSE 0)
Above Ht and f3 denote teacher and student features, (that is, 3-dimensional feature maps for VIT models and 4- 4-dimensional for ConvNets) at the output of layer and MSE represents the mcan squared error calculated MSE(X,Y) DZ y,)², for N-dimensional vectors X and Y. Normalization factor accounts for different magnitude of activations across layers of the model. This form of feature-level distillation significantly improves performance, relative to both standard retraining and regular distillation (Figure g
(left)). We note, that variants this loss were used in prior work (Sun et al. 2019 Kurtic et al. 2023
Frantar Alistarh ).22). The resulting total loss function has the following form:
Liotal Lrask Lingit + Lseat
(4)
The relative weighting of losses can be tuned, but we found uniform weighting to work well (see Figure 14
4

---

Under review as submission to TMLR
3 Experiments
3.1 Experimental setup
We compare Twar-Aware COmpression, where the objective in Equation equation ! is optimized on task- specific data, with task-agnostic Post- Training Compression (PTC), where the data is chosen uniformly random over general dataset such as ImageNet-1k and iNaturalist21, using the same number of samples for both methods.
The impact of task specialization is assessed for one-shot compression, where accuracy without further tuning is reported, and for fine-tuned accuracy, where we assess the model performance after fine-tuning in few-shot.
In the one-shot setup, one can directly compare the impact of the choice of calibration data on the solution found by the sparse solver. In the fine-tuning setting, examine whether the accuracy difference between PTC and TACO vanishes after tuning whether fine-tuning serves to bridge the performance gap. In all experiments below compare PTC and TACO after fine-tuning, and additionally consider results after only single-step compression for ImageNet subtasks.
Generalist Datasets. We consider two large image classification benchmarks: ImageNet ILSVRC12 Rus- sakovsky et al. (2015), and iNaturalist 2021 Van Horn et al. (2018). Models trained ImageNet-Ik are used in transfer learning experiments.
Generalist models. In the experiments below, we considered three models pretrained on ImageNet-1k (2022). the Transformer- CNN hybrid Swin model (Liu
specifically, DeiT-Ill-Bass following Touvron et al.
ct al. 2021 and the modern Transformer-inppicee
CNNs (ConyNeXt) (Liu et al. 2022 For iNaturalist
experiments, we adopted ViT-Base from He et al. (2022) and finetuned it on iNaturalist using the same
hyperparameters as He et al. 2022 All models are roughly of the same size ~ 90M parameters.
Task Specialization. For ImageNet, we consider specialization across subtasks of the WordNet hierarchy corresponding to ImageNet-IK,, generated via the robustness library (Engstrom et al.) 2019). We selected subtasks that contain at least 5 classes: this totals 166 possible subtasks on Image.Ntt. For iNaturalist, we specialize on various branches of the taxonomic tree corresponding to the phylum, family, order, and genus levels, in increasing order of specificity.
Training parameters and the effect of data scaling. In the experiments below, we use 10 samples per class for all tasks considered. Models are fine-tuned over 100 passes through the calibration data. Although this may seem a high number, we emphasize that the calibration set is tiny, and typically a whole pass involves a few dozen batches. The finetuning procedure takes ~10 minutes on a single NVIDIA RTX 3090 GPU. In our experiments, Section 3.7 and 32 for iNaturalist subtasks. As expected, more data and longer training improve performance of the model. Ablation w.r.t. the samples per class and calibration passes are given in Figure g
use batch size of 16 for ImageNet and transfer learning subtasks (see
Details of the pruning setup. As standard Gaale et al. 2019 we do not prune input embeddings and the classification head, since their latency cost is small. All other linear layers and convolutions are pruned to same uniform sparsity level.
3.2 The TACO Effect
To motivate the need for task-specific compression we first apply the strategy introduced above to the compression on whole ImageNet-lk dataset, shown in Figure 2 The performance drop is quite significant even at moderate sparsity (e.g., 5% relative at 70% sparsity). For high sparsity (90%), the accuracy drops are extremely significant (2 20%).
It is intuitive that one could achieve better performance for the same compression rate on more specialized task. To test this hypothesis, considered few subtasks from ImageNet taxonomy: vehicle, dog and
5

---

Under review as submission to TMLR
Relative Accuracy (%)
100
Dense Accuracy = 83.80%
100
Dense Accuracy 83.80%
90
60
80
40
70
20
PTC
PTC
50
60
70
80
90
50
60
70
80
90
Sparsity (%)
Sparsity (%)
Figure 2: (Left) Single-step compression on ImageNet-lk 1k. (Right) Finctuming ImageNet-1k.
spider subtasks with 21, 116, and sub-classes, respectively. The experiments below are performed with the DeiT-III-Base model, unless otherwise specified.
Specifically, the generalist pruning algorithm seeks to optimize performance over all 1K ImageNet classes, whereas the task- TACO sparsity solver optimizes the performance only specific categories of interest.
Vehicle Subtask
Dog Subtask
Spider Subtask
Relative Accuracy (%)
100
Dense Accuracy 93.50%
Dense Accuracy 1tttrracyy_8.50000
85.50%
110
100
Dense Accuracy 82.33%
80
70
50
PTC
20
PTC TACO
20
TACO
TACO
30
50
70
80
50
60
70
80
90
50
60
70
90
Relative Accuracy (%)
100
Dense Accuracy 93.50%
100
Dense Accuracy 95.50%
105
100
Dense Accuracy 82.33%
96
90
85
70
PTC
PTC
PTC
75
TACO
TACO
TACO
Sparsity (%)
Sparsity (%)
Sparsity (%)
Figure 3: (Top) Single-step compression on ImageNet subtasks. (Bottom) Finctuning on ImageNet sub- stasks.
We can observe in Figure 3 that the sparse solution found by TACO is noticeably better than the one produced by task-agnostic PTC, showcasing the "TACO effect" by which data-aware compression imposes task specialization in the resulting model. Importantly, the performance gap between TACO and PTC remains significant even after fine-tuning with distillation for both methods. This shows that the model specialization is preserved after fine-tuning. Specifically, one can compress the model to 75% sparsity with minor drop in accuracy, and to 90% sparsity while preserving 90% of the dense model accuracy.
3.2.1 Task Complexity and Model Compressibility
Task Complexity. In the next experiment, we wish to examine compressibility with respect to task complexity, measured as the number of task sub-classes. For this, run one-shot pruning over ImageNet
6

---

Under review as submission to TMLR
We keep only the tasks containing at least 5 output classes. We select three sparsity levels (0.6, 0.7, and 0.8) and obtain single-step pruned models using TACO. (We use the FastOBC solver in this experiment, given the large mumber-1660 of subsets.) The results of this large-scale experiment are provided in Figure ! where we examine the correlation between the accuracy of the sparse models relative to the corresponding dense one, as well as the relative accuracy drop post-compression, relative the number of classes in the target specialization dataset.
with respect to the WordNet hierarchy (Fellbaum 2010).
1.0
Sparsity
Relative Accuracy drop
Sparsity
Accuracy sparse
Sparsity 0.8
y=x
-0.15
0.7
-0.25
Sparsity 0.6
0.6
Sparsity 0.7
-0.35
Sparsity 0.8
0.5
-0.40
0.75
0.80
0.85
0,50
0.35
1.00
10a Num classes
10t
Accuracy dense
Figure 4: (Left) Accuracy on the ImageNet subtask for the dense model vs. accuracy of the sparse model. (Right) Relative accuracy drop vs. the number of classes.
In the left-hand plot, we observe a correlation between the accuracy of the uncompressed model and that of the sparse one on the given task. For most tasks, the performance drop is negligible for 60-70% sparsity and becomes pronounced al only 80% sparsity. Second, certain subtasks allow for higher compression relative to others. This correlates with the number of target classos, as a mcasure of "complexity." Specifically, the accuracy drop at a given sparsity appears to be logarithmic in the number of classes. Finally, for some subsets, the accuracy of the compressed model can be superior to the original model.
In addition, for cach subtask selected, we evaluate the PTC model and compare the performance of PTC and TACO. Due to space constraints, results are presented in the Appendix F overall, they show that TACO outperforms PTC on the majority the tasks, and that the accuracy difference becomes more pronounced with an increase in sparsity. The difference is also more significant for smaller subtasks, showing the advantage of task specialization.
3.2.2 Additional Models and Structured Sparsity
Our conclusions generalize other vision models and to struetured sparsity. Specifically, we apply the ZipLM (Kurtic et al. 2023) structured solver, which optimizes layer-wise objective and identifies the least important channels. Following Kurtic et al. (2023), we prune both attention heads and hidden dimensions in the feedforward Transformer layers, to the same sparsity level. The advantage of this approach is that structured sparsity can be accelerated both on CPU and GPU hardware, although we would expect higher accuracy drops.
Figure 5 (left) shows that we can prune the model to 30-50% sparsity with minor performance drops, and to 70% sparsity while preserving 90% of the original accuracy. Throughout, TACO outperforms PTC across sparsities, similarly to the unstructured case.
In Figure 5 (right), we evaluate TACO on two other recent vision models, ConvNext- -Bse and Swin- V2-Base, roughly of the same size and performance as DeiT-III-Base. The trend is similar across all models considered up to 80-85% sparsity, performance drops are small. DeiT-Ill-Base appears to be the easiest to compress.
3.3 TACO on iNaturalist
We further investigate task-aware compression on the challenging iNaturalist21 dataset Van Horn al. 2018), comprising 10,000 classes corresponding to diverse terrestrial and marine animal and plant species.

---

Under review as submission to TMLR
Vehicle Subtask
Vehicle Subtask
105
Relative Accuracy (%)
Dense Accuracy 99.50%
Accuracy (%)
Deff-ill-Bal Base
ConyNext-BBal
85
70
Seiin-V--Bass
60
PTC
60
TACO
75
50
30
60
70
0
50
67
7.5
80
85
90
95
Sparsity (%)
Sparsity (%)
Figure 5: (Left) Structured pruning on "Vehicle" subtask. models.
(Right) Unstructured pruning for different
In the experiments seen in Figure [ and in the Supplementary, we take categories on different levels of the taxonomic hierarchy, and compare performance between sparsity solver applied on the whole iNaturalist dataset vs task-aware compression. The setting is the same as for ImageNet: unstructured sparsity, and fine-Luning via the L2 distillation stralegy.
Mollusca
Cnidaria
Malvales
85
85
95
Accuracy (%6
80
75
85
70
70
65
PTC
PTC
PTC TACO
TACO
65
TACO
4
10
2
6
B
10
2
4
6
10
Gruiformes
Araneidae
Iris
Accuracy (%))
80
85
75
80
90
75
65
PTC
65
PTC
80
PTC
TACO
TACO
TACO
60
4
10
4
10
Sparsity (%)
Sparsity (%)
Sparsity (%)
Figure 6: Few-shot compression iNaturalist21
The results for subtasks with 10-100 classes are shown in Figure 0 showing that the impact of task-aware compression on model performance is pronounced, especially with increasing sparsity. Soc Supplementary for more details.
3.4 Speedups
Next, we examine the practical potential of our method both for CPU and GPU inference. We measure the execution lateney for processing a single image of size 224 x 224 with a DeiT-III-Baso model. (We note that larger batch processing leads higher speedups, but would be less realistic in our target deployment scenario.) Full details about the benchmarking environment provided in Supplementary D
8

---

Under review as submission to TMLR
CPU Speedups. For experiments with CPU speedups execute unstructured sparse models on the DeepSparse runtime (Kurtz et al.] 2020), which can accelerate unstructured sparsity. Latency is measured on an Apple MI processor, using 4 cores, We also quantize both model weights and activations to 8 bits, which can be done losslessly. In Figure (left), we observe that one can reduce latency by more than 4x with a minor performance drop (2%); and by 6x with a moderate (6%) accuracy drop.
Vehicle Subtask
Vehicle Subtask
Relative Accuracy (%)
105
Relative Accuracy (%)
105
100
Drnse Accurssy 93.50%
100
Dense Ascurwcy =.9..0%%
95
95
90
85
85
80
TACO 100 epochs
TACO
80
75
dense
dense
2.5
3.5
CPU Speedup,
GPU Speedup, x
Figure 7: (Left)) CPU speedups for DeiT-III-B vs relative accuracy with 8-bit quantization on the DeepSparse runtime. (Right) GPU speedups for DeiT-IlI-B vs relative accuracy with fp16 precision on the TensorRT runtime.
GPU Speedup. Next, we execute the struetured-pruned models obtained via TACO (Section 3.2.2) in the NVIDIA TensorRT runtime, an NVIDIA T4 GPU. We measure speedups and accuracies when spe- cializing on specific subtasks, shown in Figure (right). The speedups are identical on any subtask, as the classification head has a negligible cost. We observe that, through specialization, we obtain a 60% inference speedup with negligible accuracy loss, and 2.5x speedup while maintaining > 90% of the accuracy on this particular task and model combination.
3.5 Ablation Studies
Solver Impact. We have focused our experiments so far on the HybridOBC solver as the main method providing a trade-off between performance and efficiency in the few-shot scenario. In this section, we justify this choice by comparing with alternatives. Specifically, Figure (Ieft) shows that data-agnostic Magnitude pruning falls behind data-aware solvers, which leverage calibration data. Here, HybridOBC is more accurate than FastOBC, and performs on par with OBC, while being significantly faster. Specifically, on DeiT--II-- Base, Fast OBC takes 12 second, while HybridOBC takes around 12 minutes, and OBC takes more than hour.
Data Impact. The behavior of our proposed TACO approach with respect to the number of samples per class and calibration data fine-tuning passes is given in Figure g Naturally, accuracy improves with more data, but also shows diminishing returns. We find that 10 samples per class is a good balance wetw accuracy and speed of the approach.
Further, when a large amount of data is available, one can enhance the TACO effect by employing a more data- and compute-intensive Fisher-based approximation of the Hessian, tor with longer fine-tuning. See Figure 10 for comparison with the Fisher-based WoodFisher Singh & Alistarh (2020) and CAP Kuznedelev et al. 2024 pruners, and Appendix C for full results. However, in the limited-data regime with 10 samples per class, the fast HybridOBC solver works best.
Distillation Impact. Another ingredient we examine is the distillation strategy. The simplest approach would optimize the cross-entropy task loss on the calibration data; yet, this option results in suboptimal convergence as shown in Figure (right). We find that, since there only a scarce amount of data available, the model simply memorizes the samples instead of generalizing on the task. Output distillation reduces overfitting to the calibration set, but does not facilitate fast recovery. The compound loss we propose,
9

---

Under review as submission to TMLR
Vehicle Subtask
Vehicle Subtask
Accuracy (%6
Accuracy (%6)
86.00
86.48
75
86.00
83.43
70
70
78.38
69.62
64.48
Magnitude
FasDBC
0BC
HybridOBC
No XD
Output XD
Composid KD
Figure 8: (Left) Ablation of pruner choice. (Right) Ablation of distillation strategy.
Vehicle Subtask
Vehicle Subtask
85.0
Accuracy
Accuracy
30.0
86
88
70.0
10 Samples por class
100
1000
Passes
Figure 9: (Left) Accuracy on Vehicle task vs number of samples per class for fixed number of training steps for 90% sparse DeiT-IlI-Base, (Right) Accuracy vs number of passes through few-shot training set for 90% sparse DeiT-Ill-Base
Vehicle Subtask
Vehicle Subtask
90
Relative Accuracy (%)
101.0
Accuracy (%6)
80
100.0
99.0
60
HybridOBC WoodFisher
30
98.0
87.5
46
10
30
100
50
80
80
90
Samples per class
Sparsity (%)
Figure 10: (Left) Accuracy on Vehicle task number of samples per class for different pruners. (Right) Accuracy Vehicle subtask for CAP pruner in data-rich regime.
involving both the task loss, output logit distillation and intermediate activation representation, appears noticeably improve upon the former options. The evolution of validation accuracy and training loss is presented in Figure 11.
3.6 TACO on transfer learning tasks
The next test of the introduced "TACO effect" is transfer learning setting. Since ImageNet pretrained models are known to provide good initialization for "downstream" " tasks Yossiski et al. 2014 consider transfer to 6 tasks CIFAR-10, CIFAR-100, Caltech-101, Caltech-256, Oxford Pets, and Flowers 101. Firstly, we lune dense, ImageNel-prellaiineg the target compression rate, and fine-tune it via the calibration set, using the model from the previous step as a teacher. As before, we use 10 samples per cach class and 100 passes through the calibration set. The PTC variant uses generic ImageNet-1l 1k data for calibration, whereas TACO uses calibration data from the downstream task. We wish to determine whether TACO leads to better compression / aaccuray trade-off.
model on the task of interest. Then, we prune the resulting model to
10

---

Under review as submission to TMLR
Vehicle Subtask
Vehicle Subtask
0.85
Accuracy (%)
0.80
0:75
Compound
0.8
Compound
0.70
Outgut KD
No KD
0.6
Output KD
No KD
0.65
0.4
0.2
0.60
200
1000
1200
400
800
1000
1200
Step
Step
Figure 11: (Left) Validation accuracy vs step for different distillation options. (Right) Train loss accuracy vs step for different distillation options
CIFAR-10
CIFAR-100
Caltech 101
03
Accuracy (%)
92
83
90
70
TACO
TACO
TACO
Caltech 256
Oxford Pets
Flowers 102
Apcuracy
93
17.5
15.0
FTC
PTC
PTC
32.5
TACO
Compression rate, X
Compression rate, X
Compression rate, x
Figure 12: Few-shot compression on transfer learning tasks
The results in Figure 12 confirm that TACO achieves better performance than PTC in most cases. (On the X axis, the compression rate is defined as T-s for sparsity s.)
3.7 Gradual Pruning and Data-Aware Lottery Tickets
So far, our approach has been to apply TACO one-shot, initially, followed by fine-tuning. We next investi- gate its application gradually, in steps of increasing sparsity, separated by short fine-tuning. This gradual approach, formalized in Algorithm
| Cooncidees with Lottery Tickets for transfer learning as proposed by Chen et al. 2021b), and thus provides good point of comparison relative to this method.
Algorithm Overview of gradual transfer learning sctup.
1: Input: Dense model M, Sparsity target Pruning method I'
2: repeat
3: Prune 50% of remaining weights in M via method r.
4: Update to the pruned model.
5: Finetune M for 25 epochs, maintaining the sparsity mask.
6: until M's sparsity reaches some larget level a.
Return M.
Following Chen et al. 2021bb, we use a pretrained ResNNet50/ImmagNee--1 model as the base, and run TACO on a subset of 1024 samples randomly selected from the transfer task. the gradual case, each pruning step removes 50% of the remaining model weights, followed by 25 epochs of full fine-tuning on the transfer task. In the single-step case, we prune once, and fine-tune for the cumulative equivalent number of epochs. For fairness, exactly the same setting with gradual magnitude pruning (GMP) and Lottery Tickets for
11

---

Under review as submission to TMLR
Transfer (LTH) Chen et al. (2021b) methods, applying the corresponding priner at the compression step. Figure 13 shows results for 6 standard transfer datasets Kornblith et. al. (2019), for compression between 2 (50% sparsity) and 64 (98.5% sparsity). We provide dedicated plots in Appendix A
Across all tasks and sparsities, Gradual TACO achieves highest accuracy, followed by single-step TACO + finetuning. Remarkably, Gradual TACO achieves moderate accuracy loss even for 2 95% sparsity, outper- forming the GMP and LTT data-agnostic methods by large margin. The strong performance of Single-step TACO shows that specialization can occur when there is just single compression step, and the cali- bration data is not part of the pretraining task. This suggests that model specialization may also happen for tasks that are not strictly in-distribution, which bodes well for the generality of the method. The accu- racy difference between TACO and LTH/GMP shows that task-aware specialization provides much better compression for transfer learning.
4 Application to text-to-image generation
TACO can be applied to any deep-learning based computer vision task where there is potential for spe- cialization. To illustrate this, we show results for text-to-image generation on the Pokémon BLIP captions dataset Pinkney (0022) using a Stable Diffusion (SD) v1.4 model trained via the Hugging Face diffusers tutoria
We prune convolutions and linear projections in the SD UNet while keeping the image encoder dense. Since the SD UNet is much larger than ViT models considered in previous experiments (~~60M parameters) we apply the FaC sparsily solver. We use the whole task data, consisting of ~ 800 images, as a calibration set, placing us in the few-shot tuning regime. We finetune the single-step compressed model for the same amount of steps used for teacher specialization.
To quantitatively assess image generation quality, we employ the commonly used FID (Fréchet Inception Distance) metric (Heusel et al. 2018), which utilizes the Fréchet distance between features extracted by the Inception-V3 (Szegedy et al. 2016) backbone pretrained on ImageNet. This metric measures the distance between the target and generated data distributions. Specifically, using prompts from the Pokémon dataset, we generate images through dense and sparse diffusion models and calculate the FID metric based on isates from both the original dataset and the generated images. The results are presented in Table !
It can be observed that the FID score for the model with 50% sparsity is similar to that of the original model. The increase in the FID score remains moderate at 67% and 75% sparsity levels. Additionally, we provide qualitative results for eight random prompts (see Appendix G for details) in Figure 17
Table 1: FID score for Stable Diffusion v1.4 models with different sparsity.
Sparsity (%)
FID
56.91
50
53.96
67
61.45
5 Discussion
75
63.05
In this paper, we describe and analyze for the first time the effect of lask specialization in the context of model compression. Via systematic experiments, we show that task-specificity can increase the degree of compression by roughly 2x at the same level of accuracy, relative to task-agnostic techniques. Our findings are valid across one-shot compression, fine-tuning, transfer learning, and gradual pruning scenarios. They lead to significant speedups (of up to 6x) at low to moderate accuracy loss, relative to the original model, which applies to both commodity CPU and accelerator GPU hardware.
A strength of our approach is that it is generic with respect to the sparsity solver used. We have shown notable differences in performance with respect to different solvers, in particular relative to the amount of data and computation employed. Future work could investigate specialized solvers for task-aware compression, applications to other domains, such as delection or segmentation.
12

---

Under review as submission to TMLR
References
Sara Beery, Grant Van Horn, and Pietro Perona. Recognition in terra incognita. In Proceedings of the European conference on computer vision (ECCV), pp. 456 473, 2018.
Tianlong Chen, Yu Cheng, Zhe Gan, Lu Yuan, Lei Zhang, and Zhangyang Wang. Chasing sparsity in vision transformers: An end-to-end exploration. Advances in Neural Information Processing Systems, 34: 19974 19988, 2021a.
Tianlong Chen, Jonathan Frankle, Shiyu Chang, Sijia Liu, Yang Zhang, Michael Carbin, and Zhangyang Wang. The lottery tickets hypothesis for supervised and self-supe pre-training in computer vision models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp-
16306 16316, 2021b.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive lcarning of visual representations. In International conference on machine learning, pp. 1597-1607. PMLR, 2020.
Xin Dong, Shangyu Chen, and Sinno Pan. Learning to prune deep neural networks via layer-wise optimal brain surgeon. Advances in neural information processing systems, 30, 2017.
Alexcy Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale, 2020. URL
https: http://aavvvvvgroraee
11929
Logan Engstrom, Andrew Nyas, Hadi Salman, Shibani Santurkar, and Dimitris Tsipras. Robustness (python library), 2019. /ttpss///iithbb http//////////////.///////aaaaaa
Utku Evci, Fabian Pedregosa, Aidan Gomez, and Erich Elsen. The difficulty of training sparse neural networks. arXiv preprint arXiv:1966.100 2019.
Christianc Fellbaum. Wordnet. In Theory and applications of ontology: computer applications, pp. 231-243. Springer, 2010.
Elias Frantar and Dan Alistarh. SPDY: Accurate pruning with speedup guarantees. arXiv preprint arXiv:2201113000 2022.
Elins Frantar and Dan Alistarh. Sparsegpt: Massive language models can be accurately pruned in onc-shot, 2023. URL http:://aarxiv.or///aa//////3 00774
Elias Frantar, Eldar Kurtic, and Dan Alistarh. M-fac: Efficient matrix-free approximations of second-order information. Advances in Newral Information Processing Systems, 34, 2021.
Elias Frantar, Sidak Pal Singh, and Dan Alistarh. Optimal brain compression: A framework for accurate post-training quantization and pruning. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL https:://peerevvis: .etteeeeeeiieeeee//eeee
forum?1d-ksvcc00000
Trevor Gale, Erich Elsen, and Sara Hooker. The state of sparsity in deep neural networks. arXiv preprint arXiv:1902.09574, 2019.
Song Han, Huizi Mao, and William J Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding. arXiv preprint arXiw:1510000149, 2015.
Kaiming Hc, Xinlci Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pallern Recognition, pp. 16000-160000 2022.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge to local nash equilibrium, 2018.
13

---

Under review as submission to TMLR
Itay Hubara, Brian Chmicl, Moshe Island, Ron Banner, Joseph Naor, and Daniel Soudry. Accelerated sparse neural training: A provable and efficient method to find n: m transposable masks. Advances in neural information processing systems, 34:21099 21111, 2021.
Hugging Face, Inc. Hugging face's diffusers: A state-of-the-art diffusion models library. https://githe:
http://uuggingff.cc//fuussssssss 2023.
Simon Kornblith, Jonathon Shiens, and Quoc Le. Do better imagenet models transfer better? Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 2661-2671, 2019.
Eldar Kurtic, Daniel Campos, Tuan Nguyen, Elias Frantar, Mark Kurtz, Benjamin Fineran, Michael Goin, and Dan Alistarh. The optimal bert surgeo: Scalable and accurate second-order pruning for large language models. arXiv preprint arXiv:2203.0720 2022.
Eldar Kurtic, Elias Frantar, and Dan Alistarh. Ziplm: Hardwaro-aware structured pruning of language models. arviv preprint arXiv:2302.044 04089, 2023.
Mark Kurtz, Justin Kopinsky, Rati Gelashvili, Alexander Matveev, John Carr, Michael Goin, William Leiserson, Sage Moore, Bill Nell, Nir Shavit, and Dan Alistarh. Inducing and exploiting activation sparsity for fast inference deep neural networks. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 5533-5563 5543, Virtual, 13-18 Jul 2020. PMLR. URL http:://proceeddngg/ mlr. wtt.pps//19///rrt2222/. html
Denis Kuznedelev, Eldar Kurtié, Elias Frantar, and Dan Alistarh. Cap: Correlation-aware pruning for acguly- sparse vision models. Advances in Neural Information Processing Systems, 36, 2024.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hicrarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVP International Conference on Compuler Vision, pp. 10012 10022, 2021.
Zhuang Liu, Hanzi Mao, Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xic. A convnet for the 2020s. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, PD. 11976-1986,, 2022.
Markus Nagel, Rana Ali Amjad, Mart Van Baalen, Christos Louizos, and Tijmen Blankevoort. Up or down? adaptive rounding for post-training quantization. In International Conference on Machine Learning, pp. 7197-7206. PMLR, 2020.
Celil Ozkurt and Fatih Camci. Automatic traffic density estimation and vehicle classification for traffic surveillance systems using neural networks. Mathematical and Computational Applications, 14(3):187- 196, 2009.
Archit Parnami and Minwoo Lee. Learning from few examples: A summary of approaches to few-shot learning, 2022. URL https://arxiv httpss///araivoorr//bb///200000 202.aa../2.2.2119..111
Justin N. M. Pinkney. Pokemon blip captions.
pokemon-blip blip-captions/ 2022.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transfcrable visual models from natural language supervision. In International conference on machine learning, pp. 8748 8763. PMLR, 2021.
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. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJCV), 115(3):211-253 2015. doi: 10.
14

---

Under review as submission to TMLR
Sidak Pal Singh and Dan Alistarh. Woodfisher: Efficient second-order approximation for neural network compression. Advonces in Neural Information. Processing Systems, 33:18098-18 18109, 2020.
Yisheng Song, Ting Wang, Subrota Mondal, and Jyoti Prakash Sahoo. A comprehensive survey of few-
shot learning: Evolution, applications, challenges, and opportunities, 2022. URL httts://arxiv.orgg
abs/2205.06745 06743
Siqi Sun, Yu Cheng, Zhe Gan, and Jingjing Liu. Patient knowledge distillation for bert model compression. arXiv preprint Xrx: 09355, 2019.
Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed, Dragomir Anguelov, Dumitru Erhan, Vincent Vanhoucke, and Andrew Rabinovich. Going deeper with convolutions, 2014.
Christian Szegedy, Vincent Vanhoucke, Sergey loffe, Jon Shlens, and Zbigniew Wojna. Rethinking the inception architecture for computer vision. In Proceedings of the IEEE conference on compuler vision and pattern recognition, pp. 2818 2826, 2016.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou, Training data-efficient image transformers & distillation through attention. In International Conference on Machine Learning, pp. 10347-10355 10357. PMLR, 2021.
Hugo Touvron, Matthieu Cord, and Hervé Jégou. Deit ii:: Revenge of the vit, 2022. URL https: htttp:///rxii/ org/ 2ttps//224.0077
Grant Van Horn, Oisin Mac Aodha, Yang Song, Yin Cui, Chen Sun, Alex Shepard, Hartwig Adam, Pietro Perona, and Serge Belongie. The inaturalist species classification and detection dataset. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 8769 8778, 2018.
Jason Yosinski, Jeff Clune, Yoshua Bengio, and Hod Lipson. How transferable are features in deep neural networks?, 2014.
Linfeng Zhang, Jiebo Song, Anni Gao, Jingwei Chen, Chenglong Bao, and Kaisheng Ma. Be your own teacher: Improve the performance of convolutional neural networks via self distillation. In Proceedings of the IEEE/CVE International Conference on Computer Vision, pp. 3713-3722 2019.
Linfeng Zhang, Chenglong Bao, and Kaisheng Ma. Self-distillation: Towards efficient and compact neural networks. IEFE Transactions Pattern Analysis and Machine Intelligence, 44(8):4388-444 -408888443,,-0 2021.
15

---

Under review as submission to TMLR
A
Transfer learning results
Describable Textures (DTD)
Oxford- -ITT Pets
0.7
0.8
0.6
Accuracy
0.6
0.5
0.4
0.4
Gradual TACO
Gradual TACO TACO
0.3
TACO
TACO
0.2
0.2
LTH
LTH
GMP
GMP
2
16
32
64
2
8
16
32
64
Compression Rate
Compression Rate
CIFAR-10
CIFAR-100
0.85
0.96
0.80
Accuracy
0.94
0.75
0.92
0.70
Gradual TACO
Gradual TACO
TACO
TACO
0.90
LTH
0.65
LTH
GMP
GMP
8
16
32
64
2
8
16
32
64
Compression Rate
Compression Rate
Caltech-101
Caltech-256
0.8
0.8
Accuracy
0.6
0.6
0.5
0.4
0.4
Gradual TACO
Gradual TACO TACO
0.3
0.2
TACO
TACO
LTH
0.2
LTH
GMP
GMP
0.0
0.1
2
8
16
32
64
2
8
16
32
64
Compression Rate
Compression Rate
Figure 13: Experiments for model compression in a standard transfer learning setup. We compare TACO and Gradual TACO with Lottery Tickets for Transfer (LTH) and Gradual Magnitude Pruning (GMP) in terms of accuracy (Y axis) vs compression (X-axis).

---

Under review as submission to TMLR
B Few-shot tuning hyperparameters
In our few-shot tuning procedure we finetune the model via optimizing the sum of original task loss, logit and feature distillation loss. The hyperparameters adopted in training are listed in Table Learning rate 10-3,3 33
has been determined subtask and adopted in all other experiments, We apply simple augmentation pipeline Szegedy et al. (2014) with random resized crops and horizontal flips. Since distillation with intermediate teacher activations is already strong regularizer there is need for elaborate aaareetttttio/regulaaaaatto pipeline.
grid search on 11-10-5.3 3- 10 1-101,33, 10-2,1
on "Vehicle"
Table 2: Hyperparameters used in few-shot tuning.
Optimizer
Adam
Learning rate
3 10-4
LR. schedule
linear
Batch size
(32)
Num passes
100
Weight decay
0
Dropout
X
H.flip
RRC
Image size
224
Test crop ratio
0.875
Losses forming the compound distillation loss may differ significantly in magnitude and the importance for generalization performance. While the identification of optimal weighting of the losses can improve performance, we decided to keep our setup simple and tried couple values A in front of feature loss. It turned out that the results are robust to its variation within a wide range of values. However, if the coefficient in front of feature loss is too small, guidance from teacher activations may be not sufficiently strong.
Vehicle Subtask
88
87
Accuracy (%)
86
85
84
83
86.00
86.10
84.86
81
80
Areat = 0.1
Asar =
Arcat =
Figure 14: Ablation of the cocfficient in front of feature loss.
C Performance of Fisher-based pruners
The two most commonly applied approximations the Hessian for estimates of parameter importance are the layerwisc approximation Dong al. (2017): Frantar al. (2022); Frantar & Alistarh (2023 Kurtic al.
(2023) and block-diagonal Fisher approximation Singh & Alistarh (2020); Frantar et al. (2021) Kuznedelev
al. (2024) The first accounts only for importance of specific weight (group of weights) the output
2

---

Under review as submission to TMLR
of given layer, whereas the latter estimates the saliency with respect to the target task. However, it appears that methods leveraging Fisher approximation demand more samples for accurate estimate and fall behind layerwise sparsity solvers for small number of samples (Figure 15 (Left)). With the increase of amount of data for Empirical Fisher WoodFisher and CAP close the performance gap with HybridOBC and can even outperform it.
To test our method in data- and compute- rich regime we took 200 samples (20% of the training samples from ImageNet-1k) per class for Vehicle subtask and DeiT-III-Baso model and finetuned the model for 100 passes through the calibration set. One preserve 97.5% of the original model performance even at the highest 95% sparsity considered. Therefore, the proposed method works well even in the data abundant setup.
Vehicle Subtask
Vehicle Subtask
101.0
Relative Accuracy (%)
100.5
Accuracy (%)
100.0
99.5
99.0
HybridOBC WoodFisher CAP
98.5
98.0
97.5
3
10
30
100
50
60
70
80
90
Samples per class
Sparsity (%)
Figure 15: Accuracy on Vehicle subtask for CAP pruner.
D Benchmarking environment
CPU benchmarking was conducted on 4 cores of Apple M1 chip (MacBook Pro 2020). Model was compiled with DeepSparse engine (version deepsparse-nightly 1.6.0).
GPU measurements were carried out on Google Colab runtime, Nvidia T4 GPU, TensorRT (version 8. 6.6.0.1 0. 12 for CUDA Toolkit 12.0).
E Details about iNaturalist subtasks
In the main text we present results for 6 tasks from different levels species laxonomy. Below are presented some additional details about the chosen tasks.
Table 3: Additional details about iNaturalist subtasks
Dataset Mollusca
Category
Classes
phylum
169
Cnidaria Malvales
phylum
26
order
75
Gruiformes
order
29
Araneidae
family
45
Iris
genus
15
3

---

Under review as submission to TMLR
F TACO vs PTC on the ImageNet hierarchy
In order demonstrate that TACO superiority not caused by cherry-pickinn of specific ImageNet subsets we evaluated TACO and PTC compression on the whole set of ImageNet subtasks with more than 5 classes (i.e. 166 subtasks chosen). Comparison plots presented Figure 16. One can observe that TACO oulperforms PTC most of the lasks.
1.0
Sparsity 0.6
0.25
Sparsity 0.6
Sparsity 0.7
Accuracy difference
Sparsity 0.7
Accuracy TАCO
0.9
Sparsity 0.8
0.20
Sparsity 0.8
y=x
0.15
0.10
0.05
0.6
0.00
0.5
0.5
0.6
0.7
0.8
0.9
1.0
10a
10a
10a
Accuracy PTC
Num classes
Figure 16: (Left) Accuracy of PTC compressed model vs TACO the specialized task (ImageNet subset). (Right) Accuracy difference between models compressed by TACO and PTC on the specialized tasks.
G Prompts and visualization in text-to-image experiment
The prompts used for generation images 17 are the following:
à carloon bear with a ring around its neck
a very cute looking cartoon character with big eyes
a cartoon character is holding a baseball bat
drawing of green and white bird
bird with yellow eyes and a black beak
drawing of a white and black animal with blue eyes
drawing of a dragon with its mouth open
cartoon picture of giant turtle with its mouth open
Images were generated with default SD solver, 50 inference steps and efg scale of 7.5.

---

Under review as submission to TMLR
0%
50%
67%
75%
Figure 17: Generated images corresponding to random prompts from Pokémon BLIP captions. Percentages are sparsity levels.
H Sensitivity to the choice of random seed
The experiments in the main part of the text adopt single choice of random subset for specific subset. To verify the robustness of the results to the choice of random seeds we repeat the experiments for vehicle, dog and spider subtasks from ImageNet in single-step compression setting for 4 random seeds both for TACO and PTC and observed that the final result not very sensitive to the choice of random seed with TACO exhibiting smaller variance compared PTC. Results averaged across multiple runs are provided on Figure
Vehicle Subtask
Dog Subtask
Spider Subtask
Relative Accuracy (96)
Accurary 90.50%
Relative Accuracy (%)
Demse.AecaaaE. Awc.raEE.. 95.50%
Relative Accuracy (96)
100
80
80
40
50
PTC
PTC
TACO
TACO
TACO
50
80
90
70
90
50
80 Sparsily (%)
Sparsily (%)
Sparsily (%)
Figure 18: Mean and standard deviation for single-step compression ImageNet.


Reviews:
[{"review_id": "xB32eallTl", "reviewer": "Reviewer_e5gK", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The introduction of HybridOBC seems reasonable and extensive experiments illustrate the effectiveness of the work.\n2. This paper is well-organized and easy to read.\n3. This paper has done extensive experiments in verifying the efficency of the work.\n\nWeakness:\n1. According to many recent work which aims in adaping the large pre-trained models such as prompt tuning in CLIP[1], it aims in well-adapting on the downstream tasks suchas cars while maintaing the well generailized ability. If the aim is to erasing the ability of it, why not train a model from scratch. If so, there is no experiment comparing it with train from scratch.\n\n2. During the experiment, nearly all the experiments or figures are based on the TACO while there is no comparison between existing models which make the paper less convincing.\n\n\n[1] Tian X, Zou S, Yang Z, et al. ArGue: Attribute-Guided Prompt Tuning for Vision-Language Models[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024: 28578-28587.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Few-Shot Task-Aware COmpression (TACO).\n\nIn this paper, we seek to extend the reach of post-training compression to consider the requirement of model specialization. We propose a simple and versatile approach which we call Task-Aware COmpression (TACO), illustrated in Figure 1. Given a pre-trained generalist model and a “calibration set” containing a few samples for categories of interest, TACO can produce an efficient specialist model in two \nsteps.\n\nThere are two main steps in TACO including the model compression and training on the representative training images. In my view, model compression makes the most esstial novel role in the novelty of the TACO. \n\nStep1. Author split the network layer-wise, and run task-specific calibration data D through the uncompressed model. Compression is then applied to each layer by solving a layer-wise optimization problem. We also propose HybridOBC, a solver which can be seen as a best-of-both-worlds composition of OBC [1] and FastOBC [2].\n\nStep2. In order to get rid of the negative effect on the accuracy caused by the compression, the author employ a limited few-shot form of finetuning to re-calibrate the model and recover additional accuracy. We address both challenges by introducing a self-distillation-based approach where they user a self-distillation-based approach [3] to recover the accuracy of the sparsified model relative to the dense one, which does not overfit when applied over a small calibration set.\n\nTACO models provide end-to-end computational speedups 1.5-3x relative to the uncompressed models, at similar accuracy levels to dense, and 3-6x at moderate accuracy loss, on both CPUs and GPUs.\n\n\n[1] Frantar E, Alistarh D. Optimal brain compression: A framework for accurate post-training quantization and pruning\n[2] Frantar E, Alistarh D. Sparsegpt: Massive language models can be accurately pruned in one-shot[C]//International Conference on Machine Learning. PMLR, 2023: 10323-10337\n[3] Zhang L, Bao C, Ma K. Self-distillation: Towards efficient and compact neural networks[J]. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2021, 44(8): 4388-4403.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Same as above.", "broader_impact_concerns": "no", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "PHpxgfrgEZ", "reviewer": "Reviewer_rVaw", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: \n1. Authors evaluate TACO in various settings and datasets. The extensive experiments on multiple datasets, models and subtasks establish the generality and efficacy of TACO.\n2. The task-aware compression approach is novel and demonstrates potential for real-world applications in resource-constrained scenarios like identifying species in the wild. \n3. The paper has detailed ablations, quantitative and qualitative analysis to strengthen the claims. \n\nWeaknesses: \n1. In the main experiments, authors use 10 samples per class for the calibration data but failed to discuss the impact of the number of class used for the calibration dataset.\n2. This paper focuses primarily on vision classification tasks and extend a bit to image generation task, but only report a naive visualization result in supplementary material for image generation quality, without any quantitative analysis.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces TACO, a method for efficiently specializing pre-trained computer vision models for specific tasks using very small amount of task-specific data. A novel two step process was proposed, combing task-specific compression and KD based fine-tuning. The authors evaluate the proposed method in various network architectures and datasets, and further extend the evaluation task from classification to image generation. Results show TACO achieves 2x compression compared to task-agnostic method, and providing significant computational speedups both in CPU and GPU.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Please indicate what method used in Figure 2 in the caption. Is it two-step TACO? \n2. For the text-to-image generation, it is better to consider adding quantitative result for assessment. I suggest add another dataset that can be evaluate by FID metric. \n3. There are some repeat citations, like ViTs are cited twice in the same way in the introduction section.\n4. The equation $f_C^D=TACO(f,S,C,D)$ is used for defining the input of TACO but seems express same information as the text above and not relevant to other equations below.", "broader_impact_concerns": "No concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Lw6R8JKuUQ", "reviewer": "Reviewer_rnwv", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\nS1. The topic is relevant, well-motivated, and clearly presented, making the paper easy to follow in terms of writing.\n\nS2. The proposed method is straightforward and easy to understand.\n\nS3. The quantitative results are promising. The authors demonstrate the method's efficacy across multiple model architectures and datasets.\n\n\nWeaknesses:\n\nW1. Missing Related Work Section -- The related work section is absent. Given that the paper addresses multiple dimensions, it is crucial to situate the proposed approach within the context of prior research. There is substantial existing work on both few-shot learning and model compression. A detailed comparison and discussion of related approaches are necessary.\n\nW2. Lack of Experimental Comparison with Prior Work -- While the method demonstrates promising quantitative results favorable compared to the proposed Post-Training Compression (PTC) baseline, comparisons with prior work or baselines constructed using existing methods are missing. Such comparisons are essential for contextualizing the contribution of this work.\n\nW3. Sample Dependence and Lack of Error Bars -- The method relies on a small number of samples (10 images) to perform task-aware compression. This raises concerns about the method's dependence on the chosen samples. Reporting error bars in the plot visualizations is critical to account for variability and enhance the reliability of the results.\n\nW4. Severe Formatting Issues -- The paper suffers from formatting issues, particularly in the figures and references.\n- Section headings (e.g., \"Method\") are placed too close to the preceding text, rather than the content they introduce.\n- Figures are poorly formatted. For example, Figures 2 and 5 have captions that are cryptic or insufficiently detailed, making interpretation difficult without checking the text in detail.\n- Figures lack consistency in width, font size, and alignment (e.g., Figure 2 vs. Figure 3).\n- Legends overlap with content despite there being white space available (e.g., Figure 4).\n- References are inconsistent in venue naming: for instance, Hubara et al. (2022) use \"NeurIPS,\" whereas Chen et al. (2021a) use \"Advances in Neural Information Processing Systems,\" and Dong et al. (2017) use \"Conference on Neural Information Processing Systems (NeurIPS).\"\n- Some references (e.g., Szegedy et al., 2015) lack venue information.\n- Capitalization of paper titles and inclusion of page numbers are inconsistent.\n\n\nMinor Remarks:\n- Abstract: The abstract is overly long and challenging to parse. It could be shortened and made more precise to align with the paper's overall length.\n- Typo in Section 4: The phrase \"with via\" appears in the second sentence of \"Application to Text-to-Image Generation.\"\n- Figure 11: The legend is positioned in the middle of the plot. Consolidating legends into a single, well-placed instance (e.g., in the plot on the right) would improve clarity.\n- Figures 13 and 14: These figures appear before the first section of the supplemental material, disrupting the logical flow.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes Few-Shot Task-Aware Compression (TACO), a method to compress large generalist vision models into efficient task-specific models using a small amount of calibration data. The approach combines task-specific pruning and few-shot fine-tuning with knowledge distillation to maintain accuracy while achieving significant parameter reduction and computational efficiency. The authors validate their method across multiple vision architectures (e.g., ResNet, ViT, ConvNeXt) and datasets (ImageNet, iNaturalist), demonstrating the effectiveness of the approach.\nThe paper presents an interesting and relevant approach to task-aware model compression. However, issues with the presentation and missing contextualization with related work weaken the overall quality of the paper.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I expect the four identified weaknesses to be thoroughly addressed, at which point I will reassess my evaluation. Specifically, the paper requires the inclusion of a related work section, comprehensive comparisons with prior approaches, reporting of standard deviation, and a fundamental revision of its formatting.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper has the potential to be a valuable contribution to TMLR by introducing a novel and practical approach to task-aware model compression. However, it currently suffers from weaknesses in evidentiary support, benchmarking, and presentation, as highlighted by the reviewers.  "Accept with Revision" is justified because the core idea is promising, and the identified weaknesses are addressable through focused revisions.  Addressing these points will significantly strengthen the paper and make it a more convincing and impactful contribution to the field.  If the authors can successfully address these revision requirements, the paper will be a candidate for acceptance in TMLR.
Decision: Accept (poster)