Visual Transformer with Differentiable Channel Selection: An Information Bottleneck Inspired Approach

We appreciate the review and the suggestions in this review. The raised issues are addressed below.

(1) Actual Inference Time

We compare DCS-Transformer to the current state-of-the-art pruning methods for visual transformers, the results are shown in the table below. DCS-Transformer is compared to S$^2$ViTE[2], SPViT[3] and SAViT [4] on EfficientViT-B1 (r224) [1]. For S$^2$ViTE[2], SPViT[3] and SAViT [4], the pruning is performed on the ImageNet training data. After obtaining the pruned networks, we fine-tune the pruned networks using the same setting as [1]. It can be observed that DCS-EfficientViT-B1 outperforms all pruned models by at least 1.9% in top-1 accuracy with even less parameters and FLOPs. We report the actual inference time of all the model on a NVIDIA V100 GPU.

(2) Completed Figure 2

The two points for DCS-EfficientViT have been added to Figure 2 in the revised paper.

(3) Tuning Hyperparameters by Cross-Validation

Thank you for your concern about the hyperparameters. We set the initial value of the temperature $\tau$ to $4.5$ and decrease it by a factor of $0.95$ every epoch, and this is the setting used for all the experiments in this paper. In the revised paper, DCS-Transformer is trained by an improved training strategy which tunes the hyperparameters $\lambda,\eta,t_{`warm`}$ by cross-validation. The details are presented in Section B of the supplementary of the revised paper. The value of $\lambda$ is selected from $\{0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0,5\}$. The value of $\eta$ is selected from $\{0.1, 0.5, 1, 5, 10, 50, 100\}$. The value of $t_{\text{warm}}$, which is the number of warm-up epochs, is selected from $\{0.1t_{`train`}, 0.2t_{`train`}, 0.3t_{`train`}, 0.4t_{`train`}, 0.5t_{`train`}, 0.6t_{`train`}\}$, where $t_{`train`}=300$ is the total number of training epochs. We select the values of $\eta$, $\lambda$, and $t_{`warm`}$ that lead to the smallest validation loss. It is revealed that $t_{`warm`} = 90$ is chosen for all the three visual transformers in our experiments. Using the searched hyperparameters by cross-validation which is a principled method for tuning hyperparameters, we have obtained significantly improved results of DCS-Transformers shown in Table 1, and the same results are also reflected in the revised paper. For exmaple, compared to the top-1 accuracy of 79.9% of DCS-MobileViT-S (w/o IB Loss), DCS-MobileViT-S achieves a top-1 accuracy of 81.0%.

References

[1] Cai et al. "Efficientvit: Enhanced linear attention for high-resolution low-computation visual recognition." ICCV 2023.

[2] Chen, Tianlong, et al. "Chasing sparsity in vision transformers: An end-to-end exploration." NeurIPS 2021.

[3] Kong, Zhenglun, et al. "SPViT: Enabling faster vision transformers via soft token pruning." ECCV 2022.

[4] Zheng, Chuanyang, et al. "SAViT: Structure-Aware Vision Transformer Pruning via Collaborative Optimization." NeurIPS 2022.

We appreciate the review and the suggestions in this review. The raised issues are addressed below.

(1) More detailed motivation for the IB Loss and Regular Cross-Entropy Cannot Decrease the IB Loss Enough

More detailed motivation for the incorporate of information bottleneck loss is presented in Section 1 of the revised paper. It is attached below for your convenience.

Motivation

A typical transformer block can be written as $`Output` = `MLP` \left(\sigma(QK^{\top})\times V \right)$ where $Q,K,V \in \mathbb{R}^{N \times D}$ denote the query, key, and value respectively with $N$ being the number of tokens and $D$ being the input channel number. $\sigma(\cdot)$ is an operator, such as Softmax, which generates the attention weights or affinity between the tokens. We refer to $W = \sigma(QK^{\top}) \in \mathbb{R}^{N \times N}$ as the affinity matrix between the tokens. $\textup{MLP}$ (multi-layer perceptron network) generates the output features of the transformer block. There are $D$ channels in the input and output features of the MLP, and $D$ is also the channel of the attention outputs. Due to the fact that the MLP accounts for a considerable amount of FLOPs in a transformer block, the size and FLOPs of a transformer block can be significantly reduced by reducing the channels of the attention outputs from $D$ to a much smaller $\tilde D$. Our goal is to prune the attention output channels while maintaining and even improving the prediction accuracy of the original transformer. However, directly reducing the channels attention outputs, even by carefully designed methods, would adversely affect the performance of the model. In this paper, we propose to maintain or even improve the prediction accuracy of a visual transformer with pruned attention outputs channels by computing a more informative affinity matrix $W$ through selecting informative channels in the query $Q$ and the key $V$. That is, only selected columns of $Q$, which correspond to the same selected rows of $K^{\top}$, are used to compute the affinity matrix $W = \sigma(QK^{\top})$, which is refered to as channel selection for attention weights and illustrated in Figure 1a. We note that the attention outputs, which are also the input features to the MLP, is $W\times V$, and every input feature to the MLP is an aggregation of the rows of the value $V$ using the attention weights in $W$. As a result, pruning the channels of $W\times V$ amounts to pruning the channels of $V$ in the weighted aggregation. If the affinity $W$ is more informative, it is expected that a smaller number of features (rows) in $V$ contribute to such weighted aggregation, and the adverse effect of channel selection on the prediction accuracy of the transformer network is limited. Importantly, with a very informative affinity $W$, every input feature of the MLP is obtained by aggregation of the most relevant features (rows) in $V$, which can even boost the performance of visual transformers after channel selection or pruning of the attention outputs.

