什麼是NLP解碼中的Beam Search？

束搜尋（Beam Search）是一種強大的解碼演算法，廣泛應用於自然語言處理（NLP）和機器學習。它在序列生成任務中尤其重要，例如文本生成、機器翻譯和摘要。束搜尋在有效探索搜尋空間和生成高品質輸出之間取得平衡。在這篇文章中，我們將深入探討束搜尋的運作方式、它在解碼中的重要性，以及實作過程，同時探索其在現實世界中的應用和挑戰。

學習目標

了解束搜尋演算法在序列解碼任務中的概念和運作機制。

學習束寬（beam width）的重要性，以及它如何在搜尋空間中平衡探索和效率。

探索使用 Python 實作束搜尋的實際步驟。

分析束搜尋在 NLP 任務中的現實應用和挑戰。

了解束搜尋相對於其他解碼演算法（如貪婪搜尋）的優勢。

什麼是束搜尋？

束搜尋是一種啟發式搜尋演算法，用於從模型（如變壓器（Transformers）、長短期記憶網絡（LSTM）和其他序列到序列架構）中解碼序列。它通過在每一步保持固定數量（“束寬”）的最可能序列來生成文本。與貪婪搜尋僅選擇最可能的下一個標記不同，束搜尋同時考慮多個假設。這確保了最終的序列不僅流暢，而且在模型信心方面是全局最優的。

例如，在機器翻譯中，可能有多種有效的方式來翻譯一句話。束搜尋允許模型同時探索這些可能性，並跟蹤多個候選翻譯。

束搜尋是如何運作的？

束搜尋通過探索一個圖形來運作，其中節點代表標記，邊緣代表從一個標記轉換到另一個標記的概率。在每一步：

演算法根據模型的輸出邏輯（概率分佈）選擇前 k 個最可能的標記。

擴展這些標記為序列，計算它們的累積概率，並保留前 k 個序列以供下一步使用。

這個過程持續進行，直到滿足停止條件，例如達到特殊的序列結束標記或預定的長度。

束寬的概念

“束寬”決定了在每一步保留多少候選序列。較大的束寬允許探索更多序列，但會增加計算成本。相反，較小的束寬速度較快，但可能因探索有限而錯過更好的序列。

為什麼束搜尋在解碼中如此重要？

束搜尋在解碼中至關重要，原因有幾個：

改善序列質量：通過探索多個假設，束搜尋確保生成的序列是全局最優的，而不是陷入局部最優。

處理歧義：許多 NLP 任務涉及歧義，例如多種有效的翻譯或解釋。束搜尋幫助探索這些可能性並選擇最佳選擇。

效率：與全面搜尋相比，束搜尋在計算上是高效的，同時仍然探索了搜尋空間的相當一部分。

靈活性：束搜尋可以適應各種任務和取樣策略，使其成為序列解碼的多功能選擇。

束搜尋的實際實作

以下是一個束搜尋實作的實際示例。該演算法構建一個搜尋樹，評估累積分數，並選擇最佳序列：

步驟 1：安裝和導入依賴項



# 安裝 transformers 和 graphviz

!sudo apt-get install graphviz graphviz-dev

!pip install transformers pygraphviz



from transformers import GPT2LMHeadModel, GPT2Tokenizer

import torch

import matplotlib.pyplot as plt

import networkx as nx

import numpy as np

from matplotlib.colors import LinearSegmentedColormap

from tqdm import tqdm

import matplotlib.colors as mcolors

系統命令：安裝生成圖形所需的庫（graphviz）和 Python 包（transformers 和 pygraphviz）。

導入的庫：

transformers：用於加載 GPT-2 進行文本生成。

torch：用於處理張量並在模型上運行計算。

matplotlib.pyplot：用於繪製束搜尋圖形。

networkx：用於構建和管理表示束搜尋路徑的樹狀圖。

tqdm：在處理圖形時顯示進度條。

numpy 和 matplotlib.colors：用於處理數據和可視化中的顏色映射。

步驟 2：模型和標記器設置



# 加載模型和標記器

device="cuda" if torch.cuda.is_available() else 'cpu'

model = GPT2LMHeadModel.from_pretrained('gpt2').to(device)

tokenizer = GPT2Tokenizer.from_pretrained('gpt2')

model.eval()

檢測是否可用 GPU（cuda），因為它可以加速計算。如果未找到 GPU，則默認為 cpu。

從 Hugging Face 的 transformers 庫加載預訓練的 GPT-2 語言模型及其標記器。

將模型移動到適當的設備（cuda 或 cpu）。

使用 model.eval() 將模型設置為評估模式，以禁用如 dropout 等僅在訓練期間需要的功能。

步驟 3：編碼輸入文本



# 輸入文本

text = "我有一個夢想"

input_ids = tokenizer.encode(text, return_tensors="pt").to(device)

定義輸入文本“我有一個夢想”。

使用標記器將文本編碼為標記 ID，返回一個張量（return_tensors=’pt’）。

將輸入張量移動到適當的設備（cuda 或 cpu）。

步驟 4：定義輔助函數：對數概率



def get_log_prob(logits, token_id):

    probabilities = torch.nn.functional.softmax(logits, dim=-1)

    log_probabilities = torch.log(probabilities)

    return log_probabilities[token_id].item()

應用 softmax 函數將邏輯轉換為概率（詞彙的分佈）。

對這些概率取自然對數以獲得對數概率。

