def apply(self, is_train, x, x_mask=None):
x_word_dim = tf.shape(x)[1]
# (batch, x_word, key_word)
dist_matrix = self.attention.get_scores(x, x)
dist_matrix += tf.expand_dims(tf.eye(x_word_dim) * VERY_NEGATIVE_NUMBER, 0) # Mask out self
joint_mask = compute_attention_mask(x_mask, x_mask, x_word_dim, x_word_dim)
if joint_mask is not None:
dist_matrix += VERY_NEGATIVE_NUMBER * (1 - tf.cast(joint_mask, dist_matrix.dtype))
if not self.alignment_bias:
select_probs = tf.nn.softmax(dist_matrix)
else:
# Allow zero-attention by adding a learned bias to the normalizer
bias = tf.exp(tf.get_variable("no-alignment-bias", initializer=tf.constant(-1.0, dtype=tf.float32)))
dist_matrix = tf.exp(dist_matrix)
select_probs = dist_matrix / (tf.reduce_sum(dist_matrix, axis=2, keep_dims=True) + bias)
response = tf.matmul(select_probs, x) # (batch, x_words, q_dim)
if self.merge is not None:
with tf.variable_scope("merge"):
response = self.merge.apply(is_train, response, x)
return response
else:
return response
def apply(self, is_train, x, keys, memories, x_mask=None, mem_mask=None):
x_word_dim = tf.shape(x)[1]
key_word_dim = tf.shape(keys)[1]
dist_matrix = self.sim.get_scores(x, keys)
joint_mask = compute_attention_mask(x_mask, mem_mask, x_word_dim, key_word_dim)
if joint_mask is not None:
dist_matrix += VERY_NEGATIVE_NUMBER * (1 - tf.cast(joint_mask, dist_matrix.dtype))
query_probs = tf.nn.softmax(dist_matrix) # probability of each mem_word per x_word
# Batch matrix multiplication to get the attended vectors
select_query = tf.matmul(query_probs, memories) # (batch, x_words, q_dim)
if not self.q2c:
if self.query_dots:
return tf.concat([x, select_query, x * select_query], axis=2)
else:
return tf.concat([x, select_query], axis=2)
# select query-to-context
context_dist = tf.reduce_max(dist_matrix, axis=2) # (batch, x_word``s)
context_probs = tf.nn.softmax(context_dist) # (batch, x_words)
select_context = tf.einsum("ai,aik->ak", context_probs, x) # (batch, x_dim)
select_context = tf.expand_dims(select_context, 1)
if self.query_dots:
return tf.concat([x, select_query, x * select_query, x * select_context], axis=2)
else:
return tf.concat([x, select_query, x * select_context], axis=2)
def apply(self, is_train, x, keys, memories, x_mask=None, mem_mask=None):
x_word_dim = tf.shape(x)[1]
key_word_dim = tf.shape(keys)[1]
# (batch, x_word, key_word)
dist_matrix = self.attention.get_scores(x, keys)
joint_mask = compute_attention_mask(x_mask, mem_mask, x_word_dim, key_word_dim)
dist_matrix += VERY_NEGATIVE_NUMBER * (1 - tf.cast(joint_mask, dist_matrix.dtype))
if self.alignment_bias is None:
select_probs = tf.nn.softmax(dist_matrix)
else:
# Compute softmax with an additional bias term, this allows the model to 'ignore' the memories
# if needed since the sum of the weights given to each memory can be < 1.
bias = tf.exp(tf.get_variable("no-alignment-bias", initializer=tf.constant(-1.0, dtype=tf.float32)))
dist_matrix = tf.exp(dist_matrix)
select_probs = dist_matrix / (tf.reduce_sum(dist_matrix, axis=2, keep_dims=True) + bias)
# Too (batch, x_word, memory_dim)
response = tf.matmul(select_probs, memories)
if self.merge is not None:
with tf.variable_scope("merge"):
response = self.merge.apply(is_train, response, x)
return response
else:
return response
def apply(self, is_train, x, keys, memories, x_mask=None, mem_mask=None):
x_word_dim = tf.shape(x)[1]
key_word_dim = tf.shape(keys)[1]
# (batch, x_word, key_word)
dist_matrix = self.attention.get_scores(x, keys)
joint_mask = compute_attention_mask(x_mask, mem_mask, x_word_dim,
key_word_dim)
dist_matrix += VERY_NEGATIVE_NUMBER * (
1 - tf.cast(joint_mask, dist_matrix.dtype))
if self.alignment_bias is None:
select_probs = tf.nn.softmax(dist_matrix)
else:
bias = tf.exp(
tf.get_variable("no-alignment-bias",
initializer=tf.constant(-1.0,
dtype=tf.float32)))
dist_matrix = tf.exp(dist_matrix)
select_probs = dist_matrix / (
tf.reduce_sum(dist_matrix, axis=2, keep_dims=True) + bias)
# Too (batch, x_word, memory_dim)
response = tf.matmul(select_probs, memories)
with tf.variable_scope("encode_keys"):
encoded = self.encoder_layer.apply(is_train, keys, mem_mask)
return tf.concat(
[x, response, x * response, x * tf.expand_dims(encoded, 1)],
axis=2)
def apply(self, is_train, x, mask=None):
batch_size = tf.shape(x)[0]
x_word_dim = tf.shape(x)[1]
x_feature_dim = x.shape.as_list()[-1]
project_size = self.project_size
if project_size is None:
project_size = x_feature_dim // self.n_heads
if x_feature_dim % self.n_heads != 0:
raise ValueError()
mem_size = self.memory_size
if mem_size is None:
mem_size = project_size
init = get_keras_initialization(self.init)
query_proj = tf.get_variable("query_proj", (x_feature_dim, self.n_heads, project_size), initializer=init)
if self.shared_project:
key_proj = query_proj
else:
key_proj = tf.get_variable("key_proj", (x_feature_dim, self.n_heads, project_size), initializer=init)
mem_proj = tf.get_variable("mem_proj", (x_feature_dim, self.n_heads, mem_size), initializer=init)
queries = tf.tensordot(x, query_proj, [[2], [0]]) # (batch, word, n_head, project_size)
keys = tf.tensordot(x, key_proj, [[2], [0]]) # (batch, key, n_head, project_size)
if self.project_bias:
queries += tf.get_variable("query_bias", (1, 1, self.n_heads, project_size),
initializer=tf.zeros_initializer())
keys += tf.get_variable("key_bias", (1, 1, self.n_heads, project_size),
initializer=tf.zeros_initializer())
# dist_matrix = tf.matmul(queries, keys, transpose_b=True)
dist_matrix = tf.einsum("bwhd,bkhd->bwkh", queries, keys) # dots of (batch, word, key, head)
if self.scale:
dist_matrix /= tf.sqrt(float(project_size))
if self.bilinear_comp:
query_bias_proj = tf.get_variable("query_bias_proj", (x_feature_dim, self.n_heads), initializer=init)
key_bias_proj = tf.get_variable("query_bias_proj", (x_feature_dim, self.n_heads), initializer=init)
dist_matrix += tf.expand_dims(tf.tensordot(x, query_bias_proj, [[2], [0]]), 2)
dist_matrix += tf.expand_dims(tf.tensordot(x, key_bias_proj, [[2], [0]]), 1)
joint_mask = compute_attention_mask(mask, mask, x_word_dim, x_word_dim)
if joint_mask is not None:
dist_matrix += tf.expand_dims(VERY_NEGATIVE_NUMBER * (1 - tf.cast(joint_mask, dist_matrix.dtype)), 2)
dist_matrix += tf.expand_dims(tf.expand_dims(tf.eye(x_word_dim) * VERY_NEGATIVE_NUMBER, 0), 2)
if self.bias:
bias = tf.get_variable("bias", (1, 1, self.n_heads, 1), initializer=tf.zeros_initializer())
dist_matrix += bias
select_probs = tf.nn.softmax(dist_matrix) # for each (batch, word, head) probability over keys
memories = tf.tensordot(x, mem_proj, [[2], [0]]) # (batch, memory, head, mem_size)
response = tf.einsum("bwhk,bkhd->bwhd", select_probs, memories) # (batch, word, head, mem_size)
response = tf.reshape(response, (batch_size, x_word_dim, self.n_heads * mem_size)) # concat heads
if self.merge is not None:
with tf.variable_scope("merge"):
response = self.merge.apply(is_train, x, response)
return response
