首詞元截斷 LLM 取樣
從理論角度來看，LLM 所能提供的最佳回覆是透過始終挑選機率最高的詞元來取得。這種做法會讓 LLM 的輸出變得確定性，這對於許多應用來說並非理想的特性。為此，為了在保持對上下文的依循同時兼顧 LLM 的創造力，近年來已提出了不同的取樣演算法。
如今最常被使用、幾乎可說是預設的方法之一叫做 top-p：它是一種 nucleus sampling(核心取樣)，會將得分最高的詞元一路收集到總機率和達到「p」為止，然後再進行加權隨機取樣。
在這篇部落格文章中，我將探討為何我認為 nucleus sampling 可能不是最佳做法，並展示一個簡單易懂的替代方案，以避免 nucleus sampling 的問題。這個演算法仍在開發中，但現在發布是希望能激發一些討論／實作嘗試。
logits(對數勝算)中藏有黃金
儘管 LLM 的 logits 是 LLM 內部運作中少數幾個完全可以理解的部分，我通常卻很少看到有人有興趣去研究它們的特性、探索更先進的取樣方法、偵測並向使用者提示不確定性與可能的幻覺。將連續詞元的機率分布視覺化，是一個簡單而實用的練習，有助於獲得一些洞見：
在圖片中可以看到前 32 個候選詞元依機率著色（白色＝0，藍色＝1）、被選中的詞元，以及被選中詞元的排名（最高機率＝0，前一個＝1，依此類推）。
在上述範例中，Mistral 基礎模型知道 Umberto Eco（安伯托·艾可）的出生與逝世日期，因此它以佔據大部分總機率的方式，自信地標示出最有可能的詞元。其他時候模型會顯得較為困惑，原因可能是文字接續有多種表達方式，或是它對某些事實並不確定。詢問一個在訓練期間未曾學到其生日的姓名日期時，就會產生不同的詞元機率分布。
然而，一般而言我們想要避免的是選到次佳的詞元，使 LLM 的生成偏離最低 perplexity(困惑度)與幻覺最少的路徑。nucleus sampling 在這方面做得不好，因為累積詞元直到「p」為止的做法，取決於分布情況，可能會納入比首選弱得多的詞元。舉例來說，考慮上圖中「writer」這個詞元。安伯托·艾可確實是一位作家。與 writer 相關的詞元雖然得分不是非常高，但仍遠高於第二個選項。然而 p 值為 0.3 時，可能會把第二個僅有 0.01 這類低值的詞元也累積進來，並以個位數百分比的機率抽中它，從而有將生成帶往錯誤路徑的風險。
相反地，以接下來的詞元「who」為例，就有數個分數相對相近的選項，可用於產生不同的生成結果。
善用雪崩效應
這裡一個關鍵的觀察是，為了變化而挑選過弱的詞元並不划算：即使我們只利用存在少數幾個優質候選的那些生成時刻來使輸出多樣化，雪崩效應也會幫上忙：輸入的上下文會改變，LLM 的輸出因而受到擾動，效果是更有可能產生某種替代版本的文字。
First Token Cutoff(首詞元截斷)（FTC）演算法
免責聲明：此處描述的演算法尚未經過任何科學審查。我花了幾天時間試驗了幾種具有「最差詞元有界」特性的不同取樣演算法，這一個看起來在適用性、效果與可理解性之間取得了最佳平衡。
此處描述的演算法可以非正式地表述如下：
- 當 LLM 強烈偏向某個候選詞元時，就選它。
- 當存在多個可行的候選時，則產生替代選項。
- 對可能選到的最差詞元應限制在給定的範圍內。
過去的研究，例如 Tail Free Sampling(無尾取樣)，也指出應該在一小組由 LLM 產生的高品質詞元之間進行選擇。然而在 TFS 中，這樣的集合是透過計算導數來辨識：選出對應於未出現陡峭曲線、且在該曲線之後詞元品質會大幅下降的群集。
相對地，在此提出的演算法中，我們希望選擇遵循一個相對於最高分詞元 T0 更有界且更易理解的截斷，將 T0 賦予相較於所有其他詞元的特殊意義：即 LLM 在發出該詞元時的確定性程度（基本上是 perplexity 的代理指標）。因此，該演算法會拒絕所有相較於 T0 差於給定百分比的詞元。
截斷百分比的取值範圍為 0 到 1，稱為「co」。一個可行且具代表性的「co」例如可以是 0.5。
演算法的運作方式如下：
- 計算 logits 的 softmax()。
- 依機率對詞元進行排序。
- 給定 T0（最佳詞元的機率），計算所有其他詞元的比率為：r = 1 - (T[i] / T0)
- 僅選出符合 r <= co 的詞元
- 在被選中的詞元中進行加權隨機抽選。
請注意，透過這種方式，無論詞元的數值是否呈現平滑的單調遞減，我們對可納入可能性集合的詞元都有一個硬性限制。而對於其他嘗試辨識高分群集的方法，則並非如此。
實務範例
nucleus top-p sampling 在實務上看起來沒有災難性地失敗，其中一個原因是：很多時候第一個詞元的機率非常高，因此當 perplexity 較低時，加上很多時候我們碰巧不會收集並抽到低品質的詞元，生成就會沿著合理的路徑繼續。
當出現機率如下的連續詞元時，情況就會變得比較棘手：
0.25, 0.14, 0.01
在 p=0.4 的情況下，我們可能會把第三個低品質的詞元也收集進來，並以約 3% 的機率抽中它。
現在考慮使用 First Token Cutoff 且 co 值為 0.5 的情況（詞元最多可比第一個詞元差 50%）：
第二個詞元的 r 值為：
r[t1] = 1-(0.14/0.25) = 0.44 # 0.44 <= 0.5，此詞元被接受
r[t2] = 1-(0.01/0.25) = 0.96 # 0.96 > 0.5，此詞元被拒絕
輸出範例
Mistral 基礎模型（無指令微調）在提示「Sorted sets are」下使用 co=0.7 所產生的輸出。以下三個輸出皆為成功案例，並非為品質而精挑細選。
- 排序集合是 Redis 中一種強大的資料結構。它們可用於儲存已排序的清單、儲存唯一值、儲存用於排名的分數，以及儲存排序集合的已排序清單。
- 排序集合是一種非常強大的資料結構。它們讓你能以一種便於在集合中找出最高或最低值的方式來儲存資料，也讓你能對資料進行排序。這對於許多不同的任務都很有用，例如依分數為使用者排名，或在資料庫中找出最受歡迎的項目。
- 排序集合是一種非常強大的資料結構，可用於解決許多不同的問題。最常見的使用情境是儲存一份唯一元素的清單，其中每個元素都有一個關聯值。舉例來說，你可以使用排序集合來儲存家中所有人的姓名，並以他們的年齡作為關聯值。
為何可理解性很重要
取樣參數是 LLM 終端使用者或 API 使用者必須調整的少數項目之一。往往需要反覆嘗試，不過若能擁有一個具備即時、貼近現實描述與直覺的單一可調參數，將會有很大幫助。此外，「co」是一個線性參數，因此特別容易推論，相較於像 temperature 這類參數，或甚至是 top_p 的「p」——儘管同樣是線性的，卻強烈取決於 logits 分布的形狀。
未來工作
我才剛開始這方面的研究，因此會更深入地研究並評估這個演算法。比起其他，更重要的是，我非常希望看到大家對取樣演算法有更多關注，並對超越類似 top-p 的方法有更多興趣。
或許可以從 logits 分布中收集到一些有趣的資訊。舉例來說，linear probe(線性探針)很可能能夠學會判斷 LLM 的隱藏層何時正在處理某些事實性資訊。這一點加上詞元的 perplexity，或可用來向 LLM 的使用者顯示輸出的某些部分很可能是錯誤的。一般來說，將詞元的機率分布視覺化非常具有資訊量，並在某種程度上讓人能親手觸摸 LLM 的運作方式，以及每一步有哪些候選詞元。
參考實作
logits = mx.softmax(logits)
np_logits = np.array(logits) # MX -> NumPy
np_logits = np_logits.flatten()
sorted_indices = np.argsort(np_logits)
sorted_indices = sorted_indices[::-1]
co = 0.7
j = 1
t0 = np_logits[sorted_indices[0]]
while 1 - (np_logits[sorted_indices[j]] / t0) < co and j < len(np_logits):
j += 1
accepted_logits = []
for i in range(0,j):
accepted_logits.append(float(np_logits[sorted_indices[j]]))
accepted_logits = mx.array(accepted_logits)
idx = mx.random.categorical(accepted_logits)
idx = int(np.array(idx)) # Convert zero-dim array to scalar
token_id = sorted_indices[idx]致謝
這些實驗之所以能非常簡單地完成，要感謝來自 Apple 的 MLX 函式庫，以及那些持續不懈、既酷又聰明的開發者們。MLX 極其易於使用，正如它應該的那樣：畢竟 LLM 本身是難以捉摸的，但推論本身卻是一個簡單的過程。