[论文解读] Toeplitz Neural Network for Sequence Modeling

Zhen Qin, Xiaodong Han+7 more|arXiv (Cornell University)|May 8, 2023

Topic Modeling被引用 5

一句话总结

本文提出了一种名为Toeplitz神经网络（TNN）的序列建模架构，该架构用可学习的Toeplitz矩阵替代注意力机制，实现对数线性的时间-空间复杂度。通过使用轻量级的相对位置编码器和指数衰减偏差，TNN在长达14K token的序列长度上保持一致的性能表现，在Long-Range Arena等长上下文任务中优于基线模型，且显著更快。

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.

研究动机与目标

• 解决Transformer在长序列建模中自注意力机制带来的二次方复杂度瓶颈问题。
• 探究仅依赖相对位置信息是否足以实现有效的序列建模，而无需基于内容的注意力机制。
• 设计一种参数高效的架构，能够在不重新训练的情况下泛化至可变长度的序列。
• 在Long-Range Arena等长上下文基准上实现优异的性能与外推能力。

提出的方法

• 用编码token之间相对位置关系的可学习Toeplitz矩阵替代标准注意力矩阵。
• 利用基于快速傅里叶变换（FFT）的Toeplitz矩阵-向量乘法，将计算复杂度从O(n²)降低至O(n log n)。
• 引入一种轻量级相对位置编码器（RPE），在固定参数预算下生成Toeplitz系数，且与序列长度无关。
• 直接对Toeplitz矩阵应用指数衰减偏差，以在推理阶段实现对更长序列的泛化能力。
• 在token混合模块中使用门控线性单元（GLU）与门控绑定单元（GTU）以提升表征能力。
• 将TNN形式化为一个统一框架，可将Transformer、CNN和状态空间模型作为特例统一涵盖。

实验结果

研究问题

• RQ1 仅使用相对位置编码的序列模型是否能实现强大性能，而无需依赖基于内容的注意力机制？
• RQ2 对数线性复杂度的模型是否能在长上下文序列建模任务中超越二次方复杂度的Transformer？
• RQ3 一种参数高效的架构是否能无需微调即可泛化至远超其训练长度的序列长度（如14K token）？
• RQ4 在长上下文基准上，所提出的TNN与最先进模型相比，在速度、内存和准确率方面表现如何？
• RQ5 TNN是否为一个统一框架，能够涵盖如Transformer和CNN等现有架构？

主要发现

• 在WikiText-103语言建模范例上，TNN的测试困惑度达到23.98，优于基线Transformer及其他高效注意力变体。
• 在Long-Range Arena基准上，TNN实现了最高的推理速度（1K序列长度下为25.72步/秒），且在所有测试序列长度上均保持一致的性能表现。
• 即使仅在512-token序列上进行训练，TNN在推理阶段也能成功泛化至14K-token序列，平均困惑度为23.70，展现出强大的外推能力。
• 相对位置编码器（RPE）相比无RPE的TNN变体，将困惑度降低了2.47，证实其在学习位置感知表征方面的有效性。
• 以0.99的衰减率应用指数衰减可实现稳定外推，而无衰减或可学习衰减则导致性能显著下降。
• 数学上证明TNN可将Transformer、CNN和状态空间模型作为特例统一涵盖，建立了序列建模的统一视角。