本文將介紹數(shù)學(xué)推理場(chǎng)景下的首個(gè)分布外檢測(cè)研究成果。該篇論文已被 NeurIPS 2024 接收，第一作者王一鳴是上海交通大學(xué)計(jì)算機(jī)系的二年級(jí)博士生，研究方向?yàn)檎Z言模型生成、推理，以及可解釋、可信大模型。該工作由上海交通大學(xué)和阿里巴巴通義實(shí)驗(yàn)室共同完成。

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

背景與挑戰(zhàn)

分布外（Out-of-Distribution, OOD）檢測(cè)是防止深度網(wǎng)絡(luò)模型遭遇分布偏移數(shù)據(jù)時(shí)產(chǎn)生不可控輸出的重要手段，它對(duì)模型在現(xiàn)實(shí)世界中的部署安全起到了關(guān)鍵的作用。隨著語言模型的發(fā)展，復(fù)雜生成序列的錯(cuò)誤傳播會(huì)使得 OOD 數(shù)據(jù)帶來的負(fù)面影響更加嚴(yán)重，因此語言模型下的 OOD 檢測(cè)算法變得至關(guān)重要。

常規(guī)的檢測(cè)方法主要面向傳統(tǒng)生成任務(wù)（例如翻譯、摘要），它們直接計(jì)算樣本在輸入 / 輸出空間中的 Embedding 和分布內(nèi)（In-Distribution，ID）數(shù)據(jù)的 Embedding 分布之間的馬氏距離（Mahalanobis Distance）。然而，在數(shù)學(xué)推理場(chǎng)景下，這種靜態(tài) Embedding 方法遭遇了不可行性。研究團(tuán)隊(duì)可視化比較了數(shù)學(xué)推理和傳統(tǒng)文本生成任務(wù)在不同域上的輸入 / 輸出空間：

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

為了應(yīng)對(duì)這個(gè)挑戰(zhàn)，研究團(tuán)隊(duì)跳出了靜態(tài) Embedding 的方法框架，提出了一種全新的基于動(dòng)態(tài) Embedding 軌跡 的 OOD 檢測(cè)算法，稱作 “TV Score”，以應(yīng)對(duì)數(shù)學(xué)推理場(chǎng)景下的 OOD 檢測(cè)問題。

動(dòng)機(jī)與方法

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

假設(shè)語言模型有 L 層，輸出文本包含 T 個(gè) token，則第 t 個(gè)位置的 token 在第 l 層的 Embedding 輸出表示為?，F(xiàn)將每一層的平均 Embedding 稱為第 l 層的句子 Embedding 表征，則動(dòng)態(tài) Embedding 軌跡可形式化為一個(gè)遞進(jìn)的 Embedding 鏈：

2. 動(dòng)機(jī)：為什么用 Embedding 軌跡？

• 理論直覺

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

• 經(jīng)驗(yàn)分析

在初步獲取了使用 Embedding 軌跡作為測(cè)度的理論直覺后，需要繼續(xù)深入分析 ID 和 OOD 樣本的 Embedding 軌跡之間會(huì)產(chǎn)生怎樣的個(gè)性化差異。研究團(tuán)隊(duì)在 Llama2-7B 模型上統(tǒng)計(jì)了不同的 ID 和 OOD 數(shù)據(jù)集下的 Embedding 軌跡特征。其中，橫坐標(biāo)表示層數(shù)，縱坐標(biāo)表示該層與其鄰接層的 Embedding 之間的差值 2 - 范數(shù)，數(shù)值越大表示這兩個(gè)鄰接層之間的 Embedding 轉(zhuǎn)換幅度越大。通過統(tǒng)計(jì)數(shù)據(jù)得出如下發(fā)現(xiàn)：

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

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

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

因此，文章將 TV Score 的計(jì)算分為了三個(gè)步驟：

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

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

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

進(jìn)一步地，考慮到軌跡中的異常點(diǎn)可能會(huì)影響特征提取的精度，研究團(tuán)隊(duì)在此基礎(chǔ)上加入了差分平滑技術(shù) (Differential Smoothing, DiSmo):

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

• 其次，計(jì)算和之間的馬氏距離：

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

實(shí)驗(yàn)與結(jié)果

