技術文章 » 在Transformers 中使用約束波束搜尋引導文字生成

在Transformers 中使用約束波束搜尋引導文字生成

2023-06-11 12:00:09

引言

本文假設讀者已經熟悉文字生成領域波束搜尋相關的背景知識，具體可參見博文 如何生成文字: 通過 Transformers 用不同的解碼方法生成文字。

與普通的波束搜尋不同，約束 波束搜尋允許我們控制所生成的文字。這很有用，因為有時我們確切地知道輸出中需要包含什麼。例如，在機器翻譯任務中，我們可能通過查字典已經知道哪些詞必須包含在最終的譯文中; 而在某些特定的場合中，雖然某幾個詞對於語言模型而言差不多，但對終端使用者而言可能卻相差很大。這兩種情況都可以通過允許使用者告訴模型最終輸出中必須包含哪些詞來解決。

這事兒為什麼這麼難

然而，這個事情操作起來並不容易，它要求我們在生成過程中的 某個時刻 在輸出文字的 某個位置 強制生成某些特定子序列。

假設我們要生成一個句子 S，它必須按照先 \(t_1\) 再 \(t_2\) 的順序包含短語 \(p_1={ t_1, t_2 }\)。以下定義了我們希望生成的句子 \(S\):

\[S_{期望} = { s_1, s_2, …, s_k, t_1, t_2, s_{k+1}, …, s_n } \]

問題是波束搜尋是逐詞輸出文字的。我們可以大致將波束搜尋視為函數 \(B(\mathbf{s}_{0:i}) = s_{i+1}\)，它根據當前生成的序列 \(\mathbf{s}_{0:i}\) 預測下一時刻 \(i+1\) 的輸出。但是這個函數在任意時刻 \(i < k\) 怎麼知道，未來的某個時刻 \(k\) 必須生成某個指定詞？或者當它在時刻 \(i=k\) 時，它如何確定當前那個指定詞的最佳位置，而不是未來的某一時刻 \(i>k\)？

如果你同時有多個不同的約束怎麼辦？如果你想同時指定使用短語 \(p_1={t_1, t_2}\) 和 短語 \(p_2={ t_3, t_4, t_5, t_6}\) 怎麼辦？如果你希望模型在兩個短語之間 任選一個 怎麼辦？如果你想同時指定使用短語 \(p_1\) 以及短語列表 \({p_{21}, p_{22}, p_{23}}\) 中的任一短語怎麼辦？

上述需求在實際場景中是很合理的需求，下文介紹的新的約束波束搜尋功能可以滿足所有這些需求！

我們會先簡要介紹一下新的 約束波束搜尋 可以做些什麼，然後再深入介紹其原理。

例 1: 指定包含某詞

假設我們要將 "How old are you?" 翻譯成德語。它對應兩種德語表達，其中 "Wie alt bist du?" 是非正式場合的表達，而 "Wie alt sind Sie?" 是正式場合的表達。

不同的場合，我們可能傾向於不同的表達，但我們如何告訴模型呢？

使用傳統波束搜尋

我們先看下如何使用 傳統波束搜尋 來完成翻譯。

!pip install -q git+https://github.com/huggingface/transformers.git

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

tokenizer = AutoTokenizer.from_pretrained("t5-base")
model = AutoModelForSeq2SeqLM.from_pretrained("t5-base")

encoder_input_str = "translate English to German: How old are you?"

input_ids = tokenizer(encoder_input_str, return_tensors="pt").input_ids

outputs = model.generate(
    input_ids,
    num_beams=10,
    num_return_sequences=1,
    no_repeat_ngram_size=1,
    remove_invalid_values=True,
)

print("Output:\n" + 100 *'-')
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Output:
----------------------------------------------------------------------------------------------------
Wie alt bist du?

使用約束波束搜尋

但是如果我們想要一個正式的表達而不是非正式的表達呢？如果我們已經先驗地知道輸出中必須包含什麼，我們該如何 將其 注入到輸出中呢？

我們可以通過 model.generate() 的 force_words_ids 引數來實現這一功能，程式碼如下:

