#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ Created on Fri Sep 27 13:11:00 2019 Script for Multi-Head Scaled Dot-Product Attention and Context Attention layers @author: Chenkai Li """ from tf_keras import backend as K from tf_keras import initializers, regularizers, constraints, activations from tf_keras.layers import Layer class Attention(Layer): """ Adapted from https://github.com/lzfelix/keras_attention Implementation based on the work of Yang et al. "Hierarchical Attention Networks for Document Classification". [https://www.aclweb.org/anthology/N16-1174] The mathematical formulation of the model is as follows: ``` u = f(W * h + b), a_i = softmax(u_i^T * u_s), v_i = \sigma_i a_i * h_i. ``` Where h are the input tensors with shape (batch, n_timesteps, hidden_size), for instance, all hidden vectors produced by a recurrent layer, such as a LSTM and the output has shape `(batch, hidden_size)`. This layer also works with inputs with more than 3 dimensions as well, such as sentences in a document, where each input has size (batch, n_docs, n_sentences, embedding_size), outputing (batch, n_docs, embedding_size)`. # Arguments activation: The activation function f used by the layer (see [activations](../activations.md)). By default tanh is used, another common option is "linear". use_bias: Boolean, whether the layer uses a bias vector. initializer: Initializer for the `kernel` and `context` matrices (see [initializers](../initializers.md)). return_attention: If True, instead of returning the sequence descriptor, this layer will return the computed attention coefficients for each of the sequence timesteps. See Output section for details. W_regularizer: Regularizer function applied to the `kernel` weights matrix (see [regularizer](../regularizers.md)). u_regularizer: Regularizer function applied to the `context` weights matrix (see [regularizer](../regularizers.md)). b_regularizer: Regularizer function applied to the bias vector (see [regularizer](../regularizers.md)). W_constraint: Constraint function applied to the `kernel` weights matrix (see [constraints](../constraints.md)). u_constraint: Constraint function applied to the `contextl` weights matrix (see [constraints](../constraints.md)). b_constraint: Constraint function applied to the bias vector (see [constraints](../constraints.md)). # Input shape nD tensor with shape: `(batch_size, ..., timesteps, input_dim)`. The most common situation would be a 3D input with shape `(batch_size, timesteps, input_dim)`. # Outuput shape The sequence descriptor with shape `(batch_size, ..., timestamps)`. If `return_attention` is True, this layer will return the `alpha_i` weights for each timestep, and consequently its output shape will be different, namely: `(batch_size, ..., timesteps)`. """ def __init__(self, activation='tanh', initializer='glorot_uniform', return_attention=False, W_regularizer=None, u_regularizer=None, b_regularizer=None, W_constraint=None, u_constraint=None, b_constraint=None, bias=True, **kwargs): self.activation = activations.get(activation) self.initializer = initializers.get(initializer) self.W_regularizer = regularizers.get(W_regularizer) self.u_regularizer = regularizers.get(u_regularizer) self.b_regularizer = regularizers.get(b_regularizer) self.W_constraint = constraints.get(W_constraint) self.u_constraint = constraints.get(u_constraint) self.b_constraint = constraints.get(b_constraint) self.bias = bias self.supports_masking = True self.return_attention = return_attention super().__init__(**kwargs) def build(self, input_shape): # input_shape: (batch, time, amount_features) # the attention size matches the feature dimension amount_features = input_shape[-1] attention_size = input_shape[-1] self.W = self.add_weight(shape=(amount_features, attention_size), initializer=self.initializer, regularizer=self.W_regularizer, constraint=self.W_constraint, name='attention_W') self.b = None if self.bias: self.b = self.add_weight(shape=(attention_size,), initializer='zero', regularizer=self.b_regularizer, constraint=self.b_constraint, name='attention_b') self.context = self.add_weight(shape=(attention_size,), initializer=self.initializer, regularizer=self.u_regularizer, constraint=self.u_constraint, name='attention_us') super().build(input_shape) def call(self, x, mask=None): # U = tanh(H*W + b) (eq. 8) ui = K.dot(x, self.W) # (b, t, a) if self.b is not None: ui += self.b ui = self.activation(ui) # (b, t, a) # Z = U * us (eq. 9) us = K.expand_dims(self.context) # (a, 1) ui_us = K.dot(ui, us) # (b, t, a) * (a, 1) = (b, t, 1) ui_us = K.squeeze(ui_us, axis=-1) # (b, t, 1) -> (b, t) # alpha = softmax(Z) (eq. 9) alpha = self._masked_softmax(ui_us, mask) # (b, t) alpha = K.expand_dims(alpha, axis=-1) # (b, t, 1) if self.return_attention: return alpha else: # v = alpha_i * x_i (eq. 10) return K.sum(x * alpha, axis=1) def _masked_softmax(self, logits, mask): """Keras's default implementation of softmax doesn't allow masking, while this method does if `mask` is not `None`.""" # softmax(x): # b = max(x) # s_i = exp(xi - b) / exp(xj - b) # return s b = K.max(logits, axis=-1, keepdims=True) logits = logits - b exped = K.exp(logits) # ignoring masked inputs if mask is not None: mask = K.cast(mask, K.floatx()) exped *= mask partition = K.sum(exped, axis=-1, keepdims=True) # if all timesteps are masked, the partition will be zero. To avoid this # issue we use the following trick: partition = K.maximum(partition, K.epsilon()) return exped / partition def compute_output_shape(self, input_shape): """The attention mechanism computes a weighted average between all hidden vectors generated by the previous sequential layer, hence the input is expected to be `(batch_size, seq_len, amount_features)` if `return_attention` is `False`, otherwise the output should be (batch_size, seq_len).""" if self.return_attention: return input_shape[:-1] else: return input_shape[:-2] + input_shape[-1:] def compute_mask(self, x, input_mask=None): """This layer produces a single attended vector from a list of hidden vectors, hence it can't be masked as this means masking a single vector.""" return None def get_config(self): config = { 'activation': self.activation, 'initializer': self.initializer, 'return_attention': self.return_attention, 'W_regularizer': initializers.serialize(self.W_regularizer), 'u_regularizer': initializers.serialize(self.u_regularizer), 'b_regularizer': initializers.serialize(self.b_regularizer), 'W_constraint': constraints.serialize(self.W_constraint), 'u_constraint': constraints.serialize(self.u_constraint), 'b_constraint': constraints.serialize(self.b_constraint), 'bias': self.bias } base_config = super().get_config() return dict(list(base_config.items()) + list(config.items())) class ScaledDotProductAttention(Layer): """ Adapted from https://github.com/CyberZHG/keras-self-attention Implementation based on the work of Vaswani et al. "Attention Is All You Need". [https://arxiv.org/pdf/1706.03762.pdf] The attention layer that takes three inputs representing queries, keys and values. \text{Attention}(Q, K, V) = \text{softmax}(\frac{Q K^T}{\sqrt{d_k}}) V """ def __init__(self, return_attention=False, history_only=False, **kwargs): """Initialize the layer. :param return_attention: Whether to return attention weights. :param history_only: Whether to only use history data. :param kwargs: Arguments for parent class. """ super(ScaledDotProductAttention, self).__init__(**kwargs) self.supports_masking = True self.return_attention = return_attention self.history_only = history_only def get_config(self): config = { 'return_attention': self.return_attention, 'history_only': self.history_only, } base_config = super(ScaledDotProductAttention, self).get_config() return dict(list(base_config.items()) + list(config.items())) def compute_output_shape(self, input_shape): if isinstance(input_shape, list): query_shape, key_shape, value_shape = input_shape else: query_shape = key_shape = value_shape = input_shape output_shape = query_shape[:-1] + value_shape[-1:] if self.return_attention: attention_shape = query_shape[:2] + (key_shape[1],) return [output_shape, attention_shape] return output_shape def compute_mask(self, inputs, mask=None): if isinstance(mask, list): mask = mask[0] if self.return_attention: return [mask, None] return mask def call(self, inputs, mask=None, **kwargs): if isinstance(inputs, list): query, key, value = inputs else: query = key = value = inputs if isinstance(mask, list): mask = mask[1] feature_dim = K.shape(query)[-1] e = K.batch_dot(query, key, axes=2) / K.sqrt(K.cast(feature_dim, dtype=K.floatx())) e = K.exp(e - K.max(e, axis=-1, keepdims=True)) if self.history_only: query_len, key_len = K.shape(query)[1], K.shape(key)[1] indices = K.expand_dims(K.arange(0, key_len), axis=0) upper = K.expand_dims(K.arange(0, query_len), axis=-1) e *= K.expand_dims(K.cast(indices <= upper, K.floatx()), axis=0) if mask is not None: e *= K.cast(K.expand_dims(mask, axis=-2), K.floatx()) a = e / (K.sum(e, axis=-1, keepdims=True) + K.epsilon()) v = K.batch_dot(a, value) if self.return_attention: return [v, a] return v class MultiHeadAttention(Layer): """ Adapted from https://github.com/CyberZHG/keras-multi-head Implementation based on the work of Vaswani et al. "Attention Is All You Need". [https://arxiv.org/pdf/1706.03762.pdf] Multi-head scaled dot-product attention layer """ def __init__(self, head_num, activation='relu', use_bias=True, kernel_initializer='glorot_normal', bias_initializer='zeros', kernel_regularizer=None, bias_regularizer=None, kernel_constraint=None, bias_constraint=None, history_only=False, return_multi_attention=False, **kwargs): """Initialize the layer. :param head_num: Number of heads. :param activation: Activations for linear mappings. :param use_bias: Whether to use bias term. :param kernel_initializer: Initializer for linear mappings. :param bias_initializer: Initializer for linear mappings. :param kernel_regularizer: Regularizer for linear mappings. :param bias_regularizer: Regularizer for linear mappings. :param kernel_constraint: Constraints for linear mappings. :param bias_constraint: Constraints for linear mappings. :param history_only: Whether to only use history in attention layer. :param return_multi_attention: Whether to return the attention matrix. """ self.supports_masking = True self.head_num = head_num self.activation = activations.get(activation) self.use_bias = use_bias self.kernel_initializer = initializers.get(kernel_initializer) self.bias_initializer = initializers.get(bias_initializer) self.kernel_regularizer = regularizers.get(kernel_regularizer) self.bias_regularizer = regularizers.get(bias_regularizer) self.kernel_constraint = constraints.get(kernel_constraint) self.bias_constraint = constraints.get(bias_constraint) self.history_only = history_only self.return_multi_attention = return_multi_attention self.Wq, self.Wk, self.Wv, self.Wo = None, None, None, None self.bq, self.bk, self.bv, self.bo = None, None, None, None super(MultiHeadAttention, self).__init__(**kwargs) def get_config(self): config = { 'head_num': self.head_num, 'activation': activations.serialize(self.activation), 'use_bias': self.use_bias, 'kernel_initializer': initializers.serialize(self.kernel_initializer), 'bias_initializer': initializers.serialize(self.bias_initializer), 'kernel_regularizer': regularizers.serialize(self.kernel_regularizer), 'bias_regularizer': regularizers.serialize(self.bias_regularizer), 'kernel_constraint': constraints.serialize(self.kernel_constraint), 'bias_constraint': constraints.serialize(self.bias_constraint), 'history_only': self.history_only, 'return_multi_attention': self.return_multi_attention } base_config = super(MultiHeadAttention, self).get_config() return dict(list(base_config.items()) + list(config.items())) def compute_output_shape(self, input_shape): if self.return_multi_attention: if isinstance(input_shape, list): q, k, v = input_shape return (q[0], self.head_num, q[1], k[1]) return (input_shape[0], self.head_num, input_shape[1], input_shape[1]) else: if isinstance(input_shape, list): q, k, v = input_shape return q[:-1] + (v[-1],) return input_shape def compute_mask(self, inputs, input_mask=None): if isinstance(input_mask, list): return input_mask[0] return input_mask def build(self, input_shape): if isinstance(input_shape, list): q, k, v = input_shape else: q = k = v = input_shape feature_dim = int(v[-1]) if feature_dim % self.head_num != 0: raise IndexError('Invalid head number %d with the given input dim %d' % (self.head_num, feature_dim)) self.Wq = self.add_weight( shape=(int(q[-1]), feature_dim), initializer=self.kernel_initializer, regularizer=self.kernel_regularizer, constraint=self.kernel_constraint, name='%s_Wq' % self.name, ) if self.use_bias: self.bq = self.add_weight( shape=(feature_dim,), initializer=self.bias_initializer, regularizer=self.bias_regularizer, constraint=self.bias_constraint, name='%s_bq' % self.name, ) self.Wk = self.add_weight( shape=(int(k[-1]), feature_dim), initializer=self.kernel_initializer, regularizer=self.kernel_regularizer, constraint=self.kernel_constraint, name='%s_Wk' % self.name, ) if self.use_bias: self.bk = self.add_weight( shape=(feature_dim,), initializer=self.bias_initializer, regularizer=self.bias_regularizer, constraint=self.bias_constraint, name='%s_bk' % self.name, ) self.Wv = self.add_weight( shape=(int(v[-1]), feature_dim), initializer=self.kernel_initializer, regularizer=self.kernel_regularizer, constraint=self.kernel_constraint, name='%s_Wv' % self.name, ) if self.use_bias: self.bv = self.add_weight( shape=(feature_dim,), initializer=self.bias_initializer, regularizer=self.bias_regularizer, constraint=self.bias_constraint, name='%s_bv' % self.name, ) self.Wo = self.add_weight( shape=(feature_dim, feature_dim), initializer=self.kernel_initializer, regularizer=self.kernel_regularizer, constraint=self.kernel_constraint, name='%s_Wo' % self.name, ) if self.use_bias: self.bo = self.add_weight( shape=(feature_dim,), initializer=self.bias_initializer, regularizer=self.bias_regularizer, constraint=self.bias_constraint, name='%s_bo' % self.name, ) super(MultiHeadAttention, self).build(input_shape) @staticmethod def _reshape_to_batches(x, head_num): input_shape = K.shape(x) batch_size, seq_len, feature_dim = input_shape[0], input_shape[1], input_shape[2] head_dim = feature_dim // head_num x = K.reshape(x, (batch_size, seq_len, head_num, head_dim)) x = K.permute_dimensions(x, [0, 2, 1, 3]) return K.reshape(x, (batch_size * head_num, seq_len, head_dim)) @staticmethod def _reshape_from_batches(x, head_num): input_shape = K.shape(x) batch_size, seq_len, feature_dim = input_shape[0], input_shape[1], input_shape[2] x = K.reshape(x, (batch_size // head_num, head_num, seq_len, feature_dim)) x = K.permute_dimensions(x, [0, 2, 1, 3]) return K.reshape(x, (batch_size // head_num, seq_len, feature_dim * head_num)) @staticmethod def _reshape_attention_from_batches(x, head_num): input_shape = K.shape(x) batch_size, seq_len = input_shape[0], input_shape[1] return K.reshape(x, (batch_size // head_num, head_num, seq_len, seq_len)) @staticmethod def _reshape_mask(mask, head_num): if mask is None: return mask seq_len = K.shape(mask)[1] mask = K.expand_dims(mask, axis=1) mask = K.tile(mask, [1, head_num, 1]) return K.reshape(mask, (-1, seq_len)) def call(self, inputs, mask=None): if isinstance(inputs, list): q, k, v = inputs else: q = k = v = inputs if isinstance(mask, list): q_mask, k_mask, v_mask = mask else: q_mask = k_mask = v_mask = mask q = K.dot(q, self.Wq) k = K.dot(k, self.Wk) v = K.dot(v, self.Wv) if self.use_bias: q += self.bq k += self.bk v += self.bv if self.activation is not None: q = self.activation(q) k = self.activation(k) v = self.activation(v) y,a = ScaledDotProductAttention( return_attention=True, history_only=self.history_only, name='%s-Attention' % self.name, )( inputs=[ self._reshape_to_batches(q, self.head_num), self._reshape_to_batches(k, self.head_num), self._reshape_to_batches(v, self.head_num), ], mask=[ self._reshape_mask(q_mask, self.head_num), self._reshape_mask(k_mask, self.head_num), self._reshape_mask(v_mask, self.head_num), ], ) y = self._reshape_from_batches(y, self.head_num) a = self._reshape_attention_from_batches(a, self.head_num) y = K.dot(y, self.Wo) if self.use_bias: y += self.bo if self.activation is not None: y = self.activation(y) ''' if TF_KERAS: # Add shape information to tensor when using `tf.keras` input_shape = [K.int_shape(q), K.int_shape(k), K.int_shape(v)] output_shape = self.compute_output_shape(input_shape) if output_shape[1] is not None: output_shape = (-1,) + output_shape[1:] y = K.reshape(y, output_shape) ''' if self.return_multi_attention: return a return y