Learning from Offline Foundation Features with Tensor Augmentations

Authors: emir  konuk
christos  matsoukas
moein  sorkhei
phitchapha  lertsiravaramet
kevin  smith
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 13:22:48

Abstract: We introduce Learning from Offline Foundation Features with Tensor Augmentations (LOFF-TA), an efficient training scheme designed to harness the capabilities of foundation models in limited resource settings where their direct development is not feasible. LOFF-TA involves training a compact classifier on cached feature embeddings  from a frozen foundation model, resulting in up to  $37\times$ faster training and up to $26\times$ reduced GPU memory usage. Because the embeddings of augmented images would be too numerous to store, yet the augmentation process is essential for training, we propose to apply tensor augmentations to the cached embeddings of the original non-augmented images. LOFF-TA makes it possible to leverage the power of foundation models, regardless of their size, in settings with limited computational capacity. Moreover, LOFF-TA can be used to apply foundation models to high-resolution images without increasing compute.    In certain scenarios, we find that training with LOFF-TA yields better results than directly fine-tuning the foundation model.

Main paper: 
 introduction: 
 Large and expensive foundation models, designed to capture general-purpose knowledge, have become a significant focus in machine learning and computer vision research [3] . These models excel in zero-and few-shot learning [4, 23] and adapt to various domains, especially in data-scarce scenarios, through transfer learning. But adapting these models for a specific task is a resource-intensive process [35] . The cost of fine-tuning large foundation models today is already prohibitive to most individuals and organizations. As they continue to grow in size, the rising costs of foundation models risk excluding all but the wealthiest organizations. To mitigate this, parameter-efficient fine-tuning methods have been proposed. These methods incorporate rank-deficient [18] and simple affine modules [31] or learnable prompt parameters injected at the input stage [20] . Their core principle is to limit the number of parameters that need to be trained. We extend this principle to its logical conclusion by introducing no intermediate parameters or prompts, and investigate whether it is possible to completely separate the foundation model from the training process.In this study, we explore this complete separation of the resource-intensive foundation model from the training process. As seen in Figure 1 , training data is passed through the foundation model at a one-time cost, and then cached. The cached feature embeddings are later loaded and used to train a lightweight classifier. By adopting this caching strategy we can train at a significantly faster rate using less memory resources and achieve similar, 38th Conference on Neural Information Processing Systems (NeurIPS 2024). Training data is passed through a foundation model and cached. The cached embeddings are loaded and spatial tensor augmentations are applied in lieu of standard image augmentations.A lightweight classifier is trained on the cached, augmented features. This enables the use of arbitrarily large foundation models and high-resolution images at no additional cost.or even in some cases, better performance. This framework enables us to leverage the power of foundation models of any size in limited resource settings. Moreover, tasks requiring high-resolution imaging, such as medical image diagnosis can benefit from the power of foundation models without increasing computational costs. However, these benefits come at the cost of inference speed, which can be slower as a consequence of our paradigm.We call this approach "Learning from Offline Foundation Features with Tensor Augmentations", or LOFF-TA. It is a simple approach, but has not yet been explored as it is complementary to existing adaptation methods. LOFF-TA can be trivially combined with existing adaptation methods [18, 31, 20] by caching features from the adapted foundation models themselves to achieve even better performance. We conduct a series of comprehensive experiments using LOFF-TA on eleven well-known image classification benchmark datasets to demonstrate its benefits and limitations. Our key findings and contributions are outlined as follows:• We propose, LOFF-TA which decouples the training process from the resourceintensive foundation model -a classifier is trained on cached features from foundation models instead of images. • Since integrating image augmentations within LOFF-TA presents a challenge due to the tremendous storage cost of caching the embeddings of augmented images, we propose to apply spatial tensor augmentations to the cached embeddings of original images when training the compact classifier. We show that they perform nearly as well as standard image augmentations. • We show that, using LOFF-TA, it is possible to achieve similar performance to a fine-tuned foundation model at a fraction of the computational cost -training speed is accelerated up to 37×, and GPU memory usage is reduced up to 26×. Additionally, LOFF-TA allows flexibility in choosing any input image size or foundation model according to need and available resources. • Surprisingly, in some cases, LOFF-TA outperforms a fine-tuned foundation model. Despite the simplicity of our approach, our findings indicate that there are many potential benefits, in terms of both performance and economics, to training from cached foundation features. The source code used in this work can be found at https://github.com/emirkonuk/loffta. 
