def dot_product_attention(q, k, v, bias, dropout_rate=0.0):
"""Dot-product attention.
Args:
q: Tensor with shape [..., length_q, depth_k].
k: Tensor with shape [..., length_kv, depth_k]. Leading dimensions must
match with q.
v: Tensor with shape [..., length_kv, depth_v] Leading dimensions must
match with q.
bias: bias Tensor (see attention_bias())
dropout_rate: a float.
Returns:
Tensor with shape [..., length_q, depth_v].
"""
logits = tf.matmul(q, k, transpose_b=True) # [..., length_q, length_kv]
logits = tf.multiply(logits,
1.0 / math.sqrt(float(util.get_shape_list(q)[-1])))
if bias is not None:
# `attention_mask` = [B, T]
from_shape = util.get_shape_list(q)
if len(from_shape) == 4:
broadcast_ones = tf.ones([from_shape[0], 1, from_shape[2], 1],
tf.float32)
elif len(from_shape) == 5:
# from_shape = [B, N, Block_num, block_size, depth]#
broadcast_ones = tf.ones(
[from_shape[0], 1, from_shape[2], from_shape[3], 1],
tf.float32)
bias = tf.matmul(broadcast_ones,
tf.cast(bias, tf.float32),
transpose_b=True)
# Since attention_mask is 1.0 for positions we want to attend and 0.0 for
# masked positions, this operation will create a tensor which is 0.0 for
# positions we want to attend and -10000.0 for masked positions.
adder = (1.0 - bias) * -10000.0
# Since we are adding it to the raw scores before the softmax, this is
# effectively the same as removing these entirely.
logits += adder
else:
adder = 0.0
attention_probs = tf.nn.softmax(logits, name="attention_probs")
attention_probs = util.dropout(attention_probs, dropout_rate)
return tf.matmul(attention_probs, v)
def __init__(self,
vocab_size,
is_training,
input_ids,
input_mask,
segment_ids,
sample_weight=None,
reduced_size=64,
topic_size=1024,
hidden_size=768,
num_hidden_layers=12,
num_attention_heads=12,
bias=0,
scope='vae',
trainable=True,
**kwargs):
super().__init__()
# freeze parameters
config = Config(vocab_size,
hidden_size=hidden_size,
num_hidden_layers=num_hidden_layers,
num_attention_heads=num_attention_heads)
if not is_training:
config.hidden_dropout_prob = 0.0
config.attention_probs_dropout_prob = 0.0
input_shape = util.get_shape_list(input_ids, expected_rank=2)
batch_size = input_shape[0]
seq_length = input_shape[1]
# Tilda embeddings for SMART algorithm
tilda_embeddings = None
use_tilda_embedding = kwargs.get('use_tilda_embedding')
if use_tilda_embedding:
with tf.variable_scope('', reuse=True):
tilda_embeddings = tf.get_variable('tilda_embeddings')
with tf.variable_scope(scope):
with tf.variable_scope('embeddings'):
(self.embedding_output, self.embedding_table) = \
self.embedding_lookup(
input_ids=input_ids,
vocab_size=config.vocab_size,
batch_size=batch_size,
max_seq_length=seq_length,
embedding_size=config.hidden_size,
initializer_range=config.initializer_range,
word_embedding_name='word_embeddings',
tilda_embeddings=tilda_embeddings,
trainable=trainable)
self.embedding_output = self.embedding_postprocessor(
input_tensor=self.embedding_output,
batch_size=batch_size,
max_seq_length=seq_length,
hidden_size=config.hidden_size,
use_token_type=True,
segment_ids=segment_ids,
token_type_vocab_size=config.type_vocab_size,
token_type_embedding_name='token_type_embeddings',
use_position_embeddings=True,
position_embedding_name='position_embeddings',
initializer_range=config.initializer_range,
max_position_embeddings=config.max_position_embeddings,
dropout_prob=config.hidden_dropout_prob,
trainable=trainable)
with tf.variable_scope('encoder'):
# stacked transformer
attention_mask = self.create_attention_mask_from_input_mask(
input_mask, batch_size, seq_length)
self.all_encoder_layers = self.transformer_model(
input_tensor=self.embedding_output,
batch_size=batch_size,
max_seq_length=seq_length,
attention_mask=attention_mask,
hidden_size=config.hidden_size,
num_hidden_layers=config.num_hidden_layers,
num_attention_heads=config.num_attention_heads,
intermediate_size=config.intermediate_size,
intermediate_act_fn=util.get_activation(config.hidden_act),
hidden_dropout_prob=config.hidden_dropout_prob,
attention_probs_dropout_prob=\
config.attention_probs_dropout_prob,
initializer_range=config.initializer_range,
trainable=trainable)
# projection
with tf.variable_scope('projection'):
transformer_output = tf.layers.dense(
self.all_encoder_layers[-1],
reduced_size,
activation=util.gelu,
kernel_initializer=tf.truncated_normal_initializer(
stddev=config.initializer_range),
trainable=trainable)
transformer_output = tf.reshape(transformer_output,
[batch_size, -1])
input_length = tf.reduce_sum(input_mask, axis=-1)
input_length = tf.cast(input_length, tf.float32)
input_length_1d = tf.reshape(input_length, [batch_size])
input_length_2d = tf.reshape(input_length, [batch_size, 1])
broadcast_mask = tf.sequence_mask(
tf.multiply(input_length_1d, reduced_size),
seq_length * reduced_size,
dtype=tf.float32)
broadcast_mask = tf.multiply(broadcast_mask,
seq_length / input_length_2d)
transformer_output *= broadcast_mask
# latent space
miu = tf.layers.dense(
transformer_output,
topic_size,
activation='tanh',
kernel_initializer=tf.truncated_normal_initializer(
stddev=config.initializer_range),
name='miu',
trainable=trainable)
sigma = tf.layers.dense(
transformer_output,
topic_size,
kernel_initializer=tf.truncated_normal_initializer(
stddev=config.initializer_range),
name='sigma',
trainable=trainable)
self.probs['miu'] = miu
self.probs['sigma'] = sigma
with tf.variable_scope('decoder'):
with tf.variable_scope('projection'):
# reparametarization
if is_training:
noise = tf.random_normal([batch_size, topic_size])
else:
noise = tf.random_uniform([batch_size, topic_size],
minval=-bias,
maxval=bias)
decoder_input = miu + tf.exp(sigma) * noise
# projection
decoder_input = tf.layers.dense(
decoder_input,
seq_length * reduced_size,
activation=util.gelu,
kernel_initializer=tf.truncated_normal_initializer(
stddev=config.initializer_range),
trainable=trainable)
intermediate_input = tf.reshape(
decoder_input, [-1, seq_length, reduced_size])
intermediate_input = util.layer_norm(intermediate_input,
trainable=trainable)
intermediate_input = util.dropout(
intermediate_input, config.hidden_dropout_prob)
# MLP
with tf.variable_scope('intermediate'):
intermediate_output = tf.layers.dense(
intermediate_input,
4 * reduced_size,
activation=util.gelu,
kernel_initializer=util.create_initializer(
config.initializer_range),
trainable=trainable)
with tf.variable_scope('output'):
decoder_output = tf.layers.dense(
intermediate_output,
config.hidden_size,
kernel_initializer=util.create_initializer(
config.initializer_range),
trainable=trainable)
decoder_output = util.layer_norm(decoder_output,
trainable=trainable)
decoder_output = util.dropout(decoder_output,
config.hidden_dropout_prob)
self.all_decoder_layers = [intermediate_output, decoder_output]
self.all_decoder_layers = [decoder_output]
# reconstruction
with tf.variable_scope('cls/predictions'):
with tf.variable_scope('transform'):
input_tensor = tf.layers.dense(
decoder_output,
units=config.hidden_size,
activation=util.get_activation(config.hidden_act),
kernel_initializer=util.create_initializer(
config.initializer_range),
trainable=trainable)
input_tensor = util.layer_norm(input_tensor,
trainable=trainable)
output_weights = self.embedding_table
output_bias = tf.get_variable('output_bias',
shape=[config.vocab_size],
initializer=tf.zeros_initializer(),
trainable=trainable)
flatten_input_tensor = tf.reshape(input_tensor,
[-1, config.hidden_size])
logits = tf.matmul(flatten_input_tensor,
output_weights,
transpose_b=True)
logits = tf.nn.bias_add(logits, output_bias)
logits = tf.reshape(logits,
[batch_size, seq_length, config.vocab_size])
probs = tf.nn.softmax(logits, axis=-1, name='probs')
lm_log_probs = tf.nn.log_softmax(logits, axis=-1)
self.preds['preds'] = tf.argmax(probs, axis=-1)
one_hot_labels = tf.one_hot(input_ids,
depth=config.vocab_size,
dtype=tf.float32)
per_example_loss = -tf.reduce_sum(lm_log_probs * one_hot_labels,
axis=[-1])
if sample_weight is not None:
per_example_loss *= tf.expand_dims(sample_weight, axis=-1)
self.total_loss = (tf.reduce_mean(per_example_loss) +
tf.reduce_mean(tf.square(miu)) +
tf.reduce_mean(tf.exp(sigma) - sigma - 1))
self.losses['losses'] = per_example_loss
def attention_layer(self,
from_tensor,
to_tensor,
attention_mask=None,
num_attention_heads=12,
size_per_head=512,
query_act=None,
key_act=None,
value_act=None,
attention_probs_dropout_prob=0.0,
initializer_range=0.02,
do_return_2d_tensor=False,
batch_size=None,
from_max_seq_length=None,
to_max_seq_length=None,
dtype=tf.float32,
trainable=True):
def transpose_for_scores(input_tensor, batch_size, num_attention_heads,
max_seq_length, width):
output_tensor = tf.reshape(
input_tensor,
[batch_size, max_seq_length, num_attention_heads, width])
output_tensor = tf.transpose(output_tensor, [0, 2, 1, 3])
return output_tensor
# Scalar dimensions referenced here:
# B = batch size (number of sequences)
# F = from_tensor sequence length
# T = to_tensor sequence length
# N = num_attention_heads
# H = size_per_head
from_tensor_2d = util.reshape_to_matrix(from_tensor)
to_tensor_2d = util.reshape_to_matrix(to_tensor)
# query_layer = [B*F, N*H]
query_layer = tf.layers.dense(
from_tensor_2d,
num_attention_heads * size_per_head,
activation=query_act,
name='query',
kernel_initializer=util.create_initializer(initializer_range),
trainable=trainable)
# key_layer = [B*T, N*H]
key_layer = tf.layers.dense(
to_tensor_2d,
num_attention_heads * size_per_head,
activation=key_act,
name='key',
kernel_initializer=util.create_initializer(initializer_range),
trainable=trainable)
# value_layer = [B*T, N*H]
value_layer = tf.layers.dense(
to_tensor_2d,
num_attention_heads * size_per_head,
activation=value_act,
name='value',
kernel_initializer=util.create_initializer(initializer_range),
trainable=trainable)
# query_layer = [B, N, F, H]
query_layer = transpose_for_scores(query_layer, batch_size,
num_attention_heads,
from_max_seq_length, size_per_head)
# key_layer = [B, N, T, H]
key_layer = transpose_for_scores(key_layer, batch_size,
num_attention_heads,
to_max_seq_length, size_per_head)
# Take the dot product between 'query' and 'key' to get the raw
# attention scores.
# attention_scores = [B, N, F, T]
attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
attention_scores = tf.multiply(attention_scores,
1.0 / math.sqrt(float(size_per_head)))
if attention_mask is not None:
# attention_mask = [B, 1, F, T]
attention_mask = tf.expand_dims(attention_mask, axis=[1])
adder = (1.0 - tf.cast(attention_mask, dtype)) * -10000.0
attention_scores += adder
# Normalize the attention scores to probabilities.
# attention_probs = [B, N, F, T]
attention_probs = tf.nn.softmax(attention_scores, axis=-1)
# This is actually dropping out entire tokens to attend to,
# which might seem a bit unusual, but is taken from the original
# Transformer paper.
attention_probs = util.dropout(attention_probs,
attention_probs_dropout_prob)
# value_layer = [B, T, N, H]
value_layer = tf.reshape(value_layer, [
batch_size, to_max_seq_length, num_attention_heads, size_per_head
])
# value_layer = [B, N, T, H]
value_layer = tf.transpose(value_layer, [0, 2, 1, 3])
# context_layer = [B, N, F, H]
context_layer = tf.matmul(attention_probs, value_layer)
# context_layer = [B, F, N, H]
context_layer = tf.transpose(context_layer, [0, 2, 1, 3])
if do_return_2d_tensor:
# context_layer = [B*F, N*H]
context_layer = tf.reshape(context_layer, [
batch_size * from_max_seq_length,
num_attention_heads * size_per_head
])
else:
# context_layer = [B, F, N*H]
context_layer = tf.reshape(context_layer, [
batch_size, from_max_seq_length,
num_attention_heads * size_per_head
])
return (context_layer, attention_scores)
def __init__(self,
is_training,
input_tensor,
n_wide_features,
wide_features,
label_ids,
label_size=2,
sample_weight=None,
scope='cls/seq_relationship',
hidden_dropout_prob=0.1,
initializer_range=0.02,
trainable=True,
**kwargs):
super().__init__(**kwargs)
hidden_size = input_tensor.shape.as_list()[-1]
feature_size = wide_features.shape.as_list()[-1]
with tf.variable_scope('wide'):
feature_embeddings = tf.get_variable(
name='feature_embeddings',
shape=[feature_size + 1, hidden_size],
initializer=util.create_initializer(initializer_range),
trainable=trainable)
wide_output = tf.gather(feature_embeddings,
wide_features) # [B, N, H]
with tf.variable_scope('wide_and_deep'):
deep_output = tf.expand_dims(input_tensor, -1) # [B, H, 1]
attention_scores = tf.matmul(wide_output, deep_output) # [B, N, 1]
attention_scores = tf.transpose(attention_scores,
[0, 2, 1]) # [B, 1, N]
attention_scores = tf.multiply(attention_scores,
1.0 / math.sqrt(hidden_size))
feature_mask = tf.cast(
tf.sequence_mask(n_wide_features, feature_size),
tf.float32) # [B, N]
feature_mask = tf.expand_dims(feature_mask, 1) # [B, 1, N]
attention_scores += (1.0 - feature_mask) * -10000.0
attention_matrix = tf.nn.softmax(attention_scores, axis=-1)
attention_output = tf.matmul(attention_matrix,
wide_output) # [B, 1, H]
attention_output = attention_output[:, 0, :] # [B, H]
# attention_output = util.dropout(
# attention_output, hidden_dropout_prob)
input_tensor = util.layer_norm(attention_output + input_tensor,
trainable=trainable)
with tf.variable_scope(scope):
output_weights = tf.get_variable(
'output_weights',
shape=[label_size, hidden_size],
initializer=util.create_initializer(initializer_range),
trainable=trainable)
output_bias = tf.get_variable('output_bias',
shape=[label_size],
initializer=tf.zeros_initializer(),
trainable=trainable)
output_layer = util.dropout(
input_tensor, hidden_dropout_prob if is_training else 0.0)
logits = tf.matmul(output_layer, output_weights, transpose_b=True)
logits = tf.nn.bias_add(logits, output_bias)
self.preds['preds'] = tf.argmax(logits, axis=-1)
self.probs['probs'] = tf.nn.softmax(logits, axis=-1, name='probs')
log_probs = tf.nn.log_softmax(logits, axis=-1)
one_hot_labels = tf.one_hot(label_ids,
depth=label_size,
dtype=tf.float32)
per_example_loss = -tf.reduce_sum(one_hot_labels * log_probs,
axis=-1)
if sample_weight is not None:
per_example_loss = tf.cast(sample_weight,
dtype=tf.float32) * per_example_loss
thresh = kwargs.get('tsa_thresh')
if thresh is not None:
assert isinstance(
thresh,
float), ('`tsa_thresh` must be a float between 0 and 1.')
uncertainty = tf.reduce_sum(self.probs['probs'] *
tf.log(self.probs['probs']),
axis=-1)
uncertainty /= tf.log(1 / label_size)
per_example_loss = tf.cast(
tf.greater(uncertainty, thresh), dtype=tf.float32) * \
per_example_loss
self.losses['losses'] = per_example_loss
self.total_loss = tf.reduce_mean(per_example_loss)
def attention_layer(from_tensor,
to_tensor,
attention_mask=None,
num_attention_heads=1,
size_per_head=512,
query_act=None,
key_act=None,
value_act=None,
attention_probs_dropout_prob=0.0,
initializer_range=0.02,
do_return_2d_tensor=False,
batch_size=None,
from_seq_length=None,
to_seq_length=None):
'''Performs multi-headed attention from `from_tensor` to `to_tensor`.
This is an implementation of multi-headed attention based on 'Attention
is all you Need'. If `from_tensor` and `to_tensor` are the same, then
this is self-attention. Each timestep in `from_tensor` attends to the
corresponding sequence in `to_tensor`, and returns a fixed-with vector.
This function first projects `from_tensor` into a 'query' tensor and
`to_tensor` into 'key' and 'value' tensors. These are (effectively) a list
of tensors of length `num_attention_heads`, where each tensor is of shape
[batch_size, seq_length, size_per_head].
Then, the query and key tensors are dot-producted and scaled. These are
softmaxed to obtain attention probabilities. The value tensors are then
interpolated by these probabilities, then concatenated back to a single
tensor and returned.
In practice, the multi-headed attention are done with transposes and
reshapes rather than actual separate tensors.
Args:
from_tensor: float Tensor of shape [batch_size, from_seq_length,
from_width].
to_tensor: float Tensor of shape [batch_size, to_seq_length, to_width].
attention_mask: (optional) int32 Tensor of shape [batch_size,
from_seq_length, to_seq_length]. The values should be 1 or 0. The
attention scores will effectively be set to -infinity for any positions
in the mask that are 0, and will be unchanged for positions that are 1.
num_attention_heads: int. Number of attention heads.
size_per_head: int. Size of each attention head.
query_act: (optional) Activation function for the query transform.
key_act: (optional) Activation function for the key transform.
value_act: (optional) Activation function for the value transform.
attention_probs_dropout_prob: (optional) float. Dropout probability of the
attention probabilities.
initializer_range: float. Range of the weight initializer.
do_return_2d_tensor: bool. If True, the output will be of shape [batch_size
* from_seq_length, num_attention_heads * size_per_head]. If False, the
output will be of shape [batch_size, from_seq_length, num_attention_heads
* size_per_head].
batch_size: (Optional) int. If the input is 2D, this might be the batch
size of the 3D version of the `from_tensor` and `to_tensor`.
from_seq_length: (Optional) If the input is 2D, this might be the seq
length of the 3D version of the `from_tensor`.
to_seq_length: (Optional) If the input is 2D, this might be the seq length
of the 3D version of the `to_tensor`.
Returns:
float Tensor of shape [batch_size, from_seq_length,
num_attention_heads * size_per_head]. (If `do_return_2d_tensor` is
true, this will be of shape [batch_size * from_seq_length,
num_attention_heads * size_per_head]).
Raises:
ValueError: Any of the arguments or tensor shapes are invalid.
'''
def transpose_for_scores(input_tensor, batch_size, num_attention_heads,
seq_length, width):
output_tensor = tf.reshape(
input_tensor, [batch_size, seq_length, num_attention_heads, width])
output_tensor = tf.transpose(output_tensor, [0, 2, 1, 3])
return output_tensor
from_shape = util.get_shape_list(from_tensor, expected_rank=[2, 3])
to_shape = util.get_shape_list(to_tensor, expected_rank=[2, 3])
if len(from_shape) != len(to_shape):
raise ValueError(
'The rank of `from_tensor` must match the rank of `to_tensor`.')
if len(from_shape) == 3:
batch_size = from_shape[0]
from_seq_length = from_shape[1]
to_seq_length = to_shape[1]
elif len(from_shape) == 2:
if batch_size is None or from_seq_length is None or to_seq_length is None:
raise ValueError(
'When passing in rank 2 tensors to attention_layer, the values '
'for `batch_size`, `from_seq_length`, and `to_seq_length` '
'must all be specified.')
# Scalar dimensions referenced here:
# B = batch size (number of sequences)
# F = `from_tensor` sequence length
# T = `to_tensor` sequence length
# N = `num_attention_heads`
# H = `size_per_head`
from_tensor_2d = util.reshape_to_matrix(from_tensor)
to_tensor_2d = util.reshape_to_matrix(to_tensor)
# `query_layer` = [B*F, N*H]
query_layer = tf.layers.dense(
from_tensor_2d,
num_attention_heads * size_per_head,
activation=query_act,
name='query',
kernel_initializer=util.create_initializer(initializer_range))
# `key_layer` = [B*T, N*H]
key_layer = tf.layers.dense(
to_tensor_2d,
num_attention_heads * size_per_head,
activation=key_act,
name='key',
kernel_initializer=util.create_initializer(initializer_range))
# `value_layer` = [B*T, N*H]
value_layer = tf.layers.dense(
to_tensor_2d,
num_attention_heads * size_per_head,
activation=value_act,
name='value',
kernel_initializer=util.create_initializer(initializer_range))
# `query_layer` = [B, N, F, H]
query_layer = transpose_for_scores(query_layer, batch_size,
num_attention_heads, from_seq_length,
size_per_head)
# `key_layer` = [B, N, T, H]
key_layer = transpose_for_scores(key_layer, batch_size,
num_attention_heads, to_seq_length,
size_per_head)
# Take the dot product between 'query' and 'key' to get the raw
# attention scores.
# `attention_scores` = [B, N, F, T]
attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
attention_scores = tf.multiply(attention_scores,
1.0 / math.sqrt(float(size_per_head)))
if attention_mask is not None:
# `attention_mask` = [B, 1, F, T]
attention_mask = tf.expand_dims(attention_mask, axis=[1])
# Since attention_mask is 1.0 for positions we want to attend and 0.0 for
# masked positions, this operation will create a tensor which is 0.0 for
# positions we want to attend and -10000.0 for masked positions.
adder = (1.0 - tf.cast(attention_mask, tf.float32)) * -10000.0
# Since we are adding it to the raw scores before the softmax, this is
# effectively the same as removing these entirely.
attention_scores += adder
# Normalize the attention scores to probabilities.
# `attention_probs` = [B, N, F, T]
attention_probs = tf.nn.softmax(attention_scores)
# This is actually dropping out entire tokens to attend to, which might
# seem a bit unusual, but is taken from the original Transformer paper.
attention_probs = util.dropout(attention_probs,
attention_probs_dropout_prob)
# `value_layer` = [B, T, N, H]
value_layer = tf.reshape(
value_layer,
[batch_size, to_seq_length, num_attention_heads, size_per_head])
# `value_layer` = [B, N, T, H]
value_layer = tf.transpose(value_layer, [0, 2, 1, 3])
# `context_layer` = [B, N, F, H]
context_layer = tf.matmul(attention_probs, value_layer)
# `context_layer` = [B, F, N, H]
context_layer = tf.transpose(context_layer, [0, 2, 1, 3])
if do_return_2d_tensor:
# `context_layer` = [B*F, N*H]
context_layer = tf.reshape(context_layer, [
batch_size * from_seq_length, num_attention_heads * size_per_head
])
else:
# `context_layer` = [B, F, N*H]
context_layer = tf.reshape(
context_layer,
[batch_size, from_seq_length, num_attention_heads * size_per_head])
return context_layer, attention_probs