Self-attentionの計算方法と行列の形状

いまさらではあるものの、Self-attention (自己注意)の計算方法について途中の行列の形状に着目して調べてみました。

結果だけ知りたい方は、最後のまとめに進んでください。

参考文献

元論文はAttention Is All You Needです。

論文の行間を読むか参考文献を遡っていけば、具体的にどういう計算をするのか分かるのかもしれませんが、 論文の最後に、コードへのURL https://github.com/tensorflow/tensor2tensor が記載されていますので、こちらを主に参考にして、計算方法を調べていきます。

Q, K, Vはどうやって計算するの？

論文の式(1)で使われているQ=Query、K=Key、V=Valueの3つの値の計算方法を調べてみます。

論文のFig.1を見ると、入力が3つに分岐しているので、何かをどうにかして入力をQ, K, Vの3つにしていることはわかります。

まず、transformer_encode関数を見てみます。コメントによると、入力であるinputsの形状は(batch_size, input_length, 1, hidden_dim)とのことです。

関数の呼び出し直後に形状を変換していて、結局、inputsは(batch_size, input_length, hidden_dim)になっています。

その後、encoder_functionが呼び出されるのですが、これの中身は、 transformer_encoder です。この関数内の213行目から、common_attention.multihead_attentionを呼び出します。

キャッシュがなく、self-attentionの場合であれば、4650行目にて、compute_qkvが呼び出されます。 ここで、Q, K, Vが計算されているようです。

定義は

def compute_qkv(query_antecedent, ←これがcommon_kayers.layer_preprocess(x, hparams)
                memory_antecedent, ←これがNone
                total_key_depth,
                total_value_depth,
                ...

となっていて、memory_antecedentがNoneであることは、呼び出し元のtransformer_encoder:213に戻るとわかります。

memory_antecedentがNoneならquery_antecedentにしているので、self-attentionの場合、入力はquery_antecedentと考えればよさそうです。

さて、Q, K, V (コード中ではq, k, v)はcompute_attention_componentで計算されていて、入力は

antecedent: a Tensor with shape [batch, length, channels]

戻り値が

c : [batch, length, depth] tensor

となっています。filter_widthによって処理内容が異なるようですが、filter_width == 1のケースを見てみると、 4415行目で

    return common_layers.dense(
        antecedent, total_depth, use_bias=False, name=name,
        layer_collection=layer_collection)

のように書かれています。バイアスなしなので、単にMatMulの計算をしているだけになります。

Tensorflowのドキュメントによると、

Dense implements the operation: output = activation(dot(input, kernel) + bias) <中略> kernel is a weights matrix created by the layer,

なので、

depth channels depth length [Output] = length [antecedent] × channels [W]

となります。ここで、[ ]は行列を表しています。[ ]の左が行数、[ ]の上が列数です。

以上を考慮してcompute_qkvを読むと、q, k, vの形状は

q = (batch, length_q, total_key_depth)
k = (batch, length_kv, total_key_depth)
v = (batch, length_kv, total_value_depth)

となっていることがわかります。

Self-attentionの計算はどうなるの？

multihead_attentionの引数に指定するattention_typeには色々種類があるようですが、デフォルト指定されているdot_productを見てみます。

dot_product_attentionの引数の説明には、

  Args:
    q: Tensor with shape [..., length_q, depth_k].
    k: Tensor with shape [..., length_kv, depth_k]. Leading dimensions must
      match with q.
    v: Tensor with shape [..., length_kv, depth_v] Leading dimensions must
      match with q.

と書かれています。

まず、1648行目で、

logits = tf.matmul(q, k, transpose_b=True)

と計算しています。論文中の式(1)の$QK^T$の部分です。

length_kv depth_k length_kv length_q [logits] = length_q [q] × depth_k [${\rm k}^T$]

のように計算されるので、logitsの形状は (..., length_q, length_kv) となります。

式(1)の$\sqrt{d_k}$で割る部分が見当たりませんが、どこかで計算されているとして、次にsoftmaxの計算をみてみます。これは、1654行目で

weights = tf.nn.softmax(logits, name="attention_weights")

で計算されています。weightsの形状はlogitsと同じです。

ドロップアウトの処理をした後、1667行目の

return tf.matmul(weights, v)

で、式(1)の計算が完了します。これは、

depth_v length_kv depth_v length_q [Attention] = length_q [weights] × length_kv [v]

を計算していますので、この関数の戻り値の形状は(..., length_q, depth_v)となります。

dot_product_attentionのコメント部分にも

 Returns:
    Tensor with shape [..., length_q, depth_v].

と書いてあります。

以上の計算の途中で得られる[weights]の形状が(..., length_q, length_kv)となっており、 self-attentionの場合はlength_q=length_kvですので、系列長の2乗で必要になるメモリや計算量が増えていくことになります。

まとめ

コードを調べたことで、論文の式(1)の各行列の形状は、

$Q \in \mathbb{R}^{L_{\rm q} \times d_{\rm k}}$
$K \in \mathbb{R}^{L_{\rm kv} \times d_{\rm k}}$
$V \in \mathbb{R}^{L_{\rm kv} \times d_{\rm v}}$
$QK^T \in \mathbb{R}^{L_{\rm q} \times L_{\rm kv}}$
$(QK^T)V \in \mathbb{R}^{L_{\rm q} \times d_{\rm v}}$

であることが明確になりました。ここで、$L_{\rm q}$は$Q$の系列長、$L_{\rm kv}$は$K$と$V$の系列長です。 $d_{\rm k}$と$d_{\rm v}$は論文と同じです。

さらに、self-attentionの場合、$L_{\rm q} = L_{\rm kv}$ですので、単に$L$とすれば、

$Q \in \mathbb{R}^{L \times d_{\rm k}}$
$K \in \mathbb{R}^{L \times d_{\rm k}}$
$V \in \mathbb{R}^{L \times d_{\rm v}}$
$QK^T \in \mathbb{R}^{L \times L}$
$(QK^T)V \in \mathbb{R}^{L \times d_{\rm v}}$

のように$L$の添字をなくせるのですっきりします。

$Q$と$K$と$V$は、入力$X \in \mathbb{R}^{L \times C}$をバイアスなしのDenseレイヤーに通すことで得ることができます。 つまり、Denseレイヤーの重み行列をそれぞれ $W_{\rm q} \in \mathbb{R}^{C \times d_{\rm k}}$、 $W_{\rm k} \in \mathbb{R}^{C \times d_{\rm k}}$、 $W_{\rm v} \in \mathbb{R}^{C \times d_{\rm v}}$とすると、

$Q = X W_{\rm q}$
$K = X W_{\rm k}$
$V = X W_{\rm v}$

となります。ここで、$C$は、$X=\{x_1, x_2, ... ,x_i ,... x_L\}$としたときの特徴ベクトル$x_i$の次元数です。