ICML 2024 | 北大、字节提出新型双层位置编码方案，有效改善长度外推效果
ICML 2024 | 北大、字节提出新型双层位置编码方案，有效改善长度外推效果
在自然语言处理领域，如何让模型更好地理解长文本一直是研究者们关注的重点。最近，北京大学和字节跳动的研究团队提出了一种新型的双层位置编码方案（BiPE），通过将文本序列分解为模块化的段落，并为每个位置设计了段内编码和段间编码，有效改善了模型的长度外推能力。这一创新性研究已被ICML 2024接收。
研究背景
在许多场景中，文本可以有效地分解为模块化的段落，每个段落都表达一个自成一体的思想单元。例如，在自然语言中，文档通常由句子组成；在编程语言中，代码被组织成行或函数类；在数学中，证明通过一系列演绎步骤展开。
有趣的是，研究者们发现，对于不同长度的序列，每个模块段中的token数量分布通常是有界的，并且趋于大致相似。在图1中，研究者们使用了广泛使用的PG-19文本语料库进行可视化。显然，每个段落（即句子）中的token数量分布无论总序列长度如何，都保持了显著的一致性。相比之下，随着序列长度的增加，句子的数量呈线性增长。
图1：PG-19语料库中token数量分布可视化
方法
研究者们提出了一种简单但有效的用于改善长度外推的定位编码方案：BiPE。与现有的长度外推方法不同，BiPE为每个位置采用了两种不同的编码：段内编码和段间编码。
- 段内编码：标识标记在其段内的位置，帮助模型捕捉其中的语义信息。
- 段间编码：指定其所属的段，通过相对位置编码建模段间关系。
以自然语言为例，同一句子中的不同词语共享相同的段间位置编码，但具有不同的段内编码。相反，不同句子中但占据相同段间位置的词语（例如，不同句子中的第一个token）共享相同的段内编码，而具有不同的段间编码。
理论分析
研究者们利用了理论计算机科学领域广泛使用的（非确定性）有限自动机（NFA）。受分层有限自动机框架的启发，研究者们引入了一种简化模型——Bi-NFA，该模型将分层有限自动机的层级限制为两层。通过理论分析，研究者们展示了BiPE在现有位置编码方案上具有理论优势。
实验结果
研究者们在多个实验中验证了BiPE的有效性：
思维链数学推理任务：在图3中可以清楚地看到，给定相似数量的参数，基于BiPE的语言模型在此任务上始终优于其他基线方法。例如，当隐藏维度为48时，其他位置编码方法的准确率低于70%，而BiPE-ALiBi和BiPE-RoPE的准确率分别高达97%和95%。
Pile数据集预训练：在图4中，BiPE方法在长度大于训练长度的序列上始终表现出更优越的性能。例如，在序列长度为8192的PG19数据集上，BiPE-ALiBi以25.24的困惑度优于ALiBi（28.59的困惑度），领先了3.35点。
SCROLLS数据集微调测试：BiPE-RoPE和BiPE-ALiBi分别表现优于RoPE和ALiBi。例如，BiPE-RoPE以22.36的平均分数超过了RoPE（18.38的平均分数），领先了3.98分。此外，BiPE-RoPE取得了最高的平均分数，超过其他方法3分以上。
总结
研究者们提出了BiPE，一种旨在改进长度外推的新型双层位置编码方案。对于每个位置，BiPE结合了通过绝对位置编码确定其在段内位置的段内编码，以及通过相对位置编码指定其所属段的段间编码。这种双层设计很好地契合了文本数据的内在分段结构，并增强了长度外推能力。所有实验均验证了BiPE在不同文本模态任务中的长度外推能力。
论文标题：Two Stones Hit One Bird: Bilevel Positional Encoding for Better Length Extrapolation
论文链接：https://arxiv.org/abs/2401.16421
代码链接：https://github.com/zhenyuhe00/BiPE