[论文解读] KERPLE: Kernelized Relative Positional Embedding for Length Extrapolation

Ta-Chung Chi, Ting-Han Fan+2 more|arXiv (Cornell University)|May 20, 2022

Topic Modeling被引用 9

一句话总结

KERPLE 提出了一种基于核函数的相对位置嵌入框架，通过利用条件正定（CPD）核函数，将其转化为正定（PD）核函数（通过在 Softmax 归一化过程中吸收常数偏移），从而实现 Transformer 模型的上下文长度外推。该方法在 OpenWebText2、GitHub 和 ArXiv 等大规模语言建模数据集上实现了当前最优的外推性能，尤其在对数核函数变体上表现突出。

ABSTRACT

Relative positional embeddings (RPE) have received considerable attention since RPEs effectively model the relative distance among tokens and enable length extrapolation. We propose KERPLE, a framework that generalizes relative position embedding for extrapolation by kernelizing positional differences. We achieve this goal using conditionally positive definite (CPD) kernels, a class of functions known for generalizing distance metrics. To maintain the inner product interpretation of self-attention, we show that a CPD kernel can be transformed into a PD kernel by adding a constant offset. This offset is implicitly absorbed in the Softmax normalization during self-attention. The diversity of CPD kernels allows us to derive various RPEs that enable length extrapolation in a principled way. Experiments demonstrate that the logarithmic variant achieves excellent extrapolation performance on three large language modeling datasets. Our implementation and pretrained checkpoints are released at https://github.com/chijames/KERPLE.git.

研究动机与目标

• 为解决 Transformer 模型在长度外推方面的挑战，即模型必须泛化到超过其训练长度 L 的序列长度。
• 克服现有相对位置嵌入（RPE）方法在训练长度之外泛化能力不足的局限。
• 提出一种基于条件正定（CPD）核函数数学特性的系统化 RPE 设计框架，以支持外推。
• 证明可学习的 CPD 核函数能够更好地适应数据集特定的长程依赖关系，优于固定参数方法（如 ALiBi）。
• 提供一个统一框架，可泛化现有方法（如 ALiBi），并支持对具备强外推能力的 RPE 进行系统性探索。

提出的方法

• 该框架使用条件正定（CPD）核函数对相对位置差异进行核化，该核函数可推广距离度量并支持外推。
• 通过添加常数偏移，将 CPD 核函数转化为正定（PD）核函数，该偏移在自注意力机制的 Softmax 归一化过程中被隐式吸收。
• 在注意力头之间采用可学习的核函数参数化方式，使模型能够适应数据集特定的长程依赖关系。
• 该框架应用于因果语言建模任务，注意力分数由查询-键点积与核化相对位置偏置之和计算得出。
• 核函数被设计为缓慢衰减（例如对数形式），以在推理阶段处理更长序列时保持对远距离标记的关注。
• 实验基于 GPT-NeoX 代码库进行公平比较，模型在短序列（长度为 3）上进行训练，并在长序列（最长 16384 个标记）上进行评估。

实验结果

研究问题

• RQ1 能否利用条件正定（CPD）核函数设计出支持 Transformer 模型长度外推的相对位置嵌入？
• RQ2 核函数的选择（如对数函数与幂律函数）如何影响模型在推理阶段对长程依赖关系的关注能力？
• RQ3 核函数中的可学习参数是否能比固定参数方法（如 ALiBi）更好地适应数据集特定的长程依赖关系？
• RQ4 该核化 RPE 框架在多大程度上泛化了 ALiBi 等现有方法？
• RQ5 注意力有效长度（由核函数衰减速率定义）与模型在长上下文任务上的性能之间是否存在相关性？

主要发现

• 在 OpenWebText2、GitHub 和 ArXiv 三个大规模语言建模数据集上，KERPLE 的对数核变体实现了最佳的长度外推性能。
• KERPLE-log 在序列长度 16384 时保持了较低的困惑度，优于窗口注意力（window@512）和 ALiBi，后两者表现出更高的困惑度。
• 使用 KERPLE-log 训练的模型对远距离标记的关注衰减更慢，表明其具备更强的长程依赖建模能力。
• 注意力有效长度（定义为核偏置降至 -2 以下的距离）相比幂律变体，对数核函数的该值显著更长。
• 有效长度 ≤ |m−n| 的注意力头累积数量显示，与其它方法相比，KERPLE-log 在更广泛的距离范围内维持了对远距离标记的关注。
• 该方法将 ALiBi 作为特例进行泛化：ALiBi 对应于参数固定的线性 CPD 核函数，而 KERPLE 允许可学习、数据集自适应的参数。