Transformerの基本的なことをまとめてみた

Transformerが出る前まで

最初に時系列データや自然言語処理でよく使われていた再起型ニューラルネットワーク(Recurrent Neural Network)の話です。

(参照:Understanding LSTM Networks)

自然言語処理で文章を処理するケースを考えると、RNNでは文章の頭から1単語ずつ(ベクトル化した入力)を読み込んでいくことで単語間の関係を学習することができます。しかし、構造的に下記のような長期の情報の記憶が必要な場合に精度が出ない問題があります。

系列が長くなる場合、損失関数に対する系列の最初の方の重み(t=1)は下記の式で表され、重みが更新されます。(合成関数の微分はそれを構成する関数の微分の積となるため)

この時、各層の出力は活性化関数tanhを用いて下記で表されます。

微分すると…

系列が長くなると(tの値が大きくなると)その分上記の値の積が計算されますが、tanhの微分は必ず1以下になるため損失関数に対する勾配の更新幅が0になる可能性が高まります。(勾配消失がおきる)

このような長期に情報を記憶する構造を追加したのが、LSTMやGRUなどのモデルでした。

(参照:Understanding LSTM Networks)

LSTMでは、記憶セルcが追加されています。

(参照:Understanding LSTM Networks)

先程との違いは記憶セルの変換は、fとのアダマール積(要素毎の積)で計算されています。従って関数の積の微分を考えると、下記のように計算されてfの値が逆伝搬として伝わるため、勾配消失が起きにくい構造になっているのだと考えられます。

(参照：How LSTM networks solve the problem of vanishing gradients)

行列積でなく、アダマール積としている理由は、第1項を残しておきたいからなのか？と考えましたが、正直気持ちがよくわかりませんでした…。

LSTMは長期の情報を保持するという目的ではRNNから改善が見られましたが、前の層の情報を使うため並列処理ができない(計算時間がかかる)という問題点があります。

またスパム判定など全体の文章を使用して分類する場合など(系列の後の情報を使える場合)でも、前の単語の情報のみを用いて予測をするため精度が上がりにくい部分もあるようです(双方向にLSTMの構造を持つBidirectional LSTMというものもあるようですが、各層は前only or 後onlyの単語の情報しか使えないためこの後紹介するAttentionと比較すると不利なようです)

Transformerの登場

TransformerはAttention is all you needという強いタイトルの論文で公表されています。下記の図はよく見るかと思います。

最初見た時は、理解できる気がしない…と思っていましたが、基本はAttentionが複数使われているだけでLSTMのような複雑な構造をしているわけではないので、Attentionが何をやっているかを理解した後に手を動かして理解するのが良さそうです。

という事で最初にAttentionの説明になります。

Attentionは論文で下記の図で表されています。

左側がAttentionの計算処理で、右側は初期値が異なるAttentionを複数重ねている状態(Multi head attention)を表しています。

各Headの初期値が違うので、学習を進めた時に異なる情報を捉えられるようにしようというのがモチベーションのようです。Transformerでは8個のヘッドを持ちます。

左側のAttentionの計算処理で、Q,K,VはQuery, Key, Valueを表しており、ある入力xに対して、初期値が異なる重みWq, Wk, Wvの積で計算された値がQ, K, Vとなっています。

例えばx1を入力する場合(自然言語処理なら I like dogsだとしたらIの1単語)を考えると、単語を多次元の埋め込み層に変換した後(1×512)、Wq, Wk, Wv(512×64)との行列積を計算してQ, K, V(1×64)を得ます。

入力の次元(512)に対して、8個のheadを持つため512/8 = 64という計算で求められます。(実装上入力の次元がhead数で割り切れる値となるよう)

1つのAttention構造における(図の左側)計算式は下記のように表されます。

右辺のQとKの転置行列のドット積では、ある入力x1の入力と自分(x1)を含めた他の入力(x2)との類似度を算出しており、ここではある入力と他の入力との関連度を計算していると考えられます。

(引用元：The Illustrated Transformer)

重みを計算後、Q, K, Vの次元の平方根で除して正規化した後にVに掛けています。(Vも同様に、入力と重みWvの積)

