📂機械学習

論文レビュー: バダナウ・アテンション

概要と要約

• レファレンス、数式の番号、表記法などは可能な限り論文をそのまま従う。

論文『Neural Machine Translation by Jointly Learning to Align and Translate』¹はICLR 2015で発表され、著者はDzmitry Bahdanau、KyungHyun Cho、そしてYoshua Bengioである。トランスフォーマーなどに使われる現代的な🔒(26/08/07)アテンションのアイデアを初めて導入した論文として挙げられ、ここで提案された形のアテンションは第一著者の名前にちなんで、一般にバダナウ・アテンション^{Bahdanau attention}あるいはadditive attentionと呼ばれる。

当時ニューラル機械翻訳^{neural machine translation}の主流であったエンコーダ・デコーダ構造、すなわちseq2seq系列のモデルは、原文の文全体を固定された次元の一つのコンテキストベクトルに圧縮したのち、デコーダがそこから翻訳文を生成する。著者らはこの固定長ベクトルが性能向上のボトルネックであると推測する。

"In this paper, we conjecture that the use of a fixed-length vector is a bottleneck in improving the performance of this basic encoder–decoder architecture"

論文ではこれを解決するために、デコーダが翻訳文の単語を一つ生成するたびに、原文の中から関連のある部分を自ら探索できる方法を提案する。原文の$j$番目の単語の周辺を要約したベクトルを$h_{j}$とするとき、$i$番目の目標単語のためのコンテキストベクトル$c_{i}$は一つに固定されず、毎時点で次のように加重和として新たに計算される。

$$ c_{i} = \sum_{j=1}^{T_{x}} \alpha_{ij} h_{j}, \qquad \alpha_{ij} = \frac{\exp (e_{ij})}{\sum_{k=1}^{T_{x}} \exp (e_{ik})}, \qquad e_{ij} = a (s_{i-1}, h_{j}) $$

重み$\alpha_{ij}$を計算するアライメントモデル$a$は小さなニューラルネットワークであり、翻訳モデルの残りの部分とともに学習される。アライメント^alignと翻訳^translateを同時に学習するというタイトルはここから来ている。この拡張だけで英仏翻訳において既存のエンコーダ・デコーダを大きく上回り、当時最高水準であったフレーズベース^phrase-based統計翻訳システムに匹敵する性能を得たと報告している。

最近の表記法

最近使われるアテンションの表記法で当該論文を要約すると次のようになる。🔒(26/08/07)アテンションとは、クエリ行列$\mathbf{Q}$、キー行列$\mathbf{K}$、バリュー行列$\mathbf{V}$とスコア関数$f$に対して次のように定義される演算である。

$$ \operatorname{Attention} (\mathbf{Q}, \mathbf{K}, \mathbf{V}) := \mathbf{V} \operatorname{Softmax} \left( f (\mathbf{Q}, \mathbf{K}) \right) $$

バダナウ・アテンションでは、デコーダが$i$番目の単語を生成する時点のクエリはデコーダの直前の隠れ状態$s_{i-1}$一つだけである。キーとバリューは互いに区別されず、どちらもエンコーダの隠れ状態(アノテーション)である。

$$ \mathbf{q}_{i} = s_{i-1}, \qquad \mathbf{K} = \mathbf{V} = \begin{bmatrix} h_{1} & h_{2} & \cdots & h_{T_{x}} \end{bmatrix} $$

スコア関数$f$は、クエリとキーのそれぞれを線形変換して足し合わせたのち$\operatorname{Tanh}$を取る小さなニューラルネットワーク(=アライメントモデル$a$)である。

$$ \left[ f (\mathbf{q}_{i}, \mathbf{K}) \right]_{j} = a (s_{i-1}, h_{j}) = \mathbf{w}_{a}^{\mathsf{T}} \operatorname{Tanh} \left( W_{a} s_{i-1} + U_{a} h_{j} \right) $$

ここで$W_{a}, U_{a}, \mathbf{w}_{a}$はアライメントモデルの学習パラメータであり、スコアベクトル$\mathbf{w}_{a}$はバリュー$\mathbf{v}_{j}$とは無関係な別個のベクトルである。$\operatorname{Tanh}$は引数として入ってきたベクトルの各成分に双曲線正接$\tanh$を取る関数である(アライメントモデル$a$の詳しい定義は、論文の表記をそのまま従う§A.1.2にある)。

