いまさらではあるものの、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行目で、
と計算しています。論文中の式(1)の\(QK^T\)の部分です。logits = tf.matmul(q, k, transpose_b=True)
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の形状はlogitsと同じです。 ドロップアウトの処理をした後、1667行目のweights = tf.nn.softmax(logits, name="attention_weights")
で、式(1)の計算が完了します。これは、return tf.matmul(weights, v)
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\)の次元数です。
