def FProp(self, theta, inputs, paddings, class_emb):
"""Apply batch normalization.
Args:
theta: A `.NestedMap` object containing weights' values of this layer and
its children layers.
inputs: The inputs tensor. Shaped [batch, ..., dim].
paddings: The paddings tensor. Shaped [batch, ..., 1], with the same rank
as the input tensor.
class_emb: The conditioning inputs, Shaped [batch, emb_dim].
Returns:
Output after applying batch normalization, with the same shape as
'inputs'.
"""
p = self.params
batch = py_utils.GetShape(inputs)[0]
class_emb = py_utils.HasShape(class_emb, [batch, p.class_emb_dim])
if not py_utils.use_tpu():
class_emb = py_utils.with_dependencies([
py_utils.assert_less_equal(
tf.cast(class_emb, tf.int32), 1, name='one_hot_assert1'),
py_utils.assert_greater_equal(
tf.cast(class_emb, tf.int32), 0, name='one_hot_assert2'),
py_utils.assert_equal(tf.ones([batch], tf.int32),
tf.cast(tf.reduce_sum(class_emb, -1),
tf.int32),
name='one_hot_assert3'),
], class_emb)
with tf.name_scope(p.name):
norm_mean, norm_variance, beta, gamma = self.ComputeAndUpdateMoments(
theta, inputs, paddings=paddings, class_emb=class_emb)
return self._ComputeBN(inputs, paddings, gamma, beta, norm_mean,
norm_variance)
def ApplyBias():
"""Bias and update log_probs and consistent."""
def TileForBeamAndFlatten(tensor):
tensor = tf.reshape(tensor, [1, -1]) # [1, src_batch]
tensor = tf.tile(tensor,
[num_hyps_per_beam, 1
]) # [num_hyps_per_beam, src_batch]
tgt_batch = tf.shape(step_ids)[
0] # num_hyps_per_beam*src_batch
return tf.reshape(tensor, [tgt_batch])
# Consistent if step_ids == labels from previous step
# TODO(navari): Consider updating consistent only if weights > 0. Then
# re-evaluate the need for bias_only_if_consistent=True.
# Note that prev_label is incorrrect for step 0 but is overridden later
prev_label = TileForBeamAndFlatten(
tf.gather(labels, tf.maximum(time_step - 1, 0), axis=1))
is_step0 = tf.equal(time_step, 0)
local_consistence = tf.logical_or(
is_step0, tf.equal(prev_label, tf.squeeze(step_ids, 1)))
consistent = tf.logical_and(states.consistent,
local_consistence)
# get label, weight slices corresponding to current time_step
label = TileForBeamAndFlatten(
tf.gather(labels, time_step, axis=1))
weight = TileForBeamAndFlatten(
tf.gather(weights, time_step, axis=1))
if p.bias_only_if_consistent:
weight = weight * tf.cast(consistent, p.dtype)
# convert from dense label to sparse label probs
vocab_size = tf.shape(bs_results.log_probs)[1]
uncertainty = tf.constant(
1e-10,
p.dtype) # avoid 0 probs which may cause issues with log
label_probs = tf.one_hot(
label,
vocab_size,
on_value=1 - uncertainty,
off_value=uncertainty / tf.cast(vocab_size - 1, p.dtype),
dtype=p.dtype) # [tgt_batch, vocab_size]
pred_probs = tf.exp(bs_results.log_probs)
# interpolate predicted probs and label probs
weight = tf.expand_dims(weight, 1)
probs = py_utils.with_dependencies([
py_utils.assert_less_equal(weight, 1.),
py_utils.assert_greater_equal(weight, 0.)
], (1.0 - weight) * pred_probs + weight * label_probs)
return tf.log(probs), consistent
def _MaybeExpandPaddings(self, inputs, paddings):
# rank difference is at most one.
rank_diff = tf.rank(inputs) - tf.rank(paddings)
paddings = py_utils.with_dependencies([
py_utils.assert_less_equal(rank_diff, 1),
py_utils.assert_greater_equal(rank_diff, 0)
], paddings)
# Pads [1] to the end of paddings.
paddings = tf.reshape(
paddings,
tf.concat([tf.shape(paddings),
tf.tile([1], [rank_diff])], axis=0))
return paddings
def ApplyBias():
"""Bias and update log_probs and consistent."""
# Consistent if step_ids == labels from previous step
# TODO(navari): Consider updating consistent only if weights > 0. Then
# re-evaluate the need for bias_only_if_consistent=True.
# Note that prev_label is incorrrect for step 0 but is overridden
# later
prev_label = TileForBeamAndFlatten(
tf.gather(labels, tf.maximum(time_step - 1, 0),
axis=1))
is_step0 = tf.equal(time_step, 0)
local_consistence = tf.math.logical_or(
is_step0, tf.equal(prev_label, tf.squeeze(step_ids,
1)))
consistent = tf.math.logical_and(states.consistent,
local_consistence)
# get label, weight slices corresponding to current time_step
label = TileForBeamAndFlatten(
tf.gather(labels, time_step, axis=1))
weight = TileForBeamAndFlatten(
tf.gather(weights, time_step, axis=1))
if p.bias_only_if_consistent:
weight = weight * tf.cast(consistent,
py_utils.FPropDtype(p))
# convert from dense label to sparse label probs
vocab_size = tf.shape(bs_results.log_probs)[1]
label_probs = tf.one_hot(label,
vocab_size,
dtype=py_utils.FPropDtype(
p)) # [tgt_batch, vocab_size]
pred_probs = tf.exp(bs_results.log_probs)
# interpolate predicted probs and label probs
weight = tf.expand_dims(weight, 1)
probs = py_utils.with_dependencies([
py_utils.assert_less_equal(weight, 1.),
py_utils.assert_greater_equal(weight, 0.)
], (1.0 - weight) * pred_probs + weight * label_probs)
# Ensure that tf.math.log is applied to positive values.
probs = tf.maximum(probs,
tf.constant(1e-12, dtype=probs.dtype))
return tf.math.log(probs), consistent
def ComputeSparseAttention(q, k, v, sparsity_indices, paddings=None):
"""Computes attention according to a sparsity pattern.
We use the following capital letters to denote shape parameters:
B = batch size
S = length of the source sequence
T = length of the target sequence
N = number of attention heads
H = dimensions of each attention head
K = number of clusters
W = attention window (K <= S)
The 'sparsity_indices' is a tensor of integral type where the last dimension
contains W indices (W is the attention window) for each corresponding position
along S in 'k' that the query is allowed to attend to.
For example, if sparsity_indices[batch_idx, target time step, head_idx] =
[1, 7, 8], it means that token in the query attends to values with indices
1, 7, and 8, and the attention window here is 3.
The valid values in 'sparsity_indices' are [-1, S-1]. Note that the value -1
is reserved to mean paddings, distinct from the value (S-1).
For example, if W=S and 'sparsity_indices' contains range(S) on the last
dimension, this degenerates to the original full attention.
We require that 'sparsity_indices' does not contain duplicates (except for -1
to indicate paddings), but we do not require 'sparsity_indices' to be sorted.
Note that this implementation is flexible and generic but is not optimized for
time or space complexity. Please consider grouping queries that attend to the
same subset of values first for efficiency.
Args:
q: (projected) queries, [B, T, N, H];
k: (projected) keys, [B, S, N, H];
v: (projected) values, [B, S, N, H];
sparsity_indices: [B, T, N, W], where W is the attention window;
paddings: paddings for keys, [B, S] if not None.
Returns:
output: the encoded output, [B, T, N, H].
atten_probs: the attention weights, [B, T, N, S].
"""
q = tf.convert_to_tensor(q)
k = tf.convert_to_tensor(k)
v = tf.convert_to_tensor(v)
sparsity_indices = tf.convert_to_tensor(sparsity_indices)
k = py_utils.HasRank(k, 4)
_, source_length, _, dim_per_head = py_utils.GetShape(k, 4)
sparsity_indices = py_utils.HasRank(sparsity_indices, 4)
batch_size, target_length, num_heads, attention_window = py_utils.GetShape(
sparsity_indices, 4)
py_utils.assert_less_equal(
attention_window, source_length,
'The provided sparsity_indices has attention window '
' > source length. This is likely an error.')
# To prepare for gathering the relevant vectors from 'k', we prepare
# gather_idx of shape [B, T, N, W, 3] where the last dimension corresponds to
# slices in 'k' indexed by (batch index, source time step, head index),
# where the source length index comes from the original W dimension in
# 'sparsity_indices'.
seq_idx = tf.expand_dims(sparsity_indices, axis=-1)
# Overwrite the paddings -1 with valid gather indices (zeros). We will
# fix the logits with -inf in these positions later.
seq_idx = tf.where(seq_idx < 0, tf.zeros_like(seq_idx), seq_idx)
batch_idx = tf.reshape(
tf.range(0, batch_size, dtype=sparsity_indices.dtype),
[batch_size, 1, 1, 1, 1])
batch_idx = tf.tile(batch_idx,
[1, target_length, num_heads, attention_window, 1])
head_idx = tf.reshape(
tf.range(0, num_heads, dtype=sparsity_indices.dtype),
[1, 1, num_heads, 1, 1])
head_idx = tf.tile(head_idx,
[batch_size, target_length, 1, attention_window, 1])
# [B, T, N, W, 3], where last dimension is (batch index, source length index,
# head index).
gather_idx = tf.concat([batch_idx, seq_idx, head_idx], axis=-1)
# Both the gathered k and v have shape [B, T, N, W, H]
k = tf.gather_nd(k, gather_idx)
v = tf.gather_nd(v, gather_idx)
if paddings is None:
paddings = tf.zeros([batch_size, source_length])
paddings = tf.convert_to_tensor(paddings)
paddings = tf.expand_dims(paddings, axis=-1)
# [B, S, N]
paddings = tf.tile(paddings, [1, 1, num_heads])
# [B, T, N, W]
paddings = tf.gather_nd(paddings, gather_idx)
logits = tf.einsum('BTNH, BTNWH -> BTNW', q, k)
logits *= tf.math.rsqrt(tf.cast(dim_per_head, q.dtype))
very_negative_logits = (
tf.ones_like(logits) * logits.dtype.max *
tf.constant(-0.7, dtype=logits.dtype))
padded_logits = tf.where(
tf.math.logical_or(sparsity_indices < 0, paddings > 0.0),
very_negative_logits, logits)
# [B, T, N, W]
atten_probs = tf.nn.softmax(padded_logits, name='attention_weights')
atten_probs = tf.where(sparsity_indices < 0, tf.zeros_like(logits),
atten_probs)
output = tf.einsum('BTNW, BTNWH -> BTNH', atten_probs, v)
# Scatter 'atten_probs' back into the original source length.
# [B, T, N, W, 1]
batch_idx = tf.tile(
tf.range(batch_size)[:, None, None, None, None],
[1, target_length, num_heads, attention_window, 1])
# [B, T, N, W, 1]
target_seq_idx = tf.tile(
tf.range(target_length)[None, :, None, None, None],
[batch_size, 1, num_heads, attention_window, 1])
# [B, T, N, W, 1]
head_idx = tf.tile(
tf.range(num_heads)[None, None, :, None, None],
[batch_size, target_length, 1, attention_window, 1])
# seq_idx: [B, T, N, W, 1]
# [B, T, N, W, 4]
scatter_idx = tf.concat([batch_idx, target_seq_idx, head_idx, seq_idx], -1)
# [B, T, N, S]
scattered_probs = tf.scatter_nd(
scatter_idx, atten_probs,
[batch_size, target_length, num_heads, source_length])
return output, scattered_probs
def PreBeamSearchStepCallback(theta, encoder_outputs, step_ids, states,
num_hyps_per_beam, *args, **kwargs):
"""Wrapper for adding bias to _PreBeamSearchStateCallback.
Biases results.log_probs towards provided encoder_outputs.targets.
Args:
theta: a NestedMap of parameters.
encoder_outputs: a NestedMap computed by encoder.
step_ids: A tensor of shape [tgt_batch, 1].
states: A `.NestedMap` of tensors representing states that the clients
would like to keep track of for each of the active hyps.
num_hyps_per_beam: Beam size.
*args: additional arguments to _PreBeamSearchStepCallback.
**kwargs: additional arguments to _PreBeamSearchStepCallback.
Returns:
A tuple (results, out_states).
results: A `.NestedMap` of beam search results.
atten_probs:
The updated attention probs, of shape [tgt_batch, src_len].
log_probs:
Log prob for each of the tokens in the target vocab. This is of
shape
[tgt_batch, vocab_size].
out_states: a `.NestedMap` The updated states. The states relevant here
are:
time_step: A scalar indicating current step of decoder. Must be
provided and maintained by subclass.
consistent: A boolean vector of shape [tgt_batch, ] which tracks
whether each hypothesis has exactly matched
encoder_outputs.targets
so far.
"""
p = self.params
time_step = states.time_step
bs_results, out_states = self._PreBeamSearchStepCallback(
theta, encoder_outputs, step_ids, states, num_hyps_per_beam,
*args, **kwargs)
labels = encoder_outputs.targets.labels
weights = encoder_outputs.targets.weights
def TileForBeamAndFlatten(tensor):
tensor = tf.reshape(tensor, [1, -1]) # [1, src_batch]
tensor = tf.tile(
tensor,
[num_hyps_per_beam, 1]) # [num_hyps_per_beam, src_batch]
tgt_batch = tf.shape(step_ids)[
0] # num_hyps_per_beam*src_batch
return tf.reshape(tensor, [tgt_batch])
# Consistent if step_ids == labels from previous step
# TODO(navari): Consider updating consistent only if weights > 0. Then
# re-evaluate the need for bias_only_if_consistent=True.
# Note that prev_label is incorrrect for step 0 but is overridden later
prev_label = TileForBeamAndFlatten(
tf.gather(labels, tf.maximum(time_step - 1, 0), axis=1))
is_step0 = tf.equal(time_step, 0)
local_consistence = tf.logical_or(
is_step0, tf.equal(prev_label, tf.squeeze(step_ids, 1)))
out_states.consistent = tf.logical_and(states.consistent,
local_consistence)
# get label, weight slices corresponding to current time_step
label = TileForBeamAndFlatten(tf.gather(labels, time_step, axis=1))
weight = TileForBeamAndFlatten(
tf.gather(weights, time_step, axis=1))
if p.bias_only_if_consistent:
weight = weight * tf.cast(out_states.consistent, p.dtype)
# convert from dense label to sparse label probs
vocab_size = tf.shape(bs_results.log_probs)[1]
uncertainty = tf.constant(
1e-10,
p.dtype) # avoid 0 probs which may cause issues with log
label_probs = tf.one_hot(label,
vocab_size,
on_value=1 - uncertainty,
off_value=uncertainty /
tf.cast(vocab_size - 1, p.dtype),
dtype=p.dtype) # [tgt_batch, vocab_size]
pred_probs = tf.exp(bs_results.log_probs)
# interpolate predicted probs and label probs
weight = tf.expand_dims(weight, 1)
probs = py_utils.with_dependencies([
py_utils.assert_less_equal(weight, 1.),
py_utils.assert_greater_equal(weight, 0.)
], (1.0 - weight) * pred_probs + weight * label_probs)
bs_results.log_probs = tf.log(probs)
return bs_results, out_states
