通过张量化扩展注意力范围，高效处理超长文本序列。
Long Sequence Modeling with Attention Tensorization: From Sequence to Tensor Learning
- 将长序列转为张量表示，分维度做注意力以扩大感知范围。
- 在32,768上下文长度训练Llama-8B，推理时可拓展至128k，提速11倍。
- 适合需要长序列建模且关注计算效率的NLP任务。
随着对长文本数据处理需求的增长，捕捉长程依赖并保持计算效率变得愈发重要。基于注意力机制的长序列建模面临全注意力感受野有限与输入序列长程依赖不匹配的问题。本文提出通过张量化将长输入序列转化为紧凑张量表示，并在每个变换维度上进行注意力计算，从而扩展注意力感受野。所提出的张量化注意力可作为高效Transformer骨干网络，提升上下文长度，同时改善内存与时间效率。实验表明，该方法将令牌依赖编码为多跳注意力过程，等价于全注意力的克罗内克分解。大量实验显示，张量化注意力可有效适配预训练大模型。显著的是，使用张量化的Llama-8B在32,768上下文长度下训练，推理时可稳定外推至128k长度，相比FlashAttention-2的全注意力实现11倍加速。
原文摘要 · Abstract (English)
As the demand for processing extended textual data grows, the ability to handle long-range dependencies and maintain computational efficiency is more critical than ever. One of the key issues for long-sequence modeling using attention-based model is the mismatch between the limited-range modeling power of full attention and the long-range token dependency in the input sequence. In this work, we propose to scale up the attention receptive field by tensorizing long input sequences into compact tensor representations followed by attention on each transformed dimension. The resulting Tensorized Attention can be adopted as efficient transformer backbones to extend input context length with improved memory and time efficiency. We show that the proposed attention tensorization encodes token dependencies as a multi-hop attention process, and is equivalent to Kronecker decomposition of full attention. Extensive experiments show that tensorized attention can be used to adapt pretrained LLMs with improved efficiency. Notably, Llama-8B with tensorization is trained under 32,768 context length and can steadily extrapolate to 128k length during inference with $11\times$ speedup, compared to full attention with FlashAttention-2.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。