はじめに

今回はT5（Text-to-Text Transfer Transformer）という自然言語処理に関するモデルをゼロから実装してみました。
概要について理解している方は、コード部分からお読みください。

T5とは

T5（Text-to-Text Transfer Transformer）は、自然言語処理（NLP）タスクに対する汎用的なアーキテクチャであるTransformerをベースにしたモデルです。T5のアーキテクチャは比較的シンプルで、Encoder-DecoderのTransformerモデルに基づいています。大規模なテキストコーパスで事前訓練されており、多くの一般的な知識を獲得しています。これにより、特定のタスクに適応させる際の性能が向上します。

アーキテクチャ

基本的にはtransformerのエンコーダ・デコーダモデルの構造に似通っています。以下に各モジュールについて説明します。

inputs

EncoderLayerに入力するinputsは"Translate the following English text into French: Hello, how are you? target:"などの入力文をトークン化、idをEmbedding層に入力しすることで埋め込みベクトルを得ます。

Encoder

Encoderは以下のモジュールで構成された各Layerで構成されます。

Multi-HeadAttention
Add&Norm
FeedForward

Multi-HeadAttention（Encoder）

Encoder内のMulti-HeadAttentionでは、query, key, valueが全てinputsの埋め込みベクトルから作成されています。このモジュールの役割はinputsの文章内に文脈、関係性などを把握するというものです。MultiHeadということもあり、各ヘッドごとに埋め込みベクトルの次元数より少ない次元数のquery, key, valueを作成し、最後にそれら全てを結合することで、異なる視点からの特徴を捉えることが出来ます。具体的な構造は以下のようになります。Dは埋め込みベクトルの次元数、Mはヘッドの個数です。

\begin{align}
query^{(m)}_{i} = W^{(m)}_{query}h_{i} \\
key^{(m)}_{j} = W^{(m)}_{key}h_{j}\\
value^{(m)}_{j} = W^{(m)}_{value}h_{j}\\
s^{(m)}_{ij} = \frac{query^{(m)}_{i}key_{j}^{(m)T}}{\sqrt{\frac{D}{M}}}\\
a^{(m)}_{ij} = \frac{exp(s^{(m)}_{ij})}{\sum_{j^{'}}^{\frac{D}{M}}exp(s^{(m)}_{ij^{'}})}\\
o_{i} = \sum_{j=1}^{\frac{D}{M}}a^{(m)}_{ij}value^{(m)}_{j}
\end{align}

Add&Norm

このモジュールはスキップ結合を使用して、Multi-HeadAttentionと一つ前の入力を加算することで勾配爆発、消失を防ぐ役割を果たし、またLayerNormalizationを使用することでパラメータを調整しています。

FeedForward

このモジュールは単なる全結合層と活性化関数の集合です。具体的には以下のような構成となっています。fはreluなどの活性化関数を使用します。

FFN(x) = W_{2}f(W_{1}x + b_{1})+b_{2}

outputs

これは"Bonjour, comment ça va ?"といったinputsの対する回答文や翻訳文で構成されています。inputsと同様にトークン化、idをEmbedding層に入力し、埋め込みベクトルが得ます。

Decoder

Decoderは主に以下のモジュールが各Layerに構成されます。

MaskedMulti-HeadAttention
Add&Norm
Multi-HeadAttention
FeedForward

Decoderはこれらのモジュールで構成されたLayerが複数個集まったモジュールになります。

MaskedMulti-HeadAttention

このモジュールは、各位置に対してマスクを用いてその位置より未来の情報を隠し、複数のヘッドで同時に異なるアテンション重みを計算、それらを結合して出力します。

Multi-HeadAttention（Decoder）

基本的にはEncoderのMulti-HeadAttentionと同じ構造です。しかし、queryにoutputsの出力ベクトル、keyとvalueにEncoderの最終Layerから抽出した出力ベクトルを使用します。これによりoutputsと関連性の高い単語や表現をinputsから抽出することが可能となります。

相対位置埋め込み

BertやGPTといった既存のtransfomerのEncoder、decoderは文章の位置における特徴を捉えるため、埋め込みベクトルに対して絶対位置埋め込みを行っていました。しかし今回のT5では、トークン同士の距離をとらえた相対位置埋め込みを使用しています。具体的には以下のような処理を行っています。相対位置におけるスコアはスカラー値となります。