related work: 
 Foundation models [3] such as LLaMA [38] and GPT [4] have revolutionized natural language processing through large-scale training, requiring huge datasets and substantial computational resources. In a rapidly advancing competitive landscape, these models are scaling up dramatically to achieve unprecedented capabilities. In computer vision, this trend is mirrored with models like DINOv2 [35] , CLIP [37] , OpenCLIP [19] , BLIP [30] , SAM [23] and SEEM [47] .The standard approach for adapting a foundation model for a specific task involves finetuning its parameters or integrating additional trainable layers or classifiers [3, 13] . However, as the sizes of foundation models expand, these computational demands will escalate, posing significant challenges. To overcome these challenges, researchers have developed several strategies. These include fine-tuning only a subset of a model's parameters [46] and using techniques like gradient check-pointing for resource optimization [7] . Adaptation methods like [31, 17, 6, 18] introduce a limited number of tunable modules within the foundation model and incorporate learnable 'prompts' in the stem [20] . The result from these approaches is an efficiently fine-tuned, slightly bigger foundation model. Prior to deep learning, machine learning typically involved a two-stage process: first extract relevant features from data, then train a model on these features [2] . LOFF-TA, with its caching of features, echoes this latter approach [15] . As such, it complements the adaptation methods [31, 17, 6, 18, 20] and can be used in conjunction with them to improve overall performance, as we show in Section 5.5.Augmentations. Developing novel image augmentation strategies mostly rely on visual insights to design augmentations [45] . The underlying principle is the manifold hypothesis [5] and effective augmentations should not move samples too far from the image manifold. While previous work on feature augmentations relied on interpolating features of different images [42] and adding noise to these interpolations [11] to ensure the augmented features stay close to the manifold, we find that feature augmentations need not be so limited. In this work, we show that it is surprisingly beneficial to apply spatial augmentations to features, in a manner similar to image augmentations. 
methods: 
 In this work, we propose to pass training data through the foundation model, cache its features, and use them to train a lightweight classifier. We term this Learning from Offline Foundation Features (LOFF). A key challenge is LOFF's inability to incorporate image augmentations. To overcome this, we apply spatial and noise tensor augmentations directly to the pre-stored foundation features, resulting in LOFF-TA. Finally, to allow the use of high-resolution images, we pool the foundation features.Our approach embraces a straightforward yet powerful idea: to separate feature extraction using a powerful foundation model from the training process using a lightweight classifier, as depicted in 2. We use a foundation model to process training data upfront, then extract and store the output features (Step 1). These serve as rich representations of the data which can be retrieved at a later time for training one (or more) classifiers, as depicted in Step 2.In detail, given a dataset 𝐷, we generate and store a new dataset 𝑇 by applying a foundation model 𝐹 𝜙 to all images 𝑥 from 𝐷. Each sample 𝑡 in 𝑇 is a 𝑘 × 𝑑 dimensional tensor where 𝑘 is the number of tokens for each input image 𝑥 and 𝑑 is the embedding dimensionality.EQUATIONImportantly, this phase only needs to be performed once. The extracted features, 𝑇, along with the corresponding labels for each image, are cached and used as a replacement for 𝐷. The offline foundation features 𝑇 are retrieved to train a lightweight classification model 𝐹 𝜃 . Finally, we employ a standard cross-entropy loss, denoted as ℒ:EQUATIONwhere, 𝑦 is the true label, 𝑡 ∈ 𝑇 the stored features and Φ 𝑝𝑟𝑜 𝑗 is a projection module. The role of the projection module Φ 𝑝𝑟𝑜 𝑗 is to match the cached features to the dimensions expected by the classification model. Φ 𝑝𝑟𝑜 𝑗 consists of a learnable linear projection layer followed by a Layer Normalization (LN) operation [1] , which we found to improve performance.High-resolution images Many foundation models are trained to handle images larger than 224 × 224, such as [23, 35] , but GPU memory limitations on conventional hardware make it impossible for them to be fine-tuned on high-resolution images. LOFF can mitigate this issue by pooling the features before storing them as shown in Fig. 2 (Step 1). The reduced size of the pooled features allow them to be stored and used to train a classifier efficiently. We investigate both average pooling and max pooling operations [27] and assess their impact on classification performance and computational cost.Given that image augmentations are crucial for effective model training, the inability to apply them in the LOFF framework poses a significant challenge. The obstacle lies in the impracticality of storing tensors from augmented images during Step 1, which would result in prohibitive storage demands. To address this drawback, we introduce tensor augmentations 𝒜 tensor (𝑡) on the features 𝑡 ∈ 𝑇. A tensor augmentation module is applied dynamically online, before a batch of features is fed into the classifier, as depicted in 2. This method, named Learning from Offline Foundation Features with Tensor Augmentations (LOFF-TA), allows augmentations to be used for overcoming the aforementioned obstacle with image augmentations.In a standard setting, the objective for image classification is given by arg minEQUATIONwhere 𝐹 𝜃 is the model to be trained on each sample (𝑥, 𝑦) from the dataset 𝒟 and 𝒜 img indicates stochastic image augmentations. LOFF-TA changes the optimization task to become arg minEQUATIONwhere 𝒜 tensor denotes our tensor augmentation operator. The stored features 𝑡 ∈ 𝑇 and projection module Φ 𝑝𝑟𝑜 𝑗 remain the same as in LOFF.Table 1 : Main results. We train models on features extracted by DINOv2 [35] ViT-B and ViT-G models. We report the results using LOFF (no augmentations) and LOFF-TA (with tensor augmentations). We consider features extracted from 256 × 256 and 512 × 512 images (using pooling as described in 3.1). Frozen + linear and Unfrozen + linear are points of comparison consisting of a frozen/unfrozen foundation model with a linear layer trained on images directly, with image augmentations. Spatial relationships in image data are crucial for understanding the content of the image. Directly applying augmentations haphazardly to the unstructured output tokens from the foundation model may lead to undesirable results. A key aspect of our approach involves the utilization of spatial tensor augmentations during the training phase, as depicted in 2. These augmentations are chosen to consider the spatial relationships in the data, similar to how image augmentations operate. We apply spatial augmentations, denoted as 𝒜 tensor , to foundation features 𝑡 ∈ 𝑇 after a reshaping operation. These augmentations are conceptually analogous to image augmentations, treating the foundation features as if they were low-resolution, hyper-spectral images. We select a set of spatial augmentations suited for feature-level transformations, chosen to enhance the training process while maintaining the integrity of spatial relationships. We flip by mirroring the tensor on its height or width axis, resize it by upsampling or downsampling its spatial dimensions using linear interpolation. We shear the tensor using nearest neighbor interpolation and translate it by shifting along its spatial dimensions. We rotate the tensor in its spatial dimensions around its spatial center using nearest neighbor interpolation. In addition to spatial augmentations, we apply additive Gaussian noise with zero mean to the feature tensor, similar to [11] . Although we considered channel augmentations, analogous to contrast or color augmentations in images, we did not find them to be beneficial. We also note that, like image augmentations, not all tensor augmentations types are appropriate in every setting2 . 
