SCHEME: Scalable Channel Mixer for Vision Transformers

We thank the reviewer for their comments and useful feedback.

$W1$: We have included comparisons with the methods mentioned by the reviewer in Table 4. We observe that the MLP mixer methods such as gMLP and S2-MLP generally exhibit lower accuracies compared to their attention-based counterparts as typically noted in the ViT literature. While MogaNet demonstrates improved accuracy, it is significantly slower than SCHEME and other similarly accurate variants due to the additional computational overhead from multiple split and depthwise convolution operations. We also wish to clarify that since SCHEME is a generic framework, it can be readily integrated into architectures like gMLP or S2-MLP.

$W2$: Due to computational resource constraints, we are unable to scale up models beyond 60 M parameters and so we restricted our focus on models below this threshold. In the revised version, we have included additional comparisons in Table 10, incorporating more recent variants, bringing the total to 10 different ViT model evaluations. We believe this demonstrates the generalizability and scalability of SCHEME, which can be readily extended to models exceeding 60M parameters.

$W3$: Thank you for pointing out the typo. We fixed it in the revision.

We thank the reviewer for their comments and useful feedback.

$**W1**$: We apologize for the formatting issues caused by the small font size (due to the large number of comparisons presented). In the revised version, we have updated the figures with larger font sizes, and we have also included the zoomed-in versions in Figures 5 and 6. These updates will also be incorporated in the final version.

$**W2**$: Thank you for your comment. Our approach is motivated based on the observation for ViTs whose performance being highly sensitive to the expansion ratios in MLP/FFN blocks, making them ideal candidates for evaluating SCHEME. Since our proposed model builds on the Metaformer architecture, we designed and evaluated a family of models using this baseline to ensure a thorough comparison.

To address the reviewer’s concern, we have expanded our analysis to include convolution-based models such as CAFormer, which replaces attention with convolution, and CoAtNet, a hybrid model combining convolution and attention. As demonstrated in Table 10 of the appendix, the SCHEME module consistently improves the accuracy-computation tradeoff for both models, delivering higher throughput with fewer FLOPs and parameters for the same accuracy, thereby validating its broader applicability beyond ViTs.

$**W3**$: Section 4.1 (L361) outlines the hardware specifications and benchmark configurations used for all latency and throughput analyses, performed at the standard 224x224 input resolution. To ensure fair comparisons, all models were tested under identical configuration settings and in the same environment.

$**W4**$: Thank you for the suggestion. We have expanded the discussion of the differences between CCA and XCiT in the related work section (L180-190) for clarity. Below is a summary of the key distinctions.

XCiT utilizes a cross-covariance attention (XCA) operator, which is conceptually similar to a "transposed" version of self-attention and operates across feature channels. The XCiT architecture consists of three main components: the XCA operation, a local patch interaction module, and a feedforward network. The XCA mechanism computes a covariance operation across different head groups, analogous to multi-head attention.

In contrast, we introduce Channel Covariance Attention (CCA) block, which facilitates feature mixing across different channel groups. Unlike XCiT, CCA does not use "heads" or projection matrices for queries, keys, and values. Instead, it leverages the full feature set to directly compute covariance, allowing for a more holistic representation of inter-feature interactions.

Regarding the comparison with linear attention mechanisms in NLP, we would like to clarify that CCA differs in that it retains the use of softmax across the channel features, resulting in a quadratic complexity with respect to the feature dimension. Linear attention, on the other hand, reduces the complexity from O(N^2) to O(N) by expressing self-attention as a linear dot product of kernel feature maps, utilizing the associativity property of matrix products, where N is the token length.

$**W5**$: We have provided a qualitative analysis of the feature representations in Figures 9 and 10, which illustrate that the model incorporating CCA focuses more effectively on the regions where the target class is present, demonstrating the effectiveness of the proposed method. Additionally, the quantitative ablation studies presented in Tables 9 and 18 show that the feature clusters formed by CCA learn a more diverse set of features that complement each other, leading to a 1% improvement in ensemble accuracy. This further highlights SCHEME’s influence on representation learning.

$**W6**$: We included Table 10 in the appendix that shows the performance of SCHEME across 10 different ViT models including the latest variant BiFormer. Unfortunately due to the limited computational resources available in our lab, we are unable to include the comparisons to other ViT variants mentioned by the reviewer. We will include them in the camera-ready version after acceptance. We believe the proposed SCHEME module will generalize effectively to the latest variants, as demonstrated by its performance across 10 different ViTs in Table 10.

[1] Lei Zhu, Xinjiang Wang, Zhanghan Ke, Wayne Zhang, and Rynson Lau BiFormer: Vision Transformer with Bi-Level Routing Attention, CVPR 2023

We thank the reviewer for their comments and useful feedback.

