Accelerating Toeplitz Neural Network with Constant-time Inference Complexity (2023)

Abstract: Toeplitz Neural Networks (TNNs) have exhibited outstanding performance in various sequence modeling tasks. They outperform commonly used Transformer-based models while benefiting from log-linear space-time complexities. On the other hand, State Space Models (SSMs) achieve lower performance than TNNs in language modeling but offer the advantage of constant inference complexity. In this paper, we aim to combine the strengths of TNNs and SSMs by converting TNNs to SSMs during inference, thereby enabling TNNs to achieve the same constant inference complexities as SSMs. To accomplish this, we formulate the conversion process as an optimization problem and provide a closed-form solution. We demonstrate how to transform the target equation into a Vandermonde linear system problem, which can be efficiently solved using the Discrete Fourier Transform (DFT). Notably, our method requires no training and maintains numerical stability. It can be also applied to any LongConv-based model. To assess its effectiveness, we conduct extensive experiments on language modeling tasks across various settings. Additionally, we compare our method to other gradient-descent solutions, highlighting the superior numerical stability of our approach. The source code is available at https://github.com/OpenNLPLab/ETSC-Exact-Toeplitz-to-SSM-Conversion.

## introduction transformer has dominated the fields of computer vision (cv) @xcite @xcite , natural language processing (nlp) @xcite @xcite @xcite , and speech processing @xcite indicates the corresponding author (email address: zhongyiran@gmail.com). @xcite @xcite , becoming one of the best-performing approaches across different benchmarks. the core component of the transformer, the attention mechanism, has a quadratic time complexity with respect to sequence length, making it challenging to scale to long sequences and large model sizes. various methods have been proposed to address this issue, including linear attention @xcite @xcite , state space model (ssm) @xcite , toeplitz neural network (tnn) @xcite and other long-conv methods @xcite . linear attention reduces the space-time complexity of attention to linear by using a kernel trick to decompose the softmax function @xcite , but its poor performance @xcite prohibits it from being used to build large language models (llms). ssm replaces the attention operation with state space equations, resulting in log-linear training space-time complexities @xcite . however, the performance of this method in casual language modeling is often inferior @xcite and initialization-sensitive @xcite , making it unsuitable for building llms. tnn is a new class of sequence modeling methods that belongs to longconv-based methods @xcite . it models long sequences using toeplitz matrices to encode relative positional relationships. this key component allows them to effectively capture the dependencies within the sequence and make accurate predictions. it has a log-linear space-time complexity and outperforms transformers in nlp and long sequence modeling tasks @xcite . additionally, its stable training capability and insensitivity to initialization make it feasible for llms. note that the above analysis has only taken into account the training complexities for the afore-mentioned methods. however, when considering the deployment of llms, the inference complexities are also important. in decoder scenarios, i.e., casual language modeling, the time complexity of inferring the n th token in the transformer is o(n 2 d + nd 2 ), where n, d are the sequence length and the feature dimension respectively. by using the kv cache technique @xcite , the complexity can be reduced to o(nd 2 ). for linear attention, the complexity is o(dh)(h is the hidden dimension), which makes it constant with respect to the sequence length @xcite . ssm also has a constant space-time complexity of o(dh), where h is the hidden space dimension @xcite . tnn, on the other hand, has a loglinear space-time complexity of o(nd log n) in inference, which may make it challenging to handle long sequences. in this paper, we aim to accelerate the inference of tnn to constant-time complexity. we find that ssm can be thought of as a particular variation of tnn. tnn can benefit from the same inference complexity as ssm if we can convert it to ssm in inference. we show that such conversion can be viewed as an optimization problem and can be efficiently solved by a closed-form solution. specifically, given a toeplitz matrix, we first convert it to vandermoode linear system with inclusive equation reformulation (ier) and then employ the discrete fourier transform (dft) to obtain a numerical stable result. compared with gradientbased algorithms, our method is fast, training-free, and numerically stable. note that our method can be applied to other longconv-based methods @xcite as well. we conduct extensive experiments to validate the effectiveness of our method. we compare our method with gradient-based methods in terms of efficiency and errors. our method outperformed gradient-based methods significantly in efficiency while enjoying much lower error rates. we also apply our method to tnn language models and test it in real scenarios. our method has equivalent extrapolation capabilities and perplexity to the original implementation of tnn. for the number of layers, sequence length, and feature dimensions, an indepth assessment of speed and memory utilization is performed. our method clearly outperforms the original tnn inference algorithm implementation. furthermore, we demonstrate the applicability of our strategy beyond tnn by extending it to other longconv-based models. ## background and preliminary in this section, we first define sequence model inference mathematically and then briefly discuss the inference complexities of transformer and some closely related efficient sequence modeling methods such as linear transformer @xcite , ssm @xcite , and tnn @xcite . ## method the inference of tnn exhibits a time complexity of o(nd log n) and space complexity o(nd) for predicting the n-th token which poses challenges for scaling tnn to handle extremely long sequences in inference. in this section, we will present our approach to converting tnn into the form of ssm, aiming to improve generation speed and memory to a constant. ## experiments in this section, we present extensive experiments to validate our method. we first analyze the numerical stability and efficiency of our method with a comparison to a gradient-based approach. then we evaluate our method for language modeling tasks with real-world scenarios. in our inference efficiency study, we conduct an in-depth analysis of the impact of the number of layers, sequence length, and feature dimensions on the speed and memory utilization of our method. we also extend the scope of our method to other long convolutionbased methods, showcasing its versatility and generalizability.