TS‐Mixer: A lightweight text representation model based on context awareness

Abstract

AbstractLarge pre‐trained models (PTMs) have shown their powerful ability in multiple natural language processing tasks. However, using them in practical application remains a challenge due to the significant computational cost and memory requirements. In order to achieve the balance of computational cost and accuracy, MLP architecture can be used as an alternative to the self‐attention module, such as pNLP‐Mixer and Hyper‐Mixer. Experiments indicate that, MLP‐based models can attain competitive performance with low cost. They maintain the balance of computation cost and accuracy successfully, yet, this is at the expense of not being able to capture short‐range dependencies. In this paper, a novel MLP‐based model, termed TS‐Mixer, is proposed which can capture local dependencies by shifting operation. Compared with other MLP‐based models, the parameters of TS‐Mixer are decoupled from the sequence length, hence it has a smaller model size in long sequence tasks. In addition, TS‐Mixer has linear computational complexity, therefore it can be used as a lightweight alternative to the self‐attention model. Experiments show that the TS‐Mixer outperforms other MLP‐based models, which achieves higher accuracy with fewer parameters in multiple downstream tasks. Notably, compared with pre‐trained models, TS‐Mixer can reach more than 90% of their accuracy with 1% or even one thousandth of the parameters (0.174 ~ 1.2 M). These results demonstrate that TS‐Mixer can achieve a better balance between the computing resources and accuracy. Code is available at: https://github.com/wyl-privacy-project/TS-Mixer.