The idea of channel selection for the attention weights can be viewed from the perspective of Information Bottleneck (IB). Let $X$ be the input training features, $\tilde X$ be the learned features by the network, and $Y$ be the ground truth training labels for a classification task. The principle of IB is maximizing the mutual information between $\tilde X$ and $Y$ while minimizing the mutual information between $\tilde X$ and $X$. That is, IB encourages the network to learn features more correlated with the class labels while reducing their correlation with the input. Extensive empirical and theoretical works have evidenced that models respecting the IB principle enjoys compelling generalization. With channel selection for the attention weights, every feature in the attention outputs aggregates less features of the value $V$, so the attention outputs are less correlated with the training images so the IB principle is better adhered. This is reflected in Table 5 in Section C.2 of the supplementary, where a model for ablation study with channel selection for attention weights, DCS-Arch1 w/o IB Loss, enjoys less IB loss and higher top-1 accuracy than the vanilla transformer, MobileViT-S. It is noted that the model, DCS-Arch1 w/o IB Loss, only uses the regular cross-entropy loss in the retraining step, and smaller IB loss indicates that the IB principle is better respected. In order to further decrease the IB loss, we propose an Information Bottleneck (IB) inspired channel selection for the attention weights $\mathbf W$ where the learned attention weights can be more informative by explicitly optimizing the IB loss for visual transformers. Our model termed ``DCS-MobileViT-S'' in Table 5 is the visual transformer with the IB loss optimized, so that more informative attention weights are learned featuring even smaller IB loss and even higher top-1 accuracy.

We appreciate the review and the suggestions in this review. The raised issues are addressed below.

(1) More Detailed Motivation

More detailed motivation and ideas for the incorporate of information bottleneck in the proposed architecture and the rationale behind it are presented in Section 1 of the revised paper. It is attached below for your convenience.

Motivation

A typical transformer block can be written as $`Output` = `MLP` \left(\sigma(QK^{\top})\times V \right)$ where $Q,K,V \in \mathbb{R}^{N \times D}$ denote the query, key, and value respectively with $N$ being the number of tokens and $D$ being the input channel number. $\sigma(\cdot)$ is an operator, such as Softmax, which generates the attention weights or affinity between the tokens. We refer to $W = \sigma(QK^{\top}) \in \mathbb{R}^{N \times N}$ as the affinity matrix between the tokens. $\textup{MLP}$ (multi-layer perceptron network) generates the output features of the transformer block. There are $D$ channels in the input and output features of the MLP, and $D$ is also the channel of the attention outputs. Due to the fact that the MLP accounts for a considerable amount of FLOPs in a transformer block, the size and FLOPs of a transformer block can be significantly reduced by reducing the channels of the attention outputs from $D$ to a much smaller $\tilde D$. Our goal is to prune the attention output channels while maintaining and even improving the prediction accuracy of the original transformer. However, directly reducing the channels attention outputs, even by carefully designed methods, would adversely affect the performance of the model. In this paper, we propose to maintain or even improve the prediction accuracy of a visual transformer with pruned attention outputs channels by computing a more informative affinity matrix $W$ through selecting informative channels in the query $Q$ and the key $V$. That is, only selected columns of $Q$, which correspond to the same selected rows of $K^{\top}$, are used to compute the affinity matrix $W = \sigma(QK^{\top})$, which is refered to as channel selection for attention weights and illustrated in Figure 1a. We note that the attention outputs, which are also the input features to the MLP, is $W\times V$, and every input feature to the MLP is an aggregation of the rows of the value $V$ using the attention weights in $W$. As a result, pruning the channels of $W\times V$ amounts to pruning the channels of $V$ in the weighted aggregation. If the affinity $W$ is more informative, it is expected that a smaller number of features (rows) in $V$ contribute to such weighted aggregation, and the adverse effect of channel selection on the prediction accuracy of the transformer network is limited. Importantly, with a very informative affinity $W$, every input feature of the MLP is obtained by aggregation of the most relevant features (rows) in $V$, which can even boost the performance of visual transformers after channel selection or pruning of the attention outputs.