s^{(m)}_{ij} = \frac{query^{(m)}_{i}・key^{(m)T}_{j}}{\sqrt{\frac{D}{M}}} + p^{(m)}_{\lvert i-j\rvert}

この相対位置埋め込みはEncoderのMulti-HeadAttentionとDecoderのMaskedMulti-HeadAttentionに使用されています。

コード

以下が実装したコードになります。

警告
以下は個人で開発したコードです。誤った部分や、冗長な部分がある可能性があります。

Add&Norm

Add&Norm

class AddNorm(nn.Module):
    def __init__(self, batch_size, seq_length, hidden_size, check_encoder):
        super().__init__()
        self._setupAddNormModule(batch_size, seq_length, hidden_size, check_encoder)

    def _setupAddNormModule(self, batch_size, seq_length, hidden_size, check_encoder):
        if check_encoder: seq_length = seq_length[0]
        else: seq_length = seq_length[1]
        self.layer_norm = nn.LayerNorm((batch_size, seq_length, hidden_size))

    def forward(self, tokens, skipped_tokens):
        tokens += skipped_tokens
        tokens = self.layer_norm(tokens)

        return tokens

FeedForward

FeedForward

class FeedForward(nn.Module):
    def __init__(self, hidden_size, ffn_hidden_size):
        super().__init__()
        self._setupFeedForwardModule(hidden_size, ffn_hidden_size)

    def _setupFeedForwardModule(self, hidden_size, ffn_hidden_size):
        dense1 = nn.Linear(hidden_size, ffn_hidden_size)
        relu1 = nn.ReLU()
        dense2 = nn.Linear(ffn_hidden_size, hidden_size)
        self.feed_foward_module = nn.ModuleList([dense1, relu1, dense2])

    def forward(self, tokens):
        for module in self.feed_foward_module:
            tokens = module(tokens)
        return tokens

Multi-HeadAttention（マスク化も可能）

