トランスフォーマーの論文を読む⑩自己
前回は、位置エンコーディングについて解説しました。今回は、自己アテンションにフォーカスしながら論文を読み進めていきます。
まず、アテンションを使う３つの理由について考察します。さらに、ネットワーク内でどのように情報が流れるかについて、アテンション、再帰、畳み込みに関する計算量を使って比較します。
また、アテンションを使うことで得られる解釈性についても触れます。
なお、今回はこの論文を読むシリーズの最終回です。
アテンションを使う３つ理由
これまでずっと論文を読んできて、トランスフォーマーの仕組みはおおかた理解できました。ただし、アテンション機構がなぜ良いのかについては、あまり触れられていません。
セクション４「Why Self-Attention」では、自己アテンションを掘り下げ、それが再帰や畳み込みなどの仕組みよりも優れている点を述べています。なお、「自己」アテンションといっていますが、それに限定しているわけではなく、アテンション一般に関して述べています。
次のように始まります。
このセクションでは、自己アテンション層を、再帰層や畳み込み層と比較します。これらの層は、よくあるエンコーダ・デコーダの構造の中の隠れ層として、ある表現のシーケンス $${(x_1, \dots, x_n)}$$ を別の表現のシーケンス $${(z_1, \dots, z_n)}$$ にマッピングするために使われます。なお、$${x_i}$$と$${z_i}$$は、$${d}$$次元の実数$${\mathbb{R}^d}$$です。
In this section we compare various aspects of self-attention layers to the recurrent and convolutional layers commonly used for mapping one variable-length sequence of symbol representations $${(x_1, \dots, x_n)}$$ to another sequence of equal length $${(z_1, \dots, z_n)}$$ with $${x_i, z_i \in \mathbb{R}^d}$$, such as a hidden layer in a typical sequence transduction encoder or decoder.
ここでは、自己アテンション、再帰、畳み込みのどれも、エンコーダ・デコーダの中で入力シーケンスから出力シーケンスを生成するために使われているということを述べてきます。つまり、これらの層は、役割としては似たことをしているということです。よって、互いに比較の対象となります。
では、トランスフォーマーでは、なぜ自己アテンションが使われたのでしょうか。そこで、論文は望まれる３つの特性があると続けます。
自己アテンションを用いる理由として、私たちは3つの望ましい特性を考慮しています。
Motivating our use of self-attention we consider three desiderata.
まずは、２つの望ましい特性について次のように述べています。
１つ目は、各層で必要とされる計算量です。２つ目は、どれだけの計算を並列に実行できるかという点です。この並列性は、順番に処理しなければならない計算の数（つまり、並列化できない処理の数）によって測定されます。
One is the total computational complexity per layer. Another is the amount of computation that can be parallelized, as measured by the minimum number of sequential operations required.
１つ目のポイントは、層ごとの計算量が少ない方が良いということです。これは当然ですね。
２つ目のポイントは、並列化できる計算の量です。これは、たくさんの計算を同時に行うことで全体としての時間が短くなるということですが、それを測定するのに並列化できない計算の量を使います。順番に処理しなければならない計算量が少ないほど良いということです。
例えば、再帰層での計算は、シーケンス内にある$${n}$$個の埋め込みベクトルを一つずつ順番に処理する必要があります。なぜなら再帰層では、ある位置の計算が前の位置の計算結果に依存しているためです。よって、並列化ができません。
これに対して、アテンション機構や畳み込みでは並列処理が可能なので処理時間が短くなります。
情報経路の長さと学習の効率
論文は続けて、３つ目の特性を述べています。
3つ目は、ネットワーク内の遠距離の依存関係 (long-range dependencies) の経路の長さです。 多くのシーケンス変換タスクにおいて、遠距離の依存関係を学習することは重要な課題です。
The third is the path length between long-range dependencies in the network. Learning long-range dependencies is a key challenge in many sequence transduction tasks.
３つ目の特性は、シーケンス内の離れた位置にある依存関係を処理するのに、モデルがどのくらいの経路（処理の数）が必要があるのかによって測定されます。
英語では、path（経路）と書かれていますが、これを情報経路と呼びましょう。
例えば、再帰の構造では、シーケンス内の位置関係を順番に処理します。つまり、位置が離れると依存関係を処理するのに必要な処理の数が多くなります。つまり、情報経路が長くなります。
さらに、再帰の仕組みは、シーケンスの前方から後方への一方通行の処理です。もし、後方からの情報を前方へと伝えたいならば、双方向RNNなどの仕組みを使って後方からの情報を前方へと伝播する必要があります。つまり、必要な情報経路が長くなります。また、処理が複雑化します。
なお、畳み込みでは、カーネルのサイズによって、依存関係を直接計算できる範囲が限られています。よって、シーケンス内の依存関係を学習する際に、要素間の関係をすべて直接的に処理することはできません。つまり、離れた位置にある要素間の依存関係を学習するには、いくつかの処理を隔てる必要があります。これも情報経路が長くなると理解できます。
一方で、アテンション機構では、シーケンス内のすべての要素間の依存関係を前後とも直接に計算することが可能です。よって、情報経路が短いです。情報経路が短いほど処理が速く、また学習も容易（勾配消失などが起きにくい）になります。
論文でも、次のように述べています。
このような依存関係を学習する能力に影響を与える重要な要素の一つは、ネットワーク内を前向きと後ろ向きシグナルが伝わる経路の長さです。
One key factor affecting the ability to learn such dependencies is the length of the paths forward and backward signals have to traverse in the network.
ここで述べられている前向きと後ろ向きのシグナルが伝わる経路とは、次の２つのことになります。
前向きシグナル：ニューラルネットワーク内で、入力から出力へと伝わる情報のことです。いわゆる、フィードフォワードの処理によって伝えられる情報です。入力されたシーケンスの情報に基づいて、ネットワークがどのように出力を生成するかを決定するために使用されます。
後ろ向きシグナル：ニューラルネットワーク内で、出力から入力へと伝わる情報のことです。これはバックプロパゲーション（誤差逆伝播法）のことです。出力と実際の正解との誤差に基づいて、ネットワークのパラメータをどのように更新するかを決定するために使用されます。
この二つのシグナルが伝わる情報経路の長さが学習能力に大きな影響を与えると主張しているわけです。
さらに、論文は次のように続けます。
入力と出力シーケンスの任意の位置の組み合わせ間の経路が短いほど、遠距離の依存関係を学習するのが容易になります。
The shorter these paths between any combination of positions in the input and output sequences, the easier it is to learn long-range dependencies [12].
つまり、ネットワーク内の情報経路が長くなればなるほど、依存関係を学習することが難しくなることを示唆しています。これは、シグナルが伝わる間に情報が失われたり、劣化したりする可能性があるためです。
再帰の構造ではシーケンス内の位置が離れるほど情報経路が長くなり、勾配消失といった問題が生じがちです。LSTMなどによって、ある程度は改善されましたが、処理がかなり複雑になりました。
畳み込みでも、カーネルが処理できる範囲が限られているので、シーケンス内の位置が離れるほど情報経路が長くなります。
一方で、アテンション機構では、すべての要素間の関係を直接的に扱うので経路が短く学習が容易になります。よって、アテンション機構では、シーケンス内の要素の位置が離れていても、情報経路は一定です。
以上から明らかですが、異なるモデルが同じシーケンスを処理する場合、情報経路の長さが異なるので、論文では異なるネットワーク間での情報経路の最大長を比較すると述べています。
そのため、異なる種類の層で構成されたネットワークにおいて、入力と出力の任意の 2 つの位置間にある最大の経路長も比較します。
Hence, we also compare the maximum path length between any two input and output positions in networks composed of the different layer types.
異なる層による計算量の比較
情報経路の最大長に関して、より具体的に次のように述べ、表でも表現しています。
表 1 にあるように、自己アテンション層は、一定の処理ですべての位置を接続（関係を処理する）します。一方、再帰層は、$${O(n)}$$の処理を必要とします。
As noted in Table 1, a self-attention layer connects all positions with a constant number of sequentially executed operations, whereas a recurrent layer requires $${O(n)}$$ sequential operations.
下表にあるように、順番に行う処理（Sequential Operations）のオーダーを比較しています（表の真ん中）。
例えば、シーケンスの長さが$${n}$$の場合、再帰を使うと情報経路は最大で$${n}$$に比例します。つまり、処理数のオーダーが$${O(n)}$$です。しかし、アテンション機構では固定です。よって、処理数のオーダーは$${O(1)}$$となります。
また、この表では各種の層による計算の複雑さも含んでいます。
計算量で見ると、自己アテンション層は、シーケンス長$${n}$$が埋め込みベクトルの次元数$${d}$$よりも小さい場合に、再帰層よりも高速です。これは、最先端の機械翻訳モデルで使用されるword-piece（ワード・ピース）[38] やbyte-pair（バイトペア・エンコーディング） [31] などのトークン化の仕組みにおいてほとんどの場合当てはまります。
In terms of computational complexity, self-attention layers are faster than recurrent layers when the sequence length $${n}$$ is smaller than the representation dimensionality $${d}$$, which is most often the case with sentence representations used by state-of-the-art models in machine translations, such as word-piece [38] and byte-pair [31] representations.
シーケンスの長さが$${n}$$の場合、アテンション機構では、各要素間の全ての関係を処理するので、$${n \times n}$$の組み合わせの計算を行います。また、ベクトルの次元が$${d}$$なので、内積やソフトマックスの計算のオーダーは$${O(n^2 \cdot d)}$$となります。
クエリ、キー 、バリューのベクトルは、入力の埋め込みベクトルに$${d \times d}$$の重み行列を乗算（線形変換）を適用して生成されるので、各変換の計算コストのオーダーは$${O(n\cdot d^2)}$$です。
しかし、後で見るように$${n < d}$$の時には、内積やソフトマックスによる計算のオーダー$${O(n^2 \cdot d)}$$の方が大きくなります。
その一方で、再帰ではシーケンスの各要素を順番に$${n}$$個の処理します。各タイムステップで隠れ状態を更新するための線形変換（$${d \times d}$$の重み行列を乗算）が行われます。よって、計算のオーダーは$${O(n \cdot d^2)}$$となります。
この二つを比べると、$${n < d}$$の時には、層の計算量のオーダーで比較すると、アテンションの方が処理が速くなります。
例えば、ある機械翻訳モデルで、文章の平均長が 500 単語、埋め込みベクトルの次元数が 512 であるとします。この場合、シーケンス長$${n}$$は埋め込みベクトルの次元数$${d}$$よりも小さい (500 < 512) ので、アテンション層は再帰層よりも高速に処理することができます。
とは言うものの、アテンション計算量が$${n}$$の2乗に比例して増加することは、より長い文章を扱う必要がある場合に処理が遅くなる問題が生じます。
これに対して論文では、次のように述べています。
計算パフォーマンスを向上させるために、非常に長いシーケンスを扱うタスクでは、自己アテンションを入力シーケンスの特定の近傍領域に制限することが有効です。この近傍のサイズは$${r}$$で、各要素に対してその周囲$${r}$$個の要素のみを考慮に入れます。このアプローチを採用することで、最大経路長は $${O(n/r)}$$まで増加する可能性があります。この方法については、将来の研究でさらに詳しく検討する予定です。
To improve computational performance for tasks involving very long sequences, self-attention could be restricted to considering only a neighborhood of size $${r}$$ in the input sequence centered around the respective output position. This would increase the maximum path length to $${O(n/r)}$$. We plan to investigate this approach further in future work.
わざわざこのように書いているのは、長いシーケンスで計算量が２乗で増加する問題を認めているからです。
そこで、各要素に対して、$${n}$$個すべての関係を計算するのではなく、近傍の$${r}$$だけに限定すれば、計算量を減らせるという提案をしています。しかし、「将来の研究でさらに詳しく検討する予定です」と言及するのに止まっています。
なお、各出力要素が受け取る情報の範囲を狭める代わりに、情報が全ての要素に伝播するには複数のステップを経る必要が生じます。つまり、情報経路が長くなります。よって、最大経路長（Max Path Length）は元来の$${O(1)}$$から$${O(n/r)}$$にまで増加する可能性があると述べているわけです。
これは、カーネルの範囲が限定されている畳み込みのケースと似ています。論文でも次のように述べています。
カーネル幅が $${k < n}$$ の単一の畳み込み層では、入力と出力の全ての位置のペアを接続することはできません。連続するカーネルの場合、これを実現するには $${O(n/k)}$$ の畳み込み層を積み重ねる必要があります。また、拡張畳み込みの場合には $${O(\log_k (n))}$$ の畳み込み層が必要となり、ネットワーク内の任意の二点間の最長経路の長さが増加します。【18】
A single convolutional layer with kernel width $${k < n}$$ does not connect all pairs of input and output positions. Doing so requires a stack of $${O(n/k)}$$ convolutional layers in the case of contiguous kernels, or $${O(log_k (n))}$$ in the case of dilated convolutions [18], increasing the length of the longest paths between any two positions in the network.
カーネル幅がシーケンス全長より小さい場合、単一の畳み込み層では、すべての入力位置と出力位置のペアを直接接続することができません。このため、すべての入力と出力を完全に接続するには、畳み込み層を複数重ねる必要があります。なお、シーケンスにおける畳み込みは1次元の処理です（画像では2次元です）。
例えば、シーケンス長が $${n=12}$$ でカーネル幅が$${k=3}$$の場合、ストライドが１ならば、$${n-k+1}$$で、10個の新しい特徴が生成されます。つまり、畳み込み層は隣接する3つの要素を取り、それらに対してカーネルを適用し、新しい特徴を生成します。ストライドが1のため、カーネルは各ステップで1要素ずつ移動し、最終的に10個の新しい特徴を生成します。もちろん、GPUなどを使えば、実際には並列処理可能ですが、計算量は変わりません。
しかし、生成された10個の特徴だけではシーケンス全体の情報を完全には捉えきれていません。これを補うために、さらに畳み込み層を重ねることで、より広範囲の情報を捉え、特徴の抽出を深めていきます。これにより、畳み込み層を通じて情報が統合され、より高次の特徴が形成されることになります。
つまり、ある位置の周囲からの情報を徐々に集積し、より広範なコンテキストを含んだ特徴を形成していくので、シーケンス全体の情報をカバーするために、$${n/k}$$回の畳み込み層を重ねる必要があります。もちろん、ストライドを増やしたり、拡張畳み込み（dilated convolution）などを使えばより少ない回数になりますが、ここでは単純化して考えています。
続いて、論文は畳み込み層の計算量についても言及しています。
畳み込み層は一般的に再帰層よりも計算コストが高く、その差はカーネルの幅$${k}$$によるものです。しかし、チャンネル分離畳み込み（Separable convolutions）[6]を使用することで、この複雑さは大幅に減少し、$$O(k \cdot n \cdot d + n \cdot d^2)$$ まで低減されます。
Convolutional layers are generally more expensive than recurrent layers, by a factor of $${k}$$. Separable convolutions [6], however, decrease the complexity considerably, to $${O(k \cdot n \cdot d + n \cdot d^2)}$$.
通常の畳み込みの重み行列のサイズを$${d \times d}$$とすると、各畳み込み操作の計算量のオーダーは$${O(k \cdot d^2)}$$です。シーケンス全体でこの操作が$${n-k+1}$$（上記の例では10）回繰り返されるため、畳み込み層全体の計算量のオーダーは$${k \cdot O((n-k+1) \cdot d^2)}$$となります。ただし、$${k}$$はシーケンスの長さ$${n}$$に依存せず、その大きさも小さいので、単純化して$${O(k \cdot n \cdot d^2)}$$と見積もることができます。
これに対して、チャンネル分離畳み込みでは、通常の畳み込みをチャンネルごとの畳み込み（depthwise convolution）とポイントごとの畳み込み（pointwise convolution）に分けて行います。これにより、畳み込みのパラメータ数と計算コストが大幅に削減されます。
チャンネルごとの畳み込みでは、各入力チャンネルに対して独立したカーネルが適用されます。このカーネルは通常の畳み込みと比べてサイズが小さく、各チャンネル内でのみ畳み込みが行われるため、計算量は各チャンネルに$${O(k \cdot n \cdot d)}$$のオーダーになります。
ポイントごとの畳み込み（1x1畳み込みとも呼ばれる）は、チャンネルごとの畳み込みの結果を統合するために使用されます。このステップでは、畳み込みはチャンネル間で情報を混合し、出力チャンネルごとに独立した重みを適用します。よって、計算量は$${O(n \cdot d^2)}$$のオーダーです。
これら二つのステップを合わせると、チャンネル分離畳み込みの計算量は、$${O(k \cdot n \cdot d + n \cdot d^2)}$$ となります。この計算量は、通常の畳み込みの計算量 $${O(k \cdot n \cdot d^2)}$$ と比較して大幅に少なくなります。
このように、分離可能な畳み込みは計算資源が限られている環境でも効率的に畳み込み層を利用することを可能にし、特にモバイルデバイスやエッジデバイスでの応用において有利なので MobileNet で採用されました。
通常の畳み込みより計算量が少ないチャンネル分離畳み込みですが、その計算量$${O(k \cdot n \cdot d + n \cdot d^2)}$$をアテンションの計算量$${O(n^2\cdot d)}$$と比較したらどうでしょうか。アテンションは、すべての要素をカバーするので、$${k=n}$$の場合のチャンネル分離畳み込みの計算量を使って比較できます。
論文では次のように述べています。
ただし、$${k = n}$$ であっても、分離可能な畳み込みの複雑さは、自己アテンション層と位置ごとのフィードフォワード層の組み合わせと同等です。これは、私たちのモデルで採用しているアプローチです。
Even with $${k = n}$$, however, the complexity of a separable convolution is equal to the combination of a self-attention layer and a point-wise feed-forward layer, the approach we take in our model.
$${k=n}$$の場合のチャンネル分離畳み込みの計算量は、$${O(n^2 \cdot d + n \cdot d^2)}$$となり、アテンションの計算量$${O(n^2\cdot d)}$$に、トラーンスフォーマーの位置ごとのフィードフォワードの計算量$${O(n \cdot d^2)}$$を足したものと等しくなります。要するに、畳み込みのカーネルをシーケンス全体をカバーするようにした際の計算量は、トランスフォーマーのアテンションと位置ごとのフィードフォワードを合わせたものと等しいというわけです。
なぜ、畳み込みの計算量について、このように長い説明をしているのでしょうか。
上記の説明を逆にいうと、トランスフォーマーの計算は、チャンネル分離畳み込みをシーケンス全体に適用したのと同等だと主張しています。これをわざわざ言及しているのは、「長いシーケンスで計算量が２乗で増加する問題はあるけど、そんなに効率が悪いわけでもないよ」と、よりバランスの取れた視点を提供しようとしているのが伺えます。
その上で、アテンションを選ぶ理由が、単に計算コストによってだけでなく、実装の利便性や言語処理などの問題に最適であるかどうかに基づいてなされるべきであることを示唆しています。
アテンションの解釈可能性
このセクションの最後で、論文はアテンションを選ぶことによるさらなる利点を述べています。
さらなる利点として、自己アテンションはより解釈可能なモデルを提供する可能性があります。私たちはモデルからのアテンションの分布を調査し、その例を付録で提示して議論しています。これによって、個々のアテンション・ヘッドが異なるタスクを行うことを明確に学習しているのがわかります。また、それらが文の構造や意味をどのように処理しているかについての示唆が得られます。
As side benefit, self-attention could yield more interpretable models. We inspect attention distributions from our models and present and discuss examples in the appendix. Not only do individual attention heads clearly learn to perform different tasks, many appear to exhibit behavior related to the syntactic and semantic structure of the sentences.
ここで主張しているのは、アテンションによる出力が再帰や畳み込みより解釈しやすいということです。畳み込みでもある程度は、特徴量の意味を見ることができますが、アテンションでは要素間の関係を直接に計算しているのでさらに解釈しやすいという利点があります。
下図では、アテンションの分布を視覚化したものです。ここでは、making（作る）という単語が文中の各単語に対して関連性の強さを示しています。
アテンション・ヘッドは異なる色の長方形で表現され、色が濃いほど関係性が強いことを示しています。これを見ると、アテンション ヘッドの多くが、「作る...もっと（more）難しい（difficult）」といった遠距離の依存関係に注目しているのがわかります。
<EOS>は、End Of Sentenceで文章の終わりを意味します。<pad>は、使われていない位置を埋めるためだけに存在するトークンでシーケンスを固定長にするために使われます。
下図は、一つのアテンション・ヘッドからの出力を表示しています。
この図では、線が太いほど、そのアテンションの強さが大きいことを表します。一方で、線が細い、または全く描画されていない場合、モデルはそれらの単語間の関係性をあまり重要でないと判断していることを意味します。
これを見ると、モデルがどのように各単語間で複雑な関係性を識別しているかが分かります。
ちなみに、文章の意味は「法律が完璧になることは決してないが、その適用は公正であるべきですーわたしの意見では、これ（公正な適用）が不足しています」という意味です。
The Law will never be perfect, but its application should be just - this is what we are missing, in my opinion.
下図は、同じ文章において、また、its（その）という単語のアテンションに関して、もう一つのアテンション・ヘッドからの出力を追加して表示しています。
一つのアテンション・ヘッド（紫色）は、Law（法律）が its（その）という代名詞で意味として強くつながっているのが確認できます。
また、２つ目のアテンション・ヘッド（灰色）では、its が Law と application をつなげているのが分かります。its application で、「法律の適用」という意味になります。
以下の二つも異なるアテンション・ヘッドからの出力を表示しています。
正直なところ、これらが何を意味するのかは明確には分かりません（論文にも書かれていません）が、ランダムな関係性でないのは直感的に感じられます。おそらく、研究者たちも、これらの観察結果が何を意味するのかを完全には理解するのは難しいでしょう。しかし、ある程度の解釈は可能です。
緑色のラインは、より多くの単語間の広範な関連性を捉えており、このヘッドは文脈全体を捉えるように働いているようです。
これに対して、赤色のラインのヘッドは、比較的単純なパターンを示しており、このヘッドは、より直接的な単語の関連性や特定の文法的機能を捉えている可能性があるように見えます。
以上のようなアテンションの解釈性と同じような効果を再帰や畳み込みから得るのは難しいでしょう。よって、アテンションが好まれる理由の一つとなっています。
まとめ
アテンション機構を理解するためにトランスフォーマーの論文を読み出したこのシリーズも今回で終わりとなります。ここで得た知識をもとにさらにトランスフォーマーから派生したモデルの論文を読むと、より一層理解が深まるでしょう。
次回からは、OpenAIが2018年に発表した論文「Improving Language Understanding by Generative Pre-Training」を読んでいきます。
この論文では、トランスフォーマーを事前学習してから特定のタスクにファインチューニングする手法を提案しています。これがGPT（Generative Pre-trained Transformer）シリーズの最初のモデルです。その当時はGPTという略称は使われていませんでしたが、その手法は以降のGPT-2、GPT-3など更に進化したモデルにも引き継がれています。
お楽しみに！