The idea of channel selection for the attention weights can be viewed from the perspective of Information Bottleneck (IB). Let $X$ be the input training features, $\tilde X$ be the learned features by the network, and $Y$ be the ground truth training labels for a classification task. The principle of IB is maximizing the mutual information between $\tilde X$ and $Y$ while minimizing the mutual information between $\tilde X$ and $X$. That is, IB encourages the network to learn features more correlated with the class labels while reducing their correlation with the input. Extensive empirical and theoretical works have evidenced that models respecting the IB principle enjoys compelling generalization. With channel selection for the attention weights, every feature in the attention outputs aggregates less features of the value $V$, so the attention outputs are less correlated with the training images so the IB principle is better adhered. This is reflected in Table 5 in Section C.2 of the supplementary, where a model for ablation study with channel selection for attention weights, DCS-Arch1 w/o IB Loss, enjoys less IB loss and higher top-1 accuracy than the vanilla transformer, MobileViT-S. It is noted that the model, DCS-Arch1 w/o IB Loss, only uses the regular cross-entropy loss in the retraining step, and smaller IB loss indicates that the IB principle is better respected. In order to further decrease the IB loss, we propose an Information Bottleneck (IB) inspired channel selection for the attention weights $\mathbf W$ where the learned attention weights can be more informative by explicitly optimizing the IB loss for visual transformers. Our model termed ``DCS-MobileViT-S'' in Table 5 is the visual transformer with the IB loss optimized, so that more informative attention weights are learned featuring even smaller IB loss and even higher top-1 accuracy.

(2) DCS-Transformer for ViT and Swin, and DCS-Transformer for the Task of Segmentation

The proposed DCS-Transformer module can be straightforwardly applied to a broader class of visual transformers beyond the examples covered in this paper, including the vanilla ViT [1] and Swin [2]. This is because the two components of a DCS-Transformer, channel selection for attention weights and channel selection for attention outputs can be directly applied to ViT and Swin. Here we replace all the transformer blocks in ViT-S/16 [1] and Swin-T [2], obtaining DCS-ViT-S/16 and DCS-Swin-T respectively. We use the same settings for search and training as described in Section 4.1 of our paper. All the models are trained on the training set of ImageNet-1K and tested on the validation set of ImageNet-1K. The results are shown in the Table below.

In Section E of the supplementary of the revised paper, we present the results of DCS-EfficientViT-B1 for the task of instance segmentation on the COCO dataset under the same setting as [3]. Please refer to Section E for more details. The results in Table 7 of the revised paper are copied to the table below for your convenience. We report the mean bounding box Average Precision (APb) and mean mask Average Precision (APm) as well as bounding box Average Precision (APb) and mask Average Precision (APm) under IoU thresholds of 0.5 and 0.75. It can be observed that DCS-EfficientViT-B1 consistently improves the performance of segmentation across various thresholds.

(3) Comparison with Pruning Methods

We compare DCS-Transformer to the current state-of-the-art pruning methods for visual transformers, the results are shown in the table below. DCS-Transformer is compared to S$^2$ViTE[2], SPViT[3], and SAViT [4] on EfficientViT-B1 (r224) [1]. For S$^2$ViTE[2], SPViT[3] and SAViT [4], the pruning is performed on the ImageNet training data. After obtaining the pruned networks, we fine-tune the pruned networks using the same setting as [1]. It can be observed that DCS-EfficientViT-B1 outperforms all pruned models by at least 1.9% in top-1 accuracy with even less parameters and FLOPs.

We also include actual inference time of all the model on a NVIDIA V100 GPU.

References

[1] Dosovitskiy et al. "An image is worth 16x16 words: Transformers for image recognition at scale." ICLR 2021.

[2] Liu et al. "Swin transformer: Hierarchical vision transformer using shifted windows." ICCV, 2021.

[3] Liu et al. "EfficientViT: Memory Efficient Vision Transformer with Cascaded Group Attention." CVPR 2023.

[4] Cai et al. "Efficientvit: Enhanced linear attention for high-resolution low-computation visual recognition." ICCV 2023.

[5] Chen, Tianlong, et al. "Chasing sparsity in vision transformers: An end-to-end exploration." NeurIPS 2021.

[6] Kong, Zhenglun, et al. "SPViT: Enabling faster vision transformers via soft token pruning." ECCV 2022.

[7] Zheng, Chuanyang, et al. "SAViT: Structure-Aware Vision Transformer Pruning via Collaborative Optimization." NeurIPS 2022.