# -*- encoding:utf-8 -*- # Copyright (c) Alibaba, Inc. and its affiliates. """Attention layers that can be used in sequence DNN/CNN models. This file follows the terminology of https://arxiv.org/abs/1706.03762 Figure 2. Attention is formed by three tensors: Query, Key and Value. """ import tensorflow as tf from tensorflow.python.keras.layers import Layer class Attention(Layer): """Dot-product attention layer, a.k.a. Luong-style attention. Inputs are a list with 2 or 3 elements: 1. A `query` tensor of shape `(batch_size, Tq, dim)`. 2. A `value` tensor of shape `(batch_size, Tv, dim)`. 3. A optional `key` tensor of shape `(batch_size, Tv, dim)`. If none supplied, `value` will be used as a `key`. The calculation follows the steps: 1. Calculate attention scores using `query` and `key` with shape `(batch_size, Tq, Tv)`. 2. Use scores to calculate a softmax distribution with shape `(batch_size, Tq, Tv)`. 3. Use the softmax distribution to create a linear combination of `value` with shape `(batch_size, Tq, dim)`. Args: use_scale: If `True`, will create a scalar variable to scale the attention scores. dropout: Float between 0 and 1. Fraction of the units to drop for the attention scores. Defaults to `0.0`. seed: A Python integer to use as random seed in case of `dropout`. score_mode: Function to use to compute attention scores, one of `{"dot", "concat"}`. `"dot"` refers to the dot product between the query and key vectors. `"concat"` refers to the hyperbolic tangent of the concatenation of the `query` and `key` vectors. Call Args: inputs: List of the following tensors: - `query`: Query tensor of shape `(batch_size, Tq, dim)`. - `value`: Value tensor of shape `(batch_size, Tv, dim)`. - `key`: Optional key tensor of shape `(batch_size, Tv, dim)`. If not given, will use `value` for both `key` and `value`, which is the most common case. mask: List of the following tensors: - `query_mask`: A boolean mask tensor of shape `(batch_size, Tq)`. If given, the output will be zero at the positions where `mask==False`. - `value_mask`: A boolean mask tensor of shape `(batch_size, Tv)`. If given, will apply the mask such that values at positions where `mask==False` do not contribute to the result. return_attention_scores: bool, it `True`, returns the attention scores (after masking and softmax) as an additional output argument. training: Python boolean indicating whether the layer should behave in training mode (adding dropout) or in inference mode (no dropout). use_causal_mask: Boolean. Set to `True` for decoder self-attention. Adds a mask such that position `i` cannot attend to positions `j > i`. This prevents the flow of information from the future towards the past. Defaults to `False`. Output: Attention outputs of shape `(batch_size, Tq, dim)`. (Optional) Attention scores after masking and softmax with shape `(batch_size, Tq, Tv)`. """ def __init__(self, params, name='attention', reuse=None, **kwargs): super(Attention, self).__init__(name=name, **kwargs) self.use_scale = params.get_or_default('use_scale', False) self.scale_by_dim = params.get_or_default('scale_by_dim', False) self.score_mode = params.get_or_default('score_mode', 'dot') if self.score_mode not in ['dot', 'concat']: raise ValueError('Invalid value for argument score_mode. ' "Expected one of {'dot', 'concat'}. " 'Received: score_mode=%s' % self.score_mode) self.dropout = params.get_or_default('dropout', 0.0) self.seed = params.get_or_default('seed', None) self.scale = None self.concat_score_weight = None self._return_attention_scores = params.get_or_default( 'return_attention_scores', False) self.use_causal_mask = params.get_or_default('use_causal_mask', False) @property def return_attention_scores(self): return self._return_attention_scores def build(self, input_shape): self._validate_inputs(input_shape) if self.use_scale: self.scale = self.add_weight( name='scale', shape=(), initializer='ones', dtype=self.dtype, trainable=True, ) if self.score_mode == 'concat': self.concat_score_weight = self.add_weight( name='concat_score_weight', shape=(), initializer='ones', dtype=self.dtype, trainable=True, ) super(Attention, self).build(input_shape) # Be sure to call this somewhere! def _calculate_scores(self, query, key): """Calculates attention scores as a query-key dot product. Args: query: Query tensor of shape `(batch_size, Tq, dim)`. key: Key tensor of shape `(batch_size, Tv, dim)`. Returns: Tensor of shape `(batch_size, Tq, Tv)`. """ if self.score_mode == 'dot': scores = tf.matmul(query, tf.transpose(key, [0, 2, 1])) if self.scale is not None: scores *= self.scale elif self.scale_by_dim: dk = tf.cast(tf.shape(key)[-1], tf.float32) scores /= tf.math.sqrt(dk) elif self.score_mode == 'concat': # Reshape tensors to enable broadcasting. # Reshape into [batch_size, Tq, 1, dim]. q_reshaped = tf.expand_dims(query, axis=-2) # Reshape into [batch_size, 1, Tv, dim]. k_reshaped = tf.expand_dims(key, axis=-3) if self.scale is not None: scores = self.concat_score_weight * tf.reduce_sum( tf.tanh(self.scale * (q_reshaped + k_reshaped)), axis=-1) else: scores = self.concat_score_weight * tf.reduce_sum( tf.tanh(q_reshaped + k_reshaped), axis=-1) return scores def _apply_scores(self, scores, value, scores_mask=None, training=False): """Applies attention scores to the given value tensor. To use this method in your attention layer, follow the steps: * Use `query` tensor of shape `(batch_size, Tq)` and `key` tensor of shape `(batch_size, Tv)` to calculate the attention `scores`. * Pass `scores` and `value` tensors to this method. The method applies `scores_mask`, calculates `attention_distribution = softmax(scores)`, then returns `matmul(attention_distribution, value). * Apply `query_mask` and return the result. Args: scores: Scores float tensor of shape `(batch_size, Tq, Tv)`. value: Value tensor of shape `(batch_size, Tv, dim)`. scores_mask: A boolean mask tensor of shape `(batch_size, 1, Tv)` or `(batch_size, Tq, Tv)`. If given, scores at positions where `scores_mask==False` do not contribute to the result. It must contain at least one `True` value in each line along the last dimension. training: Python boolean indicating whether the layer should behave in training mode (adding dropout) or in inference mode (no dropout). Returns: Tensor of shape `(batch_size, Tq, dim)`. Attention scores after masking and softmax with shape `(batch_size, Tq, Tv)`. """ if scores_mask is not None: padding_mask = tf.logical_not(scores_mask) # Bias so padding positions do not contribute to attention # distribution. Note 65504. is the max float16 value. max_value = 65504.0 if scores.dtype == 'float16' else 1.0e9 scores -= max_value * tf.cast(padding_mask, dtype=scores.dtype) weights = tf.nn.softmax(scores, axis=-1) if training and self.dropout > 0: weights = tf.nn.dropout(weights, 1.0 - self.dropout, seed=self.seed) return tf.matmul(weights, value), weights def _calculate_score_mask(self, scores, v_mask, use_causal_mask): if use_causal_mask: # Creates a lower triangular mask, so position i cannot attend to # positions j > i. This prevents the flow of information from the # future into the past. score_shape = tf.shape(scores) # causal_mask_shape = [1, Tq, Tv]. mask_shape = (1, score_shape[-2], score_shape[-1]) ones_mask = tf.ones(shape=mask_shape, dtype='int32') row_index = tf.cumsum(ones_mask, axis=-2) col_index = tf.cumsum(ones_mask, axis=-1) causal_mask = tf.greater_equal(row_index, col_index) if v_mask is not None: # Mask of shape [batch_size, 1, Tv]. v_mask = tf.expand_dims(v_mask, axis=-2) return tf.logical_and(v_mask, causal_mask) return causal_mask else: # If not using causal mask, return the value mask as is, # or None if the value mask is not provided. return v_mask def call(self, inputs, mask=None, training=False, **kwargs): self._validate_inputs(inputs=inputs, mask=mask) q = inputs[0] v = inputs[1] k = inputs[2] if len(inputs) > 2 else v q_mask = mask[0] if mask else None v_mask = mask[1] if mask else None scores = self._calculate_scores(query=q, key=k) scores_mask = self._calculate_score_mask(scores, v_mask, self.use_causal_mask) result, attention_scores = self._apply_scores( scores=scores, value=v, scores_mask=scores_mask, training=training) if q_mask is not None: # Mask of shape [batch_size, Tq, 1]. q_mask = tf.expand_dims(q_mask, axis=-1) result *= tf.cast(q_mask, dtype=result.dtype) if self._return_attention_scores: return result, attention_scores return result def compute_mask(self, inputs, mask=None): self._validate_inputs(inputs=inputs, mask=mask) if mask is None or mask[0] is None: return None return tf.convert_to_tensor(mask[0]) def compute_output_shape(self, input_shape): """Returns shape of value tensor dim, but for query tensor length.""" return list(input_shape[0][:-1]), input_shape[1][-1] def _validate_inputs(self, inputs, mask=None): """Validates arguments of the call method.""" class_name = self.__class__.__name__ if not isinstance(inputs, list): raise ValueError('{class_name} layer must be called on a list of inputs, ' 'namely [query, value] or [query, value, key]. ' 'Received: inputs={inputs}.'.format( class_name=class_name, inputs=inputs)) if len(inputs) < 2 or len(inputs) > 3: raise ValueError('%s layer accepts inputs list of length 2 or 3, ' 'namely [query, value] or [query, value, key]. ' 'Received length: %d.' % (class_name, len(inputs))) if mask is not None: if not isinstance(mask, list): raise ValueError( '{class_name} layer mask must be a list, ' 'namely [query_mask, value_mask]. Received: mask={mask}.'.format( class_name=class_name, mask=mask)) if len(mask) < 2 or len(mask) > 3: raise ValueError( '{class_name} layer accepts mask list of length 2 or 3. ' 'Received: inputs={inputs}, mask={mask}.'.format( class_name=class_name, inputs=inputs, mask=mask)) def get_config(self): base_config = super(Attention, self).get_config() config = { 'use_scale': self.use_scale, 'score_mode': self.score_mode, 'dropout': self.dropout, } return dict(list(base_config.items()) + list(config.items()))