experimental setup: 
 To evaluate the effectiveness of LOFF-TA we benchmark over eleven datasets from various domains using different foundation models, model capacities and image resolutions.Foundation models. We employ two foundation model families: DINOv2 [35] and CLIP [37] (implemented by OpenCLIP [19] ) as the basis for our investigations. For the majority of our experiments, we utilize the ViT-B and ViT-G architectures.Classifiers. LOFF and LOFF-TA train a lightweight classifier on the features from a foundation model. While, in principle, any classifier can be used in this role, our experiments use DeiT-S [39] . The classifier is initialized using ImageNet [10] pre-trained weights, as we found empirically this gave a significant improvement over random initialization [14] . In some In our experiments, we utilize the AdamW optimizer [32] , and a batch size of 64. We incorporate a learning rate warm-up strategy and manually decrease the learning rate by a factor of 0.1 when the validation performance plateaus. For lightweight classifier in LOFF and LOFF-TA, we implement modifications to the DeiT-S architecture [39] . We remove the patchifier from the model's stem and introduce a linear projection layer followed by a normalization layer as detailed in 3.Our evaluation spans eleven image classification datasets, covering a diverse spectrum of object categories. We begin with datasets with high-resolution images, as this setting highlights new capabilities made possible by LOFF-TA. We include APTOS2019 [21] for diabetic retinopathy detection, DDSM [29] for identifying masses in mammography, ISIC [40, 8, 9] for skin lesion classification, AID [43] for aerial image classification, and NABirds [41] for fine-grained bird species classification. The resolution of these datasets varies, but we resize them to 512 × 512. We extend our evaluation to a number of standard 256 × 256 resolution benchmark datasets: Flowers102 [34] , NABirds [41] , StanfordCars [26] , StanfordDogs [22] , Oxford-III Pet [36] , Caltech-101 [12] , and SUN397 [44] . For each dataset, we report metrics appropriate to its specific evaluation criteria. We adhere to official train/validation/test splits when available, or follow [24] in their absence. Standard practice of image normalization is maintained, using the mean and variance from the training sets. 