$**W1**$: In our work, we demonstrated that SCHEME effectively improves the accuracy-latency trade-off across 10 different ViTs, as summarized in Table 10. Additionally, we showed its improvements in Object Detection and Semantic Segmentation tasks using datasets distinct from ImageNet, as detailed in Tables 5 and 6. These results collectively show the generalizability of SCHEME across diverse models, datasets, and tasks, aligning with standard benchmarks commonly used in the ViT literature.

$**W2**$: As outlined in Equations 3 and 4, SCHEME employs two hyperparameters, $g_1$ and $g_2$, to define the group structure. These parameters can be adjusted to balance computational complexity and performance. In Table 4, we presented results for configurations with $g_1 = 4, g_2 = 4$ and $g_1 = 1, g_2 = 2$, chosen to ensure comparable complexity to existing ViTs. Both configurations demonstrated improved accuracy, confirming the effectiveness of our approach. For future work, we plan to incorporate Neural Architecture Search to find the optimal architecture based on the parameter-FLOPs trade-off which is beyond the scope of the current work.

As explained in Section 3.2 (L229), splitting the feature vectors into disjoint groups without inter-group communication reduces accuracy, as shown in Table 7, where the model without CCA achieves 0.6% lower accuracy. This highlights the importance of feature mixing for achieving high performance.

$**W3**$: Regarding the motivation for CCA, as discussed in Section 4.2 (L428, L448), its design has a regularizing effect that promotes the formation of feature clusters within transformer blocks (Tables 9 and 18, Figures 9 and 10), leading to improved accuracy. We also addressed the impact of the mixing weight $\alpha$ in response to reviewer N5s3, further supporting the benefits of CCA in enhancing model performance.

We thank the reviewer for their comments and useful feedback.

$**W1**$: We apologize for the grammatical errors and formatting issues caused by the small font size (due to the large number of comparisons presented). In the revised version, we have fixed the errors, increased the font size in the figures, and we have also included the zoomed-in versions in Figures 5 and 6. These updates will also be incorporated in the final version.

$**W2**$: Table 10 in the appendix demonstrates the improvement achieved by the SCHEME module across 10 different ViTs. We have expanded the comparison to include three additional models, including original and widely used ViTs such as ViT, DeiT, and the latest ViT variant, BiFormer [1]. The table clearly shows that the SCHEME module consistently enhances accuracy (upto 1.5%) across all models while either maintaining or improving the image throughput.

Our proposed model is based on the Metaformer, and we introduce a family of models built on this baseline, comparing them extensively against it. Additionally, we compare our approach with prominent and influential ViTs in the literature, such as Swin, CS-Win, and DaViT, as shown in Table 10. However, given the large number of Vision Transformer variants in the literature (e.g., PVT, Uniformer) and the computational budget constraints for running these models, we could not include more methods in our comparison. We will include them in the revised version after acceptance. Nonetheless, we believe the proposed SCHEME module will generalize effectively to the latest variants, as demonstrated by its performance across 10 different ViTs in Table 10.

$*What is the interpretation of the gradual decay of the mixing weight ((1 -$\alpha$)) over training epochs, and how does it affect model performance or convergence?*$

The gradual decay of the mixing weight $(1 - \alpha)$ over training epochs reflects the model's shift from relying heavily on inter-group communication (enabled by the mixing) to prioritizing the independent learning of feature clusters within each group. This facilitates the emergence of more diverse and complementary features within the groups, improving the model's overall performance. We show that Channel Covariance Attention (CCA) promotes the clustering of features into naturally independent groups, improving class separability, as demonstrated in Tables 9 and 18. Experimental results show that models trained with CCA achieve higher ensemble accuracy (+1.09%/+0.6%) and exhibit superior class separability. Additionally, Table 7 highlights the accuracy improvements achieved with CCA, while Table 8 demonstrates that alternative methods are less effective compared to CCA.

[1] Lei Zhu, Xinjiang Wang, Zhanghan Ke, Wayne Zhang, and Rynson Lau BiFormer: Vision Transformer with Bi-Level Routing Attention, CVPR 2023

We thank the reviewer for their feedback on our rebuttal.

To improve readability and address quality issues, we have corrected the grammatical errors, increased the font sizes and included zoomed-in versions of the figures in the updated revision. If the reviewer could point out specific areas of the paper that could benefit from further improvements, we will gladly make additional updates and incorporate them in the final revision.

Regarding the request for additional comparisons, we are actively training the models with the backbones suggested by the reviewer and will include these results in the final version after acceptance. As highlighted in Table 10, we have already compared SCHEME with 10 different ViTs from the literature, demonstrating its effectiveness in improving the accuracy-latency tradeoff across a variety of architectures, including both attention-based and convolution-based designs. This demonstrates the versatility of SCHEME and its ability to advance the state-of-the-art across various ViT architectures. We hope this addresses the reviewer’s concerns.