この時点のアテンションの出力がまさにコンテキストベクトル$c_{i}$である。

$$ \mathbf{c}_{i} = \operatorname{Attention} (\mathbf{q}_{i}, \mathbf{K}, \mathbf{V}) = \mathbf{V} \operatorname{Softmax} \left( f (\mathbf{q}_{i}, \mathbf{K}) \right) = \sum_{j=1}^{T_{x}} \alpha_{ij} h_{j} $$

論文の表記と最近の表記の対応は次のとおりである。

論文の表記	アテンションの表記	意味
$s_{i-1}$	クエリ$\mathbf{q}_{i}$	デコーダの直前の隠れ状態
$h_{j}$	キー$\mathbf{k}_{j}$でありバリュー$\mathbf{v}_{j}$	エンコーダの$j$番目のアノテーション
$e_{ij} = a (s_{i-1}, h_{j})$	スコア$\left[ f (\mathbf{q}_{i}, \mathbf{K}) \right]_{j}$	クエリとキーの適合度
$\alpha_{ij}$	アテンション重み	スコアのソフトマックス
$c_{i}$	アテンションの出力	バリューの加重平均、コンテキストベクトル

まとめると、バダナウ・アテンションとは、(i)キーとバリューを区別せずエンコーダのアノテーションが両方の役割を兼ね、(ii)クエリがデコーディングのステップごとに一つずつ順次作られ、(iii)スコア関数が加算であるアテンションである。

1 Introduction

論文が出た当時提案されていたニューラル翻訳モデルの大部分はエンコーダ・デコーダ系列、すなわちseq2seqモデルに属する。著者らは、原文の文のすべての情報を一つの固定長ベクトルに圧縮しなければならない点がこの構造の潜在的な問題だと指摘する。この問題を解決するために、アライメントと翻訳を同時に学習するエンコーダ・デコーダの拡張を導入する。ここで言うアライメントとは、翻訳される単語が原文のどこを参照すればよいかを選択する関数のことであり、アテンションではこれをキーベクトルで計算する。翻訳とは、原文言語の単語と翻訳文言語の単語をつなぐ関数のことであり、アテンションではこれをバリューベクトルで計算する。提案されたモデルは、翻訳文の単語を一つ生成するたびに、原文の中で最も関連のある情報が集まっている位置の集合を**(ソフト)探索**^{(soft-)search}し、これらの位置に関連するコンテキストベクトルとそれまでに生成された目標単語に基づいて次の目標単語を予測する。

基本のエンコーダ・デコーダと区別される最も重要な特徴は、入力文全体を一つの固定長ベクトルにエンコードしようとしない点である。代わりに入力文をベクトルの列としてエンコードしておき、翻訳をデコードする間、そのベクトルの部分集合を適応的に^adaptively選択する。実験では、アライメントと翻訳をともに学習する提案方式が基本のエンコーダ・デコーダを有意に上回り、その改善は長い文でより顕著だが文の長さにかかわらず観察されると報告している。

2 Background: Neural Machine Translation

確率的な観点から、翻訳は与えられた原文$\mathbf{x}$に対する条件付き確率$p(\mathbf{y} | \mathbf{x})$を最大化する翻訳文$\mathbf{y}$を見つけることと同じである。すなわち次の最適化問題を解くことと同じである。

$$ \hat{\mathbf{y}} = \argmax\limits_{\mathbf{y}} p(\mathbf{y} | \mathbf{x}) $$

ニューラルネットワークに基づく機械翻訳の方式は、一般にエンコーダとデコーダの二つの部分で構成されている。エンコーダは入力された原文$\mathbf{x}$を適切な特徴ベクトル$\mathbf{c}$にエンコードし、デコーダはこれを翻訳文$\mathbf{y}$にデコードする。

$$ \mathbf{x} \overset{\text{encoding}}{\mapsto} \mathbf{c} \overset{\text{decoding}}{\mapsto} \mathbf{y} $$

2.1 RNN Encoder–Decoder