返回與給定標記對應的對數概率。

步驟 5：定義遞歸束搜尋

實現使用 GPT-2 模型的遞歸束搜尋以生成文本。



def beam_search(input_ids, node, bar, length, beams, temperature=1.0):

    if length == 0:

        return



    outputs = model(input_ids)

    predictions = outputs.logits



    # 獲取下一個標記的邏輯

    logits = predictions[0, -1, :]

    top_token_ids = torch.topk(logits, beams).indices



    for j, token_id in enumerate(top_token_ids):

        bar.update(1)



        # 計算預測標記的分數

        token_score = get_log_prob(logits, token_id)

        cumulative_score = graph.nodes[node]['cumscore'] + token_score



        # 將預測標記添加到輸入 ID 列表中

        new_input_ids = torch.cat([input_ids, token_id.unsqueeze(0).unsqueeze(0)], dim=-1)



        # 向圖中添加節點和邊

        token = tokenizer.decode(token_id, skip_special_tokens=True)

        current_node = list(graph.successors(node))[j]

        graph.nodes[current_node]['tokenscore'] = np.exp(token_score) * 100

        graph.nodes[current_node]['cumscore'] = cumulative_score

        graph.nodes[current_node]['sequencescore'] = cumulative_score / len(new_input_ids.squeeze())

        graph.nodes[current_node]['token'] = token + f"_{length}_{j}"



        # 遞歸調用

        beam_search(new_input_ids, current_node, bar, length - 1, beams, temperature)

基本情況：當長度達到 0 時停止遞歸（不再預測標記）。

模型預測：將 input_ids 傳遞給 GPT-2 以獲取下一個標記的邏輯。

選擇最佳束：使用 torch.topk() 選擇最可能的標記。

標記計分：評估標記概率以確定最佳序列。

擴展輸入：將選定的標記附加到 input_ids 以進行進一步探索。

更新圖形：通過擴展搜尋樹來跟踪進度。

遞歸調用：對每個束（束分支）重複此過程。

步驟 6：檢索最佳序列

根據累積分數找到在束搜尋過程中生成的最佳序列。



def get_best_sequence(G):

    # 找到所有葉子節點

    leaf_nodes = [node for node in G.nodes if G.out_degree(node) == 0]



    # 根據序列分數找到最佳葉子節點

    max_score_node = max(leaf_nodes, key=lambda n: G.nodes[n]['sequencescore'])

    max_score = G.nodes[max_score_node]['sequencescore']



    # 檢索從根到此節點的路徑

    path = nx.shortest_path(G, source=0, target=max_score_node)



    # 構建序列

    sequence = "".join([G.nodes[node]['token'].split('_')[0] for node in path])

    return sequence, max_score

識別所有葉子節點（沒有輸出邊的節點）。

找到最佳葉子節點（最高的序列分數）。

檢索從根節點（開始）到最佳節點的路徑。

提取並連接沿此路徑的標記以形成最終序列。

步驟 7：繪製束搜尋圖形

可視化樹狀束搜尋圖形。



def plot_graph(graph, length, beams, score):

    fig, ax = plt.subplots(figsize=(3 + 1.2 * beams**length, max(5, 2 + length)), dpi=300, facecolor="white")



    # 為每個節點創建位置

    pos = nx.nx_agraph.graphviz_layout(graph, prog="dot")



    # 根據標記分數範圍標準化顏色

    scores = [data['tokenscore'] for _, data in graph.nodes(data=True) if data['token'] is not None]

    vmin, vmax = min(scores), max(scores)

    norm = mcolors.Normalize(vmin=vmin, vmax=vmax)

    cmap = LinearSegmentedColormap.from_list('rg', ["r", "y", "g"], N=256)



    # 繪製節點

    nx.draw_networkx_nodes(graph, pos, node_size=2000, node_shape="o", alpha=1, linewidths=4,

                           node_color=scores, cmap=cmap)



    # 繪製邊

    nx.draw_networkx_edges(graph, pos)



    # 繪製標籤

    labels = {node: data['token'].split('_')[0] + f"\n{data['tokenscore']:.2f}%" \

              for node, data in graph.nodes(data=True) if data['token'] is not None}

    nx.draw_networkx_labels(graph, pos, labels=labels, font_size=10)

    plt.box(False)



    # 添加顏色條

    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)

    sm.set_array([])

    fig.colorbar(sm, ax=ax, orientation='vertical', pad=0, label="標記概率 (%)")

    plt.show()

節點代表每一步生成的標記，根據其概率進行顏色編碼。

邊連接節點，根據標記如何擴展序列。

顏色條表示標記概率的範圍。

步驟 8：主要執行



# 參數

length = 5

beams = 2



# 創建平衡樹圖

graph = nx.balanced_tree(beams, length, create_using=nx.DiGraph())

bar = tqdm(total=len(graph.nodes))



# 初始化圖形屬性

for node in graph.nodes:

    graph.nodes[node]['tokenscore'] = 100

    graph.nodes[node]['cumscore'] = 0

    graph.nodes[node]['sequencescore'] = 0

    graph.nodes[node]['token'] = text



# 執行束搜尋

beam_search(input_ids, 0, bar, length, beams)



# 獲取最佳序列

sequence, max_score = get_best_sequence(graph)

print(f"生成的文本：{sequence}")



# 繪製圖形

plot_graph(graph, length, beams, 'token')