本文將介紹數學推理場景下的首個分布外檢測研究成果。該篇論文已被 NeurIPS 2024 接收，第一作者王一鳴是上海交通大學計算機系的二年級博士生，研究方向為語言模型生成、推理，以及可解釋、可信大模型。該工作由上海交通大學和阿里巴巴通義實驗室共同完成。

• 論文題目：Embedding Trajectory for Out-of-Distribution Detection in Mathematical Reasoning
• 論文地址：https://arxiv.org/abs/2405.14039
• OpenReview: https://openreview.net/forum?id=hYMxyeyEc5
• 代碼倉庫：https://github.com/Alsace08/OOD-Math-Reasoning

背景與挑戰

分布外（Out-of-Distribution, OOD）檢測是防止深度網絡模型遭遇分布偏移數據時產生不可控輸出的重要手段，它對模型在現實世界中的部署安全起到了關鍵的作用。隨著語言模型的發展，復雜生成序列的錯誤傳播會使得 OOD 數據帶來的負面影響更加嚴重，因此語言模型下的 OOD 檢測算法變得至關重要。

常規的檢測方法主要面向傳統生成任務（例如翻譯、摘要），它們直接計算樣本在輸入 / 輸出空間中的 Embedding 和分布內（In-Distribution，ID）數據的 Embedding 分布之間的馬氏距離（Mahalanobis Distance）。然而，在數學推理場景下，這種靜態 Embedding 方法遭遇了不可行性。研究團隊可視化比較了數學推理和傳統文本生成任務在不同域上的輸入 / 輸出空間：

• 相比于文本生成，數學推理場景下不同域的輸入空間的聚類特征并不明顯，這意味著 Embedding 可能難以捕獲數學問題的復雜度；
• 更重要地，數學推理下的輸出空間呈現出高密度疊加特性。研究團隊將這種特性稱作 “模式坍縮”，它的出現主要有兩個原因：
• (1) 數學推理的輸出空間是標量化的，這會增大不同域上的數學問題產生同樣答案的可能性。例如  和 這兩個問題的結果都等于 4；
• (2) 語言模型的建模是分詞化的，這使得在數學意義上差別很大的表達式在經過分詞操作后，共享大量的 token（數字 0-9 和有限的操作符）。研究團隊量化了這一觀察，其中表示出現的所有 token 數，表示出現過的 token 種類， 表示 token 重復率，表示 token 種類在詞表中的占比，發現在一些簡單的算術場景下，token 重復率達到了驚人的 99.9%！

為了應對這個挑戰，研究團隊跳出了靜態 Embedding 的方法框架，提出了一種全新的基于動態 Embedding 軌跡 的 OOD 檢測算法，稱作 “TV Score”，以應對數學推理場景下的 OOD 檢測問題。

動機與方法

1. 定義：什么是 Embedding 軌跡？

假設語言模型有 L 層，輸出文本包含 T 個 token，則第 t 個位置的 token 在第 l 層的 Embedding 輸出表示為。現將每一層的平均 Embedding 稱為第 l 層的句子 Embedding 表征，則動態 Embedding 軌跡可形式化為一個遞進的 Embedding 鏈：

2. 動機：為什么用 Embedding 軌跡？

• 理論直覺

在數學推理場景下，輸出空間具有顯著的高密度模式坍縮特征，這使得在輸入空間相差較大的兩個起始點，通過隱藏層轉移至輸出空間后，將收斂到非常近的距離。這個 “終點收斂” 現象將增大不同樣本的 Embedding 軌跡之間產生差異的可能性，如下圖所示。該理論分析的數學建模和證明詳見論文。

• 經驗分析

在初步獲取了使用 Embedding 軌跡作為測度的理論直覺后，需要繼續深入分析 ID 和 OOD 樣本的 Embedding 軌跡之間會產生怎樣的個性化差異。研究團隊在 Llama2-7B 模型上統計了不同的 ID 和 OOD 數據集下的 Embedding 軌跡特征。其中，橫坐標表示層數，縱坐標表示該層與其鄰接層的 Embedding 之間的差值 2 - 范數，數值越大表示這兩個鄰接層之間的 Embedding 轉換幅度越大。通過統計數據得出如下發現：