Cho et al.(2014)²とSutskever et al.(2014)³が提案したフレームワークであるRNNエンコーダ・デコーダ(seq2seqモデル)を簡単に整理する。エンコーダは入力文、すなわちベクトルの列$\mathbf{x} = (x_{1}, \cdots, x_{T_{x}})$をベクトル$c$にエンコードする。最も一般的な方式はリカレントニューラルネットワーク(RNN)を使うことである。

$$ h_{t} = f (x_{t}, h_{t-1}) \tag{1} $$

$$ c = q \left( \left\{ h_{1}, \cdots, h_{T_{x}} \right\} \right) $$

ここで$h_{t} \in \mathbb{R}^{n}$は時点$t$のエンコーダの隠れ状態^{hidden state}であり、$c$は隠れ状態の列から生成されるベクトルであり、$f$と$q$はある非線形関数である。例えばSutskever et al. (2014)は$f$としてLSTMを使い、$q \left( \left\{ h_{1}, \cdots, h_{T} \right\} \right) = h_{T}$とした。デコーダは、時点$t^{\prime}$においてコンテキストベクトル$c$とそれまでに予測したすべての単語$\left\{ y_{1}, \cdots, y_{t^{\prime}-1} \right\}$が与えられたとき、次の単語$y_{t^{\prime}}$を予測するように学習される。すなわち翻訳文$\mathbf{y} = \left( y_{1}, \cdots, y_{T_{y}} \right)$に対する結合確率を、順序付けられた条件付き確率の積に分解する。

$$ p(\mathbf{y}) = \prod_{t=1}^{T} p \left( y_{t} \mid \left\{ y_{1}, \cdots, y_{t-1} \right\}, c \right) \tag{2} $$

RNNを使うと、各条件付き確率は次のようにある非線形関数$g$でモデル化される。

$$ p \left( y_{t} \mid \left\{ y_{1}, \cdots, y_{t-1} \right\}, c \right) = g (y_{t-1}, s_{t}, c) \tag{3} $$

ここで$g$は$y_{t}$の確率を出力する非線形関数であり、$s_{t}$はデコーダの隠れ状態である。既存の方法では原文$\mathbf{x}$が固定された次元のコンテキストベクトル$c$にエンコードされており、本論文ではこれが性能向上を妨げるボトルネックだと主張する。

3 Learning to Align and Translate

提案する新しいアーキテクチャは、エンコーダとして双方向RNN(§3.2)を、デコーダとしては翻訳をデコードする間に原文を探索することを模倣する構造(§3.1)を使う。

3.1 Decoder: General Description

Figure 1

新しいアーキテクチャでは、$(2)$の各条件付き確率を次のように定義する。

$$ p(y_{i} \mid y_{1}, \dots, y_{i-1}, \mathbf{x}) = g (y_{i-1}, s_{i}, c_{i}) \tag{4} $$

$s_{i}$は時点$i$のデコーダの隠れ状態で、$s_{i} = f (s_{i-1}, y_{i-1}, c_{i})$として計算される。既存のエンコーダ・デコーダではすべての$y_{i}$の予測過程で固定されたコンテキストベクトル$c$を使うのとは異なり、確率が目標単語$y_{i}$ごとに互いに異なるコンテキストベクトル$c_{i}$を使う点に注目せよ。

コンテキストベクトル$c_{i}$は、エンコーダが入力文をエンコードしたアノテーション^annotationの列$(h_{1}, \cdots, h_{T_{x}})$に依存する。($h_{j}$はエンコーダであるBRNNの二つの隠れ状態を結合したものである)各アノテーション$h_{j}$は、入力シーケンス全体の情報を含みつつ$j$番目の単語の周辺に強く集中^attentionしたベクトルであるが、どのように計算されるかは次節で扱う。コンテキストベクトルはアノテーションの線形結合として計算される。

$$ c_{i} = \sum_{j=1}^{T_{x}} \alpha_{ij} h_{j} \tag{5} $$

各アノテーション$h_{j}$の重み$\alpha_{ij}$はソフトマックスの形で計算される。

$$ \alpha_{ij} = \frac{\exp (e_{ij})}{\sum_{k=1}^{T_{x}} \exp (e_{ik})} \tag{6} $$