tokenizer = AutoTokenizer.from_pretrained("t5-base")
model = AutoModelForSeq2SeqLM.from_pretrained("t5-base")

encoder_input_str = "translate English to German: How old are you?"

force_words = ["Sie"]

input_ids = tokenizer(encoder_input_str, return_tensors="pt").input_ids
force_words_ids = tokenizer(force_words, add_special_tokens=False).input_ids

outputs = model.generate(
    input_ids,
    force_words_ids=force_words_ids,
    num_beams=5,
    num_return_sequences=1,
    no_repeat_ngram_size=1,
    remove_invalid_values=True,
)

print("Output:\n" + 100 *'-')
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Output:
----------------------------------------------------------------------------------------------------
Wie alt sind Sie?

如你所見，現在我們能用我們對輸出的先驗知識來指導文字的生成。以前我們必須先生成一堆候選輸出，然後手動從中挑選出符合我們要求的輸出。現在我們可以直接在生成階段做到這一點。

例 2: 析取式約束

在上面的例子中，我們知道需要在最終輸出中包含哪些單詞。這方面的一個例子可能是在神經機器翻譯過程中結合使用字典。

但是，如果我們不知道要使用哪種 _詞形_呢，我們可能希望使用單詞 rain 但對其不同的詞性沒有偏好，即 ["raining", "rained", "rains", ...] 是等概的。更一般地，很多情況下，我們可能並不刻板地希望 逐字母一致 ，此時我們希望劃定一個範圍由模型去從中選擇最合適的。

支援這種行為的約束叫 析取式約束 (Disjunctive Constraints) ，其允許使用者輸入一個單詞列表來引導文字生成，最終輸出中僅須包含該列表中的 至少一個 詞即可。

下面是一個混合使用上述兩類約束的例子:

from transformers import GPT2LMHeadModel, GPT2Tokenizer

model = GPT2LMHeadModel.from_pretrained("gpt2")
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")

force_word = "scared"
force_flexible = ["scream", "screams", "screaming", "screamed"]

force_words_ids = [
    tokenizer([force_word], add_prefix_space=True, add_special_tokens=False).input_ids,
    tokenizer(force_flexible, add_prefix_space=True, add_special_tokens=False).input_ids,
]

starting_text = ["The soldiers", "The child"]

input_ids = tokenizer(starting_text, return_tensors="pt").input_ids

outputs = model.generate(
    input_ids,
    force_words_ids=force_words_ids,
    num_beams=10,
    num_return_sequences=1,
    no_repeat_ngram_size=1,
    remove_invalid_values=True,
)

print("Output:\n" + 100 *'-')
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
print(tokenizer.decode(outputs[1], skip_special_tokens=True))


Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.

Output:
----------------------------------------------------------------------------------------------------
The soldiers, who were all scared and screaming at each other as they tried to get out of the
The child was taken to a local hospital where she screamed and scared for her life, police said.

如你所見，第一個輸出裡有 "screaming" ，第二個輸出裡有 "screamed" ，同時它們都原原本本地包含了 "scared" 。注意，其實 ["screaming", "screamed", ...] 列表中不必一定是同一單詞的不同詞形，它可以是任何單詞。使用這種方式，可以滿足我們只需要從候選單詞列表中選擇一個單詞的應用場景。

傳統波束搜尋

以下是傳統 波束搜尋 的一個例子，摘自之前的 博文:

與貪心搜尋不同，波束搜尋會保留更多的候選詞。上圖中，我們每一步都展示了 3 個最可能的預測詞。

在 num_beams=3 時，我們可以將第 1 步波束搜尋表示成下圖:

波束搜尋不像貪心搜尋那樣只選擇 "The dog" ，而是允許將 "The nice" 和 "The car" 留待進一步考慮 。

下一步，我們會為上一步建立的三個分支分別預測可能的下一個詞。