results: 
 In this section we show that LOFF-TA achieves competitive, sometimes superior, results compared to the baselines while significantly reducing memory usage and training time.We begin our analysis with Table 1 where we focus on the cost-benefit trade-off of using LOFF-TA, considering DINOv2 as the foundation model. Our evaluation spans five datasets where higher resolution images are known to improve performance. We measure each approach in terms of performance, throughput (TP), and memory (Mem.) footprint. The study includes a comparison with two key configurations: a baseline approach, where a linear classifier is appended to a frozen foundation model (frozen + linear), and an upperbound where the entire foundation model is fine-tuned in a typical fashion for transfer learning (unfrozen + linear). Both employ standard image augmentations. In contrast, LOFF operates without augmentations, while LOFF-TA introduces tensor augmentations. We also explore the effects of pooling on foundational features, increasing image resolution from 256 to 512, to understand how these adjustments influence the performance and efficiency.Our observations reveal several trends. Firstly, all LOFF and LOFF-TA variants (with an exception of NABirds) surpass the baseline in performance, with only a slight increase in memory usage but significantly faster training. Secondly, LOFF-TA consistently outperforms LOFF, confirming the importance of tensor augmentations. Moreover, upgrading to a larger foundation model (from ViT-B to ViT-G) doesn't alter memory or training speed. Pooling and working with higher resolution further improves performance, again without affecting memory or throughput. Remarkably, in many cases (6 out of 10), LOFF-TA with pooling exceeds the performance of the intended upper-bound model (unfrozen + linear). Finally, it's worth noting the training speed and memory efficiency of LOFF and LOFF-TA compared to fine-tuning the foundation model (unfrozen + linear); our approach offers a 37× acceleration during training, and a remarkable savings in memory as well -in fact, ViT-G is too large to fine-tune on a conventional GPU with a batch size of 64.In Table 2 , we continue our analysis in a similar fashion to Table 1 , but in this case focusing on seven standard visual object recognition datasets. Our focus remains on assessing performance, throughput (TP), and memory usage (Mem.), but with image resolution of 256. The analysis compares LOFF, LOFF-TA, frozen + linear, and unfrozen + linear, and introduces frozen + DeiT-S as a benchmark for image vs. tensor augmentations (see Section 5.3). The findings in Table 2 reinforce the patterns observed in Table 1 . LOFF-TA maintains its superiority over LOFF, highlighting the efficacy of tensor augmentations. It's important to note that in this set of experiments, the best performing LOFF-TA + Pooling configuration using high-resolution images is not considered due to the datasets' inherent resolution limit of 256. Despite this, LOFF-TA generally matches or surpasses the baseline frozen + linear in performance, while delivering a notable increase in throughput. Although frozen + linear claims the smallest memory footprint, LOFF and LOFF-TA again show significant savings compared to fine-tuning a foundation model (unfrozen + linear). Comparing foundation models of similar capacity, we notice that DINOv2 outperforms OpenCLIP in five of the seven datasets using LOFF-TA. Finally, we once again see that swapping the foundation model (between DINOv2 ViT-B, OpenCLIP ViT-B, and DINOv2 ViT-G) results in no appreciable change in throughput or memory consumption for LOFF-TA, only differences in performance.Looking at Table 2 again, we investigate the effect of image augmentations versus tensor augmentations. LOFF-TA and frozen + DeiT-S share the same architecture (a cascaded model that consists of a frozen foundation model with an appended DeiT-S), and in both cases the foundation is frozen -the only difference is LOFF-TA uses tensor augmentations while frozen + DeiT-S uses image augmentations. We also report results (Table 5 in the Appendix) when we unfreeze the foundation model in the cascaded setting for both OpenCLIP and DINOv2 for completeness. In Table 2 , we observe that image augmentations outperform tensor augmentations. As one might expect, tensor augmentations such as tensor rotation or cropping can not replicate the exact effects of image rotation or cropping since foundation models are not linear operators. Yet, the performance impact is surprisingly less than anticipated, which, considering LOFF-TA's significant computational savings, is noteworthy.In this section, we try to understand the impact of each contribution to LOFF-TA. Utilizing DINOv2 [35] , we conduct a study where components are systematically removed and then tested on the Oxford-III Pet and Caltech-101 datasets. Results are presented in Table 3 .A perhaps obvious, but critical, insight from our study is the important role played by the foundation model's CLS token in contributing to classifier performance. When this token, representing global information, is integrated into the classifier training, it proves to be a key contributor to performance gains. However, the choice of how to integrate it was not trivial. We opted to integrate the offline CLS token from the foundation model with the learned CLS token of the classifier model by summation. Other possible ways to incorporate it could be to initialize the classifier's CLS to the offline foundation CLS, or concatenate them -although our non-exhaustive experiments testing these approaches were inferior.Layer norm. We applied a Layer Norm [1] operation in the projection layer at the stem of the classifier. Similar to the CLS token ablation, we find that adding a layer norm operation before passing the foundation features to the classifier boosts performance. By centering input features, we believe the normalization plays a pivotal role similar to that in Dual Patchnorm [28] , aligning and standardizing the input.Pooling. Table 1 shows that pooling larger image features enhances performance without extra computational costs, enabling the use of large models for high-resolution tasks. In Appendix Table 6 , we compare max and average pooling, finding both perform similarly, with max pooling slightly outperforming in larger models, possibly due to better noise reduction. Thus, max pooling is recommended for larger models. While this work focuses on pooling, alternative dimensionality reduction methods like strided convolutions or bi-linear interpolation could also be effective.Augmentations. Since the cached features contain long-range spatial information that may potentially be harmed by our tensor augmentations, it is important to assess if these augmentations have a meaningful contribution to the performance of LOFF-TA. Comparing Row 3 and Row 4 of Table 3 , we see that adding spatial augmentations results in a significant boost in performance. A smaller boost is observed when Gaussian noise is added in Row 5. Replacing our augmentations with Trivial augment [33] provides a further boost to performance, although this setup was not used in our other experiments.LOFF-TA enables practitioners to use large foundation models without any modifications. A class of methods exists that modify foundation models for new tasks, so-called adaptation methods. Adaptation methods can be used in conjunction with LOFF-TA, but we also provide a comparison between them in Table 4 . The results show that standalone LOFF-TA achieves competitive performance with VPT [20] , Adaptformer [6] and SSF [31] . When LOFF-TA is combined with these methods, we observe noticeable improvements across the board, demonstrating that computationally efficient SOTA performance can be achieved by adapting foundation models and then applying LOFF-TA. 
