Transformer · Softmax · OpenModelAtlas

Decoder 6 layers › Masked attention › Dot-product attention

Softmax

no parameterstf.nn.softmax(logits, name="attention_weights") · 340

Normalize each query row over the keys.

{\htmlData{light=weights softmax}{A}}[b,i,t,s]=\frac{\exp({\htmlData{light=masked add-mask}{\tilde S}}[b,i,t,s])}{\sum_{u=0}^{T-1}\exp({\htmlData{light=masked add-mask}{\tilde S}}[b,i,t,u])}

in: $\tilde S$ [B batch, 8 heads, T queries, T keys]
out: $A$ [B batch, 8 heads, T queries, T keys]

Explanation

$b$ , $i$ , $t$ and $s$ index batch, head, query and key; $u$ runs over the keys.

Codecommon_attention.py 310–345

common_attention.py 310–345

def dot_product_attention(q,
                          k,
                          v,
                          bias,
                          dropout_rate=0.0,
                          summaries=False,
                          image_shapes=None,
                          name=None):
  """dot-product attention.

  Args:
    q: a Tensor with shape [batch, heads, length_q, depth_k]
    k: a Tensor with shape [batch, heads, length_kv, depth_k]
    v: a Tensor with shape [batch, heads, length_kv, depth_v]
    bias: bias Tensor (see attention_bias())
    dropout_rate: a floating point number
    summaries: a boolean
    image_shapes: optional tuple of integer scalars.
      see comments for attention_image_summary()
    name: an optional string

  Returns:
    A Tensor.
  """
  with tf.variable_scope(
      name, default_name="dot_product_attention", values=[q, k, v]):
    # [batch, num_heads, query_length, memory_length]
    logits = tf.matmul(q, k, transpose_b=True)
    if bias is not None:
      logits += bias
    weights = tf.nn.softmax(logits, name="attention_weights")
    # dropping out the attention links for each of the heads
    weights = tf.nn.dropout(weights, 1.0 - dropout_rate)
    if summaries and not tf.get_variable_scope().reuse:
      attention_image_summary(weights, image_shapes)
    return tf.matmul(weights, v)

Sources1

Direct evidence

Scaled scores are the product of queries and keys plus the bias; softmax over keys, dropout on the weights, then the weighted sum of values.

dot-product · tensorflow/tensor2tensor / tensor2tensor/models/common_attention.py
Pinned revision 0c66117d4297 ↗ · lines 310–345

Read code excerpt

def dot_product_attention(q,
                          k,
                          v,
                          bias,
                          dropout_rate=0.0,
                          summaries=False,
                          image_shapes=None,
                          name=None):
  """dot-product attention.

  Args:
    q: a Tensor with shape [batch, heads, length_q, depth_k]
    k: a Tensor with shape [batch, heads, length_kv, depth_k]
    v: a Tensor with shape [batch, heads, length_kv, depth_v]
    bias: bias Tensor (see attention_bias())
    dropout_rate: a floating point number
    summaries: a boolean
    image_shapes: optional tuple of integer scalars.
      see comments for attention_image_summary()
    name: an optional string

  Returns:
    A Tensor.
  """
  with tf.variable_scope(
      name, default_name="dot_product_attention", values=[q, k, v]):
    # [batch, num_heads, query_length, memory_length]
    logits = tf.matmul(q, k, transpose_b=True)
    if bias is not None:
      logits += bias
    weights = tf.nn.softmax(logits, name="attention_weights")
    # dropping out the attention links for each of the heads
    weights = tf.nn.dropout(weights, 1.0 - dropout_rate)
    if summaries and not tf.get_variable_scope().reuse:
      attention_image_summary(weights, image_shapes)
    return tf.matmul(weights, v)