Towards Better Multi-head Attention via Channel-wise Sample Permutation

## Abstract

Abstract Transformer plays a central role in many fundamental deep learning models, e.g., the ViT in computer vision and the BERT and GPT in natural language processing, whose effectiveness is mainly attributed to its multi-head attention (MHA) mechanism.
In this study, we propose a simple and novel channel-wise sample permutation (CSP) operator, achieving a new structured MHA with fewer parameters and lower complexity.
Given an input matrix, CSP circularly shifts the samples of different channels with various steps and then sorts grouped samples of each channel.
This operator is equivalent to implicitly implementing cross-channel attention maps as permutation matrices, which achieves linear complexity and suppresses the risk of rank collapse when representing data.
We replace the MHA of some representative models with CSP and test the CSP-based models in several discriminative tasks, including image classification and long sequence analysis.
Experiments show that the CSP-based models achieve comparable or better performance with fewer parameters and lower computational costs than the classic Transformer and its state-of-the-art variants.
The code is available at https://github.com/DaShenZi721/CSP .

## 1 Introduction

Transformer [] has been widely adopted in the deep learning domain.
Recent large language models like GPT [, ] and LLaMA [, ] series are built based on the Transformer and its variants, which demonstrate their remarkable abilities in natural language processing.
In the field of computer vision, Vision Transformers (ViTs) [], such as EfficientViT [, ] and SHViT [], exhibit exceptional performance and consistently push their limits.
In addition, the Transformer-based models have been designed for the complex structured data in various applications, including the Informer [] for time series broadcasting, the Transformer Hawkes process [] for continuous-time event sequence prediction, the Graphormer [] for molecular representation, the Mesh Transformer [] for 3D mesh representation, the Set-Transformer [] and Point-Transformer [] for point cloud modeling, and so on.
Although some new alternatives like Mamba [] and RWKV [] have been proposed and shown their competitiveness in some aspects, Transformer still maintains a dominant position when developing deep learning models because of its strong performance and outstanding universality.

The effectiveness of Transformer is mainly attributed to its multi-head attention (MHA) mechanism [].
However, MHA’s quadratic complexity concerning sequence length leads to a heavy, even unaffordable, computational overhead when modeling long sequences.
To improve the efficiency of MHA, many variants of Transformer introduce sparse or low-rank structures into attention maps [, , , , ] and apply algorithms friendly to GPU acceleration [, ].
At the same time, many attempts have been made to explore the mathematical reasons for the power of MHA, e.g., analyzing the representation power and rank collapse risk of MHA [, ] and revisiting attention maps through the lens of kernel theory [, ] and optimal transport [, ].
Currently, the above two research directions seem “parallel” in most situations: The acceleration methods of MHA are often empirical, but the theoretical work mainly analyzes the classic MHA, making it seldom support the rationality of the accelerated MHAs or contribute to developing a new MHA.

Figure: Figure 1: An illustration of the proposed channel-wise sample permutation operator and the equivalent implicit cross-channel attention maps.
Refer to caption: x1.png

In this study, we propose a novel Channel-wise Sample Permutation (CSP) operator, which leads to a new multi-head attention mechanism that is solid in theory and efficient in practice.
As illustrated in Figure 1, given an input matrix, CSP first shifts the samples of different channels circularly with various steps and then sorts grouped samples of each channel.
This operator is equivalent to implicitly implementing cross-channel attention maps as permutation matrices, which introduce inter- and intra-group interactions for the samples across different channels.
CSP is much simpler than the classic MHA and its existing variants.
It has no learnable parameters and can achieve linear computational complexity regarding sequence length.

The proposed CSP operator is motivated by the recent development of MHA.
In particular, the work in [, , , ] empirically demonstrate the rationality of pursuing attention maps with sparse doubly stochastic structures, which is further verified by an analytic experiment in this study.
CSP achieves permutation-based implicit attention maps that satisfy these structural properties, and thus, it has a good chance of providing a better MHA mechanism.
Moreover, such attention maps have all-one spectrums because of their permutation nature.
Based on the theoretical analysis framework provided in [], we prove that replacing MHA with CSP can suppress the risk of rank collapse when representing data.
In addition, we provide insightful understandings of the CSP operator by explaining its circular shifting and group sorting steps from the perspectives of optimal transport-based attention layer [] and channel-wise mixer [, ], respectively.

To demonstrate the usefulness of CSP, we replace the MHA of some state-of-the-art models with CSP and compare the CSP-based models with the original MHA-based ones in representative discriminative tasks, including long sequence analysis and image classification.
For each model, replacing its MHA with CSP significantly reduces the number of parameters and the computational cost while maintaining or even improving model performance.

## 2 Preliminaries and Related Work

Typically, given an input $\bm{X}\in\mathbb{R}^{N\times C}$, where $N$ indicates the length of a sequence or the size of a sample set and $C$ is the number of channels (feature dimensions), an attention head [] first obtains the value, query, and key matrices by linear maps, i.e., $\bm{V}=\bm{X}\bm{W}_{V}\in\mathbb{R}^{N\times D}$, $\bm{Q}=\bm{X}\bm{W}_{Q}\in\mathbb{R}^{N\times D}$, and $\bm{K}=\bm{X}\bm{W}_{K}\in\mathbb{R}^{N\times D}$, and then projects $\bm{V}$ as follows:

$$ $\displaystyle\begin{aligned} \text{Att}(\bm{V};\bm{Q},\bm{K}):=\bm{P}(\bm{Q}, \bm{K})\bm{V},~{}\text{where}~{}\bm{P}(\bm{Q},\bm{K})=\text{Softmax}\Bigl{(} \frac{\bm{QK}^{\top}}{\sqrt{D}}\Bigr{)}.\end{aligned}$ (1) $$

Here, we denote $\bm{V}$ as the input matrix of the head and $\bm{P}(\bm{Q},\bm{K})\in\mathbb{R}^{N\times N}$ as the attention map parametrized by $\bm{Q}$ and $\bm{K}$, respectively.
The multi-head attention mechanism applies a group of linear maps, i.e., $\theta=\{\bm{W}_{V,m},\bm{W}_{Q,m},\bm{W}_{K,m}\in\mathbb{R}^{C\times D}\}_{m=
1}^{M}$, to construct $M$ attention heads and concatenates their outputs, i.e.,

$$ $\displaystyle\begin{aligned} \text{MHA}_{\theta}(\bm{X}):=\|_{m=1}^{M}\text{ Att}(\bm{V}_{m};\bm{Q}_{m},\bm{K}_{m})\in\mathbb{R}^{N\times MD},\end{aligned}$ (2) $$

where $\bm{V}_{m}=\bm{X}\bm{W}_{V,m}$, $\bm{Q}_{m}=\bm{X}\bm{W}_{Q,m}$, and $\bm{K}_{m}=\bm{X}\bm{W}_{K,m}$ for $m=1,...,M$, and “$\|$” denotes the concatenation operation.
In practice, we set $MD=C$ for applying skip connections in the Transformer architecture, i.e., $\text{MHA}_{\theta}(\bm{X})+\bm{X}$.

The attention map in (1) has quadratic computational complexity concerning the sequence length $N$ because of its “query-key-value” (abbreviately, QKV) architecture.
Considering the high complexity per attention head, the MHA has to restrict the number of attention heads to achieve a trade-off between model capacity and computational efficiency, which may limit its representation power.

Many efforts have been made to improve the classic MHA.
SparseTrans [] and Longformer [] compute local attention maps based on the subsequences extracted by sliding windows, which leads to sparse global attention maps.
To use shorter subsequences while retaining more information, $\text{S}^{3}$Attention [] integrates global and local information by leveraging Fourier Transformation and a convolutional kernel.
Some other models sparsify the key and query matrices directly by locality-sensitive hashing (LSH) [] or ReLU [].
Besides pursuing sparse attention maps, Performer [] and Linformer [] apply low-rank attention maps.
Recently, FlashAttention and its variants [, ] further accelerate the computation of attention maps for long sequences by sophisticated I/O design, parallelism, and work partitioning.
In addition to simplifying the computation of the attention maps, some work provides new understandings of the attention mechanism.
The work in [, , ] implements attention maps as various kernel matrices.
The work in [] implements doubly stochastic attention maps by the Sinkhorn-Knopp algorithm [] and explains the computation of each attention map as a discretized Wasserstein gradient flow.

