标点符号成大模型训练神器！KV缓存狂减一半，可处理400万Tokens长序列，来自华为港大等

文字中貌似不起眼的标点符号，竟然可以显著加速大模型的训练和推理过程？

来自华为、港大、KAUST和马普所的研究者，就提出了一种新的自然语言建模视角——SepLLM。

起因是团队发现某些看似无意义的分隔符，在注意力得分中占据了不成比例的重要地位。

于是，SepLLM通过将一段文本中的信息压缩进分隔符（比如逗号，句号等）中，真的实现了加速效果，并且可以让KV缓存减少一半。

自注意力机制的平方级复杂度，给计算存储需求和训练推理速度带来了不小的挑战。

为了降低推理的复杂度，大量节约KV Cache的稀疏化方法被提出。

然而这些方法大多是基于用户的问题或者提示来筛选有用的KV Cache。

这使得如果用户再提出一个新的问题，模型回答的精度可能下降，因为包含答案信息的KV已经在上一次压缩过程中被抛弃。

除此之外，免训练方法通常无法相应地从头训练或者后训练，导致了训练和推理的流程差异性。

更重要的是现在主流的稀疏注意力改进方法，本质上更多是一种针对KV Cache存储与计算的稀疏化管理，而不是对自然语言的自然且高效的建模。

用分隔符实现自然语言高效建模

SepLLM通过将一段文本中的信息压缩进分隔符（比如逗号，句号等）中，显著加速了大型语言模型的训练和推理过程。

这一发现基于一个新颖且关键的模式：某些看似无意义的分隔符，在注意力得分中占据了不成比例的重要地位。

如下图所示，注意力可视化显示出一定的稀疏性，并且在分隔符处注意力明显更大。

由此，可以将这些自然语言中分隔符所自然分割的语义段的信息有效地压缩进分隔符中，其他tokens直接丢弃，而不会造成信息损失。

除此之外，一般一个分割符所分割的语段的长度是有限且相对均衡的，因此用分割此语段的分隔符去浓缩语段信息，可以避免类似RNN当序列过长时而出现遗忘的问题。

因为这种基于分割符的语言建模视角反映了自然语言的自然而内在的稀疏性，而不是人为用类似block/cluster等概念预先定义的稀疏性，作者认为SepLLM可以作为大语言模型的原生稀疏注意力机制和原生基线模型。

具体来说，SepLLM的基础设计包含下列三种tokens：

初始tokens：使用稀疏注意力机制时，保留初始tokens可避免生成tokens的困惑度（ppl）显著增加。
分隔符tokens：看似“无意义”的分隔符tokens在给定输入上下文中比有语义意义的tokens获得更高的注意力分数。因此假设这些分隔符可压缩其分割的文本片段信息，在免训练（training-free）的场景中，基于此策略能在许多任务上取得与原始模型相似的结果；
相邻tokens：由于语言任务通常具有局部依赖性，相邻tokens有助于形成局部平滑和连贯的上下文，所以在模型中考虑相邻tokens。

在预训练或者后训练的过程中，强迫模型当前的token只能看到前文每个片段中代表该片段的分隔符，使片段信息被强制浓缩到分隔符中。

实际上，每个分隔符（逗号、句号、分号、问号等）都是具备其特有的语义的，它们是对其分割段落的最原生和最细粒度的收尾与总结。

训练阶段，不需要将输入上下文中所有tokens对应的Query向量与所有Key向量相乘，只需乘以掩码矩阵中突出显示元素对应的Key向量；

生成阶段对KV缓存的管理较为直观，只保留初始、分隔符和相邻tokens的KV Cache。

研究者还针对Streaming场景还提出了定制的设计，包括同时维护的四个专用缓存块（初始缓存、分隔符缓存、过去窗口缓存和局部窗口缓存）及其功能，定义了四个缓存的运行时使用量和相邻tokens数量的相关变量，并详细说明了缓存系统的预设超参数。

在Streaming序列生成过程中，SepLLM会按照一定规则填充和管理这些缓存，当缓存达到一定条件时会触发压缩操作。

作者分析了KV Cache的平均使用情况，结果，SepLLM在免训练、预训练和后训练场景中都展现出了卓越的效率，首先进行一个简单总结：

训推效率提升：SepLLM在免训练、从头预训练和后训练中都展现出了卓越的效率。特别是在使用Llama-3-8B模型时，SepLLM在GSM8K和MMLU基准测试中减少了超过50%的KV缓存，同时保持了相当的性能表现。
无限长的流式处理能力：在无限长输入的流式的场景中，SepLLM能够有效处理高达400万甚至更多tokens的序列，同时保持一致的语言建模能力。
广泛的实验验证与理论分析：通过在多种任务，各种基础模型（Llama，Falcon， GPTNeoX等）和多种数据集上的广泛实验，SepLLM证明了其在不同设置下的有效性，包括免训练、预训练和后训练。除此之外，作者还提供了对SepLLM架构通用近似（Universal Approximation）的详细理论分析。