研究團(tuán)隊(duì)使用了 11 個(gè)數(shù)學(xué)推理數(shù)據(jù)集（其中 1 個(gè) ID 數(shù)據(jù)集和 10 個(gè) OOD 數(shù)據(jù)集）在兩個(gè)不同規(guī)模的語言模型（Llama2-7B 和 GPT2-XL）上進(jìn)行了實(shí)驗(yàn)。根據(jù)和 ID 數(shù)據(jù)集之間的難度差異大小，這 10 個(gè) OOD 數(shù)據(jù)集被分為兩組，分別代表 Far-shift OOD 和 Near-shift OOD。實(shí)驗(yàn)在離線檢測(cè)和在線檢測(cè)這兩個(gè)場(chǎng)景下進(jìn)行：

離線檢測(cè)場(chǎng)景：給定一組 ID 和 OOD 樣本的混合集合，檢測(cè) TV Score 對(duì)這兩類樣本的區(qū)分精度（本質(zhì)上是一個(gè)判別任務(wù)）。評(píng)估指標(biāo)采用 AUROC 和 FPR95。

• 在 Far-shift OOD 場(chǎng)景下：AUROC 指標(biāo)提高了 10 個(gè)點(diǎn)以上，F(xiàn)PR95 指標(biāo)更是降低了超過 80%；
• 在 Near-shift OOD 場(chǎng)景下：TV Score 展現(xiàn)出更強(qiáng)的魯棒性。Baseline 方法從 Far-shift 轉(zhuǎn)移到 Near-shift 場(chǎng)景后，性能出現(xiàn)明顯下降，而 TV Score 仍然保持卓越的性能。這說明對(duì)于更精細(xì)的 OOD 檢測(cè)場(chǎng)景，TV Score 表現(xiàn)出更強(qiáng)的適應(yīng)性。

在線檢測(cè)場(chǎng)景：在離線檢測(cè)場(chǎng)景中獲取一個(gè)分類閾值，之后面對(duì)新的開放世界樣本時(shí)，可以通過和閾值的大小比較自動(dòng)判定屬于 ID 或 OOD 類別。評(píng)估指標(biāo)采用 Accuracy。結(jié)果表明，TV Score 在開放世界場(chǎng)景下仍然具有十分優(yōu)秀的判別準(zhǔn)確度。

泛化性測(cè)試

研究團(tuán)隊(duì)還對(duì) TV Score 的泛化性進(jìn)行了進(jìn)一步的測(cè)試，主要分為任務(wù)泛化和場(chǎng)景泛化兩個(gè)方面：

任務(wù)泛化：測(cè)試了 OOD 場(chǎng)景下的生成質(zhì)量估計(jì)，使用 Kendall 和 Spearman 相關(guān)系數(shù)來計(jì)算 TV Score 和模型回答正確性之間的相關(guān)性。結(jié)果表明，TV Score 在該任務(wù)下仍然展現(xiàn)出了最優(yōu)性能。

場(chǎng)景泛化：研究團(tuán)隊(duì)認(rèn)為，TV Score 可以被推廣到所有輸出空間滿足 “模式坍縮” 特性的場(chǎng)景，例如多項(xiàng)選擇題，因?yàn)樗妮敵隹臻g僅包含 ABCD 等選項(xiàng)。文章選取了 MMLU 數(shù)據(jù)集，從中挑選了 8 個(gè)域的子集，依次作為 ID 子集來將剩余 7 個(gè)域作為 OOD 檢測(cè)目標(biāo)。結(jié)果表明，TV Score 仍然展現(xiàn)出良好的性能，這驗(yàn)證了它在更豐富場(chǎng)景下的使用價(jià)值。

總結(jié)

本文是 OOD 檢測(cè)算法在數(shù)學(xué)推理場(chǎng)景下的首次探索。該工作不僅揭示了傳統(tǒng)檢測(cè)算法在數(shù)學(xué)推理場(chǎng)景下的不適用性，還提出了一種全新的基于動(dòng)態(tài) Embedding 軌跡的檢測(cè)算法，可以精準(zhǔn)適配數(shù)學(xué)推理場(chǎng)景。

隨著大模型的發(fā)展，模型的應(yīng)用場(chǎng)景越來越廣泛，而這些場(chǎng)景也越來越具有挑戰(zhàn)性，早已不局限于最傳統(tǒng)的文本生成任務(wù)。因此，傳統(tǒng)安全算法在新興場(chǎng)景下的跟進(jìn)也是維護(hù)大模型在真實(shí)世界中穩(wěn)定且安全地發(fā)揮作用的不可或缺的一環(huán)。