Currently, the above accelerated or structured MHAs often lead to the performance degradation, while the theoretical understandings of MHA seldom help improve its computational efficiency in practice.
Our work attempts to bridge the gap, proposing a theoretically solid multi-head attention mechanism with low complexity and competitive performance.

## 3 Proposed Method

### 3.1 Motivation: Pursuing Sparse Doubly Stochastic Attention Maps

As shown in Section 2, many models apply various strategies to construct sparse attention maps, e.g., the locality-sensitive hashing (LSH) in [], the subsequence sampling in [, ], and the sparse activation in [].
These models achieve encouraging performance and higher efficiency than the vanilla Transformer, demonstrating sparse attention maps’ rationality.
Besides making attention maps sparse, the work in [] shows that in various discriminative tasks, the attention maps tend to be doubly stochastic automatically (i.e., $\bm{P}\in\Pi_{N}$, where $\Pi_{N}=\{\bm{A}\geq\bm{0}|\bm{A}\bm{1}_{N}=\bm{1}_{N},\bm{A}^{\top}\bm{1}_{N}
=\bm{1}_{N}$) during training,(^1^11Please refer to Section 3 in [] for more details.) and the Transformer applying doubly stochastic attention (called Sinkformer) outperforms the vanilla Transformer in image and text classification.

The above recent models show that sparse attention maps help improve the models’ computational efficiency (thus making increasing attention heads feasible), and doubly stochastic attention maps help improve the models’ discriminative power.
These phenomena imply that designing sparse doubly stochastic attention maps may lead to a better MHA mechanism and further boost model performance.
To verify this claim, we conduct an analytic experiment, replacing the attention maps in a standard ViT [] with simple permutation matrices (the doubly stochastic matrices with the strongest sparsity) and evaluating the model performance on the CIFAR-10 dataset [].
In particular, the ViT used in this experiment consists of six Transformer layers.
Each Transformer has eight attention heads (i.e., $M=8$), and each head sets $N=64$, $C=512$, and $D=64$.
For each layer, we replace the attention map of the $m$-th head with the following permutation matrix:

$$ $\displaystyle\bm{S}_{{C(m-1)}/{D}}=\begin{bmatrix}\bm{0}&\bm{I}_{{C(m-1)}/{D}} \\ \bm{I}_{N-{C(m-1)}/{D}}&\bm{0}\end{bmatrix},~{}\text{for}~{}m=1,...,M,$ (3) $$

where $\bm{I}_{N}$ indicates an identity matrix with a size $N\times N$.
Obviously, the permutation matrix $\bm{S}_{{C(m-1)}/{D}}$ corresponds to a circular shifting operator — $\bm{S}_{{C(m-1)}/{D}}\bm{V}_{m}$ means shifting the rows of $\bm{V}_{m}$ circularly with ${C(m-1)}/{D}$ steps.
Furthermore, for each layer, we can concatenate $\{\bm{V}_{m}\}_{m=1}^{M}$ to get $\bm{V}=[\bm{v}_{1},...,\bm{v}_{C}]\in\mathbb{R}^{N\times C}$ and circularly shift the channels of this matrix by applying $\bm{S}_{(c-1)\bmod{N}}\bm{v}_{c}$ for $c=1,...,C$, where “$\bmod$” is the modulo operation.
In this case, the number of attention heads, equal to the number of distinguishable permutation matrices, becomes $N$.
As shown in Table 1, even if the sparse doubly stochastic attention maps we designed are extremely simple and have no parameters, applying them with a sufficient number can still result in competitive, even better performance.
This experimental result motivates us to construct sufficiently many sparse doubly stochastic attention maps with low complexity, leading to the proposed channel-wise sample permutation operator.

**Table 1: A comparison for various MHAs and their classification accuracy (%) on CIFAR-10.**
| MHA | #Heads per layer | Parameters per layer | Top-1 Acc. | Top-5 Acc. |
| --- | --- | --- | --- | --- |
| $\|_{m=1}^{M}\bm{P}(\bm{Q}_{m},\bm{K}_{m})\bm{V}_{m}$ | $8~{}(=M)$ | $\{\bm{W}_{Q,m},\bm{W}_{K,m},\bm{W}_{V,m}\}_{m=1}^{M}$ | 81.90 | 98.85 |
| $\|_{m=1}^{M}\bm{S}_{C(m-1)/D}\bm{V}_{m}$ | $8~{}(=M)$ | $\{\bm{W}_{V,m}\}_{m=1}^{M}$ | 80.70 | 98.97 |
| $\|_{c=1}^{C}\bm{S}_{(c-1)\bmod{N}}\bm{v}_{c}$ | $64~{}(=N)$ | $\{\bm{W}_{V,m}\}_{m=1}^{M}$ | 83.84 | 99.27 |

### 3.2 Channel-wise Sample Permutation for Implicit Cross-Channel Attention

As shown in Figure 1, given an input matrix $\bm{X}\in\mathbb{R}^{N\times C}$, the CSP operator first projects $\bm{X}$ to a value matrix with the same size, i.e., $\bm{V}=\bm{X}\bm{W}=[\bm{v}_{1},...,\bm{v}_{C}]\in\mathbb{R}^{N\times C}$, where $\bm{W}\in\mathbb{R}^{C\times C}$ and $\bm{v}_{c}$ denotes the $N$ samples in the $c$-th channel.
Given $\bm{V}$, the CSP operator shifts the samples of different channels circularly with various steps and then sorts grouped samples of each channel, i.e.,

$$ $\displaystyle\begin{aligned} \text{CSP}_{\bm{W}}(\bm{X}):=\|_{c=1}^{C}\text{ GSort}_{K}(\bm{S}_{J_{c}}\bm{v}_{c})=\|_{c=1}^{C}\bm{P}_{c}\bm{v}_{c},~{}~{} \text{where}~{}\text{GSort}_{K}(\bm{v})=\|_{k=1}^{K}\text{Sort}(\bm{v}^{(k)}). \end{aligned}$ (4) $$

Here, $\bm{S}_{J_{c}}$ is the circular shifting operator defined in (3).
$\text{GSort}_{K}(\bm{v})$ denotes grouping the elements of a vector $\bm{v}$ into $K$ parts, i.e., $\bm{v}=[\bm{v}^{(1)};...;\bm{v}^{(K)}]$, and sorting each part accordingly.
When implementing the CSP operator, we take the first channel $\bm{v}_{1}$ as the reference in this study:
The circular shifting of each $\bm{v}_{c}$ is with respect to $\bm{v}_{1}$, and the group sorting permutes the elements of $(\bm{S}_{J_{c}}\bm{v}_{c})^{(k)}$ according to the element-wise order of $\bm{v}_{1}^{(k)}$, for $c=2,...,C$ and $k=1,...,K$.

The CSP operator is equivalent to implicitly implementing sparse doubly stochastic attention maps as permutation matrices, which builds interactions for the samples across different channels.
As shown in (4), we denote each attention map as $\bm{P}_{c}$.
For $\bm{v}_{1}$, $\bm{P}_{1}=\bm{I}_{N}$.
For the remaining $\bm{v}_{c}$, $\bm{P}_{c}$ can be decomposed into the following two parts:

$$ $\displaystyle\bm{P}_{c}=\bm{T}_{c}\bm{S}_{J_{c}}=\text{BlkDiag}(\{\bm{T}_{c}^{ (k)}\}_{k=1}^{K})\bm{S}_{J_{c}},~{}\text{for}~{}c=2,...,C,$ (5) $$

where $\bm{T}_{c}$ is a block-diagonal permutation matrix determined by the group sorting operation.
The $k$-th block $\bm{T}_{c}^{(k)}$ is a permutation matrix determined by the sorting within the $k$-th group, which introduces intra-group sample interactions across different channels.
The circular shifting operation introduces inter-group sampler interactions across different channels, and the ranges of the interactions are determined by the predefined shifting steps.
As a result, for arbitrary two $\bm{v}_{c}$ and $\bm{v}_{c^{\prime}}$, $\bm{P}_{c}\bm{v}_{c}$ and $\bm{P}_{c^{\prime}}\bm{v}_{c^{\prime}}$ captures their interactions determined by $\bm{P}_{c}^{\top}\bm{P}_{c^{\prime}}$.

**Table 2: A comparison for existing MHA mechanisms and CSP.**
| Model | $\text{Attention}(\bm{V};\bm{Q},\bm{K})$ | Complexity | Attention Structure |
| --- | --- | --- | --- |
| Transformer | $\text{Softmax}\left(\frac{\bm{Q}\bm{K}^{\top}}{\sqrt{D}}\right)\bm{V}$ | $\mathcal{O}(CN^{2})$ | Row-normalized |
| SparseTrans | $\text{Local2D-Softmax}\left(\frac{\bm{Q}\bm{K}^{\top}}{\sqrt{D}}\right)\bm{V}$ | $\mathcal{O}(CN^{1.5})$ | Sparse+Row-normalized |
| Longformer | $\text{Local1D-Softmax}\left(\frac{\bm{Q}\bm{K}^{\top}}{\sqrt{D}}\right)\bm{V}$ | $\mathcal{O}(CNE)$ | Sparse+Row-normalized |
| Reformer | $\text{LSH-Softmax}\left(\frac{\bm{Q}\bm{K}^{\top}}{\sqrt{D}}\right)\bm{V}$ | $\mathcal{O}(CN\log N)$ | Sparse+Row-normalized |
| CosFormer | $(\bm{Q}_{\cos}\bm{K}_{\cos}^{\top}+\bm{Q}_{\sin}\bm{K}_{\sin}^{\top})\bm{V}$ | $\mathcal{O}(\min\{CE_{QK},NE_{Q}\})$ | Sparse |
| MEGA | $f\left(\frac{\bm{Q}\bm{K}^{\top}}{\sqrt{D}}+\bm{B}\right)\bm{V}$ | $\mathcal{O}(CN^{2})\sim\mathcal{O}(CNr)$ | (Optional) Sparse+Row-normalized |
| Performer | $\phi_{r}(\bm{Q})\phi_{r}(\bm{K})^{\top}\bm{V}$ | $\mathcal{O}(CNr)$ | Low-rank |
| Linformer | $\text{Softmax}\left(\frac{\bm{Q}\psi_{r}(\bm{K})^{\top}}{\sqrt{D}}\right)\psi_<br>{r}(\bm{V})$ | $\mathcal{O}(CNr)$ | Low-rank+Row-normalized |
| Proposed | $\|_{c=1}^{C}\bm{P}_{c}\bm{v}_{c}$ | $\mathcal{O}(CN\log\frac{N}{K})\sim\mathcal{O}(CN)$ | Sparse+Doubly stochastic |

#### 3.2.1 Advantages over Existing MHAs

High computational efficiency:
Replacing MHA with CSP leads to a new variant of Transformer.
Table 2 compares the proposed model with the existing MHA-based models.
We can find that the computational complexity of CSP can be $\mathcal{O}(N\log\frac{N}{K})$ when applying QuickSort [] to implement the group sorting operation, which is much lower than the computational complexity of the existing MHAs.
When the group size is $2$, we can achieve group sorting by the simple “$\min$-$\max$” operation [, ], and the computational complexity further reduces to $\mathcal{O}(N)$.
In addition, as shown in (4), except for the projection matrix $\bm{W}$ corresponding to the value matrix, CSP does not require additional projection matrices to construct the query and key matrices.
In other words, its parameters are only one-third of the classic MHA.

A low risk of rank collapse:
Besides significantly improving computational efficiency, CSP can suppress an ordinary risk of the classic MHA, rank collapse.
In particular, we define the rank-1 estimation residual of a matrix $\bm{X}$ associated with an arbitrary matrix norm as

$$ $\displaystyle\epsilon(\bm{X})=\bm{X}-\bm{1}\hat{\bm{x}}^{\top},~{}\text{where} ~{}\hat{\bm{x}}=\sideset{}{{}_{\bm{x}}}{\operatorname*{arg\,min}}\|\bm{X}-\bm{ 1}\bm{x}^{\top}\|.$ (6) $$

In addition, for a matrix $\bm{X}=[x_{nc}]\in\mathbb{R}^{N\times C}$, we can define its $(1,\infty)$-norm as $\|\bm{X}\|_{1,\infty}=\sqrt{\|\bm{X}\|_{1}\|\bm{X}\|_{\infty}}$, where $\|\bm{X}\|_{1}=\max_{c}\sum_{n=1}^{N}|x_{nc}|$ and $\|\bm{X}\|_{\infty}=\max_{n}\sum_{c=1}^{C}|x_{nc}|$, respectively.
It has been known that $\|\epsilon(\bm{X})\|_{1,\infty}$ measures the rank collapse of $\bm{X}$ effectively, i.e., $\|\epsilon(\bm{X})\|_{1,\infty}\rightarrow 0$ means that $\bm{X}$ collapses to a rank-1 matrix.
The work in [] shows that if we construct a Transformer by stacking MHA layers without skip connections, its output matrix will lose its rank doubly exponentially with depth, i.e., $\|\epsilon(\text{MHA}_{L}\circ\cdots\circ\text{MHA}_{1}(\bm{X}))\|_{1,\infty}=
\mathcal{O}(\|\epsilon(\bm{X})\|_{1,\infty}^{3^{L}})$, where $L$ is the number of the MHA layers.

Applying CSP can suppress this risk, which is supported by the following theorem.

###### Theorem 1 .

Suppose that we construct a layer-$L$ network as $(f\circ\text{CSP})^{L}=(f_{\lambda_{L}}\circ\text{CSP}_{\bm{W}^{(L)}})\circ
\cdots\circ(f_{\lambda_{1}}\circ\text{CSP}_{\bm{W}^{(1)}})$.
For $\ell=1,...,L$, $\text{CSP}_{\bm{W}^{(\ell)}}$ is a $C$-channel CSP operator, and $f_{\lambda_{\ell}}:\mathbb{R}^{C}\mapsto\mathbb{R}^{C}$ is a $\lambda_{\ell}$-Lipschitz function.
Denote $\beta=\max_{\ell}\|\bm{W}^{(\ell)}\|_{1}$ and $\lambda=\max_{\ell}\lambda_{\ell}$.
Then, we have

$$ $\displaystyle\|\epsilon((f\circ\text{CSP})^{L}(\bm{X}))\|_{1,\infty}\leq C^{ \frac{L}{2}}(\lambda\beta)^{L}\|\epsilon(\bm{X})\|_{1,\infty},~{}\forall~{}\bm {X}\in\mathbb{R}^{N\times C}.$ (7) $$

Theorem 1 indicates a linear convergence rate of the residual.
It means that the model applying CSP avoids the rapid decay of the matrix rank.
A detailed proof is shown in Appendix A.

###### Theorem 1 .

#### 3.2.2 Implementation Details

Circular shifting:
The shifting step is crucial for the circular shifting operation, which determines the range of sample interaction.
When the sequence length is comparable to the number of channels in each layer, i.e., $N\approx C$, we can simply set the shifting step size $J_{c}=c\lceil\frac{N}{C}\rceil$ for $c=1,...,C$, so that the circular shifting operation can generate sufficient distinguishable attention heads with respect to the sequence length.
However, for long sequences, i.e., $N\gg C$, we need to set the shifting steps of different channels with high dynamics, making $C$ attention heads build diverse interactions in a long sequence.
In this study, given a Transformer-based model with $L$ layers, we consider all $L$ value matrices in these layers jointly, and set $LC$ different shifting steps based on power law, as illustrated in Figure 2.
In particular, we denote $J$ as the base shifting step.
For $c=1,...,LC$, we shift the $c$-th channel circularly with $J^{c-1}-1$ steps.
In addition, for the last channel, we require $J^{LC-1}-1\approx N-1$.
Therefore, we can set $J=\lfloor N^{{1}/{(LC-1)}}\rfloor$.

Figure: Figure 2: The shifting strategies when $N\approx C$ and $N\gg C$.
Refer to caption: x2.png

Group sorting:
Instead of merely applying the circular shifting operation (as in Section 3.1), we introduce the group sorting operation to CSP, which helps increase the number of attention heads.
Given an input matrix with size $N\times C$, the circular shifting operation constructs $\min\{N,C\}$ different attention maps, which results in repeated attention maps when $C>N$.
For the channels applying the same circular shifting steps, the group sorting operation can make their attention maps different from each other as long as the orders of their samples are inconsistent.
As a result, the group sorting helps CSP increase the number of attention heads from $\min\{C,N\}$ to $C$.

A special case of CSP:
Note that when setting $K=1$, the group sorting becomes the classic complete sorting, leading to a special case of CSP.
Given $C$ channels, the complete sorting can directly generate at most $C$ distinguishable permutation matrices/attention heads.
In addition, because of using the complete sorting, the circular shifting step of CSP becomes redundant.
In the following experiments, empirically, implementing CSP as the complete sorting often works well when modeling long sequences while the CSP combining circular shifting with group sorting helps represent visual objects.

### 3.3 Functionality and Rationality Analysis of CSP

Circular shifting works as a channel-wise mixer:
The circular shifting of CSP is similar to the channel-wise mixers used in visual representation models.
In particular, the convolution neural networks like ShuffleNet [] and its variant [] apply grouped convolution operation to reduce computational costs and increase inter-group interactions by shuffling the channels across different groups.
This shuffling strategy inspires many lightweight channel-wise mixers, e.g., the hierarchical rearrangement in Hira-MLP [], the spatial-shift module in S^2-MLP [], and the axial-shift module in AS-MLP [].
For example, given a visual feature tensor with a size $H\times W\times C$ (i.e., 2D images with $C$ channels), the axial-shift module applies horizontal and vertical shifts with zero padding to the 2D images of different channels.
The spatial-shift module first divides the input tensor into four parts by grouping its channels and then shifts the four sub-tensors along four different directions.
Both these two modules apply small shifting steps to achieve local shifting.
The circular shifting of CSP corresponds to applying these shifting modules to 1D sequences.
To capture the short-range and long-range interactions between the samples of different channels simultaneously, we apply various shifting steps to different channels and replace zero padding with circular padding.

Group sorting works as an optimal transport-based MHA:
Without causing any ambiguity, we denote $(\bm{S}_{J_{c}}\bm{v}_{c})^{(k)}$ as $\bm{v}_{c}^{(k)}$ for simplification.
It is easy to prove that the $\bm{T}_{c}^{(k)}$ in (5) is the optimal transport (OT) between $\bm{v}_{1}^{(k)}$ and $\bm{v}_{c}^{(k)}$, which can be derived by $\min_{\bm{T}\in\Pi_{N/K}}\langle-\bm{v}_{1}^{(k)}\bm{v}_{c}^{(k)\top},\bm{T}\rangle$.(^2^22See Appendix B for a detailed derivation.)
From this viewpoint, CSP achieves a new OT-based MHA mechanism.
In addition, when approximating $\bm{T}_{c}^{(k)}$ as an entropic optimal transport, we can connect CSP to the doubly stochastic attention mechanism used in Sinkformer [].
In particular, each attention head in Sinkformer derives a doubly stochastic attention map, denoted as $\bm{T}_{T,\tau}$, by the Sinkhorn-Knopp algorithm [], i.e.,

$$ $\displaystyle\begin{aligned} \bm{T}_{t,\tau}=\text{Sinkhorn}_{t}\Bigl{(}\exp \Bigl{(}\frac{\bm{QK}^{\top}}{\tau\sqrt{D}}\Bigr{)}\Bigr{)},~{}\text{and}~{} \bm{T}_{\infty,\tau}=\sideset{}{{}_{\bm{T}\in\Pi_{N}}}{\operatorname*{arg\,min }}\langle-\bm{QK}^{\top},\bm{T}\rangle+\tau\sqrt{D}H(\bm{T}).\end{aligned}$ (8) $$

Here, $\text{Sinkhorn}_{t}(\bm{A})$ means normalizing the rows and columns of a nonnegative matrix $\bm{A}$ alternatively by $t$ times, i.e., $\bm{A}^{(0)}=\bm{A}$, and $\bm{A}^{(i)}=\text{N}_{c}\circ\text{N}_{r}(\bm{A}^{(i-1)})$ for $i=1,...,t$, where $\text{N}_{c}$ and $\text{N}_{r}$ denote column-wise and row-wise normalization, respectively.
As shown in (8), the attention map corresponds to the optimal solution of an entropic optimal transport problem [] when $t\rightarrow\infty$, where $\langle\cdot,\cdot\rangle$ denotes the inner product operation, $H(\bm{T})=\langle\bm{T},\log\bm{T}\rangle$ denotes the entropy of $\bm{T}$, and its significance is controlled by $\tau>0$.

We can connect Sinkformer to CSP by modifying the attention mechanism in (8) as follows.
Given a value matrix $\bm{V}=[\bm{v}_{1},...,\bm{v}_{C}]\in\mathbb{R}^{N\times C}$, we replace the $\bm{Q}$ and $\bm{K}$ in (8) with $\bm{v}_{1}$ and $\bm{v}_{c}$, respectively, for $c=1,...,C$, and divide each vector into $K$ groups.
We can achieve a $C$-head $K$-group doubly stochastic attention mechanism by applying the Sinkhorn-Knopp algorithm to $\bm{v}_{1}^{(k)}\bm{v}_{c}^{(k)\top}$, for $c=1,...,C$ and $k=1,...,K$, i.e.,

$$ $\displaystyle\bm{T}_{c,t,\tau}=\text{BlkDiag}\Bigl{(}\Bigl{\{}\text{Sinkhorn}_ {t}\Bigl{(}\exp\Bigl{(}\frac{1}{\tau}\bm{v}_{1}^{(k)}\bm{v}_{c}^{(k)\top}\Bigr {)}\Bigr{)}\Bigr{\}}_{k=1}^{K}\Bigr{)}=\text{BlkDiag}(\{\bm{T}_{c,t,\tau}^{(k) }\}_{k=1}^{K}),$ (9) $$

where $\bm{T}_{c,t,\tau}$ is a doubly stochastic attention map with a block-diagonal structure, and $\bm{T}_{c,t,\tau}^{(k)}$ denotes a local attention map, which corresponds to computing the entropic optimal transport between $\bm{v}_{1}^{(k)}$ and $\bm{v}_{c}^{(k)}$ when $t\rightarrow\infty$, i.e., $\bm{T}_{c,\infty,\tau}^{(k)}=\operatorname*{arg\,min}_{\bm{T}\in\Pi_{N/K}}
\langle-\bm{v}_{1}^{(k)}\bm{v}_{c}^{(k)\top},\bm{T}\rangle+\tau H(\bm{T})$.

The connection between the attention map in (9) and CSP is captured by the following theorem.

###### Theorem 2 .

If $\min_{\bm{T}\in\Pi_{N/K}}\langle-\bm{v}_{1}^{(k)}\bm{v}_{c}^{(k)\top},\bm{T}\rangle$ admits a unique optimal solution, for $c=1,...,C$ and $k=1,...,K$, then for the $\bm{T}_{c,t,\tau}$ in (9), $\lim_{\tau\rightarrow 0}\bm{T}_{c,\infty,\tau}\bm{S}_{J_{c}}$ converges to the $\bm{P}_{c}$ in (5) weakly.

Theorem 2 can be derived directly based on the weak convergence of entropic optimal transport (Theorem 5.10 in []).
This theorem indicates that a Sinkformer can implement CSP approximately if it $i)$ sets the query and key matrices as the channels of the value matrix and $ii)$ applies the Sinkhorn-Knopp algorithm to the grouped samples.

