3本目！Attention Is All You Need. - 100-paper-fungoes’s diary

ではでは、いよいよかの有名なAttention Is All You Need について見ていきたいと思います。

まあ、モデルの詳しい説明はググればいっぱいあるので、そっちに置いておいて、訳の分かりにくくなるところを調べていきたいと思います。

Model Architecture

Transformerはgpt-2とは異なり、Encoder-Decoderのモデルです。

まず、Dataについて整理してみると、Encoderのinputがx = (x_1, ... x_n)でoutputが z = (z_1, ..., z_n)となっており、これをDecoderのMulti-Head-AttentionにFeedすることで、y = (y_1, .., y_n)の出力をします。

Encoder and Decoder

EncoderとDecoderは同じ数のN=6で設定されています。それぞれSublayer_1, ..., Sublayer_6まで分かれているっていうことで、EncoderのSublayer_iのoutputをDecoderのSublayer_iへ入力するみたい。

というか、ここでもResidual Connectionを採用していて、本当にResidual Networkは偉大なんだなーと思う。

でも、疑問点としてはDecoderに登場する Masked Multi-Head Attentionとやらがわからん。

実装を見てみよう！

        src_mask = get_pad_mask(src_seq, self.src_pad_idx)
        trg_mask = get_pad_mask(trg_seq, self.trg_pad_idx) & get_subsequent_mask(trg_seq)

def get_pad_mask(seq, pad_idx):
    return (seq != pad_idx).unsqueeze(-2)


def get_subsequent_mask(seq):
    ''' For masking out the subsequent info. '''
    sz_b, len_s = seq.size()
    subsequent_mask = (1 - torch.triu(
        torch.ones((1, len_s, len_s), device=seq.device), diagonal=1)).bool()
    return subsequent_mask

get_pad_maskはmask = [1, 1, 1, 1, ..., 1, 1, 0, 0, 0, ..., 0]とか出力して、tokenがPaddingかどうかのチェックを行なうやつ。というか、Batchでpaddingするときとかに使うやつ。

それでget_subsequent_maskは sequence_length * sequence_lengthの

[tex: \begin{pmatrix}
0&0&..&0&0\\
1&0&..&0&0\\
1&1&0&...&0\\
...\\
1&1&...&0&0\\
1&1&...&1&0\\
\end{pmatrix}
]

という行列ができる。そいで、これがBatch_size個重なっている。

このmaskの使われ方を見てみると両方ともMulti-Head-Attentionで使われており、

self.attention = ScaledDotProductAttention(temperature=d_k ** 0.5)

def forward(self, q, k, v, mask=None):

    d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
    sz_b, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1)

 　residual = q

        # Pass through the pre-attention projection: b x lq x (n*dv)
        # Separate different heads: b x lq x n x dv
    q = self.w_qs(q).view(sz_b, len_q, n_head, d_k)
    k = self.w_ks(k).view(sz_b, len_k, n_head, d_k)
    v = self.w_vs(v).view(sz_b, len_v, n_head, d_v)

    # Transpose for attention dot product: b x n x lq x dv
　q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)

    if mask is not None:
        mask = mask.unsqueeze(1)   # For head axis broadcasting.

    q, attn = self.attention(q, k, v, mask=mask)

    # Transpose to move the head dimension back: b x lq x n x dv
    # Combine the last two dimensions to concatenate all the heads together: b x lq x (n*dv)
    q = q.transpose(1, 2).contiguous().view(sz_b, len_q, -1)
    q = self.dropout(self.fc(q))
    q += residual

    q = self.layer_norm(q)

    return q, attn

class ScaledDotProductAttention(nn.Module):
    ''' Scaled Dot-Product Attention '''

    def __init__(self, temperature, attn_dropout=0.1):
        super().__init__()
        self.temperature = temperature
        self.dropout = nn.Dropout(attn_dropout)

    def forward(self, q, k, v, mask=None):

        attn = torch.matmul(q / self.temperature, k.transpose(2, 3))

        if mask is not None:
            attn = attn.masked_fill(mask == 0, -1e9)

        attn = self.dropout(F.softmax(attn, dim=-1))
        output = torch.matmul(attn, v)

        return output, attn

