この記事では，Pytorchで作ったseq2seq型の翻訳モデルを使って，ビームサーチによるデコーディングをします． OpenNMTやfairseqを使えば簡単に利用できるのですが，ビームサーチのためだけにこのようなフレームワークを使うのはちょっとなぁ，ということと，ビームサーチ自体はDNNに限らず様々な場面で役に立つ手法なので，この際ピュアに実装してみた，というのがこの記事です．
ちなみに一般的なseq2seqのデコードは，各タイムステップで予測したtop-1の単語を，次ステップのデコーダーの入力に使います． ビームサーチでは，このようなgreedyな条件を緩め，上位K個の予測を使って，デコードしていきます．ビームサーチをよく知らんという方は，Andrew Ngの神説明が参考になると思います． C5W3L08 Attention Model, Andrew Ng.
できたもの
- seq2seq (rnn) w/wo attentionの翻訳器のビームサーチによるデコーダー
バッチを分解して１文ずつデコードするbeam_search_decodingと（できるだけ）バッチ処理するbatch_beam_search_decodingがあります．出力は同じなので基本的に高速な後者を使うべきですが，理解のため最初の関数を書きました．
※今回始めてビームサーチを書いたのでまだまだイケてないところや変な箇所が多いと思うのでこのリポジトリは修正していく予定です．この記事では，コミット4cb1187760f3a9e7をベースに記事を書きます．
※バグや誤っているところ等，少なからずあると思うのでバグレポートしていただけると嬉しいです．
動かし方はGithubのREADMEを見てください．ビームサーチは推論時（テスト時）に動作するようになっています． github.com
解説
ここからは興味ある方へ，実装したビームサーチ部分を解説します．もっと良い方法あるかなぁと思いつつ，今回はこの方法で書きました．
beam_search_decoding (１文ずつデコード)
ビームサーチの実装は，下記の方の実装を参考にしました．OpenNMTの実装などもあるのですが，まずはシンプルな方法で実装している方を参考にしました．
この実装は，１文ずつデコードしているため実装がシンプルです（遅いけど）．また各ステップでデコードするときに，どの出力をデコードするか，を選ぶ必要があるのですが，この実装ではヒープ（heapq）を利用しています．
各デコードの結果はBeamSearchNodeという構造体にしているのですが，その状態に至ったときのスコア（対数確率の和）を負にしたものとあわせてヒープに登録しておけば，このヒープから取り出されるノードは，現状で最もスコアが高いノードになります．なのでデコード時は，このヒープからpopしていくだけで良いので高速（O(1)）で楽ですね．一方でどんどんヒープにノードが追加されるのに対して，枝刈りによるメモリ解放は行わないため，メモリ効率は良くないです．
ちなみに上位K個（K: ビームサイズ）の結果の取得は，pytorchであればtorch.topk関数で簡単に，値とその引数を取得できるので便利です．
def beam_search_decoding(decoder, enc_outs, enc_last_h, beam_width, n_best, sos_token, eos_token, max_dec_steps, device): """Beam Seach Decoding for RNN Args: decoder: An RNN decoder model enc_outs: A sequence of encoded input. (T, bs, 2H). 2H for bidirectional enc_last_h: (bs, H) beam_width: Beam search width n_best: The number of output sequences for each input Returns: n_best_list: Decoded N-best results. (bs, T) """ assert beam_width >= n_best n_best_list = [] bs = enc_outs.shape[1] # Decoding goes sentence by sentence. # So this process is very slow compared to batch decoding process. for batch_id in range(bs): # Get last encoder hidden state decoder_hidden = enc_last_h[batch_id] # (H) enc_out = enc_outs[:, batch_id].unsqueeze(1) # (T, 1, 2H) # Prepare first token for decoder decoder_input = torch.tensor([sos_token]).long().to(device) # (1) # Number of sentence to generate end_nodes = [] # starting node node = BeamSearchNode(h=decoder_hidden, prev_node=None, wid=decoder_input, logp=0, length=1) # whole beam search node graph nodes = [] # Start the queue heappush(nodes, (-node.eval(), id(node), node)) n_dec_steps = 0 # Start beam search while True: # Give up when decoding takes too long if n_dec_steps > max_dec_steps: break # Fetch the best node score, _, n = heappop(nodes) decoder_input = n.wid decoder_hidden = n.h if n.wid.item() == eos_token and n.prev_node is not None: end_nodes.append((score, id(n), n)) # If we reached maximum # of sentences required if len(end_nodes) >= n_best: break else: continue # Decode for one step using decoder decoder_output, decoder_hidden = decoder(decoder_input, decoder_hidden.unsqueeze(0), enc_out) # Get top-k from this decoded result topk_log_prob, topk_indexes = torch.topk(decoder_output, beam_width) # (1, bw), (1, bw) # Then, register new top-k nodes for new_k in range(beam_width): decoded_t = topk_indexes[0][new_k].view(1) # (1) logp = topk_log_prob[0][new_k].item() # float log probability val node = BeamSearchNode(h=decoder_hidden.squeeze(0), prev_node=n, wid=decoded_t, logp=n.logp+logp, length=n.length+1) heappush(nodes, (-node.eval(), id(node), node)) n_dec_steps += beam_width # if there are no end_nodes, retrieve best nodes (they are probably truncated) if len(end_nodes) == 0: end_nodes = [heappop(nodes) for _ in range(beam_width)] # Construct sequences from end_nodes n_best_seq_list = [] for score, _id, n in sorted(end_nodes, key=lambda x: x[0]): sequence = [n.wid.item()] # back trace from end node while n.prev_node is not None: n = n.prev_node sequence.append(n.wid.item()) sequence = sequence[::-1] # reverse n_best_seq_list.append(sequence) n_best_list.append(n_best_seq_list) return n_best_list
batch_beam_search_decoding (バッチデコード)
次に上記コードをGPUの恩恵を受けれるように，バッチ状態のまま解けるようにしたいと思います．ただここの方法は自明でなく，正直良くわからなかったです． というのも今回の実装では，各文でのビームサーチによる探索グラフ（ヒープ）は異なってきますし，探索終了タイミングもそれぞれ異なってきます． そのため私のコードでは，バッチ化できたのはRNN デコーダーに実際に投げるところです．このRNNデコーダーに投げるデータを作るため，バッチサイズ分のforループを回し，バッチデータを作っているのでこの部分は結局遅いです．また，デコードステップを回していく上で，探索を終了した文が発生しますが，データのshapeは維持されていたほうが使いやすいため，探索が終了した事例もひたすら再利用されているのが残念なところ．
def batch_beam_search_decoding(decoder, enc_outs, enc_last_h, beam_width, n_best, sos_token, eos_token, max_dec_steps, device): """Batch Beam Seach Decoding for RNN Args: decoder: An RNN decoder model enc_outs: A sequence of encoded input. (T, bs, 2H). 2H for bidirectional enc_last_h: (bs, H) beam_width: Beam search width n_best: The number of output sequences for each input Returns: n_best_list: Decoded N-best results. (bs, T) """ assert beam_width >= n_best n_best_list = [] bs = enc_last_h.shape[0] # Get last encoder hidden state decoder_hidden = enc_last_h # (bs, H) # Prepare first token for decoder decoder_input = torch.tensor([sos_token]).repeat(1, bs).long().to(device) # (1, bs) # Number of sentence to generate end_nodes_list = [[] for _ in range(bs)] # whole beam search node graph nodes = [[] for _ in range(bs)] # Start the queue for bid in range(bs): # starting node node = BeamSearchNode(h=decoder_hidden[bid], prev_node=None, wid=decoder_input[:, bid], logp=0, length=1) heappush(nodes[bid], (-node.eval(), id(node), node)) # Start beam search fin_nodes = set() history = [None for _ in range(bs)] n_dec_steps_list = [0 for _ in range(bs)] while len(fin_nodes) < bs: # Fetch the best node decoder_input, decoder_hidden = [], [] for bid in range(bs): if bid not in fin_nodes and n_dec_steps_list[bid] > max_dec_steps: fin_nodes.add(bid) if bid in fin_nodes: score, n = history[bid] # dummy for data consistency else: score, _, n = heappop(nodes[bid]) if n.wid.item() == eos_token and n.prev_node is not None: end_nodes_list[bid].append((score, id(n), n)) # If we reached maximum # of sentences required if len(end_nodes_list[bid]) >= n_best: fin_nodes.add(bid) history[bid] = (score, n) decoder_input.append(n.wid) decoder_hidden.append(n.h) decoder_input = torch.cat(decoder_input).to(device) # (bs) decoder_hidden = torch.stack(decoder_hidden, 0).to(device) # (bs, H) # Decode for one step using decoder decoder_output, decoder_hidden = decoder(decoder_input, decoder_hidden, enc_outs) # (bs, V), (bs, H) # Get top-k from this decoded result topk_log_prob, topk_indexes = torch.topk(decoder_output, beam_width) # (bs, bw), (bs, bw) # Then, register new top-k nodes for bid in range(bs): if bid in fin_nodes: continue score, n = history[bid] if n.wid.item() == eos_token and n.prev_node is not None: continue for new_k in range(beam_width): decoded_t = topk_indexes[bid][new_k].view(1) # (1) logp = topk_log_prob[bid][new_k].item() # float log probability val node = BeamSearchNode(h=decoder_hidden[bid], prev_node=n, wid=decoded_t, logp=n.logp+logp, length=n.length+1) heappush(nodes[bid], (-node.eval(), id(node), node)) n_dec_steps_list[bid] += beam_width # Construct sequences from end_nodes # if there are no end_nodes, retrieve best nodes (they are probably truncated) for bid in range(bs): if len(end_nodes_list[bid]) == 0: end_nodes_list[bid] = [heappop(nodes[bid]) for _ in range(beam_width)] n_best_seq_list = [] for score, _id, n in sorted(end_nodes_list[bid], key=lambda x: x[0]): sequence = [n.wid.item()] while n.prev_node is not None: n = n.prev_node sequence.append(n.wid.item()) sequence = sequence[::-1] # reverse n_best_seq_list.append(sequence) n_best_list.append(copy.copy(n_best_seq_list)) return n_best_list
速度比較
バッチ版のビームサーチは実装がかなり汚いですが，実行速度はどの程度恩恵があるか調べてみました．バッチ数は１２８でビーム幅は１０です． 非バッチ版の実装と比べて，翻訳結果は変わらないが，速度としては２，３倍早くなっているのがわかります（バッチサイズとビームサイズで変わりますが）． それでもまだ結構遅いなという印象ですが，まぁまぁ最初はこんなものでしょうと自分に言い訳します．
% python run.py --attention --skip_train --model_path ./ckpts/s2s-attn.pt Number of training examples: 29000 Number of validation examples: 1014 Number of testing examples: 1000 Unique tokens in source (de) vocabulary: 7855 Unique tokens in target (en) vocabulary: 5893 In: <SOS> . schnee den über laufen hunde mittelgroße zwei <EOS> for loop beam search time: 8.718 Out: Rank-1: <SOS> two medium brown dogs run across the snow . the snow . <EOS> Out: Rank-2: <SOS> two medium brown dogs run across the snow . <EOS> Out: Rank-3: <SOS> two medium brown dogs run across the snow . the snow . . <EOS> Out: Rank-4: <SOS> two medium brown dogs run across the snow . . <EOS> Out: Rank-5: <SOS> two medium brown dogs run across the snow . snow . <EOS> Batch beam search time: 2.994 Out: Rank-1: <SOS> two medium brown dogs run across the snow . the snow . <EOS> Out: Rank-2: <SOS> two medium brown dogs run across the snow . <EOS> Out: Rank-3: <SOS> two medium brown dogs run across the snow . the snow . . <EOS> Out: Rank-4: <SOS> two medium brown dogs run across the snow . . <EOS> Out: Rank-5: <SOS> two medium brown dogs run across the snow . snow . <EOS> In: <SOS> . <unk> mit tüten gehsteig einem auf verkauft frau eine <EOS> for loop beam search time: 9.654 Out: Rank-1: <SOS> a woman is selling on her Out: Rank-2: <SOS> a woman woman selling a Out: Rank-3: <SOS> a woman is her selling Out: Rank-4: <SOS> a woman is selling vegetables on a sidewalk Out: Rank-5: <SOS> a woman woman selling rice Out: Rank-6: <SOS> a woman is selling her on Out: Rank-7: <SOS> a woman is selling watermelon on a Out: Rank-8: <SOS> a woman is selling on the Out: Rank-9: <SOS> a woman is sells selling Out: Rank-10: <SOS> a woman selling selling on Batch beam search time: 3.256 Out: Rank-1: <SOS> a woman is selling on her Out: Rank-2: <SOS> a woman woman selling a Out: Rank-3: <SOS> a woman is her selling Out: Rank-4: <SOS> a woman is selling vegetables on a sidewalk Out: Rank-5: <SOS> a woman woman selling rice Out: Rank-6: <SOS> a woman is selling her on Out: Rank-7: <SOS> a woman is selling watermelon on a Out: Rank-8: <SOS> a woman is selling on the Out: Rank-9: <SOS> a woman is sells selling Out: Rank-10: <SOS> a woman selling selling on
まとめ
今回ビームサーチを生実装してみました．ヒープによる実装で楽をしましたが，何となく動作イメージは掴めてきました．スコアの正規化や効率的なバッチデコーディングにより改善点はまだまだありそうですが，今回はこんなところで．詳しい方はアドバイス求む！
