arXiv:2501.15665cs.LGcs.AI2025-01被引 2

提出StagFormer,让Transformer解码并行化,提升速度且不损失质量。

StagFormer: Time Staggering Transformer Decoding for RunningLayers In Parallel

  • 通过时间错位设计,打破层间依赖,实现解码并行执行。
  • 实验显示解码速度显著提升,生成质量与原模型持平。
  • 适合对推理速度敏感的场景,如实时对话系统。

基于Transformer的语言模型解码过程本质上是串行的,因为一个标记的嵌入必须经过网络所有层后才能生成下一个标记。本文提出新架构StagFormer(错位Transformer),沿序列轴错位执行,从而在模型深度方向实现解码并行。其核心是打破第$l$层时间步$i$的标记表示对前一层第$l-1$层时间步直至$i$的依赖,转而仅依赖前一时间步的表示。后续层仍可获取前一段“丰富”的表示,但仅限于落后一个时间步的标记位置。该方法使模型不同部分可并行运行,在模拟中实现了潜在加速且不影响生成质量。我们还探索了多种自然扩展:跨错位段共享权重更适用于内存受限场景;采用有限窗口注意力可在某些应用中进一步降低延迟;错位思想可推广至超过两个段落的结构。最后,我们展示了通过权重共享可近似实现循环模型,该变体在短文本生成中显著提升质量,同时保持延迟不变。

原文摘要 · Abstract (English)

Decoding in a Transformer based language model is inherently sequential as a token's embedding needs to pass through all the layers in the network before the generation of the next token can begin. In this work, we propose a new architecture StagFormer (Staggered Transformer), which staggers execution along the sequence axis and thereby enables parallelizing the decoding process along the depth of the model. We achieve this by breaking the dependency of the token representation at time step $i$ in layer $l$ upon the representations of tokens until time step $i$ from layer $l-1$. Instead, we stagger the execution and only allow a dependency on token representations until time step $i-1$. The later sections of the Transformer still get access to the "rich" representations from the prior section but only from those token positions which are one time step behind. StagFormer allows for different sections of the model to be executed in parallel yielding a potential speedup in decoding while being quality neutral in our simulations. We also explore many natural extensions of this idea. We present how weight-sharing across the different sections being staggered can be more practical in settings with limited memory. We explore the efficacy of using a bounded window attention to pass information from one section to another which helps drive further latency gains for some applications. We also explore the scalability of the staggering idea over more than 2 sections of the Transformer. Finally, we show how one can approximate a recurrent model during inference using weight-sharing. This variant can lead to substantial gains in quality for short generations while being neutral in its latency impact.

Transformer并行解码推理加速

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