class MultiHeadAttention(nn.Module):
    def __init__(self, batch_size, num_heads, seq_length, hidden_size, check_positional_embedding, check_mask):
        super().__init__()
        self._setupHeadQKV(num_heads, hidden_size)
        self.batch_size = batch_size
        self.num_heads = num_heads
        self.seq_length = seq_length
        self.hidden_size = hidden_size
        self.check_positional_embedding = check_positional_embedding
        self.check_mask = check_mask
        self.softmax = nn.Softmax(dim=-1)

    #ヘッド毎にquery, key, valueを設定
    def _setupHeadQKV(self, num_heads, hidden_size):
        query_module = []
        key_module = []
        value_module = []
        head_hidden_size = int(hidden_size / num_heads)

        for _ in range(num_heads):
            query_module.append(nn.Linear(hidden_size, head_hidden_size))
            key_module.append(nn.Linear(hidden_size, head_hidden_size))
            value_module.append(nn.Linear(hidden_size, head_hidden_size))

        self.query_module = nn.ModuleList(query_module)
        self.key_module = nn.ModuleList(key_module)
        self.value_module = nn.ModuleList(value_module)

    #相対位置埋め込みのスカラー値を計算
    def _outputRelativePositionalEmbeddingScalar(self, query, batch_size, seq_length, hidden_size, num_heads):
        if self.check_mask: seq_length = self.seq_length[1]
        else: seq_length = self.seq_length[0]

        embed_Module = []
        head_hidden_size = int(hidden_size / num_heads)
        position_ids = torch.tensor(list(range(seq_length)), dtype=torch.long).reshape(1, seq_length).expand(batch_size, seq_length)
        for id in range(num_heads): embed_Module.append(nn.Embedding(seq_length, head_hidden_size))
        self.embed_module = nn.ModuleList(embed_Module)
        for id in range(num_heads):
            head_query = self.query_module[id](query)
            tmp_relative_position_embedding_scalar = (head_query@(self.embed_module[id](position_ids).transpose(1, 2))).reshape(1, batch_size, seq_length, seq_length)
            if id == 0: relative_position_embedding_scalar = tmp_relative_position_embedding_scalar
            else: relative_position_embedding_scalar = torch.concat([relative_position_embedding_scalar, tmp_relative_position_embedding_scalar], dim=0)

        return relative_position_embedding_scalar
    
    #Attentionモジュールの出力
    def _outputAttention(self, query, key, value, batch_size, seq_length, hidden_size, num_heads, check_positional_embedding, check_mask, encoder_attention_mask, decoder_attention_mask):
        #どのAttentionモジュールかを判断、系列長を設定
        if check_positional_embedding:
            if check_mask:
                seq_length1 = seq_length2 = self.seq_length[1]
            else:
                seq_length1 = seq_length2 = self.seq_length[0]

        else:
            seq_length1 = self.seq_length[1]
            seq_length2 = self.seq_length[0]

        head_hidden_size = int(hidden_size / num_heads)
        encoder_attention_mask = encoder_attention_mask.reshape(batch_size, -1, 1).expand(batch_size, seq_length1, seq_length2)
        decoder_attention_mask = decoder_attention_mask.reshape(batch_size, 1, -1).expand(batch_size, seq_length1, seq_length2)
        #paddingされたベクトルを除外するmap
        padding_map = encoder_attention_mask * decoder_attention_mask
        #マスク化に使用するmap
        mask_map = torch.tensor(np.tril(np.ones((seq_length1, seq_length2))), dtype=torch.long)
        #相対位置埋め込みの計算
        if check_positional_embedding:relative_position_embedding_scalar = self._outputRelativePositionalEmbeddingScalar(query, batch_size, seq_length, hidden_size, num_heads)
        else: relative_position_embedding_scalar = torch.zeros_like(padding_map, dtype=torch.float).expand(num_heads, batch_size, seq_length1, seq_length2)

        for id in range(num_heads):
            head_query = self.query_module[id](query)
            head_key = self.key_module[id](key)
            head_value = self.value_module[id](value)
            #Attentionの計算
            if check_mask: tmp_head_attention = self.softmax(padding_map * (mask_map * (head_query@head_key.transpose(1, 2)) / torch.sqrt(torch.tensor(head_hidden_size)) + relative_position_embedding_scalar[id]))@head_value
            else:tmp_head_attention = self.softmax(padding_map * ((head_query@head_key.transpose(1, 2)) / torch.sqrt(torch.tensor(head_hidden_size)) + relative_position_embedding_scalar[id]))@head_value
            #各ヘッドの出力を結合
            if id == 0: head_attention = tmp_head_attention
            else: head_attention = torch.concat([head_attention, tmp_head_attention], dim=-1)
        output_attention = head_attention

        return output_attention

    def forward(self, query, key, value, encoder_attention_mask, decoder_attention_mask):
        output_attention = self._outputAttention(query, key, value, self.batch_size, self.seq_length, self.hidden_size, self.num_heads,
                                                  self.check_positional_embedding, self.check_mask, encoder_attention_mask, decoder_attention_mask)

        return output_attention

EncoderLayer

EncoderLayer

class EncoderLayer(nn.Module):
    #各モジュールを設定
    def __init__(self, batch_size, num_heads, seq_length, hidden_size, ffn_hidden_size):
        super().__init__()
        self.multi_head_attention = MultiHeadAttention(batch_size, num_heads, seq_length, hidden_size, check_positional_embedding=True, check_mask=False)
        self.add_norm1 = AddNorm(batch_size, seq_length, hidden_size, check_encoder=True)
        self.feed_forward = FeedForward(hidden_size, ffn_hidden_size)
        self.add_norm2 = AddNorm(batch_size, seq_length, hidden_size, check_encoder=True)
        self.dropout = nn.Dropout(p=0.1)

    def forward(self, tokens, encoder_attention_mask):
        skip1 = tokens
        multi_head_attention = self.multi_head_attention(tokens, tokens, tokens, encoder_attention_mask, encoder_attention_mask)
        add_norm1 = self.add_norm1(self.dropout(multi_head_attention), skip1)
        skip2 = add_norm1
        feed_forward = self.feed_forward(add_norm1)
        add_norm2 = self.add_norm2(self.dropout(feed_forward), skip2)
        tokens = add_norm2

        return tokens

Encoder

Encoder