• 在 20 層之前，ID 和 OOD 樣本都幾乎沒有波動；在 20 層之后，ID 樣本的 Embedding 變化幅度先增大后又被逐漸抑制，而 OOD 樣本的 Embedding 變化幅度一直保持在相對較高的范圍；
• 通過這個觀察，可以得出 ID 樣本的 “過早穩定” 現象：ID 樣本在中后層完成大量的推理過程，而后僅需做簡單的適應；而 OOD 樣本的推理過程始終沒有很好地完成 —— 這意味著 ID 樣本的 Embedding 轉換相對平滑。

3. 方法：怎么用 Embedding 軌跡？

基于上述發現，研究團隊提出了 TV Score，它可以衡量一個樣本屬于 ID 或 OOD 類別的可能性。受到靜態 Embedding 方法的啟發，文章希望通過計算新樣本的 Embedding 軌跡和 ID 樣本的 Embedding 軌跡分布之間的距離來獲取測度，但軌跡分布和軌跡距離的計算并不直觀。

因此，文章將 TV Score 的計算分為了三個步驟：

• 首先，將每一層 l 的 ID Embedding 擬合為一個高斯分布:

• 其次，對于一個新樣本，在獲取了每一層的 Embedding 后，計算它和該層高斯分布之間的馬氏距離：

• 最后，將視為新樣本的相鄰層波動率，并取所有相鄰層波動率的平均值作為該樣本的最終軌跡波動率得分：

進一步地，考慮到軌跡中的異常點可能會影響特征提取的精度，研究團隊在此基礎上加入了差分平滑技術 (Differential Smoothing, DiSmo):

• 首先，定義每一層的 k 階 Embedding 和高斯分布：

• 其次，計算和之間的馬氏距離：

• 最后，類似 TV Score 定義差分平滑后的得分：

實驗與結果

研究團隊使用了 11 個數學推理數據集（其中 1 個 ID 數據集和 10 個 OOD 數據集）在兩個不同規模的語言模型（Llama2-7B 和 GPT2-XL）上進行了實驗。根據和 ID 數據集之間的難度差異大小，這 10 個 OOD 數據集被分為兩組，分別代表 Far-shift OOD 和 Near-shift OOD。實驗在離線檢測和在線檢測這兩個場景下進行：

離線檢測場景：給定一組 ID 和 OOD 樣本的混合集合，檢測 TV Score 對這兩類樣本的區分精度（本質上是一個判別任務）。評估指標采用 AUROC 和 FPR95。

• 在 Far-shift OOD 場景下：AUROC 指標提高了 10 個點以上，FPR95 指標更是降低了超過 80%；
• 在 Near-shift OOD 場景下：TV Score 展現出更強的魯棒性。Baseline 方法從 Far-shift 轉移到 Near-shift 場景后，性能出現明顯下降，而 TV Score 仍然保持卓越的性能。這說明對于更精細的 OOD 檢測場景，TV Score 表現出更強的適應性。

在線檢測場景：在離線檢測場景中獲取一個分類閾值，之后面對新的開放世界樣本時，可以通過和閾值的大小比較自動判定屬于 ID 或 OOD 類別。評估指標采用 Accuracy。結果表明，TV Score 在開放世界場景下仍然具有十分優秀的判別準確度。

泛化性測試

研究團隊還對 TV Score 的泛化性進行了進一步的測試，主要分為任務泛化和場景泛化兩個方面：

任務泛化：測試了 OOD 場景下的生成質量估計，使用 Kendall 和 Spearman 相關系數來計算 TV Score 和模型回答正確性之間的相關性。結果表明，TV Score 在該任務下仍然展現出了最優性能。

場景泛化：研究團隊認為，TV Score 可以被推廣到所有輸出空間滿足 “模式坍縮” 特性的場景，例如多項選擇題，因為它的輸出空間僅包含 ABCD 等選項。文章選取了 MMLU 數據集，從中挑選了 8 個域的子集，依次作為 ID 子集來將剩余 7 個域作為 OOD 檢測目標。結果表明，TV Score 仍然展現出良好的性能，這驗證了它在更豐富場景下的使用價值。

總結

本文是 OOD 檢測算法在數學推理場景下的首次探索。該工作不僅揭示了傳統檢測算法在數學推理場景下的不適用性，還提出了一種全新的基于動態 Embedding 軌跡的檢測算法，可以精準適配數學推理場景。

隨著大模型的發展，模型的應用場景越來越廣泛，而這些場景也越來越具有挑戰性，早已不局限于最傳統的文本生成任務。因此，傳統安全算法在新興場景下的跟進也是維護大模型在真實世界中穩定且安全地發揮作用的不可或缺的一環。