【０から学ぶAI】第254回：Transformerモデルの実装

前回のおさらいと今回のテーマ

こんにちは！前回は、注意機構（Attention）の実装について解説しました。Attentionは、シーケンスデータ内で重要な情報に注目する技術で、特に自然言語処理（NLP）において長い文脈を捉える際に効果を発揮します。Attentionは、Transformerモデルを支える基盤技術であり、計算効率の向上や文脈情報の効果的な取り込みを可能にしました。

今回は、Transformerモデルの実装について解説します。Transformerは、Attentionを活用した革新的なアーキテクチャで、NLPの分野で幅広く使用されているモデルです。この記事では、Transformerの基本構造、Self-Attentionの応用、そして具体的な実装方法について説明します。

Transformerとは？

1. Transformerの基本概念

Transformerは、2017年にGoogleの研究者によって発表されたNLPモデルであり、従来のリカレントニューラルネットワーク（RNN）や長短期記憶（LSTM）に代わる新しいアプローチです。Transformerの特徴は、完全にAttentionメカニズムに基づいて設計されている点で、RNNのように逐次処理を行わず、シーケンス全体を並列に処理できます。

2. エンコーダ・デコーダ構造

Transformerは、エンコーダとデコーダの2つの主要な部分で構成されています。

• エンコーダ：入力シーケンスを処理して、特徴ベクトルを生成します。複数のエンコーダブロックが積み重なっており、それぞれがSelf-Attentionとフィードフォワードニューラルネットワーク（FFN）を含みます。
• デコーダ：エンコーダの出力をもとに、ターゲットシーケンスを生成します。デコーダも複数のブロックで構成されており、エンコーダの出力とターゲットのSelf-Attentionを組み合わせて予測を行います。

3. Self-Attentionとマルチヘッドアテンション

Transformerの中心には、Self-Attentionがあり、これは入力シーケンス内の単語間の関係性を捉えます。さらに、マルチヘッドアテンションでは、異なる視点から情報を捉えるために複数のSelf-Attentionを同時に実行します。

Transformerの数式と各層の説明

1. Self-Attentionの計算

Self-Attentionは、Query（Q）、Key（K）、Value（V）の3つのベクトルを用いて計算されます。具体的には、以下の式でAttentionを計算します。

[
\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) V
]

ここで、( d_k ) はKeyベクトルの次元数で、スコアのスケーリングに使用します。

2. マルチヘッドアテンション

マルチヘッドアテンションは、複数のSelf-Attentionヘッドを並行して実行し、それらを結合することで多様な情報を捉えます。

[
\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h)W^O
]

各ヘッドは異なる重み行列を使用し、並列的に計算された後、結合されて最終的な出力となります。

3. 位置エンコーディング

Transformerは並列処理のため、入力シーケンスの順序情報を保持する必要があります。これを解決するために、位置エンコーディングを利用します。位置エンコーディングは、入力ベクトルに順序情報を加える方法で、以下のように計算されます。

[
\text{PE}{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d{\text{model}}}}\right)
]
[
\text{PE}{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d{\text{model}}}}\right)
]

ここで、( pos ) は位置、( i ) は次元のインデックスを示します。

Transformerの実装例

ここでは、PythonのTensorFlowを用いて、簡単なTransformerエンコーダモデルを構築します。

1. Transformerエンコーダの実装

まず、Transformerエンコーダの基本的な構造を定義します。Self-Attentionとフィードフォワード層を含むエンコーダブロックを構築します。

import tensorflow as tf
from tensorflow.keras.layers import Layer, Dense, Dropout, LayerNormalization