class Encoder(nn.Module):
    def __init__(self, batch_size, seq_length, hidden_size, num_layer, num_heads, ffn_hidden_size):
        super().__init__()
        self._setupEncoderLayer(batch_size, seq_length, hidden_size, num_layer, num_heads, ffn_hidden_size)

    #EncoderLayerを指定回数重ねる
    def _setupEncoderLayer(self, batch_size, seq_length, hidden_size, num_layer, num_heads, ffn_hidden_size):
        encoder_layer_list = []
        for _ in range(num_layer):
            encoder_layer = EncoderLayer(batch_size, num_heads, seq_length, hidden_size, ffn_hidden_size)
            encoder_layer_list.append(encoder_layer)
        self.encoder_module = nn.ModuleList(encoder_layer_list)

    def forward(self, encoder_embedding, encoder_attention_mask):
        tokens = encoder_embedding
        for encoder_layer in self.encoder_module:
            tokens = encoder_layer(tokens, encoder_attention_mask)
        output_encoder = tokens

        return output_encoder

DecoderLayer

DecoderLayer

    
class DecoderLayer(nn.Module):
    #各モジュールを設定
    def __init__(self, batch_size, num_heads, seq_length, hidden_size, ffn_hidden_size):
        super().__init__()
        self.masked_multi_head_attention =  MultiHeadAttention(batch_size, num_heads, seq_length, hidden_size, check_positional_embedding=True, check_mask=True)
        self.add_norm1 = AddNorm(batch_size, seq_length, hidden_size, check_encoder=False)
        self.cross_multi_head_attention = MultiHeadAttention(batch_size, num_heads, seq_length, hidden_size, check_positional_embedding=False, check_mask=False)
        self.add_norm2 = AddNorm(batch_size, seq_length, hidden_size, check_encoder=False)
        self.feed_forward = FeedForward(hidden_size, ffn_hidden_size)
        self.add_norm3 = AddNorm(batch_size, seq_length, hidden_size, check_encoder=False)
        self.dropout = nn.Dropout(p=0.1)

    def forward(self, tokens, output_encoder, encoder_attention_mask, decoder_attention_mask):
        skip1 = tokens
        masked_multi_head_attention = self.masked_multi_head_attention(tokens, tokens, tokens, decoder_attention_mask, decoder_attention_mask)
        add_norm1 = self.add_norm1(self.dropout(masked_multi_head_attention), skip1)
        skip2 = add_norm1
        cross_multi_head_attention = self.cross_multi_head_attention(tokens, output_encoder, output_encoder, decoder_attention_mask, encoder_attention_mask)
        add_norm2 = self.add_norm2(self.dropout(cross_multi_head_attention), skip2)
        skip3 = add_norm2
        feed_forward = self.feed_forward(tokens)
        add_norm3 = self.add_norm3(self.dropout(feed_forward), skip3)
        tokens = add_norm3

        return tokens

Decoder

Decoder

class Decoder(nn.Module):
    def __init__(self, batch_size, seq_length, hidden_size, num_layer, num_heads, ffn_hidden_size):
        super().__init__()
        self._setupDecoderLayer(batch_size, seq_length, hidden_size, num_layer, num_heads, ffn_hidden_size)

    #DecoderLayerを指定回数重ねる
    def _setupDecoderLayer(self, batch_size, seq_length, hidden_size, num_layer, num_heads, ffn_hidden_size):
        decoder_list = []
        for _ in range(num_layer):
            decoder_layer = DecoderLayer(batch_size, num_heads, seq_length, hidden_size, ffn_hidden_size)
            decoder_list.append(decoder_layer)
        self.decoder_module = nn.ModuleList(decoder_list)

    def forward(self, decoder_embedding, output_encoder, encoder_attention_mask, decoder_attention_mask):
        tokens = decoder_embedding
        for decoder_layer in self.decoder_module:
            tokens = decoder_layer(tokens, output_encoder, encoder_attention_mask, decoder_attention_mask)
        output_decoder = tokens

        return output_decoder

T5（Encoder-Decoder）

T5