つまり、ScaledDotProductAttentionというところで、inputと同サイズのattnに対して、mask_token == 0なら、Attentionを $10^ {-9}$ にする。別の言い方をすればAttentionをほぼかけないという仕様になっている。

pytorchのmasked_fillについて見てみると、

Tensor.masked_fill_(mask, value) Fills elements of self tensor with value where mask is True. The shape of mask must be broadcastable with the shape of the underlying tensor.

と書いてある。つまり、Pytorch君が勝手にBrodcastするから、次元のばらつきは1を採用しても、次元が同じならいいよということである。 pad_maskについてはいいとして、subsequenct_maskについてはよくわかっていない。

最終的にはpad_maskのサイズは[Batch_size, 1, 1, Sequence_length]となり、subsequence_maskのサイズは[Batch_size, 1, Sequence_Length, Sequence_Length]となる。

で、maskをかける時点でのAttentionのサイズは[Batch_size, N_head, Len_Q, Len_K]となっている。 N_headでBroadcastするとして、subsequent_maskについては、 k = dec_input, q = dec_input, v=dec_inputなので、Len_Q == Len_K == Len_V == Sequence_Lengthとなっている。

まあ、サイズ的にはOKですね。

ただ、疑問点としては多分このsubsequence_maskはDecoderへの入力を順番に行うっていう意味で、Maskしてると思うんだけど、どうしてこの実装でDecoderのn番目の入力時にn+1番以降をMaskできるんだろう？

モデルの学習方法について見ていくと…。

タスクとしては

input = 私　の　こと　が　好き　です　か　？

target = あなた　が　世界　で　一番　大好き　です　。

として、EncoderにInputを代入し、Decoderには「あなた　が　世界　で　一番　大好き　です　。」を入力して、「あなた　が　世界　で　一番　大好き　です　。」を出力するように学習する。

つまり、両方に共通する「あなた　が　世界　で　一番　大好き　です　。」をそのまま出力しないようにしないといけないってことかー。 GPT-2のPre-Trainingはどうなんだろう？

追記:調べたらGPT-2のPre-trainingでも同じ方法をとってましたー。

まあ、それは置いておいて、実際に使用する際は\<BOS>だけをDecoderに突っ込んで、あとは順番に回帰して\<EOS>が出るまで出力を繰り返せばいいのか。

なんとなくだけど、Taskについてよく考えたらMaskの意味が分かるかも？

ていうか、この写真がわかりやすい。

f:id:kaya-takashiro:20210827224352j:plain — Attention

Encoderは別にAttentionを取り出すだけだから、同じような出力をする、Targetの入力のみにSubsequence_maskが適用されるのか？Seq2Seqみたいに1tokenづつ入力するわけじゃないから、あんなに学習が早いのか！納得。

Multi-Head Attention

まあ、Attentionについてはいろいろとあるから、割愛。でも、Multi-Headとかよくわからん。

実装はさっき見た通り。ただ、Tensorflowの実装のほうが見やすいのでそっちを見ると、

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

    q = self.wq(q)  # (batch_size, seq_len, d_model)
    k = self.wk(k)  # (batch_size, seq_len, d_model)
    v = self.wv(v)  # (batch_size, seq_len, d_model)

    q = self.split_heads(q, batch_size)  # (batch_size, num_heads, seq_len_q, depth)
    k = self.split_heads(k, batch_size)  # (batch_size, num_heads, seq_len_k, depth)
    v = self.split_heads(v, batch_size)  # (batch_size, num_heads, seq_len_v, depth)

    # scaled_attention.shape == (batch_size, num_heads, seq_len_q, depth)
    # attention_weights.shape == (batch_size, num_heads, seq_len_q, seq_len_k)
    scaled_attention, attention_weights = scaled_dot_product_attention(
        q, k, v, mask)

    scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])  # (batch_size, seq_len_q, num_heads, depth)

    concat_attention = tf.reshape(scaled_attention, 
                                  (batch_size, -1, self.d_model))  # (batch_size, seq_len_q, d_model)

    output = self.dense(concat_attention)  # (batch_size, seq_len_q, d_model)

    return output, attention_weights

