The architecture of Transformer - usapyoiのブログ

はじめに

Transformerのモデル構造についての備忘録。 Transformerの論文は以下から。

arxiv.org

時間があるなら、下記リンクを読んだ方が良い(より正確で、より詳細)。というより、この記事は下記リンクの翻訳、要約と解釈した方が正しい。

jalammar.github.io

なので画像をバシバシ引用する。

最上位の構成要素

Transformerは大きく分けてencodersとdecodersの2つから成る。以下は、フランス語”je suis étudiant”を入力し、英語"I am a student"を出力として得た時の図。

さらに、それぞれにはencoder, decoderが6つずつ入っている(この6という数字はマジックナンバーで、実験したら6が良さげだったから、以上の理由はない、と思う)。

ここでは、データフローに沿って処理機構を見ていこう。

Word Embedding

入力された単語はまずテンソルに変換されなければならない。そこで登場するのがThe embedding algorithmだ。詳しい話は以下のリンクを参考にして欲しいが、技術としての要点は

one-hotベクトルで表現された単語をより少ない次元数のベクトルに変換する

ということ。文脈としての近さを反映させるために、どうやらニューラルネットワークを使うらしい(例えば、動詞「育てる」は名詞「子供」と文脈的に近いが、名詞「危険」とは遠い)。

medium.com

この技術によって、入力された単語は比較的低い次元数(512 これはマジックナンバーなのだろうか？)で表現される。これを単語埋め込みベクトルと呼ぼう。

positional encoding

では早速encoderに…と行きたいところだが、transformerのencoderには入力データの時間的流れを加味する機能が抜け落ちていることに注意しよう。

例えばRNNでは再起的構造が、CNNでは畳み込みがその機能を担っていたのだが、Transformerはこれらの機能を廃したため、そのままでは単語の順番を認識できない。例えば「少女に与えられたのは、大きな銃と小さな幸せ」という文と、「少女に与えられたのは、小さな銃と大きな幸せ」という文を区別できないということだ。この文題を解決するために、positional encodingというものを考案している。

原理は単純で、先ほど生成した単語埋め込みベクトルに、以下の式で表現される値を加算するだけ。

$PE_{(pos, 2i)} = sin(\frac{pos}{10000^{\frac{2i}{d_{model}}}})$

$PE_{(pos, 2i+1)} = cos(\frac{pos}{10000^{\frac{2i}{d_{model}}}})$

ここで、 $d_{model}$ は単語埋め込みベクトルの次元数(512)で、 $pos$ は文中の単語の位置を表し、 $2i, 2i+1$ はその単語に対応するベクトルの次元に相当する。例えば、5単語目に対応するベクトルの第17次元目の値に加算されるのは、 $cos(\frac{5}{10000^{\frac{16}{512}}})$ となる。この手法では、position encodingによる値は次元によって正弦と余弦を交互に出力したものになり、なんか上手くいくらしい。

encoder

position encodingを加算された単語埋め込みベクトルはencoder(より正確には、6層のうちの最下層のencoder)に入力される。 encoderはself-attention layerとfeedforward network layerから成る。

詳細な説明は後で述べるとして、この2つのsub layerを通ったデータは、次の層のencoderに入力される。ここで、encoderのモデル構造は層の間で変わらない。違うのはパラメタと、入力データだ(最下層のみEmbeddingされたデータ、以外は一つ下の層のencoderの出力)。

self-attention layer(high level)

self-attention layerでは、とてもざっくり言うと「入力文を解釈する上での注意」が喚起される。

例えば、入力文が「昨日ゲーセンで新作のゲームをやったが、傑作だった。」という文があったとする。この時、「『傑作だった』のは何か？」と言われたら、多くの人は「ゲーム」と答えるだろう。しかし、機械は「傑作だった」のが「昨日」なのか「ゲーセン」なのか「ゲーム」なのか判断する必要がある。

この時登場するのがself-attention layerで、ある単語を読んだ時、注意すべき単語を数字で表してくれる。上記の例なら、self-attention layerは各単語に対して、「傑作だった」に対する注意度として、「昨日」は0.01、「ゲーセン」は0.1、「ゲーム」は0.8という様にスコアをつけてくれる。これにより、機械は「ああ、ここでいう『傑作』とは『ゲーム』のことなんだな」と理解できるわけだ。下の画像の例も似たようなもので、英文"The animal didn't cross the street because it was too tired"を入力した時、後半の"it"の解釈についてattention layerが評価していることを示している。

ここで、"it"の後に登場する単語にもスコアがついていることに注意。これは、self-attention layerによる評価が、まず全単語を読み、あるベクトルを生成した後に行われるからだ。詳しい話は次節へ。

self-attention layer(low level)

具体的にどうやってスコアを算出しているのだろうか？