class T5(nn.Module):
    def __init__(self, batch_size, num_embedding, seq_length, hidden_size, num_layer, num_heads, ffn_hidden_size):
        super().__init__()
        self.batch_size = batch_size
        self.seq_length = seq_length
        self.hidden_size = hidden_size
        self.num_layer = num_layer
        self.num_heads = num_heads
        self.ffn_hidden_size = ffn_hidden_size

        #Encoder, Decoderの埋め込み層を設定
        self.encoder_embedding = nn.Embedding(num_embedding, hidden_size, padding_idx=0)
        self.decoder_embedding = nn.Embedding(num_embedding, hidden_size, padding_idx=0)
        self.encoder = Encoder(batch_size, seq_length, hidden_size, num_layer, num_heads, ffn_hidden_size)
        self.decoder = Decoder(batch_size, seq_length, hidden_size, num_layer, num_heads, ffn_hidden_size)
        


    def forward(self, encoder_input_ids, decoder_input_ids, encoder_attention_mask, decoder_attention_mask):
        encoder_embedding = self.encoder_embedding(encoder_input_ids)
        output_encoder = self.encoder(encoder_embedding, encoder_attention_mask)
        decoder_embedding = self.decoder_embedding(decoder_input_ids)
        decoder_output = self.decoder(decoder_embedding, output_encoder, encoder_attention_mask, decoder_attention_mask)

        return decoder_output #下流タスクは設定していない

実行結果

以下がT5を実行した結果になります。

実行

from transformers import T5Tokenizer
from pprint import pprint
impot torch
import torch.nn as nn
import numpy as np

input = ["Translate the following English text into French: Hello, how are you? target:",
         "Translate the following English text into French: Good morning, everyone. target",
          "Translate the following English text into French: Can you help me with this? target" ]

target = ["Bonjour, comment ça va ?",
          "Bonjour à tous.",
          "Pouvez-vous m'aider avec ceci ?"]

tokenizer = T5Tokenizer.from_pretrained("t5-small")

encoder_tokenize = tokenizer(input, return_tensors="pt", padding=True)
decoder_tokenize = tokenizer(target, return_tensors="pt", padding=True)

encoder_input_ids = encoder_tokenize.input_ids
encoder_attention_mask = encoder_tokenize.attention_mask
decoder_input_ids = decoder_tokenize.input_ids
decoder_attention_mask = decoder_tokenize.attention_mask

batch_size = encoder_input_ids.size(0)
encoder_seq_length = encoder_input_ids.size(1)
decoder_seq_length = decoder_input_ids.size(1)
max_ids = torch.max(torch.max(encoder_input_ids), torch.max(decoder_input_ids)) #仮の埋め込み辞書数

kwargs = {
    "batch_size": batch_size,
    "num_embedding": max_ids + 1,
    "seq_length": (encoder_seq_length, decoder_seq_length),
    "hidden_size": 768,
    "num_layer": 12,
    "num_heads": 12,
    "ffn_hidden_size": 3072,

}
    
model = T5(**kwargs)
outputs = model(encoder_input_ids, decoder_input_ids, encoder_attention_mask, decoder_attention_mask)
print(model)
print(model(encoder_input_ids, decoder_input_ids, encoder_attention_mask, decoder_attention_mask).shape)

結果

T5(
  (encoder_embedding): Embedding(30356, 768, padding_idx=0)
  (decoder_embedding): Embedding(30356, 768, padding_idx=0)
  (encoder): Encoder(
    (encoder_module): ModuleList(
      (0-11): 12 x EncoderLayer(
        (multi_head_attention): MultiHeadAttention(
          (query_module): ModuleList(
            (0-11): 12 x Linear(in_features=768, out_features=64, bias=True)
          )
          (key_module): ModuleList(
            (0-11): 12 x Linear(in_features=768, out_features=64, bias=True)
          )
          (value_module): ModuleList(
            (0-11): 12 x Linear(in_features=768, out_features=64, bias=True)
          )
          (softmax): Softmax(dim=-1)
          (embed_module): ModuleList(
            (0-11): 12 x Embedding(18, 64)
          )
        )
        (add_norm1): AddNorm(
          (layer_norm): LayerNorm((3, 18, 768), eps=1e-05, elementwise_affine=True)
        )
        (feed_forward): FeedForward(
...
    )
  )
)
torch.Size([3, 13, 768])

無事T5のモデルを実行、出力することが出来ました。

まとめ

T5の実装を通して、Attention機構の実装の難しさを再度実感しました。相対位置埋め込みやマスク化といった処理を加えることでViTと比べて実装の難易度が何倍にも跳ね上がったように感じました。次は自然言語処理から少し離れてVision-Languageのモデルについて勉強、実装したいと思います。

参考文献

Transformerとは？AI自然言語学習の技術を解説
 Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer