SCHEME: Scalable Channel Mixer for Vision Transformers

## Abstract

Abstract Vision Transformers have achieved impressive performance in many vision tasks. While the token mixer or attention block has been studied in great detail, much less research has been devoted to the channel mixer or feature mixing block (FFN or MLP), which accounts for a significant portion of of the model parameters and computation. In this work, we show that the dense MLP connections can be replaced with a block diagonal MLP structure that supports larger expansion ratios by splitting MLP features into groups. To improve the feature clusters formed by this structure we propose the use of a lightweight, parameter-free, channel covariance attention (CCA) mechanism as a parallel branch during training. This enables gradual feature mixing across channel groups during training whose contribution decays to zero as the training progresses to convergence. In result, the CCA block can be discarded during inference, enabling enhanced performance at no additional computational cost. The resulting Scalable CHannEl MixEr (SCHEME) can be plugged into any ViT architecture to obtain a gamut of models with different trade-offs between complexity and performance by controlling the block diagonal MLP structure. This is shown by the introduction of a new family of SCHEMEformer models. Experiments on image classification, object detection, and semantic segmentation, with different ViT backbones, consistently demonstrate substantial accuracy gains over existing designs, especially for lower complexity regimes. The SCHEMEformer family is shown to establish new Pareto frontiers for accuracy vs FLOPS, accuracy vs model size, and accuracy vs throughput, especially for fast transformers of small size.

## 1 Introduction

Figure: Figure 1: Comparison of the proposed SCHEMEformer family, derived from the Metaformer-PPAA-S12 model [] with higher expansion ratios in the MLP blocks, and many SOTA transformers from the literature. The SCHEMEFormer family establishes a new Pareto frontier (optimal trade-off) for a) accuracy vs. FLOPs, b) accuracy vs model size, and c) accuracy vs, throughput. SCHEMEformer models are particularly effective for the design of fast transformers (throughput between 75 and 150 images/s) with small model size.
Refer to caption: x1.png

Vision Transformers (ViTs) [, , , ] are now ubiquitous in computer vision. They decompose an image into a set of patches which are fed as tokens to a transformer model [] of two main components: a spatial attention module, which reweighs each token according to its similarity to the other tokens extracted from the image, enabling information fusion across large spatial distances, and a channel mixer that combines the feature channels extracted from all patches using a multi-layer perceptron (MLP or FFN). A bottleneck of this model is the quadratic complexity of the attention mechanism on the number of patches. Numerous ViT variants have been proposed to address the problem, using improved attention mechanisms or hybrid architectures that replace attention or combine it with convolutions. Much less research has been devoted to the channel mixer. Most models simply adopt the two-layer MLP block of [], where channels are first expanded by a specified expansion ratio and then compressed to the original dimension. This is somewhat surprising since the mixer is critical for good transformer performance. For example, it is known that 1) pure attention, without MLPs or residual connections, collapses doubly exponentially to a rank one matrix [], 2) training fails to converge without residual connections or MLP [], and 3) replacing MLPs with more attention blocks (both spatial and channel attention) of equivalent computational complexity lowers the transformer accuracy []. All these observations indicate that the channel mixer is an indispensable ViT component.

In this work, quantify how much the channel mixer module contributes to ViT performance and investigate how to improve the trade-off between complexity and accuracy of the ViT model. We show that enhanced design of the channel mixer can lead to significant improvements in transformer performance by introducing a novel Scalable CHannEl MixEr (SCHEME) that enables the design of models with larger expansion ratios. SCHEME is a generic channel mixer that can be plugged into existing ViT variants to obtain effective scaled-down or scaled-up model versions. We replace the channel mixer of a state of the art architecture (SOTA) for low-complexity transformers, the MetaFormer-PPAA-S12 [], with SCHEME to obtain a new family of SCHEMEformer models with improved accuracy/complexity trade-off. This is illustrated in Figure 1, where the SCHEMEformer family is demonstrated to establish new Pareto frontiers for accuracy vs FLOPS, accuracy vs model size, and accuracy vs throughput, showing that SCHEME allows fine control over all these variables, while guaranteeing SOTA performance.
These properties are shown to hold for image classification, object detection and semantic segmentation tasks, as well as for different architectures, such as T2T-ViT [], CoAtNet [], Swin Transformer [], CSWin [] and DaViT [].

To develop SCHEME, we start by studying the impact of the mixer channel expansion. Transformer performance is shown to increase with expansion ration until it saturates for an expansion ratio beyond 8. However, because mixer (MLP) complexity also increases with the dimension of the intermediate representation, naive channel scaling with expansion ratios larger than 4 causes an explosion of parameters and computation, leading to models of large complexity and prone to overfit. To achieve a better trade-off between dimensionality and computation, we leverage recent findings about the increased hardware efficiency of mixers with block diagonal structure [, ]. These group the input and output feature vectors of a layer into disjoint subsets and perform matrix multiplications only within each group, as illustrated in Figure 2. We denote the resulting MLP block as Block Diagonal MLP (BD-MLP). Despite the lack of feature mixing across groups, we find that a transformer model equipped with the BD-MLP and a larger expansion ratio (to match both the parameter count and computation) achieves comparable or slightly higher accuracy than a baseline with dense MLP. This suggests that the lower accuracy of the block diagonal operations is offset by the gains of larger expansion ratios. Further analyzing the features learned by the different groups of the BD-MLP, we observe that they form feature clusters of similar ability to discriminate the target classes.
To learn better feature clusters, we seek a mechanism capable of restoring inter-group feature communication during training without increasing parameters. For this, we propose a channel attention branch that reweighs the input features of the BD-MLP according to their covariance matrix, as illustrated in Figure 2. This attention mechanism is denoted as the channel covariance attention (CCA) block. The re-weighted features are then fused with the BD-MLP output by means of a weighted residual addition with learned weights $(\alpha,1-\alpha)$.

As shown in Figure 2, SCHEME combines the sparse block diagonal structure of the BD-MLP, and the parameter-free CCA attention module, to implement a channel mixer extremely efficient in terms of parameters. Ablations of the evolution of the fusion weight $1-\alpha$ learned in the CCA branch (see Fig. 2) over training show that it gradually decays to zero during training. This happens consistently across all layers of the model and across model architectures. Hence, while the CCA is important for the formation of good feature clusters during training, it can be removed at inference without any loss, as illustrated in Table 3 and Figure 2. As a result, the model accuracy improves over simply using the BD-MLP mixer, but inference complexity does not. This leads to an extremely efficient inference setup, both in terms of parameters and FLOPs. Overall, the paper makes the following contributions,