ここで$e_{ij} = a (s_{i-1}, h_{j})$は、位置$j$の周辺の入力と位置$i$の出力がどれほどよく合うかをスコア化するアライメントモデル^{alignment model}である。スコアは、$y_{i}$を出す直前のデコーダの隠れ状態$s_{i-1}$と入力文の$j$番目のアノテーション$h_{j}$に基づく。アライメントモデル$a$は順伝播型ニューラルネットワーク^{feedforward neural network}としてパラメータ化され、システムの他のすべての構成要素とともに学習される。伝統的な機械翻訳とは異なり、アライメントをニューラルネットワークとすることでアライメントモデルと翻訳モデル全体を同時に学習できるというのが論文の説明である。

すべてのアノテーションの加重和は、可能なアライメントに対する期待アノテーション^{expected annotation}として理解できるという。$\alpha_{ij}$を目標単語$y_{i}$が原文の単語$x_{j}$にアライン(そこから翻訳)される確率と見ると、$c_{i}$はその確率に対するアノテーションの期待値である。そして論文のすぐ次の箇所でattentionという単語が言及される。

"Intuitively, this implements a mechanism of attention in the decoder. The decoder decides parts of the source sentence to pay attention to. By letting the decoder have an attention mechanism, we relieve the encoder from the burden of having to encode all information in the source sentence into a fixed-length vector."

直観的に言えば、これはデコーダにアテンション機構を実装したものである。デコーダは原文の文のどの部分に注意を向けるかを決定する。デコーダにアテンション機構を与えることで、原文の文のすべての情報を一つの固定長ベクトルにエンコードしなければならない負担をエンコーダから軽減できる。

3.2 Encoder: Bidirectional RNN for Annotating Sequences

$(1)$の普通のRNNは入力を最初の単語$x_{1}$から最後の$x_{T_{x}}$まで順に読むが、§3.1の方式では時点$i$の出力$y_{i}$が入力ベクトルの列全体を参照している。これが可能な理由は、エンコーダの構造として双方向リカレントニューラルネットワーク^{bidirectional RNN, BRNN}を使用したためである。(論文ではBiRNNと表記されている)BRNNでは、出力$y_{i}$が時点にかかわらずすべての入力ベクトル$\left\{ x_{1}, \dots, x_{T_{x}} \right\}$を参照する。

BRNNは順方向RNN $\overrightarrow{f}$と逆方向RNN $\overleftarrow{f}$で構成される。順方向RNNはデータを順序どおり($x_{1}$から$x_{T_{x}}$まで)入力として受け取り順方向の隠れ状態$\left\{ \overrightarrow{h}_{1}, \cdots, \overrightarrow{h}_{T_{x}} \right\}$を計算し、逆方向RNNはデータを逆の順序で受け取り逆方向の隠れ状態$\left\{ \overleftarrow{h}_{1}, \cdots, \overleftarrow{h}_{T_{x}} \right\}$を計算する。各単語$x_{j}$のアノテーションは、この二つを連結して^concatenate得る。

$$ h_{j} = \begin{bmatrix} \overrightarrow{h}_{j} \\ \overleftarrow{h}_{j} \end{bmatrix} $$

こうすると、アノテーション$h_{j}$は前の単語たちと後の単語たちの要約を両方含む。RNNの隠れ状態は近いデータをよりよく表現する傾向があるため、$h_{j}$は自然に$x_{j}$の周辺の単語に集中するようになるという。

A Model Architecture

実験結果に関する内容は省略し、実験に使われたモデルRNNsearchのアーキテクチャの数式を詳細に記述する。

A.2.1 Encoder

モデルはワンホットベクトルの列である原文を入力として受け取り、同じくワンホットベクトルの列である翻訳文を出力する。

$$ \mathbf{x} = (x_{1}, \dots, x_{T_{x}}), \quad x_{i} \in \mathbb{R}^{K_{x}}, \qquad \mathbf{y} = (y_{1}, \dots, y_{T_{y}}), \quad y_{i} \in \mathbb{R}^{K_{y}} $$

$K_{x}$と$K_{y}$はそれぞれ原文言語と翻訳文言語の語彙サイズであり、$T_{x}$と$T_{y}$は原文と翻訳文の長さである。エンコーダである双方向リカレントニューラルネットワークの順方向の隠れ状態は次のように計算される。

