Toeplitz Neural Network for Sequence Modeling

TL;DR

The paper tackles the quadratic bottleneck of self-attention in sequence models by introducing the Toeplitz Neural Network (TNN), which relies on relative-position Toeplitz matrices to mix tokens with log-linear complexity. A lightweight Relative Position Encoder generates the Toeplitz coefficients, decoupling parameter counts from sequence length and enabling extrapolation to very long inputs. Through stacks of GTU blocks and an exponential decay bias, TNN achieves strong performance across autoregressive and bidirectional language modeling, image modeling, and long-range tasks, often outperforming or matching state-of-the-art methods with significantly faster computation. The approach also positions itself as a generalization of Transformer, CNN, and State-space formulations, providing a unified framework for efficient sequence modeling.

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.

Figures (3)

• Figure 1: The left figure shows the training speed ($x$-axis), performances ($y$-axis), and GPU memory footprints (circle sizes) of the TNN and competing methods on Long-Range Arena benchmark. The TNN beats the competitors with a clear margin. The right figure plots the extrapolation results with different sequence lengths, where the $x$-axis denotes sequence lengths, and the $y$-axis denotes $log$ PPL. It demonstrates that regardless of the sequence length, the PPL of the TNN remains constant.
• Figure 2: Network structure overview of the proposed Toeplitz Neural Network. The proposed sequence modeling block is composed of a Gated Toeplitz Unit and a GLU shazeer2020glu and. We propose the TNO to perform token mixing with only relative position information. We use a small fully-connected network named RPE to encode relative position information.
• Figure 3: Visualization of the Toeplitz matrix used by each layer in Roberta, each element of the matrix represents the interaction between tokens. The Toeplitz matrices show similar behaviors to conventional transformer attention matrices where the diagonal concentrates the most attention.

Theorems & Definitions (7)

• Theorem 2.1
• Definition B.1
• Definition B.2
• Lemma B.3
• Lemma B.4
• proof : Proof of Lemma \ref{['circu-product']}
• proof : Proof of Theorem \ref{['toep']}