Figure: Figure 2: Proposed SCHEME channel mixer. The channel mixer of the standard transformer consists of two MLP layers, performing dimensionality expansion and reduction by a factor of $E$. SCHEME uses a combination of a block diagonal MLP (BD-MLP), which reduces the complexity of the MLP layers by using block diagonal weights, and a channel covariance attention (CCA) mechanism that enables communication across feature groups through feature-based attention. This, however, is only needed for training. The weights $1-\alpha$ decay to zero upon training convergence and CCA can be discarded during inference, as shown on the right. Experiments show that CCA helps learn better feature clusters, but is not needed once these are formed.
Refer to caption: x4.png

- •
a study of the channel mixer of ViT MLPs, showing that dense feature mixing can be replaced by sparse feature mixing of higher internal feature dimensionality for improved accuracy, without increased complexity.
- •
the SCHEME module, which combines 1) a BD-MLP to enable internal feature representations of larger dimensionality than previous MLP blocks, and 2) CCA to enable the learning of these representations without cost at inference.
- •
various models that combine SCHEME with previous transformer architectures to achieve SOTA trade-offs between accuracy and model size, FLOPS, or throughput, such as the SCHEMEformer of Figure 1. This is shown particularly effective for the design of fast transformers with small model size, of interest for edge devices, robotics, and low-power applications.
- •
Experiments on image classification, object detection, and semantic segmentation, showing consistent gains in accuracy for fixed computation and size.

## 2 Related Work

Vision Transformers:
Vision transformers advanced the SOTA in several vision tasks since [, ] successfully applied the transformer-based self-attention NLP model of [] to image generation and classification tasks. These models rely on a spatial attention mechanism, based on the matrix of dot-products between features extracted from image patches. This has quadratic complexity in the number of patches and is quite intensive. Most follow up work [, , , , , , , , , ] improved the spatial attention mechanism of ViT. DeiT [] and subsequent works [, ] introduced a distillation token to distill information, typically from a CNN teacher, without large amounts of data or compute. PvT [] proposed a progressive shrinking pyramid architecture with spatial-reduction attention that scales ViTs for dense prediction tasks beyond image classification. Swin transformers [] introduced a hierarchical shifted window attention mechanism, which reduces the complexity to linear with respect to the number of windows. HaloNet [] proposed two extensions for local ViTs [] with blocked local attention and relaxed translational equivariance for scaling ViTs. A more extensive review of ViTs is given in [].

Recently, several works have shown that spatial attention is not the critical ViT feature. Some works improved performance by relying on hybrid architectures, which augment or replace ViT layers with convolutions [, ]. Efficient transformer designs, for edge devices, frequently sacrifice spatial attention to achieve better trade-offs between FLOPs and accuracy [, ]. Other works have questioned the need for spatial attention altogether. Metaformer [] argued that the fundamental trait of ViT is the mixing of information across patches, showing that competitive results can be obtained by simply replacing attention with pooling or identity operations. Similarly, [] showed that spatial attention is not critical for vision transformers by proposing a spatial-gating MLP of comparable performance to ViT []. DaViT [] showed that spatial attention is helpful by reusing its design for channel attention, building a cascade of alternating spatial and channel attention blocks. While this improves performance, it increases the complexity of the transformer block, resulting in a model with many parameters and potentially redundant channel mixer operations.
Despite all this work on ViT architectures, little emphasis has been devoted to the channel mixer module (MLP) that follows attention. This is surprising because the mixer dominates both the parameter count and complexity (FLOPs) of the standard transformer block. [, ] modify the MLP block to mimic the inverted residual block of the MobileNetV2 [], by adding a depthwise convolution. This improved performance but increases parameter and computation costs. Switch Transformer [] replaces the FFN with a sparse mixture of experts [] that dynamically routes the input tokens. This design allows scaling models to large sizes using higher number of experts but is not effective for ViTs. In this work, we instead propose an efficient and generic channel mixer module that improves both the parameter and computational efficiency of the transformer and allows for flexible scaling of ViTs.

## 3 The SCHEME module

Figure 2 depicts the proposed SCHEME module for feature mixing in ViTs.
As shown on the left, the standard channel mixer consists of two MLP layers, which expand the dimensionality of the input features and then reduce it to the original size. Let $\mathbf{x}\in\mathbb{R}^{d\times N}$ be the matrix containing the $N$ $d$-dimensional input feature vectors extracted from $N$ image patches. The mixer computes an intermediate representation $\mathbf{z}\in\mathbb{R}^{Ed\times N}$ and an output representation $\mathbf{y}\in\mathbb{R}^{d\times N}$ according to

$$ $\displaystyle\mathbf{z}=\sigma(\mathbf{W}_{1}\mathbf{x}+\mathbf{b}_{1}{\bf 1}^ {T}_{N})$ (1) $\displaystyle\mathbf{y}=\mathbf{W}_{2}\mathbf{z}+\mathbf{b}_{2}{\bf 1}^{T}_{N}$ (2) $$

where $\mathbf{W}_{1}\in\mathbb{R}^{Ed\times N},\mathbf{W}_{2}\in\mathbb{R}^{d\times
Ed
},\mathbf{b}_{1}\in\mathbb{R}^{Ed},\mathbf{b}_{2}\in\mathbb{R}^{d},$ ${\bf 1}_{N}$ is the $N$-dimensional vector containing ones as all entries, $\sigma(.)$ is the activation function (typically GELU), and $E$ is an expansion factor, typically 4.

Figure: Table 1: MetaFormer-S12 [] ImageNet-1K validation accuracy vs MLP expansion ratio ($E$). ${*}$: results by author code.

Table 1 details the impact of the mixer on overall transformer performance by evaluating the role of the expansion factor $E$ on the performance of Metaformer-PPAA-S12 [] (Pooling, Pooling, Attention, Attention) architecture on ImageNet-1K. Classification accuracy increases from 76.0 to $81.8\%$, as $E$ ranges from $1$ to $6$, decreasing for $E=8$, which suggests overfitting. The table also shows that these gains are not trivial. The S18 model, which has no expansion ($E=1$) but more transformer layers and complexity comparable to that of the S12 model with $E=4$, has an accuracy $2.7$ points lower than the latter.
In summary, for a given computation budget, it is beneficial to trade off transformer depth for dimensionality expansion in the channel mixer. This shows that this expansion is a critical component of the transformer architecture. On the other hand, naively scaling $E$ beyond 6 severely increases parameters and computation, leading to models that over-fit and are impractical for many applications.