discussion: 
 Performance-efficiency trade-off. The aim of this paper is to consider the use of foundation models in a resource-limited setting. To do so, we propose to work with cached foundational features. The benefits of doing so are that training speed is significantly accelerated (up to 37×), and GPU memory usage is significantly reduced (up to 26×). Of course, these benefits come with trade-offs. In most cases, LOFF-TA will perform slightly worse than directly fine-tuning the foundation model, although we did observe some cases where LOFF-TA was superior. Also, the computational cost savings at training time do not translate to inference.An important benefit of LOFF-TA is that it affords flexibility, allowing practitioners to choose the most suitable foundation model for their task from a range of sizes and pretraining methods (refer to Table 2 ), including the capacity to handle high-resolution images (see Table 1 ). The caching and pooling of features makes these different choices come with equivalent computational costs during training, determined by the compact classifier.As foundation models grow, their computational demands are expected to increase, potentially creating a digital divide where advanced models are accessible only to well-resourced organizations, excluding smaller entities and individual researchers. This trend underscores the importance of developing efficient methods to utilize these models, ensuring broad access and fostering progress in various application areas. There is significant potential to impact fields like medical imaging and remote sensing, which can greatly benefit developing regions and marginalized groups. But these applications often require high-resolution image processing, which is resource-intensive. LOFF-TA facilitates wider, fairer access to sophisticated foundation models in resource-limited settings.Limitations: LOFF-TA is a training strategy meant for low resource settings. Given enough computational resources, full fine-tuning will generally outperform LOFF-TA or adaptation methods like SSF [31] . Another limitation of LOFF-TA is that at inference time it is less efficient than standalone foundation models. However, we note that due to their latency, foundation models are costly to be deployed at scale without distillation [16] (or some other measure) to limit cost which would also alleviate the limitation for LOFF-TA. 
