使用Transformers做限制集束搜尋（ConstrainedBeamSearch）的文字生成

MLNLP

(

機器學習演算法與自然語言處理

)社群是國內外知名自然語言處理社群，受眾覆蓋國內外NLP碩博生、高校老師以及企業研究人員。

社群的願景是促進國內外自然語言處理，機器學習學術界、產業界和廣大愛好者之間的交流，特別是初學者同學們的進步。

來源 | https://huggingface.co/blog/constrained-beam-search

翻譯 | 金屬成色

『簡介』

開始前，我們需要先熟悉beam search技術，詳見How to generate text: using different decoding methods for language generation with Transformers，或中文翻譯版

不像一般的beam search，constrained beam search可以對生成文字施加控制，因為很多時候，我們是明確知道生成文字之中是應該包含哪些內容的。例如在神經網路機器翻譯任務中，透過查詞典，我們明確知道生成文字應該包含的專業詞彙。有時，生成的結果是滿足語言模型的要求，但只是因為未包含部分關鍵資訊，就不能夠滿足使用者的需求。這些場景都需要讓使用者告訴模型，哪些詞彙必須被包含在生成的內容之中。

『難點在哪』

然而，做到這個並不簡單。因為這個任務需要我們在最終生成結果中的某個位置強制生成子內容。比如，我們需要生成一個句子，必須包含內容，和是按順序排列的。我們定義想要的句子如下：

難點在於beam search是在token粒度生成句子，我們可以把beam search簡化為一個函式（雖然不完全準確），也就是說，輸入從位置至的token，預測位置的token。但是這個函式是如何知道，在位置$i<k$時，強制生成的token是要等到將來生成；在位置$i=k$時，強制生成的token是要在當前生成而不是等到$i>k$時呢。</k$時，強制生成的token是要等到將來生成；在位置$i=k$時，強制生成的token是要在當前生成而不是等到$i>

如果需要多個不同的限制呢？如果想要強制包含子內容同時也要包含子內容呢？如果想要模型在這兩個子內容之間做選擇呢？如果想要強制生成子內容同時需要強制從中選出一個子內容生成呢？

以上例子是非常合理的使用者需求，下面我們透過constrained beam search來實現它們。

在這篇部落格中，我們首先快速介紹constrained beam search能做什麼，然後深入介紹實現的底層細節。

『例子一：強制詞（Forcing a Word）』

當我們翻譯How old are you?到德文時，非正式場景可以說Wie alt bist du?，正式場景可以說Wie alt sind Sie?。這需要依賴上下文，與上下文保持一致，我們如何告訴模型這樣做呢。

傳統集束搜尋（Traditional Beam Search）

下面是如何使用傳統的beam search來做文字翻譯。


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?

限制集束搜尋（Constrained Beam Search）

但是如果我們想要正式表達而不是非正式呢，如果我們已經有先驗知識，知道生成結果必須包含的子內容呢。下面是我們透過配置強制生成詞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?

就像我們看到的，我們可以透過我們想要生成的先驗知識，指導模型的生成結果。以前我們經常的做法是透過生成很多輸出結果，然後篩選符合我們要求的。下面我們會在生成階段來實現它。

『例子二：分離約束（Disjunctive Constraints）』

上面我們介紹的場景是我們明確知道必須要生成的token，例如神經網路機器翻譯任務。

但是如果我們不知道token的語法形式呢，比如輸出["raining", "rained", "rains",...]中的任何一個都可以呢。更進一步，我們經常不想要精確到一個字母都不差的詞來作為強制輸出子內容。

允許這樣配置約束的方法叫做分離約束(Disjunctive Constraints)。使用者可以輸入一組token，只要最終生成結果包含它們之中至少一個就可以。

下面是包含上面介紹的兩種限制的例子：

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。

『傳統集束搜尋（Traditional Beam Search）』

下面是一個傳統集束搜尋的例子，來自之前的部落格，或中文版

與貪心搜尋（greedy search）不同，集束搜尋保留了更長的假設序列。在上面的圖片中，我們在文字生成的每一步，展示了三個下一個token的可以取值。

num_beams=3的集束搜尋的第一步的另一種展示方式如下：

貪心搜尋會直接選擇生成The dog，而集束搜尋會允許進一步考慮The nice和The car。

在集束搜尋的下一步，我們可以考慮剛才的三條分支每個分支接下來可能的詞。

相比較貪心搜尋，雖然我們要計算更多的輸出可能性分支，但是在每步結束時仍會降至num_beams。我們不能配置num_beams太大，因為對於n步的生成，我們要計算個分支。隨著num_beams變大，分支數會快速變大，例如num_beams=10計算10步，就意味著10,000,000,000個分支。

在生成過程中，我們重複以上過程直到遇到結束條件，例如生成了<eos>，或者生成token數到達上線。在計算的每一步都會經歷，列出所有生成分支、排序、減少分支至num_beams、重複計算。

『限制集束搜尋（Constrained Beam Search）』

限制集束搜尋嘗試在生成的每一步加入想要的token。

例如我們在生成結果中，強制生成短語"is fast"。

在傳統集束搜尋中，我們會在每一個分支的下一個token計算中，選取top k個機率最高的下一個token，然後把它們都加入至考慮範圍。在限制集束搜尋中，我們仍然會這樣做，不過我們也會加入我們的強制生成token。