### 3.1 Scalable Channel Mixer (SCHEME)

Block Diagonal MLP (BD-MLP):
Block diagonal matrices have been previously used to efficiently approximate dense matrices [, ]. In CNNs, group channel operations are frequently used to design lightweight mobile models with improved accuracy-computation trade-off [, , ]. This consists of splitting the feature vectors of (1)-(2) into disjoint groups, e.g. $\mathbf{x}$ into a set of $g_{1}$ disjoint features $\{\mathbf{x}_{k}\}_{k=1}^{g_{1}}$ where $\mathbf{x}_{k}\in\mathbb{R}^{d/g_{1}\times N}$, and $\mathbf{y}$ into a set $\{\mathbf{y}_{k}\}_{k=1}^{g_{2}}$ where $\mathbf{y}_{k}\in\mathbb{R}^{Ed/g_{2}\times N}$. As illustrated in Figure 2, the MLPs of (1)-(2) are then implemented independently for each group, according to

$$ $\displaystyle\mathbf{z}_{k}=\sigma(\mathbf{W}_{1,k}\mathbf{x}_{k}+\mathbf{b}_{ 1,k}{\bf 1}^{T}_{N})$ (3) $\displaystyle\mathbf{y}_{k}=\mathbf{W}_{2,k}\mathbf{z}_{k}+\mathbf{b}_{2,k}{ \bf 1}^{T}_{N}$ (4) $$

where $\mathbf{W}_{1,k}\in\mathbb{R}^{Ed/g_{1}\times N/g_{1}},\mathbf{W}_{2,k}\in
\mathbb{R}^{d/g_{2}\times Ed/g_{2}},\mathbf{b}_{1}\in\mathbb{R}^{Ed/g_{1}},
\mathbf{b}_{2}\in\mathbb{R}^{d/g_{2}}$ and $\mathbf{z}$ is decomposed into a set $\{\mathbf{z}_{k}\}_{k=1}^{G}$ where $\mathbf{z}_{k}\in\mathbb{R}^{Ed/G\times N}$, with $G=g_{1}$ in (3) and $G=g_{2}$ in (4). Since the complexity of (3) is $g_{1}^{2}$ times smaller than that of (1) and there are $g_{1}$ groups, the complexity of the first MLP is $1/g_{1}$ times that of standard MLP. Similarly, the complexity of the second MLP is $1/g_{2}$ times that of the standard MLP. Hence, a transformer equipped with the BD-MLP and expansion factor $\frac{2g_{1}g_{2}}{g_{1}+g_{2}}E$ has identical complexity to a standard transformer of factor $E$. For example, when $g_{1}=g_{2}=g$ this allows growing the expansion factor by a factor of $g$ without computational increase.

Channel Covariance Attention (CCA):
While the introduction of groups enables accuracy gains due to the increased expansion factor by $\frac{2g_{1}g_{2}}{g_{1}+g_{2}}$, it results in sub-optimal features. This is because the features in the different groups of (3)-(4) are processed independently, i.e. there is no inter-group feature fusion. This reduces the efficiency of the BD-MLP. To enable feature mixing between all feature channels and thus induce the formation of better feature clusters, we introduce a covariance attention mechanism in a parallel branch, as illustrated in Figure 2. The input features are first transposed to obtain the $d\times d$ covariance matrix(^1^11Since the features are normalized before the mixer, i.e centered such that $\mathbf{x}\,\mathbf{1}_{N}=0$, $\mathbf{x}\mathbf{x}^{T}$ is the covariance matrix of features $\mathbf{x}$.) $\mathbf{x}\mathbf{x}^{T}$. This is then used to re-weigh the input features by their covariance with other feature channels, using

$$ $CCA(\mathbf{x})=\mbox{softmax}\left(\frac{\mathbf{x}\mathbf{x}^{T}}{\tau} \right)\mathbf{x}$ (5) $$

where the softmax operation is applied across the matrix rows and $\tau$ is a smoothing factor.
The output of the channel mixer block is the weighted summation of the BD-MLP and CCA branches

$$ $\begin{split}\textbf{y}_{out}=\alpha\textbf{y}+(1-\alpha)CCA(\textbf{x}),\end{split}$ (6) $$

where $\alpha$ is a mixing weight learned across all samples. Various other design choices are discussed in Section 4.2.

CCA as a Regularizer:
The introduction of a parameter free attention branch and a learnable weight $\alpha$ allows the model to form better feature clusters during training and gradually decay the contribution from CCA branch once the feature clusters are formed. This can be seen in Figure 4, which plots the value of the learned mixing weight $1-\alpha$ as a function of training epochs on ImageNet-1K, for all transformer layers. These plots are typical of the behavior we observed with all transformer backbones and architectures we considered. Clearly, $1-\alpha$ starts with high to intermediate values, indicating that information flows through both branches of the mixer, but decays to $1-\alpha\approx 0$ as training converges. Hence, as shown in Table 3, there is no degradation if the CCA branch is removed during inference.
This eliminates a substantial amount of computation during inference, leading to the the training and inference setup of Figure 2, where CCA is not used at inference.

We explain this behavior by conjecturing that the downside of the computational efficiency of the BD-MLP is a more difficult learning problem, due to the independent processing of channel groups. This creates symmetries in the cost function, e.g. the order of the feature groups is not important, and requires a feature clustering operation that is likely to produce more local minima. The CCA branch helps to smooth out this cost function during training, while the feature groups are not established, by allowing inter-group communication. However, once the right feature groupings are found, CCA is no longer needed, and a simple BD-MLP mixer has no loss of performance over the standard MLP. Note that the operation of (5) is basically a projection of $\mathbf{x}$ into canonical subspaces of features that are correlated in the input image. This is likely to be informative to guide the group formation, but less useful when the features are already clustered.
While this hypothesis is not trivial to test, since (5) varies from example to example, we confirmed that using CCA during training enhances class separability, which likely reduces overfitting for large expansion ratios.
See the section 4.2 for more details.

Computational Complexity:
The complexity of the BD-MLP block is controlled by the group numbers $g_{1}$, $g_{2}$ and the expansion factor $E$, with a total cost $\mathcal{O}(Ed^{2}/g_{1}+Ed^{2}/g_{2})$, where $d$ is the channel dimension. The computational cost of CCA is $\mathcal{O}(Nd^{2})$ where $N$ is the number of tokens.
Since CCA is not used during inference, it only adds to the computations during training. The SCHEME framework provides a systematic way to control the trade-off of transformer width vs depth, by controlling the block size and expansion hyperparameters.

