Press "Enter" to skip to content

加州大学伯克利分校的研究人员提出了RingAttention:一种内存高效的人工智能方法,用于减少Transformers的内存需求

在许多最先进的人工智能模型中,一种深度学习模型架构称为Transformers。它们在自然语言处理和机器学习中的各种任务中彻底改变了人工智能领域。它基于一种自注意机制,即模型在进行预测时权衡输入序列的不同部分的重要性。它们由编码器和解码器组成,用于处理输入。

然而,提高Transformers的上下文长度需要大量的工作。这是由于继承的自注意机制。自注意机制的内存成本是输入序列长度的平方,这使得将其扩展到更长的输入序列变得具有挑战性。UC伯克利的研究人员开发了一种名为环形注意力的方法来解决这个问题,基于一个简单的观察结果。他们观察到,当自注意力和前馈网络计算以块的形式进行时,序列可以分布在多个设备上并进行轻松的分析。

他们将块级注意力计算的外部循环分配给主机,每个设备管理其相应的输入块。对于内部循环,他们为所有设备计算与其指定输入块相关的块级注意力和前馈操作。他们的主机设备形成一个概念上的环,并将用于块级计算的键值块的副本发送到环中的下一个设备。他们还同时从前一个设备接收键值块。

块计算所需的时间比块传输长。团队将这些过程重叠,与标准的Transformers相比,没有额外的开销。通过这样做,每个设备只需要与块大小成比例的内存,而不受原始输入序列长度的限制。这有效地消除了个别设备所施加的内存约束。

他们的实验证明,环形注意力可以通过使其能够训练比以前的内存高效技术达到长达500倍的序列,减少Transformers对内存的需求。该方法还允许训练长度超过1亿的序列,而无需对注意力进行近似。由于环形注意力消除了个别设备所施加的内存约束,因此可以实现接近无限的上下文大小。然而,由于序列长度与设备数量成比例,因此需要许多设备。

这项研究仅涉及对该方法的有效性进行评估,没有涉及大规模训练模型。由于规模上下文长度取决于设备数量,模型的效率取决于优化;他们只研究了实现最佳计算机性能所需的低级操作。研究人员表示,他们希望将来在最大序列长度和最大计算机性能方面进行工作。接近无限的上下文的可能性开启了许多令人兴奋的机会,例如大型视频-音频-语言模型,从扩展的反馈和试错中学习,理解和生成代码库,以及使AI模型适应理解基因序列等科学数据。

Leave a Reply

Your email address will not be published. Required fields are marked *