2024 年 8 月 12 日,Zyphra 实验室宣布开发了一种新的注意力机制 —— 树状注意力(Tree Attention),提升了大型语言模型(LLM)处理超长文本的效率。大型语言模型(如 ChatGPT)的核心是 “注意力” 机制,它使模型能够关注输入文本中的重要信息。然而,随着输入文本变长,传统注意力机制的计算复杂度会呈二次方增长,导致处理速度下降。
Zyphra 实验室的研究人员通过数学推导,发现注意力计算可以转化为一个标量能量函数(scalar energy function)的梯度。基于这一发现,他们设计出能在多个 GPU 之间并行计算注意力的新算法。与现有的环状注意力(Ring Attention)等方法相比,新算法在多个 GPU 上并行计算注意力,实验中速度提升了 8 倍,同时减少了 2 倍的内存占用,通信开销也降低。这意味着 AI 模型可以更高效地处理超长文本,提升了实际应用中的性能。
研究团队还公开了相关的代码,鼓励更多的研究人员和开发者参与到这一领域的研究中来。这一举措有助于推动技术的进一步发展,为整个 AI 社区提供了资源。许多研究人员和开发者在社交媒体上表示,他们已经开始研究和测试这一新算法。
除了树状注意力机制,Zyphra 实验室还在其他领域进行了许多研究。例如,他们正在研究一种新的混合模型架构 Zamba,该架构结合了 Mamba 块和全局共享注意力层,旨在提高长距离依赖关系的学习能力。这些研究展示了 Zyphra 实验室在 AI 领域的研究方向和技术探索。