### 3.2 The SCHEMEformer family

The proposed SCHEME module enables efficient control of model complexity via the mixer hyperparameters $g_{1}$, $g_{2}$, and $E$. Table 1 shows that naively scaling down the ViT model by simply reducing $E$ causes a significant accuracy loss. The SCHEME module allows much more effective control of the accuracy/complexity trade-off, producing models of better performance for a fixed computational budget.
This is demonstrated by the introduction of a new family of models, denoted as SCHEMEformer, obtained by replacing the channel mixer of the Metaformer-PPAA [] architecture with the SCHEME module. Two such configurations are shown in Table 2 where the naming follows the convention {model-name}-{$g_{1}$$g_{2}$}-e{$E$} where the model name is skipped for brevity.

## 4 Experimental Results

Figure: Figure 3: Impact of CCA (SCHEMEformer-44-e8-S12). Left: Evolution of weight $1-\alpha$ across model layers. Right: Class separability of output features $\bf y$ (over 50 random classes of ImageNet-1K validation set) for model trained with and without CCA.
Refer to caption: x5.png

### 4.1 Comparisons to the state of the art

Image Classification:
Image classification is evaluated on Imagenet-1K, without using extra data. We report the results with single crop top-1 accuracy at $224\times 224$ input resolution. We evaluate the SCHEMEformer family of models based on the Metaformer-PPAA-S12 [] obtained by replacing the MLP of the latter with the SCHEME module. Refer to Appendix for implementation details.

We start by evaluating how this improves the trade-off between model accuracy and complexity. As discussed in Section 3.1, when $g_{1}=g_{2}$, a SCHEME transformer of expansion factor $gE$ has identical complexity to a standard transformer of expansion factor $E$. Hence, for fixed FLOPS, SCHEME allows an increase of the expansion factor by $g$. Figure 4 a) compares the performance of the Metaformer-PPAA-S12 with expansion ratios $E\in\{1,2,4,8\}$ to comparable variants of the SCHEMEformer-PPAA-S12, with SCHEME mixers of either $g=2$ (green curve) or $g=4$ (blue curve) groups. The SCHEMEformer models have a better trade-off between accuracy and FLOPS, achieving higher accuracies for all complexity levels. Among these, the one with more feature groups ($g=4$) has the best performance. While SCHEMEformer gains are observed for all FLOP levels, they are larger for lower complexity models. This makes SCHEME particularly attractive for the design of low complexity transformers, e.g. for edge devices or equivalent applications.

We next compare the SCHEMEformer family against the SOTA transformers in the literature. This is not an easy comparison, since models vary in size, FLOPS, and throughput. Because it is difficult to make any of these variables exactly the same for two different architectures, the comparison is only possible in terms of how the different architectures trade-off accuracy for any of the other factors. For a given pair of variables, e.g. FLOPS vs accuracy, the model is said to be on the Pareto frontier of the two variables if it achieves the best trade-off between the two.
Table 4 presents a comparison of the SCHEMEformer family against various SOTA transformers in the literature. Each section of the table compares a SCHEMEformer model to a group of SOTA transformers of equivalent size or complexity. Note that, in each section, the remaining models have both lower throughput and accuracy than the SCHEMEformer model. In many cases they also have more parameters and FLOPs.
Figure 1 provides a broader visualization of how SCHEME models establish Pareto frontiers for accuracy vs FLOPS, accuracy vs throughput, and accuracy vs model size (parameters). Figure 1 a) illustrates the trade-off between accuracy and FLOPS of many SOTA transformers. The dashed line connects the SCHEMEformer model results, summarizing the accuracy-FLOPs trade-off of the family. It can be seen that the SCHEME models lie on the Pareto frontier for these two objectives. This illustrates the fine control that SCHEME allows over the accuracy/complexity trade-off of transformer models. Figure 1 b) presents a similar comparison for model sizes. Like for FLOPS, the SCHEME models lie on the Pareto frontier for accuracy vs model size. In fact the two plots are quite similar, showing that in general there is a good correlation between model size and FLOPS.

Figure: Table 4: Image Classification on ImageNet-1K. Comparison with SOTA ViTs grouped by accuracy. Proposed SCHEME models use expansion ratio 8. SCHEMEformer family has higher throughput and accuracy than SOTA models.

This is not the case for throughput, which for transformers is known to not necessarily correlate with FLOPS, due to GPU parallelism. For example, Table 4 shows that, CSWin models have lower throughput despite having lower FLOPs while ViT has higher throughput despite having higher FLOPs. SCHEMEformer controls this trade-off by controlling the expansion ratio and block diagonal structure, which enables higher FLOPs utilization for a given throughput [].
We compute throughput on a NVIDIA-Titan-X GPU with a batch size of 1 averaged over 1000 runs. While comparisons could be made for larger batch sizes, we consider the setting for live/streaming applications, where speed is most critical. In these applications, the concern is usually inference throughput, which requires batch size of 1. Figure 1 c) illustrates the trade-off between accuracy and throughput of various models, including very fast ResNet models of low accuracy. It can bee seen that the SCHEMEformer again achieves the best trade-off between these two variables, thus lying on the Pareto frontier for accuracy vs throughput. Its performance is particularly dominant in the range of throughputs between 75 and 150 images/sec, where it significantly outperforms the other methods. These results demonstrate how the SCHEME module endows transformer designers with the ability to produce multiple models at different points of the Pareto frontiers of accuracy vs model size, FLOPS, or throughput.

SCHEME with other ViTs: Figure 4 shows the Accuracy-Throughput curves for various SCHEME models obtained by replacing the MLP blocks of popular ViT architectures. It shows that SCHEME improves the accuracy for a fixed throughput, the throughput for a fixed accuracy, or both. The gains can be substantial, e.g. about $1\%$ accuracy gain (constant throughput) for the fastest transformer or speed gains of up to $20\%$ at constant accuracy. This shows that SCHEME benefits various ViT backbones, not just the MetaFormer.

