0. TL; DR

本文介绍了一种新型的卷积混合架构，旨在通过结合卷积和注意力机制的优势，提升大规模语言模型的效率和性能。该架构通过定制化的卷积操作符和硬件感知算法设计，实现了在大规模参数（400亿参数）下的高效训练，相比优化后的Transformer模型，训练速度提升了1.2到2.9倍。此外，该架构在处理字节级数据时表现出色，为基因组学等领域的基础模型开发提供了新的可能性。

1. 背景介绍

近年来，语言模型的架构改进主要集中在以下几个方面：

注意力机制的优化：通过减少kv缓存的大小，如GQA、MQA、MLA、滑动窗口和线性注意力等方法，提升模型效率。
数值稳定性与量化改进：引入预归一化、SwiGLU、QK归一化等技术，增强模型对异常值的鲁棒性和量化能力。
模型容量与长文本建模：通过RoPE、MoE等技术提升模型在长文本建模中的表现。

尽管这些方法在一定程度上提升了模型性能，但在大规模语言模型训练中，能够实现一致性能提升的改进仍然有限。

混合架构通过引入新的输入依赖操作符（如卷积和注意力的组合），优化其组合，以实现质量和效率的双重提升。然而，现有的混合架构（如基于线性注意力或状态空间模型的混合架构）在短序列和宽模型的常见预训练场景中，通常比Transformer慢，且在长序列上表现不佳。

本文提出了一种新的卷积混合架构，通过结合不同类型的卷积操作符（短、中、长）和硬件感知算法设计，实现了在大规模参数下的高效训练。

2. 方法介绍

本文提出的卷积混合架构基于输入依赖卷积操作符（Hyena），其基本结构如下：

\[q_{\alpha}^{t} = T_{\alpha}^{tt'}(x_{\beta}^{t'}W_{\beta}^{\alpha}) \\ k_{\alpha}^{t} = H_{\alpha}^{tt'}(x_{\beta}^{t'}U_{\beta}^{\alpha}) \\ v_{\alpha}^{t} = K_{\alpha}^{tt'}(x_{\beta}^{t'}P_{\beta}^{\alpha}) \\ y_{\alpha}^{t} = (q_{\beta}^{t} G_{\beta}^{tt'}k_{\beta}^{t'}v_{\beta}^{t'} )M_{\beta}^{\alpha}\]

其中，$T, H, K, G$ 是Toeplitz矩阵，对应卷积操作符的滤波器；$W, U, P, M$ 是密集矩阵，用于参数化。该架构通过定制不同类型的卷积操作符，实现对不同子任务（如回忆、压缩、多查询回忆和模糊回忆）的优化。

本文设计了三种卷积操作符，这些卷积操作符通过分层组合，形成了完整的卷积混合架构。

Hyena-SE（短显式卷积）：使用短的显式滤波器，适用于短序列和局部多标记回忆任务。通过硬件感知实现，Hyena-SE在短序列上表现出色，且在硬件上具有高吞吐量。
Hyena-MR（中等正则化卷积）：使用中等长度的显式滤波器，适用于跨数百标记的高效建模。通过简单的指数衰减正则化，Hyena-MR在长序列上表现出色。
Hyena-LI（长隐式卷积）：使用长的隐式滤波器，适用于跨整个序列的信息聚合。Hyena-LI通过线性组合的实指数参数化，能够实现常数内存的递归参数化。

3. 实验分析

本文在大规模参数（400亿参数）和海量数据（超过9万亿标记）上验证了卷积混合架构的有效性。实验使用了H100 GPU集群，模型宽度为4096。实验中，作者对比了卷积混合架构与Transformer、线性注意力和状态空间模型等其他架构的性能。

主要实验结论：

训练速度提升：卷积混合架构在400亿参数规模下，相比优化后的Transformer模型，训练速度提升了1.2到2.9倍。在H100 GPU上，单个操作符的吞吐量比线性注意力和状态空间模型高出两倍。
硬件感知优化：通过定制化的卷积算法和上下文并行策略，卷积混合架构在硬件上实现了高效的计算。例如，通过两阶段阻塞算法和FFT卷积，进一步提升了模型的吞吐量。
上下文扩展能力：通过位置插值和调整基频等技术，卷积混合架构能够扩展到长达100万的上下文长度，且在扩展过程中保持了良好的性能。