self-attention layerはまず、"Queries" "Keys" "Values"という3つのベクトルを生成する。これらのベクトル( ${\bf q}, {\bf k}, {\bf v}$ とする)は、単語データ( ${\bf x}$ とする　何度も言うが、最下層ではEmbeddingされたデータ、以外は一つ下の層のencoderの出力)と、対応する3つの行列( ${\bf W^Q}, {\bf W^K}, {\bf W^V}$ とする)との行列積によって算出される。入力データは各単語ごとに生成されるため、これらのベクトルも各単語ごとに生成される。なお、 ${\bf q}, {\bf k}, {\bf v}$ の次元数は64に設定されている。これもマジックナンバーだが、計算上都合が良いそうだ。

次に、単語 $x_i$ に対する単語 $x_j$ の評価値(の元となる値)を、 ${\bf q}_i$ と ${\bf k}_j$ との内積によって表現する。イメージとしては、単語 $x_i$ に対する評価値を、 $j$ を1つずつ動かしていって精査する感じだ。

単語'Thinking'に対する評価値を生成する様子　評価元と評価先が同じ場合もある

この評価値をそれぞれ8で割っている。8はベクトル ${\bf k}$ の次元数64の平方根であるが、なぜ64なのか、そしてなぜ平方根なのかについて、理論的な説明はなされていないようだ。どうやらそうしたら性能が上がったから、というさっき見たような理由らしい。 ※2023/04/23 補足：8である必要はないが、少なくとも割る事によって、後述のsoftmax関数によって値が0と1に二極化することを防いでいるらしい。

その後、それら評価値をsoftmax関数にかけている。分類問題などで馴染み深い関数だが、軽く説明しておくと総和が1になるようにスケーリングをする関数だ。これにより、評価値の総和が1になり、注意度がそのまま解釈すべく単語である確率と見做せるようになる。

softmax関数により、総和が1となり、確率と見做せるようになるのは分類問題では馴染み深い

これらの確率が、このself-attention layerによる各単語に対するスコアとなる。

では、このスコアからどのように出力を求めているのか？

この疑問は簡単で、各単語が持つベクトル ${\bf v}$ に先ほどのスコアをスカラー倍し、その総和を出力としている。つまり、 ${\bf v}_j$ にスコア ${\rm softmax}({\bf q}_i \cdot {\bf k}_j / 8)$ を乗じて、 $j$ を制御変数とした総和を取っている。

では、ここまでのまとめとして、self-attention layerの出力を数式で表現してみよう。 $N$ 語の単語を含む入力データに対して、 $i$ 番目の単語に対するself-attention layerの出力 ${\bf z}_i$ は、

${\bf z}_i = \sum_{j = 1}^{N} {\rm softmax}(\frac{{\bf q}_i \cdot {\bf k}_j}{\sqrt{d_k}(=8)}){\bf v}_j$

と表現できる。

行列による表現

当たり前のことだが、実際は上記のような計算が逐次的に行われている訳ではなく、行列同士の積によって一気に計算が行われている。先ほどの計算式を、行列を用いたものに変換してみよう。

まずは、計算に用いる行列の定義から始める。入力する単語データを各行にまとめた行列を ${\bf X}$ とする。この時、行数は単語データの数、列数は各単語データの次元数(恐らく論文では512のはずだ)になる。また、この入力単語行列と ${\bf W^Q}, {\bf W^K}, {\bf W^V}$ との行列積によって得られる行列を、それぞれ ${\bf Q}, {\bf K}, {\bf V}$ とする。定義から明らかだが、 ${\bf Q}, {\bf K}, {\bf V}$ の行数は単語データの数、列数は内部ベクトルの次元数(つまり64)となる。この辺の話は説明が難しいから、自分で書いて確かめよう。

最後に、出力される行列を $Z$ とする。この行列の形状も ${\bf Q}, {\bf K}, {\bf V}$ と同じで、行数は単語データの数、列数は内部ベクトルの次元数だ。

入力単語行列の列数は512(図では4マス)、内部ベクトルの次元数は64(3マス)に設定されている

これらを用いると、上の計算式は、以下のように書き表せる。

${\bf Z} = {\rm softmax}(\frac{{\bf Q} {\bf K}^{T}}{\sqrt{d_k}(=8)}){\bf V}$

ここで、ベクトルの内積を行列で計算するために、 ${\bf K}$ を転置していることに注意。

論文の計算式は以下のようになっている。

${\rm Attention}({\bf Q}, {\bf K}, {\bf V}) = {\rm softmax}(\frac{{\bf Q} {\bf K}^{T}}{\sqrt{d_k}}){\bf V}$

ニュアンスはわかったんじゃないかな。

Multi-Head Attention layer

ではMulti-Head Attention layerが登場する。これは先ほどのSelf-Attention layerを多層化したものだ。つまり、同じ入力データに対して、複数の(8、例に漏れずマジックナンバー) ${\bf W^Q}, {\bf W^K}, {\bf W^V}$ が評価を行っている。

さて、これらのself-attention layerがそれぞれの評価を下し、今ここに8つの出力 ${\bf Z}_0, {\bf Z}_1, ..., {\bf Z}_7$ がある。もちろんこのまま出力することはできない。次元数が異なってしまうからだ。ではどうするか？

答えは行列である。最終的な出力を算出する行列 ${\bf W}^{O}$ とし、各層の出力を結合したもの( ${\bf Z}_{con}$ とでも置こう)に乗ずる。つまりこうだ。

${\bf Z}_{con} = [{\bf Z}_{0}; {\bf Z}_{1}; ...; {\bf Z}_{7}$ ]

${\bf Z} = {\bf Z}_{con}{\bf W}^{O}$

Attention layerの説明は以上。最後に全体像と、論文の計算式を載せる。

論文の計算式:

${\rm MultiHead}({\bf Q}, {\bf K}, {\bf V}) = {\rm Concat}({\rm head}_{1}, ..., {\rm head}_{h}){\bf W}^{O}$

${\rm where} \quad{\rm head}_{i} = {\rm Attention} ({\bf Q}{\bf W}_{i}^{Q}, {\bf K}{\bf W}_{i}^{K}, {\bf V}{\bf W}_{i}^{V})$

Position-wise feed forward neural Network

Attention層の次は全結合層を通る。Position-wizeとは、単に各単語ごとに独立したニューラルネットワークであること(但し、重みは共通)を意味する。つまり、式としては同じだが、他の単語からの影響がない層であるということ。

${\rm FFN}({\bf x}) = {\rm max}(0, {\bf x}{\bf W}_{1} + b_{1}){\bf W}_2 + b_{2}$

式の通り、この層は2層のネットワークで、活性化関数として ${\rm RELU}$ を採用しているようだ。なお、中間層の次元数は2048らしい。マジックナンバーである。

residual connection

先に述べたattention layerやPosition-wise FFNにはresidual connection(残差接続)が備わっている。下の層の情報がある程度残りながら上の層に渡されるようだ。なお、接続後はnormalize(正規化)される。

残差接続は勾配消失を防ぐためと習ったが、この場合でも情報喪失が起きるのだろうか？

decoder

さてdecoderの説明に移る訳だが、「まだ半分かよ…」と思う人もいるかもしれない。安心してほしい。確かに入出力は多少違うが、基本的構造や技術はencoderで紹介済みだ。ここでは、encoderとの違いについてフォーカスして述べたいと思う。

恐らくもう論文に掲載されているモデル構造を見ても言いたいことがわかるのではないだろうか？

左側がencoder、右側がdecoderである。大きな相違点として、encoderは2層構造であるのに対し、decoderは3層構造であることが挙げられる。2つのMulti-Head Attention layerの役割は以下の通り。

Masked Multi-Head Attention(1番下)…直前の出力系列(最下層のdecoderなら更にWord embeddingとpositional encodingをしたもの)を入力とし、評価する層。encoderの時では、現在の単語からすべての単語への評価を算出すると述べたが、今回の場合、現在の位置以降の単語に対する評価を行ってはいけない(カンニングになるので)。よって、以降の単語に対する評価は $-\infty$ とマスキングしている。だからMasked。
encoder-decoder Attention(下から2番目)…一つ下のAttention layerの出力と、最上層のencoderの出力を受け取る層。具体的には、Attention layerから $q$ を、encoderから $k, v$ を受け取る。つまり、この層には3つのベクトル $q, k, v$ を自前で生成する能力がないことに注意。この辺の説明がみんな雑で、自分もよくわかってない。

decoderもまた6層あるので、この作業を6度繰り返すことになる。 decoderの説明は(なんと！)以上。最後の最後に行われる処理だけ見ていこう。

The final Linear and Softmax layer

長い長い旅を終え、入力単語はついにdecoderから出てきた。しかし、decoderの出力はベクトルであり、単語ではない。人間が理解できる形にするために、単語に直してやる必要がある。ここで登場するのが、最後のLinear layerとSoftmax layerだ。

まず、Linear layerは出力されたベクトルをより高次元の空間に射影する。要するに高次元ベクトルに変換する。この時の次元数は、モデルが知っている単語の数による。例えば、日本語を英語に翻訳するタスクで、モデルが覚えている英単語の数が10000であるとすると、この層は1000次元のベクトルに変換する。

この説明から、Softmax layerの役割について見当がついたのではないだろうか？Softmax layerでは、ベクトルを正規化し、総和を1にする。これにより、ベクトルの各要素を生起確率と解釈することができるため、もっとも高い確率を示している要素に対応する単語を出力すれば良い。

要するにこの作業は、ニューラルネットワークにおける分類問題の最後の作業と全く同じである。

終わりに

モデル構造については以上だ。自分もつい数日前にTaransformerについて手を出し始めたばかりであるから、正確性に欠ける情報を垂れ流してしまっているかもしれない。もしそうだったら指摘してね。

参考文献

原論文 arxiv.org

英語サイト jalammar.github.io

本 Transformerによる自然言語処理 | Denis Rothman, 黒川利明 |本 | 通販 | Amazon