Semantic Segmentation:
Table 5 compares the semantic segmentation performance of two SCHEMEformer models using the semantic FPN framework [] to various SOTA models of similar complexity, on ADE20K. Since we do not not have access to the throughput of most models, we report only parameter sizes and FLOPS. In each section of the table, the remaining models have comparable or larger FLOPs and model size but lower accuracy than the corresponding SCHEMEformer. For example, SCHEMEformer-44-e8-S12 achieves 40.9% mIoU, which is 5.2/3.7 points higher than the PvT-Tiny/PoolFormer-S12, which both have comparable size and FLOPs. Similarly, the S24 model outperforms PoolFormer-M36 using only 41% of its parameters. To demonstrate the applicability of SCHEME to larger models, we also present a comparison of the SCHEME version of the DaViT-Tiny using the UperNet framework []. While the DaViT-Tiny is already the best model in the table, the use of the SCHEME mixer improves its performance by an additional 0.8 points.

Object Detection:
Table 6 compares SCHEMEformer models to models of similar complexity on the COCO-17 object detection and instance segmentation benchmark, for both Retinanet and Mask-RCNN detection heads. Again, the SCHEMEformer models outperform most other models of the same or smaller size. The only exception is the PoolFormer-S24, which slightly outperforms ($0.1$ points) the comparable SCHEME-former-44-S24, for the RetinaNet head. However, with the stronger Mask R-CNN head, the SCHEMEformer-44-S24 beats the PoolFormer-S24 by $0.8$ points.
For the top performing models in the bottom third of each section of Table 6, SCHEME-DaViT-Tiny improves over DaViT-Tiny by an additional 0.7% and 0.9%, for RetinaNet and MaskRCNN heads respectively, while maintaining a comparable throughput.

Figure: Table 6: COCO-17 Object Detection and Instance Segmentation. All backbones pretrained on ImageNet-1K (1x learning schedule). ($AP^{b}$, $AP^{m}$): (bounding box AP, mask AP). $T$: throughput (images/sec).

### 4.2 Ablation Studies

Contribution of BD-MLP and CCA branch:
Table 7 shows an ablation of the contribution of the BD-MLP and CCA branches. Starting from the Metaformer-PPAA-11-e2-S12, with expansion ratio $E=2$ and dense MLP, we replace the channel mixer by SCHEME to obtain the SCHEMEformer-PPAA-44-e8-S12 ($g_{1}=g_{2}=4$ and $E=8$), which maintains the number of parameters and FLOPs constant. The SCHEME model with only BD-MLP improves on the baseline by 0.2%. The addition of the CCA branch provides an additional gain of 0.6%, showing the gains of better feature clusters. Since CCA is not used at inference, its gains are free in terms of additional parameters/FLOPs.

Regularizing effect of CCA: Fig. 4 shows the evolution of the weight $1-\alpha$ of the CCA branch in (6) during training. While initially large, it gradually decays to zero as training progresses. This holds for all network layers. Hence, CCA can be discarded at inference. Fig. 1 in supplementary plots the weights $1-\alpha$ upon training convergence, for the family of SCHEMEformer-PPAA-44-e8 models, confirming that the weights are indeed very close to zero across all layers.

Effect of large expansion ratios:
The left of Figure 4 shows the effect of simultaneously increasing the expansion ratio $E$ and adjusting groups to realize different models of similar size and complexity. All models are based on the Metaformer-PPAA-S12. For fixed parameters/FLOPs, SCHEMEformer models achieve a gain of 1.0% to 1.4% over the baseline by increasing $E$ from 1 to 4. The gains increase with larger expansions and saturate at larger FLOPs. This shows that higher internal feature dimensions are important for obtaining better accuracy with smaller ViT models.

Alternative designs of Channel Mixer:
We investigate whether alternative choices to the CCA branch could accomplish this goal more effectively. Table 8 compares models that replace CCA with other feature mixing operations: the channel shuffling operation of ShuffleNet [], a squeeze and excitation [] network (SENet), a single layer of convolution, and a dynamic version of CCA (DyCCA), where the weight $\alpha$ of (6) is predicted dynamically, using GCT attention []. CCA obtains the best result. While CCA is computationally heavier than some of these alternatives, it is not needed at inference, as shown in Table 3. This is not true for the alternatives, which produce much more balanced weights $\alpha$ after training convergence, and cannot be discarded at inference without performance drop. We conjecture that, because the alternatives have learnable parameters, the network learns to use them to extract complementary features, which must be used at inference.

Feature Clustering: We conjectured above that CCA helps training because it facilitates feature clustering into naturally independent groups that do not require inter-group communication. We tested this hypothesis by studying the intermediate feature vectors ${\bf v}_{l}$ (obtained after max-pooling the features $\bf y$ of (4)) extracted from four randomly selected layers $l$ of two ImageNet pretrained models, trained with and without CCA. We split the features into the 4 groups used in the model ${\bf v}_{l,g},g\in\{1,\ldots 4\}$ and concatenated the features of all layers in the same group. This produced four vectors ${\bf u}_{g}=\mbox{concat}(\{v_{l,g}\}_{l})$ containing the features of each group $g$ extracted throughout the network. A linear classifier was then learned over each vector ${\bf u}_{g}$.
Table 9 shows the top-1 accuracy per feature group and model. To evaluate whether groups learn different class clusters, we also average the outputs from the four group classifiers to obtain the final accuracy. Without CCA, i.e. no group communication during training, the network produces feature groups individually more predictive of the image class, but less predictive when combined. This suggests that there is redundancy between the features of the different groups. By introducing inter-group communication, CCA enables the groups to learn more diverse features, that complement each other. Fig. 4 shows the class separability of the intermediate features $\bf y$ of (4) of a randomly chosen layer of the SCHEMEformer-PPAA-44-e8-S12. Class separability was measured as in [], with a final value obtained by averaging the class separability across all classes. The model trained with CCA has higher class separability than that without it. This confirms that CCA is helpful in forming feature clusters that increase class separability during training.
Conversely, we tested if CCA is helpful when channel shuffling is inserted in between the two mixer MLP layers, which destroys the group structure. This variant of the SCHEMEformer-PPAA-44-e8-S12 model achieved an accuracy of 79.1% for both training with and without CCA (Table 7, row 2). This shows that CCA is not helpful when feature groups are mixed. Similarly, CCA did not provide any gains when applied to the standard MLP branch with full feature mixing. These results suggest that CCA indeed helps to form the independent feature clusters needed to achieve the computational efficiency of channel groups without performance degradation.
Table 9 shows that, despite the higher accuracy of the individual group features of the model trained without CCA, the model trained with CCA has 0.6% higher ensemble accuracy.

## 5 Limitation and Conclusion

Although SCHEME module improves the accuracy-throughput curves of popular ViTs, it incurs a slight overhead in memory (see appendix section A.2.2) during training due to the channel attention operation which limits the batch size on smaller GPUs. However, Figure 4 shows that the contribution of the CCA branch becomes negligible after 150 epochs. Therefore, the CCA branch can be removed beyond this point, significantly reducing the training overhead.
In this work, we proposed the SCHEME module for improving the performance of ViTs. SCHEME leverages a block diagonal feature mixing structure to enable MLPs with larger expansion ratios, a property that is shown to improve transformer performance, without increase of model parameters or computation. It uses a weighted fusion of a BD-MLP branch, which abstracts existing MLPs with block diagonal structure, and a parameter-free CCA branch that helps to cluster features into groups during training. The CCA branch was shown to improve training but not be needed at inference. Experiments showed that it indeed improves the class separability of the internal feature representation of the BD-MLP branch, helping create feature clusters that are informative of the image class. The standard transformer MLP was replaced with the SCHEME module to obtain a new family of SCHEMEformer models with improved performance for classification, detection, and segmentation, for fixed parameters and FLOPs with favorable latency. SCHEME was shown to be effective for various ViT architectures and to provide a flexible way to scale models, always outperforming models with smaller MLP expansion ratios having the same complexity.

## Appendix A Appendix

### A.1 Implementation Details

#### A.1.1 Hyperparameter Settings

Table A.1.1 shows the detailed hyperparameter settings of the family of SCHEMEformer models reported in the main paper.

Figure: Table 10: Hyperparameter Settings for the family of SCHEMEformer models trained on ImageNet-1K dataset.

Figure: (a) SCHEMEformer-44-e8-S12
Refer to caption: extracted/5625602/images/weights_s12.png

Figure: Table 11: Ablation study on the formation of feature clusters in the BD-MLP branch of the SCHEME module.

### A.2 Ablation Studies

#### A.2.1 Plot of the learned 1 − α 1 𝛼 1-\alpha 1 - italic_α weights in the SCHEME module for SCHEMEformer models

The learned weights $1-\alpha$ for the family of SCHEMEformer models is shown in Fig. 5. Interestingly, the learned weights coarsely approximate the shape of a gaussian distribution. The learned weights reach a peak value in the middle layers of the network and drop to zero for all the other layers. The middle of the network typically correspond to the initial few layers of the third stage of the model that contains the maximum number of transformer blocks for all the models shown in Fig. 5. We conjecture that the weights for these layers have not fully converged and that feature mixing can still be useful for these layers and so training for more epochs will allow the $1-\alpha$ weights of these layers to converge to zero. To test this hypothesis, we trained the SCHEMEformer-PPAA-12-e8-S12 model for an additional 200 epochs beyond the standard 300 epochs and observed that the peak value of the $1-\alpha$ weights decreased further by 0.11 as compared to the model trained for 300 epochs and the accuracy improved by 0.4%. The weight norm (of all layers) decreased from 0.18 for 300 epoch model to 0.05 for 500 epoch model. This confirms our hypothesis and training SCHEMEformer models for larger epochs can further improve the accuracy.

