Cross-Layer Cache Aggregation for Token Reduction in Ultra-Fine-Grained Image Recognition

TL;DR

The paper tackles the high computational cost of Vision Transformer–based Ultra-Fine-Grained Image Recognition (UFGIR) by introducing Cross-Layer Cache Aggregation (CLCA), a plug-and-play framework that combines a Cross-Layer Aggregation (CLA) head with a Cross-Layer Cache (CLC) to mitigate information loss during token reduction. CLCA enables very low token keep rates (as low as $0.10$) while maintaining or improving accuracy, by aggregating intermediate features from multiple encoder groups and recovering lost information via the CLC. Extensive experiments across 5 UFGIR datasets, 9 backbones, 7 TR schemes, 5 keep rates, and 2 image sizes (over 2000 runs) demonstrate robust improvements and favorable FLOPs-accuracy trade-offs. The approach is modular and compatible with diverse TR methods, offering a practical path to scalable, high-accuracy UFGIR.

Abstract

Ultra-fine-grained image recognition (UFGIR) is a challenging task that involves classifying images within a macro-category. While traditional FGIR deals with classifying different species, UFGIR goes beyond by classifying sub-categories within a species such as cultivars of a plant. In recent times the usage of Vision Transformer-based backbones has allowed methods to obtain outstanding recognition performances in this task but this comes at a significant cost in terms of computation specially since this task significantly benefits from incorporating higher resolution images. Therefore, techniques such as token reduction have emerged to reduce the computational cost. However, dropping tokens leads to loss of essential information for fine-grained categories, specially as the token keep rate is reduced. Therefore, to counteract the loss of information brought by the usage of token reduction we propose a novel Cross-Layer Aggregation Classification Head and a Cross-Layer Cache mechanism to recover and access information from previous layers in later locations. Extensive experiments covering more than 2000 runs across diverse settings including 5 datasets, 9 backbones, 7 token reduction methods, 5 keep rates, and 2 image sizes demonstrate the effectiveness of the proposed plug-and-play modules and allow us to push the boundaries of accuracy vs cost for UFGIR by reducing the kept tokens to extremely low ratios of up to 10\% while maintaining a competitive accuracy to state-of-the-art models. Code is available at: \url{https://github.com/arkel23/CLCA}

Figures (4)

• Figure 1: Accuracy vs FLOPs for DeiT3 touvron_deit_2022 with EViT liang_evit_2021 token reduction method with (CLCA) and without (Base) our proposed modules on the SoyGene dataset yu_benchmark_2021 as a function of different keep rates (10, 25, 50, 70 and 100%) and image sizes (in parenthesis).
• Figure 2: From top to bottom, on the left is the overview for a ViT with the proposed Cross-Layer Cache Aggregation (CLCA). A learnable CLS and Cross-Layer Register (CLR) tokens are attached to a sequence of local features corresponding to image patches. This sequence is passed through a series of ViT Encoder Groups equipped with the proposed Cross-Layer Cache (CLC). Each Encoder Group is composed of a series of transformer encoder blocks followed by a reduction block that drops tokens at the end. After each encoder block we cache the spatial Global Average Pooling (GAP) of the local features, the tokens surrounded by the green dashed box, and the last token in the sequence (the CLR). After the reduction block we recover intermediate features from previous layers from the CLC, empty it, and then cache the current layer outputs. We repeat the process for each encoder group. Thereafter, the CLS token outputs of each Encoder Group are processed by the Cross-Layer Aggregation (CLA) Classification Head. These CLS tokens are first normalized using a BatchNorm layer, then forwarded through a Depth-Wise Convolution layer that combines cross-layer features for each channel separately. A second BatchNorm and activation function are employed for learning non-linear interactions between intermediate features. Finally, the vector of features is passed through a Point-Wise Convolution that outputs the classification predictions.
• Figure 3: Distribution of accuracies for multiple keep rates (25, 50, and 70%) on 9 different pretrained backbones as a function of the TR selection scheme (and the baseline without TR) on SoyLocal and SoyGene datasets with and without proposed CLCA.
• Figure 4: Ablation showing max magnitude of gradient across all layers of the model during training for DeiT3 B-16 touvron_deit_2022 on SoyLocal dataset with our proposed modules and the corresponding top-1 accuracy.