雖然我們 考查 了明顯多於 num_beams 個候選詞，但在每步結束時，我們只會輸出 num_beams 個最終候選詞。我們不能一直分叉，那樣的話， beams 的數目將在 \(n\) 步後變成 \(\text{beams}^{n}\) 個，最終變成指數級的增長 (當波束數為 \(10\) 時，在 \(10\) 步之後就會變成 \(10,000,000,000\) 個分支！)。

接著，我們重複上述步驟，直到滿足中止條件，如生成 <eos> 標記或達到 max_length 。整個過程可以總結為: 分叉、排序、剪枝，如此往復。

約束波束搜尋

約束波束搜尋試圖通過在每一步生成過程中 _注入_所需詞來滿足約束。

假設我們試圖指定輸出中須包含短語 "is fast" 。

在傳統波束搜尋中，我們在每個分支中找到 k 個概率最高的候選詞，以供下一步使用。在約束波束搜尋中，除了執行與傳統波束搜尋相同的操作外，我們還會試著把約束詞加進去，以 看看我們是否能儘量滿足約束。圖示如下:

上圖中，我們最終候選詞除了包括像 "dog" 和 "nice" 這樣的高概率詞之外，我們還把 "is" 塞了進去，以儘量滿足生成的句子中須含 "is fast" 的約束。

第二步，每個分支的候選詞選擇與傳統的波束搜尋大部分類似。唯一的不同是，與上面第一步一樣，約束波束搜尋會在每個新分叉上繼續強加約束，把滿足約束的候選詞強加進來，如下圖所示:

組 (Banks)

在討論下一步之前，我們停下來思考一下上述方法的缺陷。

在輸出中野蠻地強制插入約束短語 is fast 的問題在於，大多數情況下，你最終會得到像上面的 The is fast 這樣的無意義輸出。我們需要解決這個問題。你可以從 huggingface/transformers 程式碼庫中的這個 問題 中瞭解更多有關這個問題及其複雜性的深入討論。

組方法通過在滿足約束和產生合理輸出兩者之間取得平衡來解決這個問題。

我們把所有候選波束按照其 滿足了多少步約束分到不同的組中，其中組 \(n\) 裡包含的是 滿足了 \(n\) 步約束的波束列表 。然後我們按照順序輪流選擇各組的候選波束。在上圖中，我們先從組 2 (Bank 2) 中選擇概率最大的輸出，然後從組 1 (Bank 1) 中選擇概率最大的輸出，最後從組 0 (Bank 0) 中選擇最大的輸出; 接著我們從組 2 (Bank 2) 中選擇概率次大的輸出，從組 1 (Bank 1) 中選擇概率次大的輸出，依此類推。因為我們使用的是 num_beams=3，所以我們只需執行上述過程三次，就可以得到 ["The is fast", "The dog is", "The dog and"]。

這樣，即使我們 強制 模型考慮我們手動新增的約束詞分支，我們依然會跟蹤其他可能更有意義的高概率序列。儘管 The is fast 完全滿足約束，但這並不是一個有意義的短語。幸運的是，我們有 "The dog is" 和 "The dog and" 可以在未來的步驟中使用，希望在將來這會產生更有意義的輸出。

圖示如下 (以上例的第 3 步為例):

請注意，上圖中不需要強制新增 "The is fast"，因為它已經被包含在概率排序中了。另外，請注意像 "The dog is slow" 或 "The dog is mad" 這樣的波束實際上是屬於組 0 (Bank 0) 的，為什麼呢？因為儘管它包含詞 "is" ，但它不可用於生成 "is fast" ，因為 fast 的位子已經被 slow 或 mad 佔掉了，也就杜絕了後續能生成 "is fast" 的可能性。從另一個角度講，因為 slow 這樣的詞的加入，該分支 滿足約束的進度 被重置成了 0。

最後請注意，我們最終生成了包含約束短語的合理輸出: "The dog is fast" ！

起初我們很擔心，因為盲目地新增約束詞會導致出現諸如 "The is fast" 之類的無意義短語。然而，使用基於組的輪流選擇方法，我們最終隱式地擺脫了無意義的輸出，優先選擇了更合理的輸出。