根據模型計算下一個詞最高的機率是dog和nice，同時我們也把強制生成tokenis也放入考慮分支中，從而儘可能生成我們想要的短語is fast。

在下一步中，去除候選分支的策略與傳統集束搜尋相似，但和第一步一樣，限制集束搜尋會在每個分支上，再次加入強制生成token。

Banks

在討論下一步前，我們需要考慮上面步驟中不想要的結果。在強制生成想要的短語is fast時，大多時候，我們得到的是不符合邏輯的輸出，例如The is fast。這實際上是一個較複雜問題。在huggingface/transformers的request issue中有深入討論這個問題的複雜性。

Banks透過平衡強制限制生成和模型機率生成，解決上面的問題。

Bank 表示在分支中已經有了個token滿足了強制限制。在透過機率排序所有分支後，我們做輪序排程選擇（round-robin selection）。在以上的例子中，我們從Bank 2中選擇機率最高的一個輸出分支；然後從Bank 1中選擇機率最高的一個輸出分支；從Bank 0中選擇機率最高的一個輸出分支。接下來從Bank 2中選擇機率第二高的分支、Bank 1中機率第二高的分支、Bank 0中機率第二高的分支，以此類推。因為我們配置的num_beams=3，我們只保留三個分支，所以留下了["The is fast", "The dog is", "The dog and"]，分別對應機率最高的Bank 2、Bank 1、 Bank 0。

透過這種方法，雖然我們強制模型考慮我們想要加入的token的分支，但是我們仍然保持了高機率的輸出序列分支，從而使得輸出結果更有意義。儘管"The is fast"完全滿足我們的強制限制，但它是不符合常識的短語。幸運的是，我們還有"The dog is"、"The dog and"分支可以在後面的步驟中繼續計算，它們很有希望會輸出更符合常識的結果，進而在BANK 2的排序中替換掉"The is fast"。

步驟三的例子如下圖所示：

注意，"The is fast"分支的下一個token預測，不再需要加入強制限制token了，因為強制限制token已經完全滿足了。同時注意分支如"The dog is slow"或"The dog is mad"，它們雖然包含了限制詞"is"，但是在"is"後面加入了"slow"。因此只能重新開始生成"is fast"，所以它們從Bank 1回到了Bank 0。

最終，在Bank 2上我們得到了"The dog is fast"，即滿足了強制限制的短語，又滿足較高的輸出機率，即符合常識。

我們剛才擔心的，為了強制限制token導致不符合常識語義的短語"The is fast"已經在輪序排程選擇（round-robin selection）中被排除掉了，因為它只在Bank 2中排到最後一名，如上圖所示。

『關於Constraint Classes、Custom Constraints的更多資訊』

主要流程簡要如下。在每一步，我們要求模型考慮滿足限制token的分支，同時也要考慮那些不滿足限制，有著高機率的分支，直到找到即高機率又滿足限制短語的分支。

儘管在model.generate()函式中我們有了force_words_ids來控制強制生成，但我們可以做一個更好的實施設計。我們把每個限制設計成一個限制物件，它們在集束搜尋過程中，分別記錄下一個限制生成的token，如下所示：

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?

你可以自己定義一個限制類，放入限制列表中，從而設計自己的定製限制。只需要自己建立滿足介面需要的限制子類即可。可以從這裡獲得更多關於限制的定義資訊。

關於定製限制，有一些獨特的想法（還未實現，也許你可以試試），例如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."
]

或者使用者不關心兩個token之間有幾個token，只是使用OrderedConstraint。

『總結』

限制集束搜尋（Constrained beam search）可以讓使用者優雅的引入外部知識做文字生成。早先，是很難控制模型依照這樣的限制規則的。

限制生成結果必須包含短語
一些短語是有可選列表，一些是不可選的
短語生成在指定的位置的

現在我們透過多個限制物件（Constraint objects）的子類可以充分控制文字生成效果。這些新功能來自一些論文：

Guided Open Vocabulary Image Captioning with Constrained Beam Search
Fast Lexically Constrained Decoding with Dynamic Beam Allocation for Neural Machine Translation
Improved Lexically Constrained Decoding for Translation and Monolingual Rewriting
Guided Generation of Cause and Effect

就像剛才介紹的，有很多研究論文正在探索引用外部知識（如知識庫、知識圖譜）輔助指導深度神經網路大模型輸出結果，而限制集束搜尋會是另一種有效的實現目標的方法。

技術交流群邀請函

△長按新增小助手

掃描二維碼新增小助手微信

請備註：姓名-學校/公司-研究方向

（如：小張-哈工大-對話系統）

即可申請加入自然語言處理/Pytorch等技術交流群

關於我們

MLNLP社群 ( 機器學習演算法與自然語言處理 ) 是由國內外自然語言處理學者聯合構建的民間學術社群，目前已經發展為國內外知名自然語言處理社群，旗下包括 萬人頂會交流群、AI臻選匯、AI英才匯 以及 AI學術匯 等知名品牌，旨在促進機器學習，自然語言處理學術界、產業界和廣大愛好者之間的進步。

社群可以為相關從業者的深造、就業及研究等方面提供開放交流平臺。歡迎大家關注和加入我們。