つまり、Multi-Headとはd_modelのサイズの行列をnum_headsとdepthの大きさに分けて、Attentionを計算することみたい。

何故？

Attentionの分割をしても、別にinputの分割をするわけではないし、ただattention内の重みの影響が互いに少なくなったくらいの影響しかないのかと思った。というか、なんか、それぞれのHeadの重みが異なるっていうのがなんか信じられない。

調べてみると

This allows the model to attend to different “representation sub-spaces” at different positions, akin to using different filters to create different features maps in a single layer in a CNN.

ということで、CNNのFilterと同じような作用をしているってことらしいんだけど、CNNでつかうFilterのmapとか固定値なのに、QueryもKeyもValueも全部可変値やん！とか突っ込み入れたくなる。なんか、Mulit-Headにすると精度は上がるけど、実際はBlackboxっぽい。

EncoderのAttentionについて Q = K = V = EncoderInputとなってる。本当に　self-Attentionだ。 DecoderのAttentionについては、あのsubsequent_maskを突っ込むやつは Q = K= V = Decoder_Inputで、もう一方のEncoderのOutputを突っ込むやつは Q = Decoder_output, K =Encoder_output, V = Encoder_outputとなっている。

QとKで入力文と出力文の関係を見つけ出す！っていうのはいいんだけど、実際Valueはどういった役割してんのかはわからん。QとKの積をsoftmaxかけるってことはそれぞれのHead内で、なんの確率を出してんだ？

たぶん、これが本当の意味でのAttentionの確率っていうことかな。つまり、どこのEncoder_inputの単語に着目するのかっていう。

それをValue=Encoder_inputに掛けることでDecoderの予測値を出すのか？

そういえば、実装を見てみると、

     self.trg_word_prj = nn.Linear(d_model, n_trg_vocab, bias=False)

    def forward(self, src_seq, trg_seq):

        src_mask = get_pad_mask(src_seq, self.src_pad_idx)
        trg_mask = get_pad_mask(trg_seq, self.trg_pad_idx) & get_subsequent_mask(trg_seq)

        enc_output, *_ = self.encoder(src_seq, src_mask)
        dec_output, *_ = self.decoder(trg_seq, trg_mask, enc_output, src_mask)
        seq_logit = self.trg_word_prj(dec_output)
        if self.scale_prj:
            seq_logit *= self.d_model ** -0.5

        return seq_logit.view(-1, seq_logit.size(2))

っていうことだから、出力は[Batch_size, Decoder_sequence_Length, Vocab_size]ってことかな？

まあ、それは置いておいて、DecoderにおけるDecoderのSelf-AttentionとEncoder-Decoder Attentionの組み合わせ方は以下の通り。

  def call(self, x, enc_output, training, 
           look_ahead_mask, padding_mask):
    # enc_output.shape == (batch_size, input_seq_len, d_model)

    attn1, attn_weights_block1 = self.mha1(x, x, x, look_ahead_mask)  # (batch_size, target_seq_len, d_model)
    attn1 = self.dropout1(attn1, training=training)
    out1 = self.layernorm1(attn1 + x)

    attn2, attn_weights_block2 = self.mha2(
        enc_output, enc_output, out1, padding_mask)  # (batch_size, target_seq_len, d_model)
    attn2 = self.dropout2(attn2, training=training)
    out2 = self.layernorm2(attn2 + out1)  # (batch_size, target_seq_len, d_model)

    ffn_output = self.ffn(out2)  # (batch_size, target_seq_len, d_model)
    ffn_output = self.dropout3(ffn_output, training=training)
    out3 = self.layernorm3(ffn_output + out2)  # (batch_size, target_seq_len, d_model)

    return out3, attn_weights_block1, attn_weights_block2