何故平方根なのかということに関しては不明でしたが…、qとkの要素が平均0、分散1の独立正規分布に従うとするとqとkの積を取る1要素は(x1y1+x2y2…xdkydkとあらわされるので)分散がdk倍になります。(V[X + Y] = V[X]+V[Y]　(X,Yが独立の時))

分散が大きくなるとソフトマックス関数に入力した際に、マイナス方向に小さいと出力(勾配)が非常に小さくなる要素が増える可能性が高まるため、分散が大きくならないようにしたいという気持ちでのスケーリングだと思われます。

平均0分散1の配列をスケーリングした場合の分布

Attention maskとは何か？

エンコーダー側のみを用いて時系列データ予測を行う場合や、翻訳時にデコーダー側で前の単語の翻訳結果を使用する場合など、学習時に将来のデータを利用しないように(leakageしないように)対応する将来の系列情報から得られる重み(attention weight)にmaskをする必要(重みを0にする必要)があります。

これは、重みを計算するQ・Ktの積の結果が全て-infになるように(重みが0になるように)処理が行われます。

maskは上記のように、上三角行列の形で重みを0にする要素の位置を与えます。その後、マスク箇所の要素が0となった行列と、Vの積を取ることで、学習させない要素を除いた出力を得ることができます。

Positional Encodingとは何か？

AttentionはLSTMのような系列の順番に入力をするのではなく、全ての入力に対して内積を取る形で学習を行うので、系列の順番をモデルに教える必要があります。

位置情報の与え方として単純に入力のembeddingに足し合わせるAbsolute positional encodingという方法と、相対的な位置情報(3番目の単語にとっては、1番目の単語は2つ前にあるなど)を与えるRelative positional encodingという方法があります。

(引用元：The Illustrated Transformer)

論文では、図のように入力に位置毎に値が一意に決まるを足し合わせるAbsolute positional encodingが紹介されています。

こちらは、下記の式で表されています。

次元の偶数奇数別々に正弦波関数と余弦波関数を用いて、位置を表す値が計算されています。これは、posからk離れた位置情報をposとkの各位位置の値の線形和で表せるようにして、相対的な位置(pからk離れている入力であるということ)を学習しやすくなるためという説明がありました。

しかし、最近の研究でAbsolute positional encodingによる位置情報の付与より、相対位置を考慮したRelative positional encodingを行う方が精度が良いという論文もありました。(Self-Attention with Relative Position)

こちらは、最初のレイヤーのみでなく、各レイヤーに相対位置を学習させるパラメータを付与するようです。(正しく理解できているか怪しい…)

さらに、2020年の研究では、そもそも入力のembeddingと位置情報のembeddingは違う性質のものなので足し合わせるのはおかしいということで、別々にパラメータを学習させた方が精度が良いという論文(RETHINKING POSITIONAL ENCODING IN LANGUAGE PRE-TRAINING)が公表されていました。

入力と位置情報を足し合わせた場合のQueryとKeyのドット積の計算は、(Q=(input_emb+position)WqとK=(input_emb+position)Wkの積になるので)、input-input, input-position, position-input, position-positionの4つの項に分解できます。

大規模データで学習させた後の各項のドット積を可視化すると、下記のようにinputとpositionからなる2,3項目(左から2,3番目)の値が小さいため位置情報と単語の関係は弱く(ある単語が最初にあろうが、最後にあろうがタスクを解く上であまり関係ない)、これらを一緒に学習すると非効率になるというのが、主張でした。

上記は、自然言語処理の場合のようなので、株価などの時系列データを扱う場合は、また違ってくるのかもしれません。

最後に

Attention何もわかっていなかったので(今も正しく理解できているのかわかりませんが…)、とても勉強になりました。

直近ではkaggleのRiiid! Answer Correctness Predictionというコンペで、Transformerを少し変えたモデルが使われており、実装が公開されているので、手を動かしてみると理解が深まる気がします。また、こちらのハーバード大学が公開している実装と解説もわかりやすくて助かりました。

仕事やコンペで使ってみよう…。