arXiv:2410.12982cs.LGcs.AI2024-10ICLR被引 3

让长序列卷积模型推理速度接近线性,大幅加速生成

Flash Inference: Near Linear Time Inference for Long Convolution Sequence Models and Beyond

  • 用分块计算减少内存移动,共享计算资源
  • 实现110倍的局部推理加速,端到端提速7.8倍
  • 适合需要高效长序列生成的场景,如超长文本

尽管变压器在序列生成模型中占据核心地位,其计算开销仍随序列长度呈二次增长。虽然一些次二次架构(如长卷积序列模型,LCSM)在训练时解决了该问题,但推理阶段仍为二次复杂度。本文提出一种方法,将LCSM的精确推理时间降至准线性 $O(L\ ext{log}^2L)$,识别出实现此加速的关键性质,并构建通用框架加以利用。该方法受松弛多项式插值启发,采用分块策略,有效降低内存访问量并共享计算。同时,该设计几乎实现了位置混合层的完全并行化。我们以Hyena为例进行了实证验证,结果表明:位置混合部分推理速度提升110倍,端到端推理性能提升达7.8倍。

原文摘要 · Abstract (English)

While transformers have been at the core of most recent advancements in sequence generative models, their computational cost remains quadratic in sequence length. Several subquadratic architectures have been proposed to address this computational issue. Some of them, including long convolution sequence models (LCSMs), such as Hyena, address this issue at training time but remain quadratic during inference. We propose a method for speeding up LCSMs' exact inference to quasilinear $O(L\log^2L)$ time, identify the key properties that make this possible, and propose a general framework that exploits these. Our approach, inspired by previous work on relaxed polynomial interpolation, is based on a tiling which helps decrease memory movement and share computation. It has the added benefit of allowing for almost complete parallelization across layers of the position-mixing part of the architecture. Empirically, we provide a proof of concept implementation for Hyena, which gets up to $7.8\times$ end-to-end improvement over standard inference by improving $110\times$ within the position-mixing part.

长序列生成推理加速卷积模型

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。