$$ \overrightarrow{h}_{i} = \begin{cases} (1 - \overrightarrow{z}_{i}) \odot \overrightarrow{h}_{i-1} + \overrightarrow{z}_{i} \odot \overrightarrow{\tilde{h}}_{i} & \text{if } i > 0 \\ 0 & \text{if } i = 0 \end{cases} $$

ここで$\odot$はベクトルの成分ごとの積を意味するアダマール積である。論文では関数の合成を意味する$\circ$で表記されているので注意が必要である。各項は次のとおりである。

$$ \begin{align*} \overrightarrow{\tilde{h}}_{i} &= \tanh \left( \overrightarrow{W} \overline{E} x_{i} + \overrightarrow{U} \left[ \overrightarrow{r}_{i} \odot \overrightarrow{h}_{i-1} \right] \right) \\ \overrightarrow{z}_{i} &= \sigma \left( \overrightarrow{W}_{z} \overline{E} x_{i} + \overrightarrow{U}_{z} \overrightarrow{h}_{i-1} \right) \\ \overrightarrow{r}_{i} &= \sigma \left( \overrightarrow{W}_{r} \overline{E} x_{i} + \overrightarrow{U}_{r} \overrightarrow{h}_{i-1} \right) \end{align*} $$

$\overline{E} \in \mathbb{R}^{m \times K_{x}}$は単語の埋め込み行列であり、$\overrightarrow{W}, \overrightarrow{W}_{z}, \overrightarrow{W}_{r} \in \mathbb{R}^{n \times m}$と$\overrightarrow{U}, \overrightarrow{U}_{z}, \overrightarrow{U}_{r} \in \mathbb{R}^{n \times n}$は重み行列であり、$m$は埋め込みの次元、$n$は隠れユニットの個数である。$x_{i}$が1-of-$K$ベクトルであるため、$\overline{E} x_{i}$は埋め込み行列から該当単語の列を抜き出す演算である。

逆方向の隠れ状態$\overleftarrow{h}_{1}, \cdots, \overleftarrow{h}_{T_{x}}$も、入力を逆の順序で読みながら同じ方式で計算される。このとき埋め込み行列$\overline{E}$は順方向と逆方向のRNNが共有するが、重み行列は共有しない。二つの隠れ状態を連結すると§3.2で述べたアノテーションになる。

$$ h_{i} = \begin{bmatrix} \overrightarrow{h}_{i} \\ \overleftarrow{h}_{i} \end{bmatrix} \in \mathbb{R}^{2n} $$

A.2.2 Decoder

デコーダの隠れ状態$s_{i}$は、エンコーダが出したアノテーションから次のように計算される。

$$ s_{i} = (1 - z_{i}) \odot s_{i-1} + z_{i} \odot \tilde{s}_{i} $$

$$ \begin{align*} \tilde{s}_{i} &= \tanh \left( W E y_{i-1} + U \left[ r_{i} \odot s_{i-1} \right] + C c_{i} \right) \\ z_{i} &= \sigma \left( W_{z} E y_{i-1} + U_{z} s_{i-1} + C_{z} c_{i} \right) \\ r_{i} &= \sigma \left( W_{r} E y_{i-1} + U_{r} s_{i-1} + C_{r} c_{i} \right) \end{align*} $$

$E$は目標言語の単語埋め込み行列であり、$W, W_{z}, W_{r} \in \mathbb{R}^{n \times m}$、$U, U_{z}, U_{r} \in \mathbb{R}^{n \times n}$、$C, C_{z}, C_{r} \in \mathbb{R}^{n \times 2n}$は重み行列である。

さて、デコーダから翻訳文がどのように出力されるかを最初のステップから見よう。初期の隠れ状態は、逆方向RNNの最後の隠れ状態、すなわち原文を末尾から読んできた要約から次のように計算される。$W_{s} \in \mathbb{R}^{n \times n}$である。

$$ s_{0} = \tanh \left( W_{s} \overleftarrow{h}_{1} \right) $$

コンテキストベクトル$c_{1}$は以下のように計算される。

