即使最強大的 LLM 也難以透過 token 索引來關注句子等概念，現在有辦法了。
最近兩天，馬斯克和 LeCun 的口水戰妥妥成為大家的看點。這兩位 AI 圈的名人你來我往，在推特（現為 X）上相互拆對方臺。
LeCun 在宣傳自家最新論文時，也不忘手動 @ 一把馬斯克，並意味深長地囑咐道：「馬斯克，我們這項研究用來改善你家的 Grok 也沒問題。」
LeCun 宣傳的這篇論文題目為《 Contextual Position Encoding: Learning to Count What’s Important 》，來自 Meta 的 FAIR。
罵戰歸罵戰，這篇論文的重要性不言而喻。短短 24 小時之內就成為了 AI 領域最熱門的論文之一。它有望解決如今大模型（LLM）最讓人頭疼的問題。
論文地址：https://arxiv.org/pdf/2405.18719
總的來說，該研究提出了一種新的用於 transformer 的位置編碼方法 CoPE（全稱 Contextual Position Encoding），解決了標準 transformer 無法解決的計數和複製任務。傳統的位置編碼方法通常基於 token 位置，而 CoPE 允許模型根據內容和上下文來選擇性地編碼位置。CoPE 使得模型能更好地處理需要對輸入資料結構和語義內容進行精細理解的任務。文章透過多個實驗展示了 CoPE 在處理選擇性複製、計數任務以及語言和編碼任務中相對於傳統方法的優越性，尤其是在處理分佈外資料和需要高泛化能力的任務上表現出更強的效能。
CoPE 為大型語言模型提供了一種更為高效和靈活的位置編碼方式，拓寬了模型在自然語言處理領域的應用範圍。
有網友表示，CoPE 的出現改變了在 LLM 中進行位置編碼的遊戲規則，此後，研究者能夠在一個句子中精確定位特定的單詞、名詞或句子，這一研究非常令人興奮。
這篇論文主要講了什麼，我們接著看。
許多常見的資料來源（例如文字、音訊、程式碼）都是順序序列（ordered sequences）。在處理此類序列時，順序（ordering）資訊至關重要。對於文字，位置資訊不僅對於解碼單詞之間的含義至關重要，而且在其他尺度（例如句子和段落級別）上都是必需的。
作為當前大型語言模型 (LLM) 的主要支柱 Transformer 架構，依賴於注意力機制，而這種機制本身就缺乏順序資訊，因此，需要一種額外的機制來編碼資料的位置資訊。
先前有研究者提出了位置編碼（PE，Position encoding），該方法透過為每個位置分配一個嵌入向量，並將其新增到相應的 token 表示中來實現這一點。然而，當前的位置編碼方法使用 token 計數來確定位置，因此無法推廣到更高層次如句子。
為了將位置與更具有語義的單元（如單詞或句子）聯絡起來，需要考慮上下文。然而，使用當前的位置編碼方法無法實現這一點，因為位置定址是獨立於上下文計算的，然後再與上下文定址合併。
Meta 認為，位置與上下文定址的這種分離是問題的根本所在，因此他們提出了一種新的 PE 方法，即上下文位置編碼（ CoPE ），將上下文和位置定址結合在一起。
方法介紹
CoPE 首先使用上下文向量確定要計數的 token。具體來說，給定當前 token 作為查詢向量，接著使用先前 token 的鍵向量計算一個門值（gate value）。然後彙總這些門值，以確定每個 token 相對於當前 token 的相對位置，如圖 1 所示。
與 token 位置不同，上下文位置可以取分數值，因而不能具有指定的嵌入。相反，該研究插入賦值為整數值的嵌入來計算位置嵌入。與其他 PE 方法一樣，這些位置嵌入隨後被新增到鍵向量中，因此查詢向量可以在注意力操作中使用它們。由於上下文位置可能因查詢和層而異，因此該模型可以同時測量多個單元的距離。
在 CoPE 中，位置是透過上下文相關的方式來測量的，而不是簡單的 token 計數。該方法的工作原理是首先決定在使用上下文向量測量距離時應包含哪些 token。因此，對每個查詢 q_i 和鍵 k_j 對計算門值
其中 j < i 且 σ 是 sigmoid 函式。門值為 1 表示該鍵將被計入位置測量中，而 0 表示將被忽略。例如，要計算 token i 和 j 之間的句子，僅對於諸如 “.” 之類的句子分隔 token，門值應為 1。門以查詢為條件，如果需要，每個查詢可以有不同的位置測量。軟門控函式（soft gating function）允許微分，以便可以透過反向傳播來訓練系統。
然後，該研究透過新增當前 token 和目標 token 之間的門值來計算位置值。
值得注意的是，如果門值始終為 1，則 p_ij = i − j + 1 ，並且恢復基於 token 的相對位置。因此，CoPE 可以被視為相對 PE 的泛化。然而，一般來說，p_ij 可以是特定單詞或單詞型別（如名詞或數字）的計數、句子的數量或 Transformer 認為在訓練期間有用的其他概念。
與 token 位置不同，位置值 p_ij 不限於整數，並且因為 sigmoid 函式的原因可以採用小數值。這意味著不能像相對 PE 中那樣使用嵌入層將位置值轉換為向量。
首先，該研究為每個整數位置 p ∈ [0, T] 分配一個可學習的嵌入向量 e [p]，那麼位置 p_ij 的嵌入將是兩個最接近的整數嵌入的簡單插值。
最後，計算類似於如下等式的注意力權重。
然而，在實踐中，計算和儲存向量 e [p_ij ] 需要使用額外的計算和記憶體。該研究透過首先計算所有整數位置 p 的乘法，再對結果值進行插值來提高效率：
如下方程 (4) 所示，p_ij 的最大值是上下文大小 T，這意味著需要 T + 1 個位置嵌入（包括位置 0）。然而，如果門被稀疏啟用（例如計算句子），則可以用更少的位置覆蓋整個上下文 T。因此，該研究透過設定，使得最大可能位置 p_max < T。
CoPE 的多頭擴充套件非常簡單，因為每個頭都會獨立執行自己的 CoPE。頭之間的鍵和查詢向量是不同的，這意味著它們可以實現不同的位置測量。
實驗結果
Flip-Flop 任務
Liu 等人 [2024] 提出了 Flip-Flop 語言建模任務，以揭示 Transformer 模型無法在長距離輸入序列上進行穩健推理。
結果如表 2（左）所示。結果表明，CoPE 優於現有方法，使模型不僅可以學習分佈內任務，還可以推廣到 OOD 序列 —— 這是現有 PE 方法無法提供的屬性。
選擇性複製任務
Gu 和 Dao [2023] 提出的選擇性複製任務需要上下文感知推理才能進行選擇性記憶。
表 2（右）中給出的結果顯示，在分佈內測試集上，新方法 CoPE 可以解決該任務，而其他方法則無法解決。同樣的，CoPE 在密集和稀疏 OOD 測試集上都具有更好的泛化能力。空白 token 的存在使得找到下一個要複製的 token 變得更加困難，但 CoPE 只能計算非空白 token，因此更加穩定。在每個步驟中，它可以簡單地複製距離為 256（非空白）的非空白 token。重複此操作 256 次將複製整個非空白序列。
計數任務
計數比簡單地回憶上一個例項更具挑戰性，因為它需要在一定範圍內更均勻的注意力。
結果見表 3 和圖 2。具有相對 PE 的基線模型很難學習此任務，尤其是當有多個變數需要跟蹤時。絕對 PE 的表現更差。最佳表現來自 CoPE，在 1 個變數的情況下獲得滿分。對於 OOD 泛化，相對 PE 表現出較差的泛化能力，而 CoPE 的泛化能力非常好，如表 4 所示。有關這些實驗的標準差，請參見附錄表 9。
語言建模
為了在語言建模任務上測試新方法，研究人員使用了 Wikitext-103 資料集，該資料集包含從 Wikipedia 中提取的 1 億個 token。
表 5（左）中比較了不同的 PE 方法：絕對 PE 表現最差，CoPE 優於相對 PE，與相對 PE 結合使用時效果更佳。這表明，即使在一般語言建模中，CoPE 也能帶來改進。
接下來，作者測試了 CoPE 推廣到比訓練上下文更長的上下文的效果。
結果如圖 3 所示。相對 PE 推廣到更長的上下文效果不佳。相比之下，相對上限版本的表現要好得多。然而 CoPE 的表現仍然優於它，當測試上下文比訓練上下文長得多時，差距會擴大（見圖 3 右）。
如圖 4 所示，作者展示了使用 sep-keys 訓練的模型的注意力圖示例（gate 是用分離的鍵計算的）。注意力圖僅根據位置構建（它們必須與上下文注意力相乘才能得到最終的注意力），這能讓我們更好地瞭解 CoPE 正在做什麼。作者還進行了歸一化，以便每個查詢的最大注意力權重始終為 1。首先，我們可以看到位置明顯具有上下文相關性，因為無論它們的相對位置如何，注意力都傾向於落在特定的 token 上。
仔細觀察這些 token 會發現，注意力主要集中在最後一段（左）或部分（右）上。為清楚起見，實際的段落和部分邊界用黑色加號標記。在 CoPE 中，這是可能的，因為一個注意力頭可以計數段落，而另一個注意力頭計數部分，然後它可以只關注位置 0。
程式碼建模
作者透過對程式碼資料進行評估來進一步測試 CoPE 的能力。與自然語言相比，程式碼資料具有更多的結構，並且可能對上下文學習更敏感。
結果總結在表 5（右）中。CoPE 嵌入的困惑度比絕對 PE 和 RoPE 分別提高了 17% 和 5%。將 RoPE 和 CoPE 嵌入結合在一起可以改善 RoPE，但不會比所提出的嵌入方法帶來任何改進。