attention_layer.py 6.78 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Implementation of multiheaded attention and self-attention layers."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
21
22
import math

23
import tensorflow as tf
24
from official.nlp.modeling import layers
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50


class Attention(tf.keras.layers.Layer):
  """Multi-headed attention layer."""

  def __init__(self, hidden_size, num_heads, attention_dropout):
    """Initialize Attention.

    Args:
      hidden_size: int, output dim of hidden layer.
      num_heads: int, number of heads to repeat the same attention structure.
      attention_dropout: float, dropout rate inside attention for training.
    """
    if hidden_size % num_heads:
      raise ValueError(
          "Hidden size ({}) must be divisible by the number of heads ({})."
          .format(hidden_size, num_heads))

    super(Attention, self).__init__()
    self.hidden_size = hidden_size
    self.num_heads = num_heads
    self.attention_dropout = attention_dropout

  def build(self, input_shape):
    """Builds the layer."""
    # Layers for linearly projecting the queries, keys, and values.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
51
    size_per_head = self.hidden_size // self.num_heads
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
52
53
54
55
56
57
58

    def _glorot_initializer(fan_in, fan_out):
      limit = math.sqrt(6.0 / (fan_in + fan_out))
      return tf.keras.initializers.RandomUniform(minval=-limit, maxval=limit)

    attention_initializer = _glorot_initializer(input_shape[-1],
                                                self.hidden_size)
59
60
    self.query_dense_layer = layers.DenseEinsum(
        output_shape=(self.num_heads, size_per_head),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
61
        kernel_initializer=attention_initializer,
62
63
64
65
        use_bias=False,
        name="query")
    self.key_dense_layer = layers.DenseEinsum(
        output_shape=(self.num_heads, size_per_head),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
66
        kernel_initializer=attention_initializer,
67
68
69
70
        use_bias=False,
        name="key")
    self.value_dense_layer = layers.DenseEinsum(
        output_shape=(self.num_heads, size_per_head),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
71
        kernel_initializer=attention_initializer,
72
73
        use_bias=False,
        name="value")
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
74
75

    output_initializer = _glorot_initializer(self.hidden_size, self.hidden_size)
76
77
78
    self.output_dense_layer = layers.DenseEinsum(
        output_shape=self.hidden_size,
        num_summed_dimensions=2,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
79
        kernel_initializer=output_initializer,
80
81
        use_bias=False,
        name="output_transform")
82
83
84
85
86
87
88
89
90
    super(Attention, self).build(input_shape)

  def get_config(self):
    return {
        "hidden_size": self.hidden_size,
        "num_heads": self.num_heads,
        "attention_dropout": self.attention_dropout,
    }

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
91
92
93
  def call(self, query_input, source_input, bias, training, cache=None,
           decode_loop_step=None):
    """Apply attention mechanism to query_input and source_input.
94
95

    Args:
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
96
97
98
99
100
      query_input: A tensor with shape [batch_size, length_query, hidden_size].
      source_input: A tensor with shape [batch_size, length_source,
        hidden_size].
      bias: A tensor with shape [batch_size, 1, length_query, length_source],
        the attention bias that will be added to the result of the dot product.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
101
102
103
      training: A bool, whether in training mode or not.
      cache: (Used during prediction) A dictionary with tensors containing
        results of previous attentions. The dictionary must have the items:
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
104
105
106
107
            {"k": tensor with shape [batch_size, i, heads, dim_per_head],
             "v": tensor with shape [batch_size, i, heads, dim_per_head]}
        where i is the current decoded length for non-padded decode, or max
        sequence length for padded decode.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
108
109
      decode_loop_step: An integer, step number of the decoding loop. Used only
        for autoregressive inference on TPU.
110
111

    Returns:
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
112
      Attention layer output with shape [batch_size, length_query, hidden_size]
113
    """
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
114
    # Linearly project the query, key and value using different learned
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
115
116
117
118
119
    # projections. Splitting heads is automatically done during the linear
    # projections --> [batch_size, length, num_heads, dim_per_head].
    query = self.query_dense_layer(query_input)
    key = self.key_dense_layer(source_input)
    value = self.value_dense_layer(source_input)
120
121
122

    if cache is not None:
      # Combine cached keys and values with new keys and values.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
123
124
125
126
      if decode_loop_step is not None:
        cache_k_shape = cache["k"].shape.as_list()
        indices = tf.reshape(
            tf.one_hot(decode_loop_step, cache_k_shape[1], dtype=key.dtype),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
127
            [1, cache_k_shape[1], 1, 1])
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
128
129
130
131
        key = cache["k"] + key * indices
        cache_v_shape = cache["v"].shape.as_list()
        indices = tf.reshape(
            tf.one_hot(decode_loop_step, cache_v_shape[1], dtype=value.dtype),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
132
            [1, cache_v_shape[1], 1, 1])
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
133
134
135
136
        value = cache["v"] + value * indices
      else:
        key = tf.concat([tf.cast(cache["k"], key.dtype), key], axis=1)
        value = tf.concat([tf.cast(cache["v"], value.dtype), value], axis=1)
137
138

      # Update cache
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
139
140
      cache["k"] = key
      cache["v"] = value
141

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
142
143
    # Scale query to prevent the dot product between query and key from growing
    # too large.
144
    depth = (self.hidden_size // self.num_heads)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
145
    query *= depth ** -0.5
146
147

    # Calculate dot product attention
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
148
    logits = tf.einsum("BTNH,BFNH->BNFT", key, query)
149
    logits += bias
150
151
152
153
    # Note that softmax internally performs math operations using float32
    # for numeric stability. When training with float16, we keep the input
    # and output in float16 for better performance.
    weights = tf.nn.softmax(logits, name="attention_weights")
154
155
    if training:
      weights = tf.nn.dropout(weights, rate=self.attention_dropout)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
156
    attention_output = tf.einsum("BNFT,BTNH->BFNH", weights, value)
157

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
158
159
    # Run the outputs through another linear projection layer. Recombining heads
    # is automatically done --> [batch_size, length, hidden_size]
160
161
162
163
164
165
166
    attention_output = self.output_dense_layer(attention_output)
    return attention_output


class SelfAttention(Attention):
  """Multiheaded self-attention layer."""

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
167
168
169
170
  def call(self, query_input, bias, training, cache=None,
           decode_loop_step=None):
    return super(SelfAttention, self).call(
        query_input, query_input, bias, training, cache, decode_loop_step)