MSCFFN : A New FFN with Multi-Space Cross to Accelerate Transformer (2023)

Abstract: Transformer models have achieved impressive success in various natural language processing tasks. But it is also limited used in some areas and the heavy computation complexity is one of the main limitations. Many model structures have been proposed to reduce the computation complexity and some are really effective. The previous research can be divided into two categories. One is to use more effective training and inference strategies and the other is focused on how to replace the standard self-attention mechanism with linear attention method. Differently, we revisit the design in Transformer and find that the feed forward network (FFN) is also computationally expensive, especially when the hidden dimension is large. In this paper, we propose a new FFN structure, named MSCFFN, which splits the large matrix space to several small space to reduce the computation complexity and uses the Multi-Space Cross method to ensure the accurate result. To the best of our knowledge, this is the first time to redesign FFN to accelerate Transformers. We experimentally validate the effectiveness of the proposed method on the Long-Range Arena benchmark. And the results show MSCFFN can achieve a faster speed with a similar or even better accuracy.

## introduction since transformer models @xcite were proposed, especially the bert @xcite with pre-trained language models, they have achieved state-of-the-art results in various domains of natural language processing, such as machine translation @xcite ), text classification @xcite , natural language inference @xcite , question answering @xcite and so son. the effectiveness of transformer comes from the multi-head self-attention mechanism (mhsa), fully connected feed-forward network (ffn) and the deep layers. the mhsa helps model to learn the relation between tokens in different position of a sequence, the ffn helps enlarge the network space and enhance the non-linear capability, and the deep layers are helpful to build a more complex sentence relationship. however, these operations also cause the high computation complexity, and it is hard to use for user with low computation resource. to tackle this essential obstacle, there has been much prior literature on improving the efficiency of transformers. the main research consists of two parts, one is to use more effective strategies. for example, pyramid-bert @xcite uses the core-set based token selection method to reduce complexity, and pcee-bert @xcite accelerates bert using patient and confident early exiting method. the other is focused on how to reduce the computation complexity of the standard self-attention, which is quadratic complexity correlated with sequence length, to linear attention method. such as performer @xcite , longformer @xcite , cosformer @xcite , flowformer @xcite and so on. although these methods have been proven to be very effective in reducing computation complexity, they almost don't discuss the influence of hidden dimension to computation complexity. in fact, transformer models are also quadratic complexity correlated with hidden dimension. when the hidden dimension is large, it is computationally prohibitive to train or inference @xcite . actually, to achieve better results, bert xlarge uses 2048 of hidden dimension and albert xxlarge @xcite uses 4096 of hidden dimension, which is really computationally expensive for a common researcher. for the transformer models, the main unit with quadratic complexity of hidden dimension (d) is ffn which has the complexity of o(8d 2 ) in vanilla transformer, because the dimension of inner-layer is usually set to 4d for a better result. therefore, in this paper, we focus on how to reduce the computation complexity of ffn. according to the characteristic of quadratic complexity related to hidden dimension, if we shrink the hidden dimension to 1 n d, the computation complexity will be reduced to 1 n 2 of the original. therefor, we firstly split the original large matrix space to several matrix spaces, and we apply the different weight parameters for each small matrix to conduct the feed-forward network separately. obviously, the representation ability of n small matrices with d 2 n 2 which is the space of weight parameters is much less than it of d 2 , and it is easy to cause the effect decline. in order to solve this problem, we propose a method of space cross mechanism, and we name the final model mscffn. through the cross of different small space, we can enlarge the matrix space and obtain a stronger representation capability. furthermore, benefitting from the information from different representation subspaces, a larger coefficient of inner-layer and the deeper network of ffn, mscffn can achieve a better accuracy with a faster speed. we conduct experimental study on the long-range arena benchmark @xcite , and the results demonstrate that our mscffn is effective in speeding up and performs better than baselines. ## related work as mentioned in section 1, the most common methods to reduce the computation complexity are as follows: effective strategies: with this method, they don't change the structure of mhsa and ffn, but change the strategies of layers to train or tokens to train. for example, @xcite proposed funnel-transformer which gradually compresses the sequence of hidden states to a shorter one and hence reduces the computation cost. @xcite proposed pyramid-bert which used a coreset based token selection method to successively shorten sequence length and void expensive computation resource. pcee-bert @xcite was proposed by @xcite and it achieves the goal of speedup through exiting early when enough numbers of consecutive intermediate layers are confident about their prediction. groupbert @xcite uses grouped transformations to reduce the density of fully connected layers. sparse mixtures of experts (sparse moes) reduce the computation cost through conditional computation which only activates a subset of the overall model, such as @xcite . linear attention method: this method is to resolve the quadratic complexity of vanilla selfattention mechanism and try to reduce the complexity of o(n 2 ) to o(nlog(n)) or even o(n). for example, locally-sensitive hashing (lsh) attention @xcite ) applies a multi-round hashing scheme to compute dot-product attention. longformer utilizes a sliding windows to learn the token relation in a local window. performer introduces an unbiased linear estimation of the attention matrices. cosformer proposes a linear attention by using the kernel methods with a non-linear re-weighting scheme. ## experiments in this section, we testify the effectiveness and generality of our proposed model mscffn on the dataset of long-range arena (lra). lra is a benchmark which contains five different tasks: long sequence listops @xcite , byte-level text classification @xcite , document retrieval @xcite , image classification on sequence of pixels @xcite , and pathfinder @xcite . their length of sequences ranges from 1k to 16k. lra is a suitable benchmark for testing the quality of efficient transformer variants, especially for the ability of long input sequence. we chose this dataset to verify that mscffn is general for different tasks and is efficient in different sequence length. we utilize directly the official method of data processing, experiment protocol and training parameters foot_0 to ensure the fair comparison. and we use mscffn directly to replace the original ffn in four base models. the experiment is conducted on 4 v100 gpus. to ensure result reliability, we conducted five separate experiments for each task with different seed values, averaging them for the final reported result. the experimental results are shown in to further prove the effectiveness of our multispace cross method, we conduct ablation studies whose results shown in we also conduct experiment in efficiency analysis, which is shown in ## conclusions in this paper, we propose a new ffn structure named mscffn, which can replace the original ffn in transformer models and speedup model training and inference no matter how long or short the input sequence is. in mscffn, we reduce the computation complexity through splitting the matrix space to several small space. benefiting from multi-space cross mechanism, a larger coefficient of inner-layer and a deeper network, mscffn can achieve comparable or even better performances than the vanilla ffn. the experiment results on five tasks of lra demonstrate the effectiveness and generality of our proposed model mscffn. furthermore, the proposed model is probably not only used for ffn in transformer models, but also for other model structures which contain ffn.