class MultiHeadAttention(Layer):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        self.num_heads = num_heads
        self.d_model = d_model
        assert d_model % self.num_heads == 0

        self.depth = d_model // num_heads
        self.wq = Dense(d_model)
        self.wk = Dense(d_model)
        self.wv = Dense(d_model)
        self.dense = Dense(d_model)

    def split_heads(self, x, batch_size):
        x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))
        return tf.transpose(x, perm=[0, 2, 1, 3])

    def call(self, v, k, q, mask):
        batch_size = tf.shape(q)[0]

        q = self.wq(q)
        k = self.wk(k)
        v = self.wv(v)

        q = self.split_heads(q, batch_size)
        k = self.split_heads(k, batch_size)
        v = self.split_heads(v, batch_size)

        scaled_attention_logits = tf.matmul(q, k, transpose_b=True) / tf.math.sqrt(tf.cast(self.depth, tf.float32))
        if mask is not None:
            scaled_attention_logits += (mask * -1e9)

        attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)
        output = tf.matmul(attention_weights, v)
        output = tf.transpose(output, perm=[0, 2, 1, 3])
        concat_attention = tf.reshape(output, (batch_size, -1, self.d_model))

        return self.dense(concat_attention)

class TransformerEncoderBlock(Layer):
    def __init__(self, d_model, num_heads, dff, dropout_rate=0.1):
        super(TransformerEncoderBlock, self).__init__()
        self.mha = MultiHeadAttention(d_model, num_heads)
        self.ffn = tf.keras.Sequential([
            Dense(dff, activation='relu'),
            Dense(d_model)
        ])
        self.layernorm1 = LayerNormalization(epsilon=1e-6)
        self.layernorm2 = LayerNormalization(epsilon=1e-6)
        self.dropout1 = Dropout(dropout_rate)
        self.dropout2 = Dropout(dropout_rate)

    def call(self, x, training, mask):
        attn_output = self.mha(x, x, x, mask)
        attn_output = self.dropout1(attn_output, training=training)
        out1 = self.layernorm1(x + attn_output)

        ffn_output = self.ffn(out1)
        ffn_output = self.dropout2(ffn_output, training=training)
        out2 = self.layernorm2(out1 + ffn_output)

        return out2

2. エンコーダを用いたモデルの構築

上記のエンコーダブロックを積み重ね、テキスト分類タスクに応用します。

class TransformerEncoder(tf.keras.Model):
    def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size, dropout_rate=0.1):
        super(TransformerEncoder, self).__init__()
        self.d_model = d_model
        self.num_layers = num_layers

        self.embedding = tf.keras.layers.Embedding(input_vocab_size, d_model)
        self.pos_encoding = positional_encoding(1000, d_model)
        self.enc_layers = [TransformerEncoderBlock(d_model, num_heads, dff, dropout_rate) for _ in range(num_layers)]
        self.dropout = Dropout(dropout_rate)

    def call(self, x, training, mask):
        seq_len = tf.shape(x)[1]
        x = self.embedding(x)
        x *= tf.math.sqrt(tf.cast(self.d_model, tf

.float32))
        x += self.pos_encoding[:, :seq_len, :]

        x = self.dropout(x, training=training)
        for i in range(self.num_layers):
            x = self.enc_layers[i](x, training, mask)

        return x

Transformerのメリット

1. 並列計算が可能

従来のRNNの逐次処理と異なり、Transformerはシーケンス全体を同時に処理するため、計算速度が大幅に向上します。

2. 長いシーケンスでも性能が向上

Self-Attentionにより、長いシーケンスでも重要な情報を保持しやすくなり、長距離依存関係の問題が緩和されます。

まとめ

今回は、Transformerモデルの実装について、その基本構造とSelf-Attentionの役割、具体的な実装例を解説しました。TransformerはNLPの最先端技術の一つであり、多くのタスクで驚異的な性能を発揮しています。次回は、Transformerを応用したBERTのファインチューニングについて解説します。

次回予告

次回は、BERTのファインチューニングについて解説します。事前学習済みモデルを特定のタスクに適用する方法を学びましょう。

---

注釈

1. フィードフォワードニューラルネットワーク（FFN）：入力に対して順方向の計算のみを行うニューラルネットワーク。
2. 位置エンコーディング：シーケンスの順序情報をTransformerに取り込むための手法。
3. Self-Attention：入力シーケンス内の単語同士の関連性を捉えるAttentionの一種。