arXivDaily arXiv每日学术速递 周一至周五更新
北美退伍军人节(Veterans Day),12 日无更新
arXiv 2310.01889cs.CL
收藏

面向近无限上下文的基于分块Transformer的Ring Attention

Ring Attention with Blockwise Transformers for Near-Infinite Context

UC Berkeley(加州大学伯克利分校)

Hao Liu, Matei Zaharia, Pieter Abbeel

AI总结:

提出Ring Attention方法,通过分块计算和通信重叠实现近无限上下文,支持数百万token且无额外开销。

摘要:

Transformer已成为许多最先进AI模型的首选架构,在广泛的AI应用中展现出卓越性能。然而,Transformer带来的内存需求限制了其处理长序列的能力,从而在复杂环境中利用视频、动作及其他长格式序列和模态时构成了挑战。我们提出了一种新方法,即基于分块Transformer的Ring Attention,它利用自注意力和前馈网络的分块计算,将长序列分布到多个设备上,同时将键值块的通信与分块注意力的计算完全重叠。我们的方法能够训练和推理长度达到设备数量倍于先前内存高效Transformer可处理序列的序列,且无需借助近似或产生额外的通信与计算开销。在语言建模和强化学习任务上的大量实验证明了该方法在支持数百万token上下文大小及提升性能方面的有效性。

英文摘要:

Transformers have emerged as the architecture of choice for many state-of-the-art AI models, showcasing exceptional performance across a wide range of AI applications. However, the memory demands imposed by Transformers limit their ability to handle long sequences, thereby posing challenges in utilizing videos, actions, and other long-form sequences and modalities in complex environments. We present a novel approach, Ring Attention with Blockwise Transformers (Ring Attention), which leverages blockwise computation of self-attention and feedforward to distribute long sequences across multiple devices while fully overlapping the communication of key-value blocks with the computation of blockwise attention. Our approach enables training and inference of sequences that are up to device count times longer than those achievable by prior memory-efficient Transformers, without resorting to approximations or incurring additional communication and computation overheads. Extensive experiments on language modeling and reinforcement learning tasks demonstrate the effectiveness of our approach in allowing millions of tokens context size and improving performance.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