トランスフォーマーの論文を読む⑧埋込

前回に引き続き、論文を読み進めていきます。今回は、埋め込みにフォーカスします。

埋め込みに関しては、エンコーダ・ブロックやデコーダ・ブロックの外側の話になります。下図の上部の「線形層」、「ソフトマックス」と書かれている部分や、下部の「入力文章の埋め込み」、「出力文章の埋め込み」と書かれている部分です。

では、論文のセクション3.4「Embeddings and Softmax」（埋め込みとソフトマックス）を読んでいきましょう。

トークンから埋め込みへ

セクション3.4「Embeddings and Softmax」は次にように始まります。

言語モデルでは、トークン（番号）を埋め込みベクトルに変換するのは、よくある手順です。$${d_\text{model}}$$次元のベクトルがエンコーダやデコーダで使われるというのは、論文の中でもすでに述べられていることではありますが、それが入力シーケンスや出力シーケンスの埋め込みでも同じであることがわかります。

「学習された埋め込みを使用して」というのは、既存の埋め込み（例えば、Word2Vec）などではなく、トランスフォーマーの訓練を通して埋め込みが学習されるということです。

入力トークンから埋め込みベクトルへの変換が簡単なのは分かります。エンコーダはこの埋め込みベクトルに対して処理を施していきます。しかし、デコーダが処理して出力した埋め込みベクトルはどのようにして、次のトークンの予想へと変換されるのでしょうか。

埋め込みから確率分布へ

論文を読み続けましょう。

ここで述べられているのは、次の二つの処理が行われて、デコーダの出力を次のトークンの確率分布へと変換しているということです。

• 学習された線形変換

• ソフトマックス関数

学習された線形変換に対して「通常の」と形容しているのは、特別な処理をしているわけではないということを意味します。つまり、出力された埋め込みベクトルに線形層を適用して、その後にソフトマックス関数を使って確率分布に変換しているだけと言っています。

このことは、デコーダの上部からの処理をみると明らかです。

さて、ソフトマックス関数で確率分布に変換するというのは、全てのトークンに対して確率を付与するということです。つまり、ボキャブラリとして扱えるトークンの総数が10,000個だとすると、10,000個のトークンそれぞれに対して確率を計算します。また、確率として解釈できるようにするために、その全ての確率を足し合わせると１になる必要があります。このような用途ではソフトマックス関数がよく使われます。

以上から理解できるのは、線形変換が埋め込みベクトルを$${d_\text{model}}$$次元からトークンの総数と同じ数の要素を持つベクトルへと変換しているということです。そしてソフトマックス関数を適用すれば、ボキャブラリとして扱える全てのトークンにわたる確率分布を計算できます。

埋め込みの重みの共有

論文は次のように続きます。

「２つの埋め込み層」とは、入力側と出力側の埋め込みのことを指します。図で言うと、下の方にある赤い箱のことを指します。

「ソフトマックス前の線形変換」とは、前述した線形層のことです。

これらが「同じ重み行列を共有」していると言っているのですが、どういうことでしょうか。また、「これは[30]に類似しています」とあるので、参照されている文献から必要な情報を見つける必要があります。

この参照された論文では、LSTMを使った機械翻訳モデルで、同じ重みを以下の三つに対して使う実験をしました。

• エンコーダ：入力文章の埋め込み

• デコーダ：出力文章の埋め込み（デコーダへの入力となる）

• デコーダ：ソフトマックス前の線形層

このように埋め込みの重みを共有することで、必要なパラメータ数を大幅に削減することができる上、共有しない場合と同等のパフォーマンスを発揮しました。

なお、この手法では、異なる言語のトークンを同じ埋め込みで扱うことになります。例えば、フランス語から英語への翻訳の場合、トークンのボキャブラリとしてフランス語と英語の両方が含まれるものが使われます。

注目するべきは、デコーダのソフトマックス前の線形層の重みが埋め込みと同じであることで、デコーダから出力された埋め込みベクトルがどの全てのトークンの埋め込みベクトルとの内積が計算されていることになります。この結果がソフトマックス関数を通ることで、どのトークンに近いのかが確率分布として出力されるわけです。

