More than 3 years have passed since last update.

Transformerにおける相対位置エンコーディングを理解する。

Last updated at 2021-10-02Posted at 2021-09-28

自然言語処理を中心に近年様々な分野にて成功を納めているTransformerでは、入力トークンの位置情報をモデルに考慮させるために「positional encoding（位置エンコーディング）」と呼ばれる処理が必要になります。

この位置エンコーディングは通常、トークン系列の開始時点を基準とした絶対的な位置情報を埋め込みます。
しかし、タスクによってはトークンの絶対的な位置ではなくトークン間の相対的な位置関係をモデルに考慮させたい場合もあると思います。
そういった需要に応えるため、トークン間の相対的な位置を考慮した位置エンコーディング技法（relative position representation）がShawら(2018)により初めて提案され、Chengら(2018)により若干の定義変更とアルゴリズムの効率化がなされました。

本記事は、従来のSelf-Attention機構と位置エンコーディングをおさらいした後、後発のChengらにより提案された相対位置エンコーディングの考え方とアルゴリズムを解説します。

なお、私の確認した既存実装は以下です。

https://github.com/jason9693/MusicTransformer-pytorch (pytorch)
https://github.com/gwinndr/MusicTransformer-Pytorch (pytorch)
tensorflowではtensor2tensorに実装済みとのことです。

Self-Attentionと位置エンコーディングのおさらい

本題に入る前に、従来のSelf-Attentionと位置エンコーディングについて簡単におさらいします。
ここでは簡単のため、Multi-Head Attentionのヘッド数は考慮しません。

Self-Attention

Self-AttentionはTransformerの根幹をなす機構であり、入力トークンの特徴量を他のトークンとの「関係性」を加味しながら変換をしていく仕組みです。
例として、以下の簡単なトークン列をSelf-Attentionで変換する仕組みを図示していきます。

まず、以下に図示するように、固定次元の単語ベクトルへとembeddingされた各トークンを「Query」「Key」「Value」と呼ばれる3種類のベクトルへと線形層で写像します。
Query, Key, Valueそれぞれについて全トークン分をスタックした行列を$Q, K, V$と表記します。

（簡単のため助詞を無視しています）
ここで、あるトークンのQuery, Key, Valueそれぞれの役割を簡単に説明します。

Value: そのトークンから得られる「値」です。
Query: その名の通り「問い合わせ」であり、関連の強い別のトークン（自分自身も含む）を知り、そのValueを得るために存在します。
Key: 自分自身や別のトークンのQueryから問い合わせを受ける「索引」です。問い合わせてきたQueryからの関連性を計算し、それに応じて自身に結びついたValueを返すために存在します。

QueryからKeyへの関連性の強さは内積により表され、その導出は次に示すように行列積$QK^T$を用いて全クエリについて並列に計算されます。
見ていただくと分かるとおり、$QK^T$は全てのトークンの組み合わせ$(i,j)$について「トークン$i$からトークン$j$への関連の強さ」を規定しています。

得られた$QK^T$はsoftmax関数にかけられ、同じ形の行列$A$に変換されます。
これにより各トークン$i$について別のトークンへの関連の強さは"割合"に変換され、その割合に応じてValueの重み付き和が計算されます。
これが、トークン$i$の次の特徴ベクトルとなります。

（絶対）位置エンコーディング

上で説明したSelf-Attention機構ではトークンの位置情報を加味できる仕組みが見当たりません。
このことは、仮に元の文章が「可愛い私の猫は太い」でも「太い猫は私の可愛い」でも全く同じ特徴表現が各トークンについて得られてしまうことを示し、トークンの位置関係に意味がある場合は好ましくありません。
この問題を解決するために、Transformerの原著論文にて「位置エンコーディング」が提案されました。
これは、次に図示するようにトークンの絶対的な位置に応じた固有の位置ベクトルを加算するものです。

トークンの位置を$pos$、ベクトルの次元を$i$として、加算されるベクトルは以下のように定義されます。

相対位置エンコーディング

本題です。
最初にも述べましたが、本記事では後発のChengらにより提案された相対位置エンコーディングを解説します。定義がShawらのものと若干違うため注意してください。

モチベーション

実は、通常の位置エンコードを用いた場合でも、絶対的な位置情報と同時にある程度は相対的な位置情報もモデルは学習できると言われています。
それでも完全に相対位置のみを考慮した位置エンコーディングを用いるモチベーションは、大きく分けて以下の2つになります。

① データの定性的性質から相対位置情報の有用性が明らか

近年では自然言語以外の様々な系列データにTransformerが適用されていますが、その全てのデータにおいて絶対位置情報が重要であるとは限りません。
例えばChengらの論文では音楽の譜面データに対してTransformerを適用しており、その性質について以下のように述べています。

Music has multiple dimensions along which relative differences arguably matter more than their absolute values; the two most prominent are timing and pitch.
（音楽には、絶対的な値よりも相対的な違いが間違いなく重要な複数の次元があり、最も顕著なのはタイミングとピッチの2つです。）

つまり音楽では譜面の絶対的な時刻よりも、相対的な位置関係が重要であると述べています。
私の個人的な考えでは、このように相対位置が重要な系列データは他にもたくさんあると思われます。
例えば音声・画像・映像などの物理世界由来の時空間データは、基本的に相対的な位置関係が重要であり、畳み込みニューラルネットワークが成功している理由の一つもここにあるのではないかと思われます。

② トレーニングデータに現れなかった長い系列データをデコードしたい。

Chengらは論文中の実験で、1. 相対位置エンコード付きTransformer、2. 通常の位置エンコード付きTransformer、3. LSTMの3種類のモデルについて、トレーニングデータに含まれていた譜面の2倍の長さを持つ譜面をデコードさせました。
その結果が以下の図になります。
上の行から順にモデル1, 2, 3が生成した譜面です。

通常の位置エンコードを用いたTransformerはトレーニングデータの譜面長を超えた直後に崩壊を来たしていますが、相対位置エンコード付きTransformerはデコードする譜面長にかかわらず一貫して安定したデータを生成していることがわかります。

実験的に示されたものではありますが、この事実は相対位置エンコードを採用する大きな理由になるのではないでしょうか。

手法

Chengらの相対位置エンコーディングは、以下のように行列$QK^T$に対して新たな行列$S^{rel}$を加算することで実現されます。

行列$QK^T$は、全てのトークンの組み合わせ$(i,j)$について「トークン$i$からトークン$j$への関連の強さ」を規定するものであるとおさらいしました。
対して$S^{rel}$は、全てのトークン$i$について、「トークン$i$と、そこからの距離に基づいて決定される$QK^T$の補正値」です。
この補正値は単純に距離に基づいて決まるのではなく、各トークンそれぞれについて距離に対する補正値が計算されます。
例えば、上の例では、「『猫』から距離1の関係性調整」と「『可愛い』から距離1の関係性調整」の値は別物です。

これらの値は、Self-Attention Layerが内部に学習可能パラメータとして保持する距離embeddingとQueryとの内積として決定されます。
以下は「『猫』から距離1の関係性調整」を決定する計算を図示したものです。
距離embeddingはTransformer内部の各layer・各headについてそれぞれ固有のものを持ちます。