Hyena成下一代Transformer？StripedHyena-7B开源：最高128k输入，训练速度提升50%

文章来源：新智元

Hyena处理长序列输入比FlashAttention速度高100倍！最新发布的StripedHyena模型可能成下一代AI架构的新标准？

图片来源：由无界 AI生成

最近几年发布的AI模型，如语言、视觉、音频、生物等各种领域的大模型都离不开Transformer架构，但其核心模块「注意力机制」的计算复杂度与「输入序列长度」呈二次方增长趋势，这一特性严重限制了Transformer在长序列下的应用，例如无法一次性处理一整本书，或是处理千兆像素级别的图像。

即便强如GPT-4也难以摆脱这种缺陷。

最近，Together Research开源了一个全新的语言模型StripedHyena，采用了针对「长上下文」的新架构，可以处理高达128k个token的长上下文，并且改进了Transformer架构在训练和推理上的性能，为目前的主流架构提供了一种可选方案。

开源链接：https://github.com/togethercomputer/stripedhyena

StripedHyena也是「首个」在短上下文和长上下文评估中，以相同模型尺寸，实现了与最佳开源Transformer模型性能相匹敌的模型：在OpenLLM基准任务上与Llama-2, Yi和Mistral 7B实现了相当的性能，并且在长上下文摘要上表现更出色。

StripedHyena是一种混合架构，由多头、分组查询注意力和排列在Hyena块中的门控卷积组成，不同于传统的decoder-only的Transformer：通过将卷积表示为状态空间模型（SSM，state-space model）或截断滤波器，在Hyena块中进行常数内存解码。

实验结果显示，StripedHyena在32k个token、64k个token和128k个token序列的端到端训练中比传统transformer快30%、50%和100%以上。

SH 7B的另一个优点是，与Transformer相比，自回归生成期间的内存占用减少了50%以上；在Transformers中，每个层的键和值在预填充阶段被缓存，以避免重新计算并加快增量解码。

Hyena块

现有的基于低秩和稀疏近似的次二次（subquadratic）方法需要与稠密的注意层相结合才能匹配Transformer，也就是说二者在表达能力上存在差距。

也就是说，注意力机制在语言处理中只利用了其二次方能力的一小部分，所以研究问题在于，是否存在一个次二次算子，在大规模训练时的性能可以与注意力机制相匹敌？

今年2月，来自斯坦福大学和蒙特利尔大学（Mila and Université de Montréal）的研究人员提出了一个次二次下降的注意力替代品Hyena：在对数千到数十万个token的序列进行召回和推理任务时，Hyena比依赖于状态空间和其他隐式和显式方法的运算符提高了50多个点的准确性，匹配基于注意力的模型。 

论文链接：https://arxiv.org/abs/2302.10866

研究人员在标准数据集（WikiText 103和The Pile）的语言建模上设置了一个新的免密集注意力架构，达到了Transformer的质量，在序列长度为2k时所需的训练计算减少了20%；在序列长度为8k时，Hyena算子的速度是高度优化注意力的两倍，在序列长度为64k时，速度快100倍。

研究人员对高效次二次primitive进行组合，如元素乘法（门控）和长卷积（即滤波器大小与输入一样长的卷积），最终从实验结果中得到了肯定的答案。

根据最近在机制可解释性方面所做的工作（如召回recall和归纳induction），研究人员制定了一系列有针对性的推理任务，以提炼出注意力与其性能相关的三个特性，以及与现有次二次元方法之间的质量差距：

1. 数据控制（Data control）

注意力机制实现了一种富有表现力的数据控制线性算子，在单个块中对整个线性函数族进行编码。

2. 次线性参数缩放（sublinear parameter scaling）

将注意力层的参数数量与序列长度脱钩，允许Transformers在注意力层之间的其他地方分配更多参数，例如前馈神经网络（FFN）。

3. 无限制的上下文（unrestricted context）