バイト対符号化

なお、上述の論文では、バイト対符号化（Byte Pair Encoding、BPE）という手法でトークンを生成します。英語、ドイツ語、フランス語などではBPEを使うことでかなりのトークンに重なりが生じるので、埋め込みのパラメータ数を減らすことに貢献しています。BPEや類似したトークン化の手法はトランスフォーマーでも使われています。

BPEを使用することで、これらの言語において共通の語根ごこんや部分語（サブワード、subword）を効果的に特定し、再利用することが可能になります。

たとえば、「run」のような語根ごこんは多くの関連単語で共有されます。例えば、「runner」、「running」などの単語は、BPEによって「run」という共通のサブワードに分割される可能性があります。これにより、これらの単語間で「run」というサブワードの埋め込みが共有され、モデルが単語の意味をより効率的に学習し、未知の単語や新しい単語形式にも対応できるようになります。また、この手法は語彙の多様性を保ちながら、モデルのパラメータ数を削減するのにも役立ちます。

埋め込みのスケーリング

セクション3.4の最後はこう結ばれています。

$${d_{\text{model}}}$$ これまで何度か登場しましたが、モデルの埋め込みベクトルの次元数を示します。

なお、ここで埋め込み層と言っているのは、エンコーダで使われる埋め込み（入力文章の埋め込み）とデコーダで使われる埋め込み（デコーダへの入力に対する埋め込み）です。これらが共有される場合は、同じ埋め込み行列を指します。

また、ソフトマックス前の線形層の重みに対しては行われません。ここでは、内積を計算しどのトークンの埋め込みに似ているのかを計算しているからです。

さて、このように $${(\sqrt{d_{\text{model}}}}$$ でスケーリングをする理由は、論文には書かれていません。一説には、埋め込みベクトルへと追加される位置エンコーディングの値との差をつけるためというものがあります。

TensorFlowのKerasを使ってTransformerを実装するサンプルコードがありますが、その中に次のようなコードがあります。

class PositionalEmbedding(tf.keras.layers.Layer):
  def __init__(self, vocab_size, d_model):
    super().__init__()
    self.d_model = d_model
    self.embedding = tf.keras.layers.Embedding(vocab_size, d_model, mask_zero=True) 
    self.pos_encoding = positional_encoding(length=2048, depth=d_model)

  def compute_mask(self, *args, **kwargs):
    return self.embedding.compute_mask(*args, **kwargs)

  def call(self, x):
    length = tf.shape(x)[1]
    x = self.embedding(x)
    # This factor sets the relative scale of the embedding and positonal_encoding.
    x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
    x = x + self.pos_encoding[tf.newaxis, :length, :]
    return x

その中のコメントで埋め込みベクトルを$${(\sqrt{d_{\text{model}}}}$$ でスケーリングしているところに次のようなコメントがあります。

なので、これがもっともらしい理由でありそうですが、論文の著者が明確にしていないので100％とは言えません。ネットで調べると、必要ないのではないかと考える人もいます。

実際、こちらのKerasによるトランスフォーマーの実装では、このようなスケーリングは行われていません。

class TokenAndPositionEmbedding(layers.Layer):
    def __init__(self, maxlen, vocab_size, embed_dim):
        super().__init__()
        self.token_emb = layers.Embedding(input_dim=vocab_size, output_dim=embed_dim)
        self.pos_emb = layers.Embedding(input_dim=maxlen, output_dim=embed_dim)

    def call(self, x):
        maxlen = ops.shape(x)[-1]
        positions = ops.arange(start=0, stop=maxlen, step=1)
        positions = self.pos_emb(positions)
        x = self.token_emb(x)
        return x + positions

ただし、上述の実装では、オリジナルのトランスフォーマーで使われた位置エンコーディングではなく、位置埋め込みを使っているのでちょっと異なります。位置埋め込みでは、位置エンコーディング自体を学習するのでスケーリングが必要ないのでしょう。

ちょっと、先走って位置エンコーディングに触れてしまいましたが、それについては次回に解説します。

次回予告

次回は、次のセクションである3.5「Positional Encoding」を読み進めます。

お楽しみに！