Toeplitz Neural Network for Sequence Modeling

TL;DR

This paper introduces a Toeplitz neural network that models sequences efficiently using a relative position encoded Toeplitz matrix, enabling long sequence processing with improved performance and speed over transformers.

Contribution

The paper proposes a novel Toeplitz matrix-based sequence modeling approach with a relative position encoder, reducing complexity and enabling extrapolation to much longer sequences.

Findings

Achieves better performance than competitors on multiple tasks.

Significantly faster sequence processing.

Handles sequences up to 14K tokens in inference.

Abstract

Sequence modeling has important applications in natural language processing and computer vision. Recently, the transformer-based models have shown strong performance on various sequence modeling tasks, which rely on attention to capture pairwise token relations, and position embedding to inject positional information. While showing good performance, the transformer models are inefficient to scale to long input sequences, mainly due to the quadratic space-time complexity of attention. To overcome this inefficiency, we propose to model sequences with a relative position encoded Toeplitz matrix and use a Toeplitz matrix-vector production trick to reduce the space-time complexity of the sequence modeling to log linear. A lightweight sub-network called relative position encoder is proposed to generate relative position coefficients with a fixed budget of parameters, enabling the proposed Toeplitz neural network to deal with varying sequence lengths. In addition, despite being trained on 512-token sequences, our model can extrapolate input sequence length up to 14K tokens in inference with consistent performance. Extensive experiments on autoregressive and bidirectional language modeling, image modeling, and the challenging Long-Range Arena benchmark show that our method achieves better performance than its competitors in most downstream tasks while being significantly faster. The code is available at https://github.com/OpenNLPLab/Tnn.