conclusion: 
 Foundation models have significantly impacted the community and will likely become more crucial as they grow in size and capability. However, adapting these models for specific tasks is increasingly challenging due to their high resource demands. Revisiting the classical machine learning paradigm of perceive, then reason is pragmatic in this context. Using foundation models as powerful feature extractors allows us to benefit from their rich representations, while at the same time mitigating computational costs by training a lightweight classifier on their cached features. LOFF-TA achieves a low memory footprint and high throughput during training, regardless of the chosen foundation model or input image size. It is particularly beneficial for high-resolution image processing when pooling is applied, allowing LOFF-TA to outperform linear classifiers using frozen foundation models and compete with or surpasses fine-tuned foundation models.  
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: We provided a societal impact statement addressing potential positive and negative impacts of the study.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). 
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 study is about efficient training strategies and it does not pose a high risk for misuse.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 
 
 Appendix: We evaluate the model's performance in the cascaded setting (foundation model + DeiT-S classifier) when the whole model is fine-tuned. We present our findings in Table 5 , comparing the use of DINOv2 and OpenCLIP as foundation models. Once again, we find that DINOv2 outperforms OpenCLIP, in most cases. When DINOv2 is unfrozen, it either surpasses or matches the performance of its frozen alternative (see Table 2 ), though the margin of improvement is rather small. OpenCLIP shows a similar trend, but its performance seems more dependent on the dataset it is evaluated on. Overall, unfrozen foundation models appear to outperform their frozen alternatives. However, their computational cost makes them prohibitive for larger models or image sizes.Table 5 : Fine-tuning a foundation model + DeiT-S classifier. We append a DeiT-S classifier after a DINOv2 [35] or a OpenCLIP [19] foundation model and the whole cascaded model. In our experiments, we consistently observe a performance improvement when applying tensor augmentations. However, the need for task-specific and domain-appropriate augmentation strategies should be emphasized. For example, in scene classification using the SUN397 dataset, applying vertical flips to images when training DeiT-S (pretrained on ImageNet) led to decreased accuracy (-2%), likely due to the unrealistic expectation of upside-down building facades during testing. One might expect a similar effect with tensor augmentations, since the feature space preserves the spatial orientations of objects (see Figure 4 ). Interestingly we see a much smaller performance drop when we apply vertical flips as tensor augmentations on this dataset (-0.2%). This suggests LOFF-TA's tensor augmentations exhibit greater resilience to improper augmentations compared to the image domain. This observation merits further exploration to understand its underlying mechanisms and consequences.Intriguingly, LOFF-TA demonstrates benefits from spatial tensor augmentations despite their potential conflict with the spatial information present in feature tokens. Feature tokens in foundation models carry positional data, informed by positional embeddings and attention mechanisms. Spatial augmentations, such as horizontal flips, disrupt this positional context, yet our experiments, especially with Trivial Augment, show a notable performance enhancement (see Table 3 ). This paradox suggests that these disruptions may actually bolster the classifier's ability to learn robust features, thereby improving classification outcomes.The precise dynamics of this phenomenon remain unclear, presenting an exciting avenue for future research, particularly in the realm of auto augment strategies for foundation features. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).It was not possible to train ViT-G on a single GPU with batch size of 64. Instead we report the memory footprint across 8 NVIDIA Quadro RTX 8000 using distributed training.We omit vertical flips for the SUN397 dataset[44]. 


Decision: The reviewers agree that the manuscript "Learning from Offline Foundation Features with Tensor Augmentations" (LOFF-TA) presents a valuable and innovative contribution to the field. The proposed approach, which enables efficient training of classifiers using cached embeddings from a foundation model, shows significant potential, particularly in resource-constrained settings. LOFF-TA not only substantially reduces computational and memory requirements but also demonstrates competitive performance across a variety of benchmarks.

The reviewers commend the authors for clarifying the suitability of LOFF-TA for high-resolution images, highlighting that the method’s computational efficiency and reduced memory usage make it especially advantageous in these scenarios. Furthermore, the reviewers appreciate the authors' commitment to extending their work to include experiments on 3D volumes and to visualizing the latent feature spaces. These additions will undoubtedly strengthen the paper’s contributions.

The reviewers encourage the authors to incorporate the proposed revisions and additional experiments into the final manuscript. Doing so will enhance the clarity of the work and further validate the claims made.

In summary, the paper presents a promising approach with the potential to significantly impact the field, especially in environments with limited computational resources. The reviewers recommend accepting this paper, with the expectation that the authors will implement the proposed improvements.

