Sliceformer: Make Multi-head Attention as Simple as Sorting in
  Discriminative Tasks
Shen Yuan Hongteng Xu
As one of the most popular neural network modules, 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. The
effectiveness of the Transformer is often attributed to its multi-head
attention (MHA) mechanism. In this study, we discuss the limitations of MHA,
including the high computational complexity due to its ``query-key-value''
architecture and the numerical issue caused by its softmax operation.
Considering the above problems and the recent development tendency of the
attention layer, we propose an effective and efficient surrogate of the
Transformer, called Sliceformer. Our Sliceformer replaces the classic MHA
mechanism with an extremely simple ``slicing-sorting'' operation, i.e.,
projecting inputs linearly to a latent space and sorting them along different
feature dimensions (or equivalently, called channels). For each feature
dimension, the sorting operation implicitly generates an implicit attention map
with sparse, full-rank, and doubly-stochastic structures. We consider different
implementations of the slicing-sorting operation and analyze their impacts on
the Sliceformer. We test the Sliceformer in the Long-Range Arena benchmark,
image classification, text classification, and molecular property prediction,
demonstrating its advantage in computational complexity and universal
effectiveness in discriminative tasks. Our Sliceformer achieves comparable or
better performance with lower memory cost and faster speed than the Transformer
and its variants. Moreover, the experimental results reveal that applying our
Sliceformer can empirically suppress the risk of mode collapse when
representing data. The code is available at
\url{https://github.com/SDS-Lab/sliceformer}.