Feature Clustering: We conjectured above that CCA helps training because it facilitates feature clustering into naturally independent groups that do not require inter-group communication. We tested this hypothesis by studying the intermediate feature vectors ${\bf v}_{l}$ (obtained after max-pooling the features $\bf y$ of Eqn. (4) of the paper extracted from four randomly selected layers $l$ of two ImageNet pretrained models, trained with and without CCA respectively. We split the features into the 4 groups used in the model ${\bf v}_{l,g},g\in\{1,\ldots 4\}$ and concatenated the features of all layers in the same group. This produced four vectors ${\bf u}_{g}=\mbox{concat}(\{v_{l,g}\}_{l})$ containing the features of each group $g$ extracted throughout the network. A linear classifier was then learned over each vector ${\bf u}_{g}$.
Table A.1.1 shows the top-1 accuracy per feature group and model. To evaluate whether groups learn different class clusters, we also average the outputs from the four group classifiers to obtain the final accuracy. Without CCA, i.e. no group communication during training, the network produces feature groups that tend to be individually more predictive of the image class, but less predictive when combined. This suggests that there is some amount of redundancy between the features extracted by the different groups. By introducing inter-group communication, CCA enables the groups to learn more diverse sets of features, that complement each other.

Fig. 4b in the main paper showed the class separability of the intermediate features $\bf y$ of Eqn. (4) of the paper of a randomly chosen layer of the SCHEMEformer-PPAA-44-e8-S12 model.
Conversely, we tested if CCA is helpful when channel shuffling is inserted in between the two mixer MLP layers, which destroys the group structure. This variant of the SCHEMEformer-PPAA-44-e8-S12 model achieved an accuracy of 79.1% for both training with and without CCA. This shows that CCA is not helpful when feature groups are mixed. Similarly, CCA did not provide any gains when applied to the standard MLP branch with full feature mixing. These results suggest that CCA indeed helps to form the independent feature clusters needed to achieve the computational efficiency of channel groups without performance degradation.

Figure: Table 12: Training Overhead of CCA. CCA adds only a small overhead in GPU memory and training time.

Figure: Table 13: Ablation study of longer training for SCHEMEformer-PPAA-44-e8-S12.

#### A.2.2 Training Overhead of CCA

Table 12 compares the training GPU memory and throughput for SCHEME mixer with and without using CCA. CCA improves the accuracy by 0.6% with only a slight increase in the GPU memory (+12.5%) and training time (+16.7%). Further, CCA is not needed during inference thereby providing gains for ”free” without additional computational cost at inference.

#### A.2.3 Training for longer epochs

Table 13 shows the comparison of training for longer epochs for SCHEMEformer with the baseline Metaformer model. We train for an additional 200 epochs from the standard 300 epochs. SCHEMEformer-PPAA-44-e8-S12 trained for 300 epochs even outperforms the baseline model trained for 500 epochs. On continuing the training from 300 to 500 epochs, SCHEMEformer continues to improve the performance without saturation suggesting that it is beneficial to train with CCA for longer epochs. 300 to 500 epochs is a much larger increase of training time (67%) than the 16.7% increase in training time required by SCHEME mixer (see Table 12).

Figure: Table 14: Impact of removing CCA branch during inference.

Figure: Table 15: Ablation study of removing CCA for COCO-17 Object Detection and Instance Segmentation. Removing CCA at inference does not impact the AP values as they are identical to the model using CCA at inference. $AP^{b}$ and $AP^{m}$ denote bounding box AP and mask AP, respectively. Backbone models denote SCHEMEformer-PPAA-44-e8- variants.

#### A.2.4 Impact of removing CCA at inference

Table A.2.3 shows the impact of removing CCA at inference, for various backbones. While the number of FLOPS decreases, the top-1 accuracy changes very little ($\approx 0.02$ difference). Hence, there is no advantage in using CCA at inference. This is unlike training, where the use of CCA makes a non-negligible difference, as shown in Table 6 of the main paper.

Table A.2.3 and A.2.3 show the results of removing CCA at inference for object detection and semantic segmentation models, respectively. For both tasks, the results are identical to the model using CCA showing that the CCA also generalizes to downstream tasks.

Figure: Table 17: Ablation study on the effect of larger expansion ratios in BD-MLP block of SCHEME on ImageNet-1K validation dataset.

Figure: Table 18: Ablation study on the formation of feature clusters in the BD-MLP branch of the SCHEME module. We train a linear classifier on top of the four feature groups extracted from the final MLP mixer of the transformer block of the network. The model trained with CCA forms feature clusters that learn diverse and complementary set of features that can obtain 1% higher validation accuracy than the model trained without CCA.

Figure: Figure 6: GRAD-CAM [] visualization for a few validation samples on ImageNet-1K dataset for SCHEMEformer-PPAA-44-e8-S12 model and comparison with other competing methods.
Refer to caption: x9.png

#### A.2.5 CCA

In the main paper, feature groups across different layers of the model was used to demonstrate the learning of feature clusters by CCA. Here, the feature from the final layer of the model is only considered. Table A.2.4 shows the top-1 accuracy per feature group and model. The average of the outputs from the four group classifiers is reported in the final column of the table. The effect is more pronounced when using a single layer feature with +1.09% accuracy difference between the model with and without CCA. This further reinforces that by introducing inter-group communication, CCA enables the groups to learn more diverse sets of features, that complement each other.

#### A.2.6 Larger expansion ratios

Table A.2.4 shows additional results of using larger expansion ratios with SCHEME mixer (illustrated in Fig. 3 of the main paper) for the same number of parameters and FLOPs using the Metaformer-PPAA-S12 baseline model. We observe that SCHEMEformer-PPAA consistently outperforms the baseline for larger expansions ratios with larger gains at lower FLOPs. The performance saturates as the model size and FLOPs increases.

### A.3 Qualitative Analysis: Grad-CAM [ 29 ] Visualization

Fig. 6 shows the results of class activation maps for SCHEMEformer-PPAA-44-e8-S12 model for a few examples from the validation set of ImageNet-1K dataset. The stronger heatmap responses around the salient features of an object (e.g., body of a bird, cat) shows that the model ignores the background and attends to more discriminative spatial regions. Fig. 6 also shows the qualitative comparison with a few existing methods such as ResNet-50, DeiT-S, Poolformer etc. of similar complexity. It demonstrates that SCHEMEformer-PPAA-44-e8-S12 attends to the complete object class and less spurious features showing that it is better than the competing methods.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 1
questions: 1. 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?
soundness: 2
strengths: 1. SCHEME introduces a clever way of reducing complexity by leveraging block diagonal MLPs, which directly targets the computational bottleneck of standard transformers.
2. The addition of the Channel Covariance Attention (CCA) during training adds flexibility and improves feature clustering without increasing inference complexity, making it an efficient regularization tool.
3. The approach is shown to be effective across multiple transformer backbones (MetaFormer, T2T, Swin), demonstrating its adaptability in various settings.
summary: The paper proposes SCHEME (Scalable CHannEl MixEr), a channel mixing mechanism for Vision Transformers (ViTs). The method focuses on replacing the dense MLP layers in the transformer’s channel mixing block with a block diagonal MLP (BD-MLP) structure, allowing for larger expansion ratios and improving computational efficiency. Additionally, the authors introduce a Channel Covariance Attention (CCA) mechanism to enable inter-group communication, which is discarded after training to reduce inference complexity. The SCHEME mechanism is integrated into multiple Vision Transformer architectures and is evaluated across various benchmarks, showing improved accuracy while maintaining efficiency.
weaknesses: 1. The manuscript contains several grammatical errors and formatting issues, such as the phrase 'In result' in the abstract, as well as figures that are too small to clearly convey the details, which fails to meet the standards of ICLR.
2. Experiments focuses heavily on comparisons with a few models like MetaFormer, lacking comparisons with state-of-the-art transformer-based models like PVT, UniFormer or TransNeXt.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. I am concerned about the generalization of this method to other data sets.
2. Some details are missing. For Block Diagonal MLP, how to split the feature vectors of (1)-(2) into disjoint groups? What criteria is it based on？What's the motivation for designing the Channel Covariance Attention (CCA)? It didn't show more details.
3. How is the design like channel mixer relate to improving performance for example? It seems inadequate that none of were seriously discussed in the manuscript.
soundness: 3
strengths: 1. The writing is easy to read and clearly explains everything in the paper.
2. The experimental result is good compared to the previous works. Empirically, the method seems to offer strong accuracy, compared to existing methods with similar architectures.
summary: This paper study the channel mixer of ViT MLPs, showing that dense feature mixing can be replaced by sparse feature mixing of higher internal feature dimensionality for improved accuracy, without increased complexity.  The work presents a range of experiments that sufficiently support its claims.
weaknesses: 1. I am concerned about the generalization of this method to other data sets.
2. Some details are missing. For Block Diagonal MLP, how to split the feature vectors of (1)-(2) into disjoint groups? What criteria is it based on？What's the motivation for designing the Channel Covariance Attention (CCA)? It didn't show more details.
3. How is the design like channel mixer relate to improving performance for example? It seems inadequate that none of were seriously discussed in the manuscript.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: Please first deal with the major concerns above. There are some minor issues below.

1. The learnable weight $\alpha$ enables CCA as a regularizer during training, discarded for efficiency during inference. However, this approach may create potential conflicts in optimization objectives. Why not consider using more established techniques, such as Structural Re-parameterization, to achieve structure changes between training and inference?

2. Clarification is needed on the smoothing factor $\tau$ — is it a learnable parameter or a constant?

3. There are a few typographical errors, such as a double 'of' in Line 041 of the abstract.
soundness: 2
strengths: 1. The paper is well-structured, providing a clear and accessible overview of the motivation and methodology behind the proposed approach. The logical flow enables readers to easily understand both the problem addressed and the design of the proposed solution.

2. The methodology section is well-developed, with good supporting arguments for each component. The progression from MLP to BD-MLP, followed by the integration of CCA, is clear, well-motivated.

3. The use of CCA as a regularizer for BD-MLP to facilitate interaction across channel groups during training, while discarding it during inference to streamline computational efficiency, is an innovative and effective approach.

4. The proposed SCHEME module and its various configurations demonstrate impressive performance across a range of tasks.
summary: This paper studies the channel-mixer of modern vision architectures, and propose the SCHEME module, a new channel-mixer that split features to multiple sub channel groups and project these channel groups to lager dimensionality for learning richer internal representation. A channel covariance attention is designed to make information exchange among sub channel groups during training time. The experiments show that the proposed SCHEME outperforms the traditional channel-mixers (e.g., FFN/MLP) in several tasks.
weaknesses: 1. The labels and coordinates in Figures 1, 3, and 4 are too small, making them challenging to interpret. For a basic academic paper, readability of figure details is essential.

2. The effectiveness of the proposed SCHEME is primarily validated on MetaFormer. While the authors briefly present results on other backbones like Swin and DaViT in Figure 4 (yet difficult to interpret). SCHEME’s generalizability is limited by testing primarily on ViT-based architectures, lacking ablation on modern ConvNet architectures (e.g., ConvNeXt [1], FasterNet [2], and InceptionNeXt [3]) that also use FFN/MLP modules.

3. Although SCHEME shows theoretical complexity advantages, its on-device efficiency is underexplored. Only limited throughput comparisons are presented in Tables 4 and 8. Since an on-device benchmark has been implemented, including detailed on-device results across ablation and comparison studies would substantiate claims about SCHEME’s efficiency over conventional FFNs. Additionally, details on the benchmark configurations, hardware, and input shapes are necessary to ensure fair comparisons.

4. The CCA component functions as a channel mixer and resembles linear attention mechanisms in NLP and prior work in XCiT [4]. A discussion of the differences between CCA and these methods, alongside proper citations, would clarify the novelty and address potential ethical concerns regarding prior art.

5. While SCHEME’s performance is strong and well-motivated, the contribution could be strengthened with deeper analysis. SCHEME integrates Block Diagonal MLP and CCA, elements with roots in existing literature. A more detailed explanation of how SCHEME specifically influences learned representations would underscore its contribution.

6. The comparison methods are outdated and do not represent state-of-the-art accuracy-efficiency models. Including comparisons with recent approaches like FasterViT [5], FastViT [6], and MobileOne [7] would enhance the experimental validity.

7. SCHEME is tested only on small model scales (<60M), which aligns with its efficiency-focused motivation. However, a benchmark of on-device speed on mobile devices is necessary to validate the proposed method’s real-world application, as many recent studies have done [6,7].



[1] Liu Z, Mao H, Wu C Y, et al. A convnet for the 2020s[C]//Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022: 11976-11986.

[2] Chen J, Kao S, He H, et al. Run, don't walk: chasing higher FLOPS for faster neural networks[C]//Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2023: 12021-12031.

[3] Yu W, Zhou P, Yan S, et al. Inceptionnext: When inception meets convnext[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024: 5672-5683.

[4] Ali A, Touvron H, Caron M, et al. Xcit: Cross-covariance image transformers[J]. Advances in neural information processing systems, 2021, 34: 20014-20027.

[5] Hatamizadeh A, Heinrich G, Yin H, et al. Fastervit: Fast vision transformers with hierarchical attention[J]. arXiv preprint arXiv:2306.06189, 2023.

[6] Vasu P K A, Gabriel J, Zhu J, et al. FastViT: A fast hybrid vision transformer using structural reparameterization[C]//Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023: 5785-5795.

[7] Vasu P K A, Gabriel J, Zhu J, et al. Mobileone: An improved one millisecond mobile backbone[C]//Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2023: 7907-7917.

---

confidence: 4
contribution: 4
initial_rating: 5
presentation: 3
questions: See `Weaknesses`.
soundness: 4
strengths: 1. Channel mixer design, the focus of this paper, is a more than interesting topic for the backbone design.
2. Comprehensive vision experiments (*e.g.*, classification, detection, segmentation) and analysis (*e.g.* `Paragraph 3`in `Sec 4.2` for channel covariance attention) were conducted to validate the proposed method's performance.
summary: This paper concentrate on the  channel mixer or feature mixing block design and propose a block diagonal MLP structure with a lightweight, parameter-free, channel covariance attention (CCA) module to improve the feature clusters formed. Replacing the vanilla MLP with proposed SCHEME module, various experiments were conducted to validate its performance.
weaknesses: 1. With regard to the design of channel mixer, there exists many previous research, such as gMLP[1], channel aggregation module in moganet[2], S2-MLP[3]. In addition to comparing with models using traditional MLP, these should also be taken into account.
2. More scaling experiments are needed. In this paper, the biggest model was 58M in `Tab.1`, which may could only match the `base-size` model in general. To comprehensively validate the method's scaling capability in backbone, `Large` or even `X-Large` model are also need.
3. Typos for the dimension of $W_1$ in `Line 196,221`, $N$ -> $d$.

I would be more than willing to reconsider the rating based on the authors' response.

Refs:

[1] Liu, Hanxiao et al. “Pay Attention to MLPs.” Neural Information Processing Systems (2021).

[2] Li, Siyuan et al. “MogaNet: Multi-order Gated Aggregation Network.” International Conference on Learning Representations (2022).

[3] Yu, Tan et al. “S2-MLP: Spatial-Shift MLP Architecture for Vision.” 2022 IEEE/CVF Winter Conference on Applications of Computer Vision (WACV) (2021): 3615-3624.

Meta-review: The paper proposes a scalable channel mixing mechanism for Vision Transformers (ViTs). Its main components are block diagonal MLP (BD-MLP) and channel covariance attention (CCA). The proposed method is integrated into multiple Vision Transformer architectures and is evaluated across various benchmarks, showing improved accuracy while maintaining efficiency. The main concerns raised by the reviewers were the marginal novelty, the quality of writing, and the limited experimental comparisons and analyses. In particular, all reviewers commonly pointed out the need for more extensive experiments with recent architectures and benchmarks. The authors addressed some of the corners during discussions but failed to address the issue of limited experiments fully. After post-rebuttal discussion, two reviewers remain rejects (one reject and one borderline reject) and while other two reviewers gave accepts (borderlines); there was no strong support from the reviewer. AC agrees with reviewers (in particular, N5s3 and vWvm) that the generalizability of the proposed method needs to be validated on more recent architectures and also larger models, which is lacking in the current version. AC thus recommends rejection.

Decision: Reject