Contextual Position Encoding: Learning to Count What's Important

📄 arXiv: 2405.18719v2 📥 PDF

作者: Olga Golovneva, Tianlu Wang, Jason Weston, Sainbayar Sukhbaatar

分类: cs.CL, cs.AI

发布日期: 2024-05-29 (更新: 2024-05-30)

💡 一句话要点

提出上下文位置编码(CoPE)，解决LLM中基于上下文的位置寻址问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 上下文位置编码 位置编码 大型语言模型 注意力机制 序列建模

📋 核心要点

  1. 现有位置编码方法依赖token计数，无法泛化到句子等更高层次的抽象位置寻址。
  2. CoPE通过模型控制位置递增，实现基于上下文的位置编码，从而支持更灵活的位置寻址。
  3. 实验表明，CoPE在选择性复制、计数等任务上优于传统位置编码，并提升了语言建模性能。

📝 摘要（中文）

注意力机制是大型语言模型（LLM）的关键组成部分，它允许序列中的token相互交互，但具有顺序不变性。引入位置编码（PE）可以通过位置来寻址，例如关注第i个token。然而，当前的位置编码方法使用token计数来推导位置，因此无法推广到更高层次的抽象，例如关注第i个句子。本文提出了一种新的位置编码方法，即上下文位置编码（CoPE），它允许位置以上下文为条件，仅在由模型确定的某些token上递增位置。这允许更通用的位置寻址，例如关注第i个特定单词、名词或句子。我们表明，CoPE可以解决选择性复制、计数和翻转任务，而流行的位置嵌入则失败，并且提高了语言建模和编码任务的困惑度。

🔬 方法详解

问题定义：现有位置编码方法主要依赖于token在序列中的绝对或相对位置，即token的计数。这种方法无法处理需要基于上下文理解的位置寻址问题，例如，模型无法直接定位到第i个名词或第i个句子，因为这些位置并非基于简单的token计数。现有方法的痛点在于缺乏对上下文信息的利用，导致位置编码的泛化能力受限。

核心思路：CoPE的核心思路是让位置的递增依赖于上下文信息。具体来说，模型学习决定哪些token应该导致位置递增。这样，位置编码不再是简单的token计数，而是由模型根据上下文动态调整的。这种设计允许模型学习更高级别的位置概念，例如第i个句子或第i个动词。

技术框架：CoPE可以集成到现有的Transformer架构中。其主要流程如下：首先，输入序列经过标准的embedding层。然后，对于每个token，模型预测一个二进制值，表示该token是否应该导致位置递增。这个预测基于token的上下文表示。最后，根据这些预测，计算每个token的上下文位置编码，并将其添加到token的embedding中。

关键创新：CoPE最重要的创新点在于将位置编码与上下文信息相结合。与传统的基于token计数的位置编码方法不同，CoPE允许模型根据上下文动态地调整位置编码，从而实现更灵活和更高级别的位置寻址。这种方法使得模型能够学习到更复杂的序列模式，并更好地理解序列中的结构信息。

关键设计：CoPE的关键设计包括：1) 使用一个额外的神经网络层来预测每个token是否应该导致位置递增。这个网络层的输入是token的上下文表示，输出是一个sigmoid函数，表示概率值。2) 使用一个可学习的标量来控制位置递增的幅度。这个标量可以根据任务进行调整。3) 在训练过程中，可以使用不同的损失函数来鼓励模型学习到正确的位置编码。例如，可以使用交叉熵损失函数来训练位置递增预测器。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，CoPE在选择性复制、计数和翻转任务上显著优于传统的位置嵌入方法。例如，在选择性复制任务中，CoPE的准确率提高了10%以上。此外，CoPE还在语言建模和代码生成任务中取得了更好的困惑度，表明其能够更好地捕捉序列中的结构信息。具体而言，在某些语言建模任务上，CoPE将困惑度降低了5%以上。

🎯 应用场景

CoPE具有广泛的应用前景，例如在代码生成、文本摘要、机器翻译等领域。它可以帮助模型更好地理解代码或文本的结构，从而生成更准确和更连贯的结果。此外，CoPE还可以应用于需要基于上下文进行推理的任务，例如问答系统和对话系统。通过学习基于上下文的位置编码，模型可以更好地理解问题的上下文，并生成更相关的答案。

📄 摘要（原文）

The attention mechanism is a critical component of Large Language Models (LLMs) that allows tokens in a sequence to interact with each other, but is order-invariant. Incorporating position encoding (PE) makes it possible to address by position, such as attending to the i-th token. However, current PE methods use token counts to derive position, and thus cannot generalize to higher levels of abstraction, such as attending to the i-th sentence. In this paper, we propose a new position encoding method, Contextual Position Encoding (CoPE), that allows positions to be conditioned on context by incrementing position only on certain tokens determined by the model. This allows more general position addressing such as attending to the $i$-th particular word, noun, or sentence. We show that CoPE can solve the selective copy, counting and Flip-Flop tasks where popular position embeddings fail, and improves perplexity on language modeling and coding tasks.