Reviews: review: paper_summary: paper_summary: This paper proposed a framework to efficiently use foundation features for online serving. The idea (offline feature extraction from foundation model and online inference with lightweight model) is simple but interesting/effective.
strength_weakness: strength_weakness: 1. Can you please list the classification accuracy of the raw visual foundation models (VFM)? If the performance of the raw VFM is high, we can directly output the probability score and save them in a database instead of extracting features offline and then do online inference with a followup lightweight model.
2. I'm concerned the potential application scenarios of the proposed scheme is limited, since it may only be able to be leveraged on image classification scenarios.
3. When the authors claim efficiency improvement, the baseline is somewhat unfair, because the proposed methods actually need an offline feature extraction step which would also consume computation resources.
4. The tensor augmentation is actually limited to several spatial augmentations, while there are many for raw image classification, for example, color-jittering and mixup..
questions: questions: 1. I agree the tensor augmentation is reasonable for image embeddings (such as flip, rotate, translate and resizing). Is there any way to expand the the tensor augmentation w.r.t spatial operators to the video embeddings?
limitations: limitations: N/A
paper_summary: paper_summary: The authors propose to store low dimensional representations of images that are obtained after passing them through pretrained (foundation) models. Only afterwards augmentation strategies ("tensor transformations") for training the downstream classifier are employed. This has the benefit that the foundation model is not part of the training processm, which is computationally cheaper. More importantly, some dataset specific optimization is obtained, without the need of finetuning the complex / expensive foundation model itself.  The authors present a set of relevant tensor transformation and compare against the baseline.
strength_weakness: strength_weakness: * It would have been interesting to study tensor augmentations in a more systematic / more formal way
* I am not sure why this approach would be particular suited for high-resolution images. The authors describe how foundation model optimized for large images can be used to deal with small images via the offline eature / tensor augmentation approach. This implies that the presented method is particulary well suited deal with_low-resolution_ images, but not the opposite, isn't it?
* Using a pretrained (foundation) model as a feature extractor is not necessarily an overly innovative idea.
questions: questions: * See above
limitations: limitations: All very well done.
paper_summary: paper_summary: This paper introduced a training scheme considering offline foundation features with tensor augmentations LOFF-TA, focusing on limited resource settings. They basically trained a classifier on cached features from frozen foundation models with an augmentation process to these cached embedded features. Their proposed training mechanism is claimed to speed up the training process by $37\times$ with appx. $26\times$ reduction of GPU memory usage. The authors validated their approach on eleven image classification benchmarks and showed similar performances as standard image augmentations and fine-tuned foundation models.
strength_weakness: strength_weakness: **W1. Adaptation of spatial augmentations.** For high-resolution images, performing augmentations on cached embedded features highly depends on the size of the latent space from where the offline foundation features are being cached. Typically in the medical image domain where one needs to deal with very high-resolution images/volumes (e.g., $256\times256\times256$), it is necessary to consider the topological features when we generate/augment images from raw data that should not change the image class. In that case, where the authors are performing spatial augmentations on the latent features, there might be scenarios where the topological features get distorted from the flip or the classes change from other spatial augmentation operations. Though I understand that spatial augmentations might be beneficial in terms of natural images, these kinds of augmentations might suffer from experimenting on high-dimensional medical images. 


**W2. LOFF-TA generalizability on high dimensional 3D volumes.** I appreciate the authors for experimenting with diverse 2D datasets ranging from natural to medical images. However, in the medical domain, we typically follow 3D volumes for complex disease diagnoses such as classifying tumors, neurodegenerative disease, myocardial infarction, etc., which are considered to be high dimensional, e.g., $128^3$ or $256^3$, where it is pretty complex to preserve anatomical information by augmenting on low-dimensional latent features. Experimenting with this type of application on datasets like OASIS, ADNI, ACDC, Synapse multi-organ, etc. would further strengthen the argument about offline tensor augmentations on high-dimensional medical images.


**W3. Pictorial latent space depictions.** To validate the proposed model's efficiency compared to image augmentations, the authors might want to visualize the latent feature spaces on both image/tensor augmentations. This would further ensure that offline latent tensors preserve the important spatial features similar to the latent image features. After that, augmenting the offline tensors would make more sense in the carried-out image analysis tasks.
questions: questions: Please see the weaknesses section. I tried to discuss all the findings and questions there. Following, please find some of my concerns and suggestions the author might want to rebut in the rebuttal

1. The validation of their model on high-dimensional medical images is missing. The author might want to show very brief experiments on how their proposed approach performs on 3D medical classification tasks. 

2. How the offline tensor augmentations preserve the anatomical features compared to the online tensor/image augmentations? Consider the small anatomical features such as ventricles that are very important when we experiment on finding the presence of neurodegenerative disease from 3D brain MRIs.
limitations: limitations: The authors carefully addressed the limitations and societal impact of their work.
paper_summary: paper_summary: The paper discusses leveraging the large foundational models in resource-scare settings to fine-tune an image classification task. The authors propose saving the features from the frozen foundational model and then using them to train/fine-tune a smaller classifier model. They introduce tensor augmentations to improve the method's robustness and demonstrate the advantages of reduction in GPU memory usage and improved training speed and performance on different open-source dataset tasks. The performance of the proposed method is comparable to the full fine-tuning of the foundations model along with a linear layer.
strength_weakness: strength_weakness: - Table 1 results. Question about the inference throughput. The table shows a considerable reduction compared to the original model with linear. I assume that inference is done end-to-end, i.e., the image passed through the Foundational Model + the extra bits of the proposed LOFF-TA/LOFF pipeline. While the improved training speed is appreciable, the drop in inference is much less desirable, especially when deployed in scenarios that require faster throughput.  So although the resource saving is highly beneficial for the training scenarios, unfortunately for inference the whole system still needs to be loaded onto the GPU because caching is not possible. I would appreciate the author's insight on this. 

