在 Hugging Face Transformers 中使用受限束搜索(Constrained Beam Search)引導文本生成

Hugging Face 已在 transformers 函式庫中實作了受限束搜索,使用戶能夠透過強制包含特定單詞或短語來對模型輸出進行精確控制。這項功能對於神經機器翻譯等任務至關重要,例如在這些任務中,字典查詢可以規定強制性術語,或者需要根據上下文強制執行特定的正式程度(例如,德語中的正式與非正式)。

受限生成的挑戰

將特定的子序列強制加入文本生成任務並非易事,因為束搜索(beam search)是逐個 token 進行操作的。標準的解碼函數會根據當前的序列預測下一個 token,但本身沒有內在機制來判斷所需的 token 是否應該出現在當前步驟或未來的某個步驟。

此外,管理多個約束條件——例如同時要求兩個不同的短語,或允許模型從多個可能的選項中選擇一個短語——會為解碼過程增加顯著的複雜性。

受限束搜索的功能

受限束搜索允許用戶將先驗知識注入生成階段,而不是在生成後過濾輸出。該實作支持兩種主要的約束類型:

強制單詞約束

使用 model.generate() 中的 force_words_ids 關鍵字參數,用戶可以指定必須出現在輸出中的確切 token。例如,在「How old are you?」的英德翻譯中,強制使用單詞「Sie」可以確保模型生成正式的翻譯(「Wie alt sind Sie?」)而不是非正式的翻譯(「Wie alt bist du?」)。

析取約束(Disjunctive Constraints)

析取約束允許模型透過包含所提供列表中的至少一個單詞來滿足要求。這在多種單詞形式(例如,「raining」、「rained」、「rains」)皆可接受時非常有用。在混合約束的情景中,模型可以被強制逐字包含一個特定單詞,同時從相關術語的靈活列表中選擇一個選項。

技術實作:Banks 與輪詢選擇(Round-Robin Selection)

受限束搜索的工作原理是在生成的每一步注入所需的 token,以使序列更接近於滿足約束條件。然而,盲目地強制加入 token 通常會導致不合邏輯的結果(例如,「The is fast」)。為了解决這個問題,Hugging Face 利用了「Banks」系統來平衡約束的履行與語言學概率。

Bank 系統

束(Beams)根據其滿足約束條件的進度被分類到不同的 banks 中:

  • Bank $n$:包含在滿足約束方面已取得 $n$ 步進展的 beams。

輪詢選擇

為了防止模型僅選擇受限程度最高(但可能不合邏輯)的 beams,該算法會在不同的 banks 之間進行輪詢選擇。如果 num_beams=3,系統會先從 Bank 2 選擇最可能的輸出,然後是 Bank 1,最後是 Bank 0。這確保了即使在模型被推向約束條件時,高概率且合理的序列仍能被保留。

如果一個 beam 偏離了約束條件(例如,在約束為「is fast」時生成了「slow」),其進度將被重置,並移回 Bank 0。

透過 Constraint Class 擴展約束

雖然 force_words_ids 是主要的用戶介面,但後端利用了一個 Constraint 對象來追蹤進度並建議下一個 token。開發者可以透過繼承 Constraint 抽象接口類別來創建自定義約束。

提到的潛在未來擴展包括:

  • OrderedConstraints:確保約束按特定順序被滿足。
  • TemplateConstraints:強制模型遵循具有待填充空格的特定結構模板。

研究基礎

這項功能基於多篇專注於詞彙受限解碼和引導生成的研究論文,包括關於開放詞彙圖像描述(open vocabulary image captioning)和神經機器翻譯的工作。

Sources

相關