$$ c_{1} = \sum_{j=1}^{T_{x}} \alpha_{1j} h_{j}, \qquad \alpha_{1j} = \frac{\exp (e_{1j})}{\sum_{k=1}^{T_{x}} \exp (e_{1k})}, \qquad e_{1j} = v_{a}^{\mathsf{T}} \tanh \left( W_{a} s_{0} + U_{a} h_{j} \right) $$

初期値$y_{0}$が文の始まりを知らせるあるトークンとして与えられているとしよう。デコーダの隠れ状態$s_{1}$は以下のように計算される。

$$ \begin{align*} z_{1} &= \sigma \left( W_{z} E y_{0} + U_{z} s_{0} + C_{z} c_{1} \right) \\ r_{1} &= \sigma \left( W_{r} E y_{0} + U_{r} s_{0} + C_{r} c_{1} \right) \\ \tilde{s}_{1} &= \tanh \left( W E y_{0} + U \left[ r_{1} \odot s_{0} \right] + C c_{1} \right) \\ s_{1} &= (1 - z_{1}) \odot s_{0} + z_{1} \odot \tilde{s}_{1} \end{align*} $$

デコーダの状態$s_{1}$とコンテキストベクトル$c_{1}$、直前に生成された単語$y_{0}$が与えられたとき、次の単語の確率ベクトル(確率分布)は次のように定義される。

$$ p (\cdot \mid s_{1}, y_{0}, c_{1}) = \operatorname{softmax} \left( W_{o} t_{1} \right) \in \mathbb{R}^{K_{y}} $$

$$ t_{1} = \begin{bmatrix} \max \left\{ \tilde{t}_{1, 1}, \tilde{t}_{1, 2} \right\} & \cdots & \max \left\{ \tilde{t}_{1, 2l-1}, \tilde{t}_{1, 2l} \right\} \end{bmatrix}^{\mathsf{T}} \in \mathbb{R}^{l} $$ $$ \tilde{t}_{1} = U_{o} s_{0} + V_{o} E y_{0} + C_{o} c_{1} $$

ここで$\tilde{t}_{i,k}$は$\tilde{t}_{i}$ベクトルの$k$番目の成分である。$W_{o} \in \mathbb{R}^{K_{y} \times l}$、$U_{o} \in \mathbb{R}^{2l \times n}$、$V_{o} \in \mathbb{R}^{2l \times m}$、$C_{o} \in \mathbb{R}^{2l \times 2n}$は重み行列である。ベクトル$\tilde{t}_{i} \in \mathbb{R}^{2l}$の成分を二つずつ組にして大きい方だけ残すのが隠れ層一つのマックスアウト層であり、その出力$t_{i} \in \mathbb{R}^{l}$に$W_{o}$を掛けて語彙サイズ$K_{y}$次元のスコアを作ったのちソフトマックスで正規化すると確率になる。ここで値が最も大きいインデックスの単語を選択すると、それがまさに$y_{1}$になる。

$$ y_{1} = \operatorname*{argmax}_{1 \le w \le K_{y}} \left[ p (\cdot \mid s_{1}, y_{0}, c_{1}) \right]_{w} $$

以後この過程を繰り返すと翻訳文$\left\{ y_{1}, \dots, y_{T_{y}} \right\}$を得る。

A.1.2 Alignment Model

アライメントモデルは、長さがそれぞれ$T_{x}$、$T_{y}$の文のペアごとに$T_{x} \times T_{y}$回計算されなければならない点を考慮して設計しなければならない。論文では計算量を減らすために、隠れ層が一つの単層パーセプトロンを使う。

$$ a (s_{i-1}, h_{j}) = v_{a}^{\mathsf{T}} \tanh \left( W_{a} s_{i-1} + U_{a} h_{j} \right) $$

ここで$W_{a} \in \mathbb{R}^{n \times n}$、$U_{a} \in \mathbb{R}^{n \times 2n}$、$v_{a} \in \mathbb{R}^{n}$は重み行列である。$U_{a} h_{j}$は$i$に依存しないので、文ごとに一度だけ事前に計算しておけば計算コストを最小化できる。$\tanh$の中でクエリ$s_{i-1}$の線形変換とキー$h_{j}$の線形変換が足し合わされるこの形が、先ほどバダナウ・アテンションがadditiveアテンションとも呼ばれると述べた理由である。