つまり、自分自身にAttentionをた結果に対して、Encoderの結果からAttentionを作成し、そのAttentionをEncoderの結果に掛けているのだね。マジで頭の中がこんがらがる（笑）。

Point-wise Feedforward-Network

Attention以外にも全結合層を入れてる。表現力を高めているのかな？なんか、Transformerってこうして見てみると、Seq2seqよりもCNNに近い気がする。

Positinal Encoding

やってまいりました。一番わけのわかんないやつ。 ReccurentもCNNもない（つまり、ベースはFNNだけ）のモデルがTransformerであり、Tokenの相対位置or絶対位置を入れるために、入れたとある。

よくわからんから、実装を見てみる。

def get_angles(pos, i, d_model):
  angle_rates = 1 / np.power(10000, (2 * (i//2)) / np.float32(d_model))
  return pos * angle_rates

def positional_encoding(position, d_model):
  angle_rads = get_angles(np.arange(position)[:, np.newaxis],
                          np.arange(d_model)[np.newaxis, :],
                          d_model)

  # 配列中の偶数インデックスにはsinを適用; 2i
  angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])

  # 配列中の奇数インデックスにはcosを適用; 2i+1
  angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])

  pos_encoding = angle_rads[np.newaxis, ...]

  return tf.cast(pos_encoding, dtype=tf.float32)

これをグラフに直すとこうなるらしい。

f:id:kaya-takashiro:20210827215839j:plain — positionalEncoding

グラフのpositionがtokenの位置で、DepthがEmbedding次の次元？んなわけないか。

    self.embedding = tf.keras.layers.Embedding(input_vocab_size, d_model)
    self.pos_encoding = positional_encoding(maximum_position_encoding, 
                                            self.d_model)

実装の中では上記のように定義していて

    # 埋め込みと位置エンコーディングを合算する
    x = self.embedding(x)  # (batch_size, input_seq_len, d_model)
    x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
    x += self.pos_encoding[:, :seq_len, :]

っていう感じで使用している。つまり、d_modelのサイズ = positonal_encodingのdepth, seq_len = positional_encodingのpositionってことか。

つまり、各Depthにおいて、token間の距離 = positionの相対差が一定になっているんじゃね？問題は各Depth間の関係がわからん。

あっ、条件は2つだけか！ 1.各positionのVectorは重複しない一意なものであること。 2.同じ距離のPostion間のVectorは一定であること。

なんか、これを満たしてるっぽい気がする。

実験してみる。

    encoding = positional_encoding(10000, 1024)

    print(encoding[0] - encoding[1])
    print(encoding[100] - encoding[101])

    print(np.linalg.norm(encoding[0] - encoding[1]))
    print(np.linalg.norm(encoding[100] - encoding[101]))

    print(np.linalg.norm(encoding[0] - encoding[100]))
    print(np.linalg.norm(encoding[1000] - encoding[1100]))

Result
[-8.41470985e-01  4.59697694e-01 -8.31705202e-01 ...  5.37303912e-09
 -1.01815172e-04  5.18316468e-09]
[-9.58391428e-01 -2.96859975e-02  2.35073541e-01 ...  1.07996133e-06
 -1.01809842e-04  1.04179791e-06]
5.208103571044113
5.208103571044114
24.015006646071527
24.01500664607152

確かに距離が一定である！すげー。マジで。

つまり、word2vecでの king - man = queen - woman と同じようにVectorがマップされるのか。Wow!

次回はBERT: Pre-training of Deep Bidirectional Transformers for Language Understandingっていう論文に行こう！

そうしたら、Dialogue系の論文探すか…。