**Table 3: The comparison for various models on the number of parameters ($\times 10^{6}$) and classification accuracy (%). The best result on each dataset is bold, and the second best result is underlined.**
| Model | Attention | CIFAR-10 | CIFAR-100 | ImageNet-1k |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| #Param. | Top-1 | Top-5 | #Param. | Top-1 | Top-5 | #Param. | Top-1 | Top-5 |  |  |
| ViT | MHA | 9.52 | 81.90 | 98.85 | 9.65 | 53.30 | 79.97 | 22.05 | 76.53 | 92.81 |
| Circular Shifting | 6.38 | 83.84 | 99.27 | 6.50 | 58.38 | 84.26 | 18.50 | 75.64 | 92.42 |  |
| Group Sorting | 6.38 | 79.41 | 99.03 | 6.50 | 51.47 | 79.67 | 18.50 | 64.77 | 85.28 |  |
| CSP (Ours) | 6.38 | 84.81 | 99.35 | 6.50 | 59.16 | 84.76 | 18.50 | 76.66 | 93.05 |  |
| 9.52 | 85.02 | 99.37 | 9.65 | 59.23 | 85.09 | 22.05 | 77.14 | 93.23 |  |  |

Figure: Figure 3: The singular spectrums of the output matrices achieved on ImageNet-1k.
Refer to caption: x3.png

###### Theorem 2 .

## 4 Experiments

To demonstrate the effectiveness and efficiency of CSP, we conduct comprehensive comparative and analytic experiments in two representative discriminative tasks, image classification and long sequence analysis.
The implementation details can be found in Appendix C.

### 4.1 Image Classification

We conduct comparative experiments and ablation studies on three image datasets, including CIFAR-10, CIFAR-100 [], and ImageNet-1k [].
For each dataset, we treat the classic ViT as the baseline and replace its MHA layers with $i)$ circular shifting, $ii)$ group sorting, and $iii)$ the proposed CSP operator, respectively.
Here, the circular shifting and the group sorting are two simplified CSP operators that help analyze the contributions of different CSP modules.
Table 3 shows these models’ size and classification accuracy.
Applying CSP and its simplified variants can reduce the model size significantly without the query and key matrices.
The circular shifting operator achieves competitive performance in all three datasets.
In addition, although the standalone group sorting operator results in performance degradation, combining it with the circular shifting operator, i.e., the proposed CSP, can achieve the best performance.
These observations are consistent with the experimental results achieved by mixer-MLP models [, ]: $i)$ Simple channel-wise interactions can replace the dense and smoothed attention maps and lead to promising model performance, and $ii)$ the shifting operator is crucial for CSP in computer vision tasks because it fully leverages the local similarity nature of the image.
Moreover, when we increase the number of channels per layer and make the model size comparable to the original ViT, we can further boost the performance of the CSP-based models and achieve the best performance.

In Figure 3, we illustrate the singular spectrums of the output matrices achieved by different methods on ImageNet-1k.
The spectrums achieved by the circular shifting and CSP operators decay much more slowly than the spectrum achieved by MHA.
This observed phenomenon serves as a strong validation of the theoretical result in Theorem 1, providing further evidence that the representation model using permutation-based attention maps indeed carries a lower risk of rank collapse compared to the classic MHA-based model.

### 4.2 Long Range Arena Benchmark

Long Range Arena (LRA) is a benchmark designed to evaluate models for long sequence analysis [], which consists of six discriminative tasks, including ListOps [], byte-level text classification [], byte-level document retrieval [], and three sequentialized image classification tasks, i.e., CIFAR-10 [], PathFind, and Path-X [].(^3^33Given a set of gray-level images, each of which plots two points and several curves, PathFind aims to recognize whether there exists a path connecting the points in each image. Path-X is a more challenging version of Pathfind because of applying high-resolution images.)
Each image is formulated as a long sequence of pixels in the three image classification tasks.
We first replace the MHA of the classic Transformer with CSP and compare it with other variants of Transformer.
As shown in Figure 4 and the first part of Table 4, the Transformer using CSP outperforms other models on both performance and computational efficiency.
It achieves the highest average score and the fastest training speed among all the models, and its memory cost is comparable to the most efficient variant of Transformer.
For long sequence modeling, we simply implement CSP as the complete sorting operator in this experiment, which can capture the long-range interactions between the samples with the highest flexibility.

Figure: Figure 4: The performance and efficiency of various models on the LRA benchmark. The disk area indicates the memory cost of each method.
Refer to caption: x4.png

Besides improving the classic Transformer, we further plug CSP into the state-of-the-art attention-based model, MEGA [], and analyze its impacts on the model performance.
As evidenced in the second part of Table 4, the MEGA with dense attention maps currently outperforms all other methods, including those based on the state space model (SSM), such as S5 [] and SPADE [], on the LRA benchmark.
When MEGA applies chunked attention maps, its performance degrades slightly but its computational complexity can reduce from $\mathcal{O}(CN^{2})$ to $\mathcal{O}(CNr)$, where $r$ is the chunk size.
When replacing the attention mechanism of MEGA with the proposed CSP operator, the complexity of the CSP-based MEGA becomes linear and thus comparable to that of the chunked MEGA and the SSM-based models.
At the same time, the CSP-based MEGA is better than the chunked MEGA in the overall performance.
These results serve as compelling evidence, demonstrating the practical rationality of CSP.

**Table 4: Results (%) of various methods on the LRA benchmark. The first group contains the classic Transformer and its variants, and the second group contains the state-of-the-art methods on LRA. The best result on each dataset is bold, and the second best result is underlined.**
| Type | Model | ListOps | Text | Retrieval | Image | PathFind | Path-X | Avg. |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| MHA | Transformer [] | 36.37 | 64.27 | 57.46 | 42.44 | 71.40 | FAIL | 54.39 |  |
| LocalAttention [] | 15.82 | 52.98 | 53.39 | 41.46 | 66.63 | FAIL | 46.06 |  |  |
| LinearTrans [] | 16.13 | 65.90 | 53.09 | 42.34 | 75.30 | FAIL | 50.55 |  |  |
| Reformer [] | 37.27 | 56.10 | 53.40 | 38.07 | 68.50 | FAIL | 50.67 |  |  |
| Sinkformer [] | 30.70 | 64.03 | 55.45 | 41.08 | 64.65 | FAIL | 51.18 |  |  |
| SparseTrans [] | 17.07 | 63.58 | 59.59 | 44.24 | 71.71 | FAIL | 51.24 |  |  |
| SinkhornTrans [] | 33.67 | 61.20 | 53.83 | 41.23 | 67.45 | FAIL | 51.29 |  |  |
| Linformer [] | 35.70 | 53.94 | 52.27 | 38.56 | 76.34 | FAIL | 51.36 |  |  |
| Performer [] | 18.01 | 65.40 | 53.82 | 42.77 | 77.05 | FAIL | 51.41 |  |  |
| Synthesizer [] | 36.99 | 61.68 | 54.67 | 41.61 | 69.45 | FAIL | 52.88 |  |  |
| Longformer [] | 35.63 | 62.85 | 56.89 | 42.22 | 69.71 | FAIL | 53.46 |  |  |
| BigBird [] | 36.05 | 64.02 | 59.29 | 40.83 | 74.87 | FAIL | 55.01 |  |  |
| Cosformer [] | 37.90 | 63.41 | 61.36 | 43.17 | 70.33 | FAIL | 55.23 |  |  |
| Transformer using CSP (Ours) | 37.65 | 64.60 | 62.23 | 48.02 | 82.04 | FAIL | 58.91 |  |  |
| Type | Model | Complexity | ListOps | Text | Retrieval | Image | PathFind | Path-X | Avg. |
| CNN | CCNN [] | $\mathcal{O}(CN^{2})$ | 43.60 | 84.08 | FAIL | 88.90 | 91.51 | FAIL | 68.02 |
| SSM | ETSMLP [] | $\mathcal{O}(CN)$ | 62.55 | 88.49 | 86.72 | 75.34 | 91.66 | 93.78 | 83.09 |
| S4 [] | 58.35 | 76.02 | 87.09 | 87.26 | 86.05 | 88.10 | 80.48 |  |  |
| S5 [] | 62.15 | 89.31 | 91.40 | 88.00 | 95.33 | 98.58 | 87.46 |  |  |
| SPADE [] | 59.70 | 87.55 | 90.13 | 89.11 | 96.42 | 94.22 | 86.19 |  |  |
| MHA | MEGA [] | $\mathcal{O}(CN^{2})$ | 63.14 | 90.43 | 91.25 | 90.44 | 96.01 | 97.98 | 88.21 |
| MEGA-chunk [] | $\mathcal{O}(CNr)$ | 58.76 | 90.19 | 90.97 | 85.80 | 94.41 | 93.81 | 85.66 |  |
| MEGA using CSP (Ours) | $\mathcal{O}(CN)$ | 61.85 | 90.27 | 90.09 | 87.42 | 93.74 | 91.98 | 85.89 |  |

## 5 Conclusion & Future Work

We have proposed a novel channel-wise sample permutation operator, leading to a simple but effective surrogate of existing multi-head attention mechanisms.
In theory, we demonstrate that the proposed CSP operator overcomes the rank collapse problem of the classic MHA because of implementing sparse doubly stochastic attention maps as permutation matrices.
In addition, we explain the operator from the perspective of channel-wise mixer and optimal transport-based attention.
For representative MHA-based models, replacing their MHA layers with the CSP operator helps improve their performance in discriminative tasks and reduce their computational cost at the same time.
In summary, our work provides a promising solution to developing a better multi-head attention mechanism, demonstrating the usefulness of discrete algorithms like shifting and sorting in model design.

Limitations and Future Work.
Currently, the design of CSP is motivated by pursuing sparse doubly stochastic attention maps, which restricts its application to discriminative tasks — the attention maps of Transformer decoder in generative tasks are lower-triangular, so that imposing the doubly stochastic constraint on the attention maps results in trivial identity matrices.
For the Transformers in generative tasks [, , ], how to achieve effective and efficient attention maps by simple algorithms is still an open problem, which is left as our future work.
In addition, we implement our method based on Pytorch at the current stage.
To maximize the computational efficiency of our method, we plan to refactor its underlying code and optimize its I/O, parallelism, and partitioning strategies as FlashAttention [, ] did.

## Appendix A The Proof of Theorem 1

A Single Channel: Given a CSP’s output matrix, we can derive a residue for each channel as

$$ $\displaystyle\epsilon(\bm{P}_{c}\bm{X}\bm{w}_{c})=\bm{P}_{c}\bm{X}\bm{w}_{c}- \bm{1}\hat{a}_{c},~{}\text{for}~{}c=1,...,C.$ (10) $$

where $\hat{\bm{a}}=[\hat{a}_{c}]$ and

$$ $\displaystyle\hat{a}_{c}=\sideset{}{{}_{a}}{\operatorname*{arg\,min}}\|\bm{P}_ {c}\bm{X}\bm{w}_{c}-\bm{1}a\|=\sideset{}{{}_{a}}{\operatorname*{arg\,min}}\| \bm{X}\bm{w}_{c}-\bm{1}a\|.$ (11) $$

Then, we have

$$ $\displaystyle\begin{aligned} \|\epsilon(\bm{P}_{c}\bm{X}\bm{w}_{c})\|=\|\bm{P} _{c}\bm{X}\bm{w}_{c}-\bm{1}\hat{a}_{c}\|=\|\bm{X}\bm{w}_{c}-\bm{1}\hat{a}_{c} \|\leq\|(\bm{X}-\bm{1}\bm{x}^{\top})\bm{w}_{c}\|\leq\|\bm{w}_{c}\|\|\epsilon( \bm{X})\|,\end{aligned}$ (12) $$

where the second equation is based on the permutation invariance of the matrix norm, the first inequation is based on (11), and the second inequation is based on the sub-multiplicativity (or called consistency) of the matrix norm.

A Single CSP:
Considering all $C$ heads and specifying the matrix norm to be $1$-norm and $\infty$-norm, respectively, we have

$$ $\displaystyle\begin{aligned} \|\epsilon(\text{CSP}_{\bm{W}}(\bm{X}))\|_{1}&=\| (\|_{c=1}^{C}\bm{P}_{c}\bm{X}\bm{w}_{c})-\bm{1}\hat{\bm{a}}^{\top}\|_{1}\\ &=\max_{c}\|\bm{P}_{c}\bm{X}\bm{w}_{c}-\bm{1}\hat{a}_{c}\|_{1}\\ &\leq(\max_{c}\|\bm{w}_{c}\|_{1})\|\epsilon(\bm{X})\|_{1}\\ &=\|\bm{W}\|_{1}\|\epsilon(\bm{X})\|_{1}.\end{aligned}$ (13) $$

$$ $\displaystyle\begin{aligned} \|\epsilon(\text{CSP}_{\bm{W}}(\bm{X}))\|_{\infty }&=\|(\|_{c=1}^{C}\bm{P}_{c}\bm{X}\bm{w}_{c})-\bm{1}\hat{\bm{a}}^{\top}\|_{ \infty}\\ &\leq\sideset{}{{}_{c=1}^{C}}{\sum}\|\bm{P}_{c}\bm{X}\bm{w}_{c}-\bm{1}\hat{a}_ {c}\|_{\infty}\\ &\leq\sideset{}{{}_{c=1}^{C}}{\sum}\|\bm{w}_{c}\|_{\infty}\|\epsilon(\bm{X})\| _{\infty}\\ &\leq\sideset{}{{}_{c=1}^{C}}{\sum}\sideset{}{{}_{n=1}^{N}}{\sum}|w_{nc}|\| \epsilon(\bm{X})\|_{\infty}\\ &\leq C(\max_{c}\|\bm{w}_{c}\|_{1})\|\epsilon(\bm{X})\|_{\infty}\\ &=C\|\bm{W}\|_{1}\|\epsilon(\bm{X})\|_{\infty}.\end{aligned}$ (14) $$

Combining the above two inequations, we have

$$ $\displaystyle\begin{aligned} \|\epsilon(\text{CSP}_{\bm{W}}(\bm{X}))\|_{1, \infty}\leq\sqrt{C}\|\bm{W}\|_{1}\|\epsilon(\bm{X})\|_{1,\infty}.\end{aligned}$ (15) $$

A Single CSP followed by a Lipschitz function. Given a Lipschitz function $f_{\lambda}:\mathbb{R}^{C}\mapsto\mathbb{R}^{C}$, we apply it to each row of a CSP’s output matrix.
For the residual of $f_{\lambda}\circ\text{CSP}_{\bm{W}}(\bm{X})$, we have

$$ $\displaystyle\begin{aligned} \|\epsilon(f_{\lambda}\circ\text{CSP}_{\bm{W}}( \bm{X}))\|&=\|f_{\lambda}\circ\text{CSP}_{\bm{W}}(\bm{X})-\bm{1}\hat{\bm{y}}^{ \top}\|\\ &\leq\|f_{\lambda}\circ\text{CSP}_{\bm{W}}(\bm{X})-\bm{1}f_{\lambda}^{\top}( \hat{\bm{a}})\|\\ &\leq\lambda\|\text{CSP}_{\bm{W}}(\bm{X})-\bm{1}\hat{\bm{a}}^{\top}\|\\ &=\lambda\|\epsilon(\text{CSP}_{\bm{W}}(\bm{X}))\|,\end{aligned}$ (16) $$

where $\hat{\bm{y}}=\operatorname*{arg\,min}_{\bm{y}}\|f_{\lambda}\circ\text{CSP}_{
\bm{W}}(\bm{X})-\bm{1}\bm{y}^{\top}\|$ and $\hat{\bm{a}}$ is the vector associated to $\epsilon(\text{CSP}_{\bm{W}}(\bm{X}))$.

Stacking $L$ CSP Operators:
We can recursively leverage the above results and derive the following inequation:

$$ $\displaystyle\begin{aligned} \|\epsilon((f\circ\text{CSP})^{L}(\bm{X}))\|_{1, \infty}\leq&\lambda_{L}\sqrt{C}\|\bm{W}^{(L)}\|_{1}\|\epsilon((f\circ\text{CSP })^{L-1}(\bm{X}))\|_{1,\infty}\\ \leq&C^{\frac{L}{2}}\Bigl{(}\sideset{}{{}_{\ell=1}^{L}}{\prod}\lambda_{\ell}\| \bm{W}^{(\ell)}\|_{1}\Bigr{)}\|\epsilon(\bm{X})\|_{1,\infty}\\ \leq&C^{\frac{L}{2}}(\lambda\beta)^{L}\|\epsilon(\bm{X})\|_{1,\infty},\end{aligned}$ (17) $$

where $\beta=\max_{\ell}\|\bm{W}^{(\ell)}\|_{1}$ and $\lambda=\max_{\ell}\lambda_{\ell}$.
When the $f$’s are the identity map, we have $\|\epsilon(\text{CSP}^{L}(\bm{X}))\|_{1,\infty}\leq C^{\frac{L}{2}}\beta^{L}\|
\epsilon(\bm{X})\|_{1,\infty}$.

## Appendix B The OT-based Explanation of Cross-Channel Sorting

For convenience, denote the group size $N/K$ by $G$.
For $\bm{v}_{1}^{(k)},\bm{v}_{c}^{(k)}\in\mathbb{R}^{G}$
The optimal transport distance between $\bm{v}_{1}^{(k)}$ and $\bm{v}_{c}^{(k)}$ can be defined as the following linear programming problem:

$$ $\displaystyle W(\bm{v}_{1}^{(k)},\bm{v}_{c}^{(k)}):=\sideset{}{{}_{\bm{T}\in \Pi_{G}}}{\min}\langle\bm{D},\bm{T}\rangle,$ (18) $$

where $\bm{D}=(\bm{v}_{1}^{(k)}\odot\bm{v}_{1}^{(k)})\bm{1}_{G}^{\top}+\bm{1}_{G}(\bm
{v}_{c}^{(k)}\odot\bm{v}_{c}^{(k)})^{\top}-2\bm{v}_{1}^{(k)}\bm{v}_{c}^{(k)\top}$ is the squared Euclidean distance matrix, and $\odot$ denotes the Hadamard product.
Denote the optimal solution of (18) as $\bm{T}^{*}$.
Because $\bm{T}\in\Pi_{G}$, we have

$$ $\displaystyle\begin{aligned} \bm{T}^{*}&=\sideset{}{{}_{\bm{T}\in\Pi_{G}}}{ \operatorname*{arg\,min}}\langle\bm{D},\bm{T}\rangle\\ &=\sideset{}{{}_{\bm{T}\in\Pi_{G}}}{\operatorname*{arg\,min}}\langle(\bm{v}_{1 }^{(k)}\odot\bm{v}_{1}^{(k)})\bm{1}_{G}^{\top}+\bm{1}_{G}(\bm{v}_{c}^{(k)} \odot\bm{v}_{c}^{(k)})^{\top}-2\bm{v}_{1}^{(k)}\bm{v}_{c}^{(k)\top},\bm{T} \rangle\\ &=\sideset{}{{}_{\bm{T}\in\Pi_{G}}}{\operatorname*{arg\,min}}\langle\bm{v}_{1} ^{(k)}\odot\bm{v}_{1}^{(k)},\bm{T}\bm{1}_{G}\rangle+\langle\bm{v}_{c}^{(k)} \odot\bm{v}_{c}^{(k)},\bm{T}\bm{1}_{G}\rangle-2\langle\bm{v}_{1}^{(k)}\bm{v}_{ c}^{(k)\top},\bm{T}\rangle\\ &=\sideset{}{{}_{\bm{T}\in\Pi_{G}}}{\operatorname*{arg\,min}}\underbrace{ \langle\bm{v}_{1}^{(k)}\odot\bm{v}_{1}^{(k)},\bm{1}_{G\times G}\rangle+\langle \bm{v}_{c}^{(k)}\odot\bm{v}_{c}^{(k)},\bm{1}_{G\times G}\rangle}_{\text{A Constant}~{}C_{0}}-2\langle\bm{v}_{1}^{(k)}\bm{v}_{c}^{(k)\top},\bm{T}\rangle \\ &\Leftrightarrow\sideset{}{{}_{\bm{T}\in\Pi_{G}}}{\operatorname*{arg\,min}} \langle-\bm{v}_{1}^{(k)}\bm{v}_{c}^{(k)\top},\bm{T}\rangle.\end{aligned}$ (19) $$

In addition, because $\bm{v}_{1}^{(k)}$ and $\bm{v}_{c}^{(k)}$ are 1D vectors, their OT distance can be computed by aligning the elements of $\bm{v}_{c}^{(k)}$ to align to those of $\bm{v}_{1}^{(k)}$, which corresponds to the sorting operation, i.e.,

$$ $\displaystyle W(\bm{v}_{1}^{(k)},\bm{v}_{c}^{(k)})=\|\bm{v}_{1}^{(k)}-\bm{T}_{ c}^{(k)}\bm{v}_{c}^{(k)}\|_{2}^{2}=\underbrace{\langle\bm{v}_{1}^{(k)},\bm{v}_ {1}^{(k)}\rangle+\langle\bm{v}_{c}^{(k)},\bm{v}_{c}^{(k)}\rangle}_{=C_{0}}-2 \langle\bm{v}_{1}^{(k)}\bm{v}_{c}^{(k)\top},\bm{T}_{c}^{(k)}\rangle,$ (20) $$

where $\bm{T}_{c}^{(k)}$ is the permutation matrix.
Therefore, as long as $W(\bm{v}_{1}^{(k)},\bm{v}_{c}^{(k)})$ has a unique optimal transport, $\bm{T}_{c}^{(k)}=\bm{T}^{*}$.

## Appendix C Implementation Details

### C.1 Image Classification

The detailed hyperparameter setups are presented in Table 5.
Both training and testing are conducted on 8 NVIDIA GeForce RTX 4080 SUPER GPUs.

**Table 5: The hyperparameters of ViT using CSP on image classification tasks.**
| Dataset | CIFAR-10 | CIFAR-100 | ImageNet-1k |
| --- | --- | --- | --- |
| #Groups $K$ | 32 | 128 | 98 |
| Shifting step | Linear | Linear | Linear |
| Batch Size | 64 | 64 | 256 |
| Epochs | 100 | 100 | 300 |
| Learning Rate | 1E-04 | 1E-04 | 5E-04 |
| LR scheduler | cosine | cosine | cosine |
| Optimizer | Adam | Adam | AdamW |
| Dropout Rate | 0.1 | 0.1 | 0.1 |
| Hidden Dims | 512 | 512 | 386 |
| Num. Layers | 6 | 6 | 12 |
| Pooling Type | mean | mean | mean |
| #Param. | 6.46M | 6.50M | 18.50M |

### C.2 Long Range Arena Benchmark

We strictly follow the LRA benchmark []’s default data processing and experimental design.
The detailed hyperparameter setups are presented in Table 6 and Table 7.
For Image task, we only apply the circular shifting operation.
Both training and testing are conducted on 8 NVIDIA RTX A6000 GPUs.

**Table 6: The hyperparameters of Transformer using CSP on LRA.**
| Dataset | ListOps | Text | Retrieval | Image | PathFind |
| --- | --- | --- | --- | --- | --- |
| #Groups $K$ | 1 | 1 | 1 | 1 | 1 |
| Shifting step | — | — | — | — | — |
| Batch Size | 32 | 32 | 8 | 256 | 256 |
| Train steps | 5000 | 20000 | 5000 | 35156 | 125000 |
| Learning Rate | 5E-02 | 5E-02 | 5E-02 | 8E-03 | 1E-03 |
| LR scheduler | sqrt | sqrt | sqrt | cosine | cosine |
| Optimizer | Adam | Adam | Adam | Adam | Adam |
| Weight Decay | 1E-01 | 1E-01 | 1E-01 | 0 | 0 |
| Hidden Dims | 512 | 256 | 128 | 128 | 64 |
| Num. Layers | 4 | 4 | 4 | 4 | 6 |
| Pooling Type | cls | cls | cls | cls | cls |

**Table 7: The hyperparameters of MEGA using CSP on LRA.**
| Dataset | ListOps | Text | Retrieval | Image | PathFind | Path-X |
| --- | --- | --- | --- | --- | --- | --- |
| #Groups $K$ | 1 | 1 | 1 | — | 512 | 8192 |
| Shifting step | — | — | — | Linear | Linear | Exp |
| Batch Size | 64 | 25 | 8 | 50 | 64 | 60 |
| Epochs | 60 | 50 | 40 | 200 | 200 | 100 |
| Learning Rate | 1E-03 | 4E-03 | 6E-03 | 1E-02 | 3E-02 | 1E-02 |
| LR scheduler | linear | linear | linear | linear | linear | linear |
| Optimizer | Adam | Adam | Adam | Adam | Adam | Adam |
| Weight Decay | 1E-02 | 1E-02 | 4E-02 | 2E-02 | 1E-02 | 1E-02 |
| Dropout Rate | 0.1 | 0.1 | 0.1 | 0.0 | 0.1 | 0.5 |
| Hidden Dims | 160 | 256 | 256 | 1024 | 256 | 128 |
| Num. Layers | 6 | 4 | 6 | 8 | 6 | 4 |
| Pooling Type | mean | mean | mean | mean | mean | mean |

In Figure 4, we compare Transformer using CSP with other baselines based on JAX [].
These models are trained on 4 NVIDIA GeForce RTX 3090 GPUs.
The detailed settings are as follows:
The length of the sequence is 3K.
The x-axis corresponds to the number of training steps per second.
The y-axis corresponds to the average score (%) on the LRA benchmark.
The peak memory usage of each model is represented as the area of the corresponding circle.
For a better comparison, the values (GB) of the top-2 models are shown.

## Peer Reviews

confidence: 2
contribution: 3
initial_rating: 6
presentation: 2
questions: I have a few questions about this paper.
1. In Tab. 3, do you use ViT-base as backbone for evaluation?
2. When the group size increases, will the sorting algorithm decline the speed in the real application in Cuda?
soundness: 3
strengths: 1. The proposed operator achieves comparable performance to current Transformer variants with fewer parameters and lower complexity.
2. The experimental comparisons effectively demonstrate performance compared to current state-of-the-art Transformers.
summary: This submission presents a novel channel-wise sample permutation operator (CSP operator) to achieve fewer parameters and lower complexity compared traditional multi-head attention mechanism. Specifically, CSP operator circularly shifts the samples of different channels with various steps and then sorts grouped samples of each channel. The experiments on several discriminative tasks demonstrate the proposed approach achieves comparable or better performance than the state-of-the-art Transformer variants.
weaknesses: 1. The experiments are majorly evaluated on discriminative tasks.
2. According to Tab. 4, the improvement compared to previous approaches is limited.
3. The paper focus on the shift operator. I think it is necessary to discuss with previous shift operator, including those approaches in channel shift or spatial shift. e.g. TSM: Temporal Shift Module for Efficient and Scalable Video Understanding on Edge Device.

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 2
questions: see weakness
soundness: 3
strengths: 1. Paper writing is clear and the methodology is more accurately described
2. The CSP module was used with good results, obtaining promising results in several tests.
summary: In this paper, in order to improve the performance of Multi-head Attention, Channel-wise Sample Permutation is proposed , this method can reduce the number of Multi-head Attention parameters which reduces Multi-head Attention time complexity, and at the same time, it can reduce the risk of rank collapse and improve the stability of the training, the method is tested on several public datasets, and promising results are obtained.
weaknesses: 1. The motivation of the modeling is clearer in the paper, and it can be understood that CSP needs to construct intergroup and intragroup permutations to achieve feature mixing, but there is some confusion in the paper's explanation of the model, and the linkage with the illustration (e.g., Fig. 1) is weak, which makes it more difficult to understand. For example, the specific definition of the substitution matrix $T_{c}^{k}$ and K in eq(5); also, unlike the definition in eq(3), $J_{c}$ seems to be independent of the number of heads, does it imply that shifting matrices are shared between different heads?
2. It seems that when $J_{c}$ is given, which parts of the whole model can be mixed between groups is determined; in other words, such a pre-determined mixing restricts the fusion between the features of the samples at any moment; moreover, unlike linear weighting, permutation makes each moment sample end up retaining only its own features or those of a certain other moment, and although the authors have demonstrated on some tests the algorithm's Even though the authors have demonstrated the effectiveness of the algorithm on some tests, the generalizability of the algorithm is questionable.
3. The authors have optimized the complexity of the model, but there are still many empirical settings in the model, such as the definition of shifting steps, the choice of K, and so on. These empirical modeling further introduces concerns about the utility of the model.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: The idea of this paper is good. However, the proposed approach is not experimentally verified. I suggest the authors to provide more experimental analysis to demonstrate the advantage of the proposed approach.
soundness: 3
strengths: - The presentation of this paper is good. The authors clearly explain the motivation of this paper and give specific implementation details on the proposed method.

- The authors provide theoretical analysis to show the advantages of the proposed approach.

- Experimental results show that the proposed approach receives better results than previous attention mechanisms on CIFAR, ImageNet classification and long sequence analysis.
summary: This paper presents a new attention mechanism based on the multi-head self-attention. It introduces a channel-wise sample permutation operation, which circularly shifts the samples of different channels with various steps and then performs a sorting operation for the  grouped samples of each channel. The authors theoretically analyze the advantages of the proposed attention mechanism over the original self-attention. Experiments show that Transformers with the proposed attention mechanism performs well on several tasks, including image classification and long sequence analysis.
weaknesses: - It seems that the proposed approach performs better than other ViT variants as shown in Table 3. However, the authors did not compare the proposed approach with ViTs with other types of attention mechanisms, like the ones shown in Table 2.

- According to the appendix, the training recipe used in image classification is not new. It seems that the results in image classification is much less than vision transformers, like DeiT. Many popular training strategies are not used. Have the authors used stronger training recipes to conduct the image classification experiments?

- There are no analysis experiments to show how the proposed approach works. For example, how does the shifting step affect the model performance? Should the shifting step be set larger or smaller? There is no related experiments to demonstrate this.

- I do not think the experiments in this paper are comprehensive. The authors only show results on image classification and long sequence analysis. I suggest the authors to provide more results on more vision and language tasks to show the advantages of the proposed approach.

- The method text in Fig. 4 is too small. It is really difficult to identify the performance of different methods.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 1
questions: - If this model is trained to sequences of length $L_{1}$, can it process sequences of $L_{2}$ where $L_{2} >> L_{1}$? Similarly, can this model be generalized to multiple $L$ settings, such as in the auto-regressive setting?
- Experimental details are not sufficient. When comparing to other methods, how are the experiments being controlled? Are every models using same number of layers? Are they parameter matched? For instance, how much is the number of parameters for each in table 4? In addition, I believe M2 [B] should be also added and SSM part should have S4D [C] and S4 with updated scores reported in [C].
- In Table 3, #Param of MHA is noted as 9.52 while CSP has two variants of 6.38 and 9.52. How are these #param matched? Specifically, the paper suggests that Q, K, V of MHA can be simplified to V only. Considering that given $D := d_model$, a transformer layer is $12D^2$, where $4D^2$ from MHA and $8D^2$ is from FFN, the proposed method should have $10D^2$ thanks to removing Q K projections in MHA. However, 6.38 is too small compared to 9.52, so I am not sure how these numbers are achieved.
- Why is Transformer so slow in Figure 4? Is it measured using FlashAttention?
- Does CSP also work on other tasks such as MQAR [B]?

[A] Fu et al. Monarch Mixer: A Simple Sub-Quadratic GEMM-Based Architecture\
[B] Arora et al. Zoology: Measuring and Improving Recall in Efficient Language Models\
[C] Gu et al. On the Parameterization and Initialization of Diagonal State Space Models
soundness: 2
strengths: Similar to ShuffleNet, this paper presents a new channel mixing method that shuffles channels of a given set of tokens then performs token-wise projection. This mechanism can remove the parameters necessary for explicit token mixing such as Attention, MLP-mixer, and SSM. The authors show that not only it leads to efficient number of parameters, but also have a mathematical strength over Attention that it is robust from a rank collapse. This paper demonstrates that this efficient approach performs on par, and even surpass previous methods  on some benchmarks.
summary: This paper presents a new module that aims to replace self-attention in Transformer layers. Specifically, this authors make an extreme use of channel shift operations to remove the quadratic computational overhead in Transformers as well as reducing the number of parameters. This method can be viewed as similar to MLP-mixer, which replaces the trainable channel mixing matrix of MLP-mixer to a permutation matrix that's driven according to their proposed algorithm. The algorithm has subquadratic theoretical complexity and achieves competitive results on a number of benchmarks.
weaknesses: - I mainly concern regarding the uncertainty of whether this method can be generally applicable. For example, as the authors mention in Sec 3.2.2, it is intractable to handle an input with a large number of tokens (e.g., long sequences) using this module, because of the limited number of channel dimension. The authors suggest a solution that incorporate multiple number of layers as a whole for channel shifting, but having two separate algorithms for different settings (N = C and N >> C) involves additional tuning of the model.

- Although I understand it could be due to lack of compute resources, the experiments are done in very light-weight configurations. Therefore, I am not sure if this method could be extended to a bit more general cases (e.g., 110M).

- I believe an ablation study of comparing to previous channel mixing methods is missing. For instance, I believe ShuffleNet is very similar to this paper in that it first shuffles channels then performs a token-wise projection. How much is this method better than the previous ones?

- Even though the main approach of this paper is very simple, it is confusing to understand due to the presentation of the paper. I believe it's mainly because of the ordering of the sections. Also, as Transformers indeed do have a residual connection, what are clear practical benefits of having a mixer that prevents from the rank collapse even if without a residual connection instead of Attention? Since the paper uses a lot of space for explaining mathematical analysis of the proposed module, I believe there should be a practical benefit of having such mathematical properties.