近年来，Transformer已成为许多最先进的人工智能系统的核心架构，并在各种AI应用中展现出了令人印象深刻的性能，例如：ChatGPT、GPT-4、ChatGLM、PanGu等。然而，Transformer的内存需求限制了其处理长序列的能力，当面对长序列依赖任务时具有一定的挑战性。
针对该挑战，伯克利大学提出了一种独特的方法：环注意力（Ring Attention），在保持内存效率的同时处理更长的输入序列，与最先进的技术相比长出数倍，最高可实现了超过100M的上下文大小。
Paper：https://browse.arxiv.org/pdf/2310.01889.pdf
Code：coming soon~
Transformer架构主要由self-attention和位置前馈机制组成，这些组件有助于捕获输入Tokens之间的依赖关系，并通过高度并行计算实现可扩展性。
由于Transformer自注意力架构，其内存成本与输入序列长度呈二次方关系，这使得扩展到更长的输入序列具有挑战性。然而，能够支持长序列的Transformer对于人工智能模型应用至关重要，例如，可以将其应用到长篇书籍、高分辨图像、长视频、代码库的处理。并且目前也已经出现了支持长序列的语言模型，例如：ChatGPT可支持16K上下文长度、GPT-4可支持32K上下文长度、MPT可支持65K上下文长度、Claude可支持100K上下文长度。
为了降低Transformer的内存成本，使其能够应用到更长的序列任务上来。研究人员们发现自注意力中的softmax矩阵计算无需在完整矩阵情况下计算，这促进了自注意力和前馈的块式计算的发展，该类方法虽然内存减少了，每层输出的存储仍存在挑战。由于后续层的自注意力依赖于前一层的输出，如果不对中间过程输出做存储，计算成本会成倍增加。
因为如果不做存储就必须为每个序列元素重新计算每个输出，这对于较长的序列来说是不切实际的。从内存需求的角度来看，即使处理批量大小为 1 的情况，对于隐藏大小为 1024 的普通模型来说，处理 1 亿个Token也需要超过1000GB的内存。这远远大于当前 GPU 和 TPU 的容量。
为了能够拓展Transformer可支持的序列长度。作者提出以分块方式执行自注意力和前馈网络计算，跨多个主机设备分布序列维度，从而实现并发计算和通信，由于该方法将环中主机设备之间的键值块通信与块计算重叠，因此将其命名：环注意(Ring Attention)。
具体地，该方法在主机设备之间构建注意力计算块的外循环，每个主机设备具有一个查询块，并通过键值块遍历主机设备环，以逐块的方式进行注意力和前馈网络计算。当计算注意力时，每个主机将键值块发送到下一个主机，同时从前一个主机接收键值块。这里作者使用与原始 Transformer 相同的模型架构，但重新组织了计算。具体如下图所示：
下表展示了不同模型和加速器集群上的设备内存可支持的最大上下文长度。Ring Attention 大大优于基线Transformer，并且比最先进的技术长出数倍，最高可实现了超过100M的上下文大小。
本文提出了Ring Attention方法，用于降低Transformer的内存需求。该方法允许上下文长度与设备数量线性扩展，同时保持性能，消除了个别设备的内存瓶颈。Ring Attention 大大优于基线Transformer，并且比最先进的技术长出数倍，最高可实现了超过100M的上下文大小。
虽然该方法在Transformer模型的上下文长度方面达到了最新水平，但也存在一定的限制：1）由于计算预算限制，本文实验侧重于评估所提出方法的有效性，而没有大规模训练模型；2）虽然Ring Attention可以将上下文长度与设备数量线性扩展并保持性能，但为了实现最佳计算性能，还需要优化低级操作。
联系客服
