Decoder 6 layers › Cross-attention › Prepare Q, K, V
Take keys
×no parametersk · tf.split(combined, …) · common_attention.py 411
The first features of the fused projection are the keys.
- in
- [B batch, S source positions, 1024 k, v]
- out
- [B batch, S source positions, 512 d]
Codecommon_attention.py 348–421
def multihead_attention(query_antecedent,
memory_antecedent,
bias,
total_key_depth,
total_value_depth,
output_depth,
num_heads,
dropout_rate,
summaries=False,
image_shapes=None,
name=None):
"""Multihead scaled-dot-product attention with input/output transformations.
Args:
query_antecedent: a Tensor with shape [batch, length_q, channels]
memory_antecedent: a Tensor with shape [batch, length_m, channels]
bias: bias Tensor (see attention_bias())
total_key_depth: an integer
total_value_depth: an integer
output_depth: an integer
num_heads: an integer dividing total_key_depth and total_value_depth
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.
Raises:
ValueError: if the key depth or value depth are not divisible by the
number of attention heads.
"""
if total_key_depth % num_heads != 0:
raise ValueError("Key depth (%d) must be divisible by the number of "
"attention heads (%d)." % (total_key_depth, num_heads))
if total_value_depth % num_heads != 0:
raise ValueError("Value depth (%d) must be divisible by the number of "
"attention heads (%d)." % (total_value_depth, num_heads))
with tf.variable_scope(
name,
default_name="multihead_attention",
values=[query_antecedent, memory_antecedent]):
if memory_antecedent is None:
# self attention
combined = common_layers.conv1d(
query_antecedent,
total_key_depth * 2 + total_value_depth,
1,
name="qkv_transform")
q, k, v = tf.split(
combined, [total_key_depth, total_key_depth, total_value_depth],
axis=2)
else:
q = common_layers.conv1d(
query_antecedent, total_key_depth, 1, name="q_transform")
combined = common_layers.conv1d(
memory_antecedent,
total_key_depth + total_value_depth,
1,
name="kv_transform")
k, v = tf.split(combined, [total_key_depth, total_value_depth], axis=2)
q = split_heads(q, num_heads)
k = split_heads(k, num_heads)
v = split_heads(v, num_heads)
key_depth_per_head = total_key_depth // num_heads
q *= key_depth_per_head**-0.5
x = dot_product_attention(
q, k, v, bias, dropout_rate, summaries, image_shapes)
x = combine_heads(x)
x = common_layers.conv1d(x, output_depth, 1, name="output_transform")
return xSources1
Direct evidence
Self-attention projects queries, keys and values with one 1×1 convolution; attention to a memory projects queries from the input and keys and values from the memory; heads are split, queries scaled by depth^-0.5, attended, combined and projected out.
multihead · tensorflow/tensor2tensor / tensor2tensor/models/common_attention.py
Pinned revision 0c66117d4297 ↗ · lines 348–421
Read code excerpt
def multihead_attention(query_antecedent,
memory_antecedent,
bias,
total_key_depth,
total_value_depth,
output_depth,
num_heads,
dropout_rate,
summaries=False,
image_shapes=None,
name=None):
"""Multihead scaled-dot-product attention with input/output transformations.
Args:
query_antecedent: a Tensor with shape [batch, length_q, channels]
memory_antecedent: a Tensor with shape [batch, length_m, channels]
bias: bias Tensor (see attention_bias())
total_key_depth: an integer
total_value_depth: an integer
output_depth: an integer
num_heads: an integer dividing total_key_depth and total_value_depth
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.
Raises:
ValueError: if the key depth or value depth are not divisible by the
number of attention heads.
"""
if total_key_depth % num_heads != 0:
raise ValueError("Key depth (%d) must be divisible by the number of "
"attention heads (%d)." % (total_key_depth, num_heads))
if total_value_depth % num_heads != 0:
raise ValueError("Value depth (%d) must be divisible by the number of "
"attention heads (%d)." % (total_value_depth, num_heads))
with tf.variable_scope(
name,
default_name="multihead_attention",
values=[query_antecedent, memory_antecedent]):
if memory_antecedent is None:
# self attention
combined = common_layers.conv1d(
query_antecedent,
total_key_depth * 2 + total_value_depth,
1,
name="qkv_transform")
q, k, v = tf.split(
combined, [total_key_depth, total_key_depth, total_value_depth],
axis=2)
else:
q = common_layers.conv1d(
query_antecedent, total_key_depth, 1, name="q_transform")
combined = common_layers.conv1d(
memory_antecedent,
total_key_depth + total_value_depth,
1,
name="kv_transform")
k, v = tf.split(combined, [total_key_depth, total_value_depth], axis=2)
q = split_heads(q, num_heads)
k = split_heads(k, num_heads)
v = split_heads(v, num_heads)
key_depth_per_head = total_key_depth // num_heads
q *= key_depth_per_head**-0.5
x = dot_product_attention(
q, k, v, bias, dropout_rate, summaries, image_shapes)
x = combine_heads(x)
x = common_layers.conv1d(x, output_depth, 1, name="output_transform")
return x