關於 Constraint 類的更多資訊及自定義約束

我們總結下要點。每一步，我們都不斷地糾纏模型，強制新增約束詞，同時也跟蹤不滿足約束的分支，直到最終生成包含所需短語的合理的高概率序列。

在實現時，我們的主要方法是將每個約束表示為一個 Constraint 物件，其目的是跟蹤滿足約束的進度並告訴波束搜尋接下來要生成哪些詞。儘管我們可以使用 model.generate() 的關鍵字引數 force_words_ids ，但使用該引數時後端實際發生的情況如下:

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, PhrasalConstraint

tokenizer = AutoTokenizer.from_pretrained("t5-base")
model = AutoModelForSeq2SeqLM.from_pretrained("t5-base")

encoder_input_str = "translate English to German: How old are you?"

constraints = [
    PhrasalConstraint(
        tokenizer("Sie", add_special_tokens=False).input_ids
    )
]

input_ids = tokenizer(encoder_input_str, return_tensors="pt").input_ids

outputs = model.generate(
    input_ids,
    constraints=constraints,
    num_beams=10,
    num_return_sequences=1,
    no_repeat_ngram_size=1,
    remove_invalid_values=True,
)

print("Output:\n" + 100 *'-')
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Output:
----------------------------------------------------------------------------------------------------
Wie alt sind Sie?

你甚至可以定義一個自己的約束並將其通過 constraints 引數輸入給 model.generate() 。此時，你只需要建立 Constraint 抽象介面類的子類並遵循其要求即可。你可以在 此處 的 Constraint 定義中找到更多資訊。

我們還可以嘗試其他一些有意思的約束 (尚未實現，也許你可以試一試！) 如 OrderedConstraints 、 TemplateConstraints 等。目前，在最終輸出中約束短語間是無序的。例如，前面的例子一個輸出中的約束短語順序為 scared -> screaming ，而另一個輸出中的約束短語順序為 screamed -> scared 。 如果有了 OrderedConstraints， 我們就可以允許使用者指定約束短語的順序。 TemplateConstraints 的功能更小眾，其約束可以像這樣:

starting_text = "The woman"
template = ["the", "", "School of", "", "in"]

possible_outputs == [
   "The woman attended the Ross School of Business in Michigan.",
   "The woman was the administrator for the Harvard School of Business in MA."
]

或是這樣:

starting_text = "The woman"
template = ["the", "", "", "University", "", "in"]

possible_outputs == [
   "The woman attended the Carnegie Mellon University in Pittsburgh.",
]
impossible_outputs == [
  "The woman attended the Harvard University in MA."
]

或者，如果使用者不關心兩個詞之間應該隔多少個詞，那僅用 OrderedConstraint 就可以了。

總結

約束波束搜尋為我們提供了一種將外部知識和需求注入文字生成過程的靈活方法。以前，沒有一個簡單的方法可用於告訴模型 1. 輸出中需要包含某列表中的詞或短語，其中 2. 其中有一些是可選的，有些必須包含的，這樣 3. 它們可以最終生成至在合理的位置。現在，我們可以通過綜合使用 Constraint 的不同子類來完全控制我們的生成！

該新特性主要基於以下論文:

與上述這些工作一樣，還有許多新的研究正在探索如何使用外部知識 (例如 KG (Knowledge Graph) 、KB (Knowledge Base) ) 來指導大型深度學習模型輸出。我們希望約束波束搜尋功能成為實現此目的的有效方法之一。

感謝所有為此功能提供指導的人: Patrick von Platen 參與了從 初始問題 討論到 最終 PR 的全過程，還有 Narsil Patry，他們二位對程式碼進行了詳細的反饋。

本文使用的圖示來自於 Freepik - Flaticon。

英文原文: https://hf.co/blog/constrained-beam-search

原文作者: Chan Woo Kim

譯者: Matrix Yao (姚偉峰)，英特爾深度學習工程師，工作方向為 transformer-family 模型在各模態資料上的應用及大規模模型的訓練推理。

審校/排版: zhongdongy (阿東)