- The effectiveness of these augmentations depends heavily on the type of features extracted from the foundation model. Specifically, different layers within a neural network capture varying level of information: earlier layers focus on the spatial structure of the image, while deeper layers learn more abstract and semantic features. The paper needs to clarify what layer of the ViT features were extracted and cached. Further, there is a strong dependency on this choice of layers, and I don't see any ablations on this choice. 

- It's interesting that pooling actually improves the results. It's surprising because I would assume pooling reduces some of the features when reducing the dimension, but it somehow actually helps. Can the authors comment on this? It's a little counterintuitive, in my opinion. 

- Authors propose a $\phi_{proj}$ adapter to match the cached features to the dimensions expected in the classifier model. Is the layer trained end-to-end along with the classifier?
questions: questions: See above
limitations: limitations: There is a lack of clarity in some of the choices in the design. I have a few questions on some of the results and the missing ablation experiments.

Comments: comment: There may be a misunderstanding regarding the results.

> But I don't agree with the fact that LOFF-TA is necessarily a better learner than a simple linear.

The opposite is what we have shown empirically in our main results, in Table 1. LOFF-TA is consistently better than a simple linear classifier, i.e. the first rows of our tables.

> The way I see it if you put a few a few more trainable ResNet blocks beyond the last frozen layer, would it probably match the accuracy LOFF-TA?

We already give an answer to this criticism in the “Frozen+DeiT-S” rows in our Table 2 (not with ResNet layers, but with DeiT layers). LOFF-TA results are generally on par with such a setup, but LOFF-TA is far less computationally demanding. More importantly, this setup still prevents training with high resolution images due to prohibitive computational costs.

> Yes the augmentations are a nice addition in a unique feature space. But it would justify changing my rating. I would like to keep my rating but I definitely think this is interesting work the authors should build upon.

We thank the reviewer. However, based on the stated weaknesses and our discussion, we do not think their comments justify their rating. We kindly ask them to reconsider if our paper warrants a “3” which according to the guidelines indicates “technical flaws”, “weak evaluation”, “inadequate reproducibility” or “incompletely addressed ethical considerations”.
We are struggling to understand what the reviewer’s criticisms are.

> the paper should better thoroughly consider the different types of features, i.e., spatial detailed features and rich semantic features.

We do consider both features, i.e. spatial features are the features we augment (Section 3.3), and rich semantic features (we assume the reviewer is referring to the CLS features), are also included during training with a subsection (lines 257-265) detailing how different inclusion strategies might impact performance.

> the paper should formerly discuss the augmentation types which can be used in TA and which cannot be used now.

There must be a typo in the reviewer’s comment. We believe the reviewer means “the augmentation types that could be used in pixel space that cannot be used in tensor space”. We discussed these issues in Lines 158-162, Section 5.3, and Appendix C.
Thanks for the authors' response. I agree that it is also very valuable to address image classification tasks. However, if it is only for image classification, the paper should better thoroughly consider the different types of features, i.e., spatial detailed features and rich semantic features. Also, the paper should formerly discuss the augmentation types which can be used in TA and which cannot be used now.
I thank the reviewers for their comments and addressing some of my comments on the weakness. I appreciate the clarifying comments. I would like to follow up on one of their comments: 

> Therefore, for practitioners planning to use foundation models, such as in medical imaging, latency should not be a primary concern from the start. It's important to recognize that in such applications, the challenge isn't inference throughput but rather the capacity to learn a model that can handle high-resolution images effectively. 

> Moreover, if we solely focus on real-time applications, foundation models typically require distillation or similar techniques to enhance inference latency for feasible deployment. Should distillation be employed, our approach would match the latency of a standalone foundation model.

Agree that accuracy might be more important than inference latency in medical applications. But I don't agree with the fact that LOFF-TA is necessarily a better learner than a simple linear. The way I see it if you put a few a few more trainable ResNet blocks beyond the last frozen layer, would it probably match the accuracy LOFF-TA? i.e. extracting out the improvements by the tensor augmentation + projection layer, versus a ResNet block + projection block. I feel like they are similar learners. Yes the augmentations are a nice addition in a unique feature space. But it would justify changing my rating.  I would like to keep my rating but I definitely think this is interesting work the authors should build upon.