def layer_norm(x, dim, epsilon=1e-6, name="layer_prepostprocess"): """Layer normalization over dimension dim. Args: x: a mtf.Tensor whose shape contains dim. dim: a mtf.Dimension epsilon: a floating point number name: a string. variable scope. Returns: a mtf.Tensor with same shape as x. """ with tf.variable_scope(name + "/layer_norm"): scale = mtf.get_variable(x.mesh, "layer_norm_scale", mtf.Shape([dim]), initializer=tf.ones_initializer(), activation_dtype=x.dtype) bias = mtf.get_variable(x.mesh, "layer_norm_bias", mtf.Shape([dim]), initializer=tf.zeros_initializer(), activation_dtype=x.dtype) reduced_shape = x.shape - dim mean = mtf.reduce_mean(x, output_shape=reduced_shape) variance = mtf.reduce_mean(mtf.square(x - mean), output_shape=reduced_shape) norm_x = (x - mean) * mtf.rsqrt(variance + epsilon) return norm_x * scale + bias
def batch_norm(x, is_training, momentum, epsilon=1e-9, name=None): """Batch normalization. Args: x: a mtf.Tensor whose shape contains [batch_dim, ..., dim] is_training: a boolean, whether mode is training. momentum: a floating point number, specifying batch norm decay value. epsilon: a floating point number. name: a string. variable scope. Returns: a mtf.Tensor with same shape as x. """ with tf.variable_scope(name, default_name="batch_norm", values=[x]): batch_dim = x.shape.dims[0] reduced_shape = x.shape - batch_dim scale = mtf.get_variable(x.mesh, "batch_norm_scale", mtf.Shape([batch_dim]), initializer=tf.ones_initializer(), activation_dtype=x.dtype) bias = mtf.get_variable(x.mesh, "batch_norm_bias", mtf.Shape([batch_dim]), initializer=tf.zeros_initializer(), activation_dtype=x.dtype) moving_mean = mtf.get_variable( x.mesh, "moving_mean", reduced_shape, initializer=tf.random_normal_initializer(stddev=1.0), activation_dtype=x.dtype, trainable=False) moving_variance = mtf.get_variable(x.mesh, "moving_variance", reduced_shape, initializer=tf.ones_initializer(), activation_dtype=x.dtype, trainable=False) # At training time, calculate mean and variance and normalize across batch # dim. if is_training: mean = mtf.reduce_mean(x, output_shape=reduced_shape) variance = mtf.reduce_mean(mtf.square(x - mean), output_shape=reduced_shape) norm_x = (x - mean) * mtf.rsqrt(variance + epsilon) # Update running mean and running variance. moving_mean = mtf.assign( moving_mean, momentum * moving_mean + (1 - momentum) * mean) moving_variance = mtf.assign( moving_variance, momentum * moving_variance + (1 - momentum) * variance) else: # At eval and test time, use the running mean and variance. norm_x = (x - moving_mean) * mtf.rsqrt(moving_variance + epsilon) return norm_x * scale + bias
def dense(x, output_dim, reduced_dims=None, expert_dims=None, use_bias=True, activation=None, master_dtype=tf.float32, slice_dtype=tf.float32, name=None): """Dense layer doing (kernel*x + bias) computation. Args: x: a mtf.Tensor of shape [..., reduced_dims]. output_dim: a mtf.Dimension reduced_dims: an optional list of mtf.Dimensions of x to be reduced. If omitted, we reduce the last dimension. expert_dims: an optional list of mtf.Dimension which represent different experts. Different experts get different weights. use_bias: a boolean, whether to add bias. activation: an optional function from mtf.Tensor to mtf.Tensor master_dtype: a tf.dtype slice_dtype: a tf.dtype name: a string. variable scope. Returns: a mtf.Tensor of shape [..., output_dim]. """ if expert_dims is None: expert_dims = [] if reduced_dims is None: reduced_dims = x.shape.dims[-1:] w_shape = mtf.Shape(expert_dims + reduced_dims + [output_dim]) output_shape = mtf.Shape( [d for d in x.shape.dims if d not in reduced_dims] + [output_dim]) with tf.variable_scope(name, default_name="dense"): stddev = mtf.list_product(d.size for d in reduced_dims)**-0.5 w = mtf.get_variable( x.mesh, "kernel", w_shape, initializer=tf.random_normal_initializer(stddev=stddev), master_dtype=master_dtype, slice_dtype=slice_dtype, activation_dtype=x.dtype) y = mtf.einsum([x, w], output_shape) if use_bias: b = mtf.get_variable(x.mesh, "bias", mtf.Shape(expert_dims + [output_dim]), initializer=tf.zeros_initializer(), activation_dtype=x.dtype) y += b if activation is not None: y = activation(y) return y
def multihead_self_attention_incremental(query_antecedent, prev_k, prev_v, step_num, name="multihead_attention"): """Incremental self-attention (one decode step). In order to use only one variable containing the four weight matrices packed together, we insist that the query and memory antecedents have the same dimensionality (io_channels) and that the keys and values have the same dimensionality (kv_channels). Args: query_antecedent: a mtf.Tensor with shape [batch..., io_channels] prev_k: mtf.Tensor with shape [batch..., heads, memory_length, kv_channels] prev_v: mtf.Tensor with shape [batch..., heads, memory_length, kv_channels] step_num: mtf Scalar with dtype tf.int32 name: an optional string. Returns: y: A mtf.Tensor with shape [batch..., io_channels] new_k: mtf.Tensor with shape [batch..., heads, memory_length, kv_channels] new_v: mtf.Tensor with shape [batch..., heads, memory_length, kv_channels] Raises: ValueError: if the dimensions do not match. """ batch_dims = query_antecedent.shape.dims[:-1] io_channels = query_antecedent.shape.dims[-1] heads, memory_length, kv_channels = prev_k.shape.dims[-3:] with tf.variable_scope(name, default_name="multihead_attention"): q_var, k_var, v_var, o_var = multihead_attention_vars( query_antecedent.mesh, heads, io_channels, kv_channels, query_antecedent.dtype) memory_antecedent = query_antecedent q = mtf.einsum([query_antecedent, q_var], mtf.Shape(batch_dims + [heads, kv_channels])) k = mtf.einsum([memory_antecedent, k_var], mtf.Shape(batch_dims + [heads, kv_channels])) v = mtf.einsum([memory_antecedent, v_var], mtf.Shape(batch_dims + [heads, kv_channels])) k = prev_k + mtf.multiply( k, mtf.one_hot(step_num, memory_length), output_shape=prev_k.shape) v = prev_v + mtf.multiply( v, mtf.one_hot(step_num, memory_length), output_shape=prev_v.shape) mask = mtf.to_float( mtf.greater( mtf.range(query_antecedent.mesh, memory_length, dtype=tf.int32), step_num)) * -1e9 o = dot_product_attention(q, k, v, mask) y = mtf.einsum([o, o_var], query_antecedent.shape) return y, k, v
def masked_local_attention_1d_incremental(x, prev_k, prev_v, step_num, master_dtype, slice_dtype, name=None): """Incremental local self-attention (one decode step). Incremental version of masked_local_attention_1d() Args: x: a mtf.Tensor with shape [batch..., io_channels] prev_k: mtf.Tensor with shape [batch..., heads, window_length, kv_channels] prev_v: mtf.Tensor with shape [batch..., heads, window_length, kv_channels] step_num: mtf Scalar with dtype tf.int32 master_dtype: a tf.dtype slice_dtype: a tf.dtype name: an optional string. Returns: y: A mtf.Tensor with shape [batch..., io_channels] new_k: mtf.Tensor with shape [batch..., heads, window_length, kv_channels] new_v: mtf.Tensor with shape [batch..., heads, window_length, kv_channels] Raises: ValueError: if the dimensions do not match. """ batch_dims = x.shape.dims[:-1] io_channels = x.shape.dims[-1] heads, window_length, kv_channels = prev_k.shape.dims[-3:] with tf.variable_scope(name, default_name="multihead_attention"): q_var, k_var, v_var, o_var = multihead_attention_vars( x.mesh, heads, io_channels, kv_channels, master_dtype, slice_dtype, x.dtype) q = mtf.einsum([x, q_var], mtf.Shape(batch_dims + [heads, kv_channels])) k = mtf.einsum([x, k_var], mtf.Shape(batch_dims + [heads, kv_channels])) v = mtf.einsum([x, v_var], mtf.Shape(batch_dims + [heads, kv_channels])) current_position = mtf.equal( mtf.range(x.mesh, window_length, dtype=tf.int32), mtf.mod(step_num, window_length.size)) k = mtf.where(current_position, k, prev_k, output_shape=prev_k.shape) v = mtf.where(current_position, v, prev_v, output_shape=prev_v.shape) o = dot_product_attention(q, k, v, mask=None) y = mtf.einsum([o, o_var], x.shape) return y, k, v
def multihead_encdec_attention_incremental(query_antecedent, q_var, o_var, k, v, mask, name="multihead_attention"): """Incremental attention over encoder (one decode step). In order to use only one variable containing the four weight matrices packed together, we insist that the query and memory antecedents have the same dimensionality (io_channels) and that the keys and values have the same dimensionality (kv_channels). memory_dims is a subset of query_dims Args: query_antecedent: a mtf.Tensor with shape query_dims + [io_channels] q_var: a mtf.Tensor with shape [heads, io_channels, kv_channels] o_var: a mtf.Tensor with shape [heads, io_channels, kv_channels] k: memory_dims + [heads, memory_length, kv_channels] v: memory_dims + [heads, memory_length, kv_channels] mask: mask Tensor (see attention_mask()) name: an optional string. Returns: A mtf.Tensor with shape [batch, qlen, io_channels] """ heads, _, kv_channels = k.shape.dims[-3:] query_dims = query_antecedent.shape.dims[:-1] with tf.variable_scope(name, default_name="multihead_attention"): q = mtf.einsum([query_antecedent, q_var], mtf.Shape(query_dims + [heads, kv_channels])) o = dot_product_attention(q, k, v, mask) return mtf.einsum([o, o_var], query_antecedent.shape)
def my_gather(tensor): return mtf.gather(tensor, top_beam_index, beam_dim, output_shape=mtf.Shape([ double_beam if d == beam_dim else d for d in tensor.shape.dims ]))
def gather(tensor, name): with tf.name_scope(prefix + name): output_shape = mtf.Shape([ beam_dim if d == old_beam_dim else d for d in tensor.shape.dims ]) return mtf.gather(tensor, topk_indices, old_beam_dim, output_shape=output_shape)
def dot_product_attention(q, k, v, mask, dropout=0.0, dropout_broadcast_dims=None, extra_logit=None): """Dot-product attention. Args: q: Tensor with shape [...., length_q, depth_k]. Typically leading dimensions are [batch, heads]. 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. mask: mask Tensor (see attention_mask()) dropout: a float. dropout_broadcast_dims: an optional list of mtf.Dimension extra_logit: an optional scalar or tensor Returns: Tensor with shape [..., length_q, depth_v]. """ length_kv = k.shape.dims[-2] logits_shape = mtf.Shape(q.shape.dims[:-1] + [length_kv]) logits = mtf.einsum([q, k], logits_shape) if mask is not None: logits += mask weights = mtf.softmax(logits, length_kv, extra_logit=extra_logit) if dropout != 0.0: weights = mtf.dropout(weights, 1.0 - dropout, noise_shape=weights.shape - dropout_broadcast_dims) depth_v = v.shape.dims[-1] outputs_shape = mtf.Shape(q.shape.dims[:-1] + [depth_v]) outputs = mtf.einsum([weights, v], outputs_shape) return outputs
def multihead_attention_vars(mesh, heads, io_channels, kv_channels, master_dtype, slice_dtype, activation_dtype): """Create Parameters for Multihead Attention. Args: mesh: a Mesh heads: a Dimension io_channels: a Dimension kv_channels: a Dimension master_dtype: a tf.dtype slice_dtype: a tf.dtype activation_dtype: a tf.dtype Returns: q_var: a Tensor with shape [heads, io_channels, kv_channels] k_var: a Tensor with shape [heads, io_channels, kv_channels] v_var: a Tensor with shape [heads, io_channels, kv_channels] o_var: a Tensor with shape [heads, io_channels, kv_channels] """ qkvo = mtf.Dimension("qkvo", 4) qk_stddev = (io_channels.size**-0.5) * (kv_channels.size**-0.25) v_stddev = io_channels.size**-0.5 o_stddev = (io_channels.size * heads.size)**-0.5 def qkvo_initializer(shape, dtype=None, partition_info=None, verify_shape=None): del partition_info, verify_shape return tf.random_normal(shape, dtype=dtype) * tf.reshape( tf.cast([qk_stddev, qk_stddev, v_stddev, o_stddev], dtype or tf.float32), [4, 1, 1, 1]) var = mtf.get_variable(mesh, "qkvo", mtf.Shape([qkvo, heads, io_channels, kv_channels]), initializer=qkvo_initializer, master_dtype=master_dtype, slice_dtype=slice_dtype, activation_dtype=activation_dtype) q_var, k_var, v_var, o_var = mtf.unstack(var, qkvo) return q_var, k_var, v_var, o_var
def dense_relu_dense(x, hidden_channels, dropout=0.0, dropout_broadcast_dims=None, name=None): """Hidden layer with ReLU activation followed by linear projection. The output has the same number of channels as the input. Args: x: a mtf.Tensor hidden_channels: a mtf.Dimension - channels in the hidden layer dropout: an optional float dropout_broadcast_dims: an optional list of mtf.Dimension name: an optional string Returns: a mtf.Tensor with the same shape as x. """ with tf.variable_scope(name, default_name="dense_relu_dense"): io_channels = x.shape.dims[-1] stddev = (hidden_channels.size * io_channels.size)**-0.25 io = mtf.Dimension("io", 2) w = mtf.get_variable( x.mesh, "kernel", mtf.Shape([io, io_channels, hidden_channels]), initializer=tf.random_normal_initializer(stddev=stddev), activation_dtype=x.dtype) wi, wo = mtf.unstack(w, io) h = mtf.relu(mtf.einsum([x, wi])) if dropout != 0.0: h = mtf.dropout(h, 1.0 - dropout, noise_shape=h.shape - dropout_broadcast_dims) return mtf.einsum([h, wo])
def multihead_attention(query_antecedent, memory_antecedent, mask, kv_channels, heads, dropout=0.0, dropout_broadcast_dims=None, name="multihead_attention"): """Multihead scaled-dot-product attention with input/output transformations. In order to use only one variable containing the four weight matrices packed together, we insist that the query and memory antecedents have the same dimensionality (io_channels) and that the keys and values have the same dimensionality (kv_channels). Args: query_antecedent: a mtf.Tensor with shape [<batch_dims>, query_length, io_channels] memory_antecedent: a mtf.Tensor with shape [batch, memory_length, io_channels] (optional) mask: mask Tensor (see attention_mask()) kv_channels: a mtf.Dimension (the size of the key and value vectors) heads: a mtf.Dimension (the number of heads) dropout: a floating point value dropout_broadcast_dims: an optional list of mtf.Dimension name: an optional string. Returns: A mtf.Tensor with shape [batch, query_length, io_channels] Raises: ValueError: if the dimensions do not match. """ batch_dims = query_antecedent.shape.dims[:-2] query_length, io_channels = query_antecedent.shape.dims[-2:] with tf.variable_scope(name, default_name="multihead_attention", values=[query_antecedent, memory_antecedent]): q_var, k_var, v_var, o_var = multihead_attention_vars( query_antecedent.mesh, heads, io_channels, kv_channels, query_antecedent.dtype) if memory_antecedent is None: memory_antecedent = rename_length_to_memory_length( query_antecedent, query_length.name) memory_batch_dims = memory_antecedent.shape.dims[:-2] memory_length, memory_channels = memory_antecedent.shape.dims[-2:] if memory_batch_dims != batch_dims: raise ValueError("memory batch must equal query batch") if memory_channels != io_channels: raise ValueError("memory channels must equal query channels") q = mtf.einsum([query_antecedent, q_var], mtf.Shape(batch_dims + [heads, query_length, kv_channels])) k = mtf.einsum([memory_antecedent, k_var], mtf.Shape(batch_dims + [heads, memory_length, kv_channels])) v = mtf.einsum([memory_antecedent, v_var], mtf.Shape(batch_dims + [heads, memory_length, kv_channels])) o = dot_product_attention(q, k, v, mask, dropout, dropout_broadcast_dims) return mtf.einsum([o, o_var], mtf.Shape(batch_dims + [query_length, io_channels]))
def masked_local_attention_1d(query_antecedent, memory_antecedent, kv_channels, heads, block_length=128, name=None): """Attention to the source position and a neighborhood to the left of it. The sequence is divided into blocks of length block_size. Attention for a given query position can only see memory positions less than or equal to the query position, in the corresponding block and the previous block. Args: query_antecedent: a mtf.Tensor with shape [batch, query_length, io_channels] memory_antecedent: a mtf.Tensor with shape [batch, memory_length, io_channels] (optional). Currently, memory_length must have the same size as query_length, but a different name. kv_channels: a mtf.Dimension (the size of the key and value vectors) heads: a mtf.Dimension (the number of heads) block_length: an integer, representing receptive fields for attention. name: an optional string. Returns: a Tensor of shape [batch, query_length, io_channels] Raises: ValueError: if channels or depth don't match. """ with tf.variable_scope(name, default_name="multihead_attention", values=[query_antecedent, memory_antecedent]): batch, query_length, io_channels = query_antecedent.shape.dims q_var, k_var, v_var, o_var = multihead_attention_vars( query_antecedent.mesh, heads, io_channels, kv_channels, query_antecedent.dtype) if memory_antecedent is None: memory_antecedent = rename_length_to_memory_length( query_antecedent, query_length.name) memory_batch, memory_length, memory_channels = memory_antecedent.shape.dims if memory_batch != batch: raise ValueError("memory batch must equal query batch") if memory_channels != io_channels: raise ValueError("memory channels must equal query channels") # Get query q, keys k and values v. q = mtf.einsum([query_antecedent, q_var], mtf.Shape([batch, heads, query_length, kv_channels])) k = mtf.einsum([memory_antecedent, k_var], mtf.Shape([batch, heads, memory_length, kv_channels])) v = mtf.einsum([memory_antecedent, v_var], mtf.Shape([batch, heads, memory_length, kv_channels])) # Let's assume for now we don't have padding and the block length equally # divides the memory length. block_length = (query_length.size if query_length.size < block_length * 2 else block_length) blength = mtf.Dimension("block_length", block_length) mlength = mtf.Dimension("mem_block_length", block_length) num_blocks = mtf.Dimension("num_blocks", query_length.size // block_length) q = mtf.reshape( q, mtf.Shape([batch, heads, num_blocks, blength, kv_channels])) k = mtf.reshape( k, mtf.Shape([batch, heads, num_blocks, mlength, kv_channels])) v = mtf.reshape( v, mtf.Shape([batch, heads, num_blocks, mlength, kv_channels])) # compute attention for the first query block. def first_block_attention(): """Compute attention for the first block.""" first_q = mtf.slice(q, 0, 1, num_blocks.name) first_k = mtf.slice(k, 0, 1, num_blocks.name) first_v = mtf.slice(v, 0, 1, num_blocks.name) first_output = dot_product_attention(first_q, first_k, first_v, mask=None) return first_output # Attention for first block, since query_length = key_length. first_output = first_block_attention() # Concatenate two adjacent blocks to compute the overlapping memory block. def local(x): """Helper function to get memory blocks.""" prev_block = mtf.slice(x, 0, num_blocks.size - 1, num_blocks.name) cur_block = mtf.slice(x, 1, num_blocks.size - 1, num_blocks.name) local_block = mtf.concat([prev_block, cur_block], mlength.name) return local_block local_k = local(k) local_v = local(v) # Calculate the causal mask to avoid peeking into the future. We compute # this once and reuse it for all blocks since the block_size is known. mlength = local_k.shape.dims[3] mask = attention_bias_local_block(query_antecedent.mesh, blength, mlength) # Remove the first block from q since we already computed that. tail_q = mtf.slice(q, 1, num_blocks.size - 1, num_blocks.name) tail_output = dot_product_attention(tail_q, local_k, local_v, mask=mask) # Now concatenate the first and rest of the blocks. final_output = mtf.concat([first_output, tail_output], num_blocks.name) final_output = mtf.reshape( final_output, mtf.Shape([batch, heads, query_length, kv_channels])) return mtf.einsum([final_output, o_var], mtf.Shape([batch, query_length, io_channels]))
def local_self_attention_spatial_blocks(query_antecedent, kv_channels, heads, memory_w_dim=None, mask_right=False, name=None): """Attention to the source position and a neighborhood to the left or right. The sequence is divided into blocks of length block_size. Attention for a given query position can only see memory positions less than or equal to the query position, in the corresponding block and the previous block. Args: query_antecedent: a mtf.Tensor with shape [batch, num_h_blocks, num_w_blocks, h_dim, w_dim, io_channels] must have the same size as query_length, but a different name. kv_channels: a mtf.Dimension (the size of the key and value vectors) heads: a mtf.Dimension (the number of heads) memory_w_dim: mtf Dimension, for the memory width block. mask_right: bool, flag specifying whether we mask out attention to the right for the decoder. name: an optional string. Returns: a Tensor of shape [batch, num_h_blocks, num_w_blocks, h_dim, w_dim, io_channels] Raises: ValueError: if channels or depth don't match. """ with tf.variable_scope(name, default_name="multihead_attention", values=[query_antecedent]): w_dim, io_channels = query_antecedent.shape.dims[-2:] batch, num_w_blocks = query_antecedent.shape.dims[:2] q_var, k_var, v_var, o_var = multihead_attention_vars( query_antecedent.mesh, heads, io_channels, kv_channels, query_antecedent.dtype) # Rename dimensions for the memory height and width. memory_antecedent = mtf.rename_dimension(query_antecedent, w_dim.name, memory_w_dim.name) # Call einsum over the query and memory to get query q, keys k and values v. q = mtf.einsum([query_antecedent, q_var], mtf.Shape( [batch, heads, num_w_blocks, w_dim, kv_channels])) k = mtf.einsum([memory_antecedent, k_var], mtf.Shape( [batch, heads, num_w_blocks, w_dim, kv_channels])) v = mtf.einsum([memory_antecedent, v_var], mtf.Shape( [batch, heads, num_w_blocks, w_dim, kv_channels])) # Halo exchange for memory blocks. if memory_w_dim is not None: k, v = local_1d_halo_exchange(k, v, num_w_blocks, w_dim, memory_w_dim, mask_right) # Calculate the causal mask to avoid peeking into the future. We compute # this once and reuse it for all blocks since the block_size is known. mask = None if mask_right: mask = attention_bias_local_block(query_antecedent.mesh, w_dim, memory_w_dim) output = dot_product_attention(q, k, v, mask=mask) return mtf.einsum([output, o_var], mtf.Shape([batch, num_w_blocks, w_dim, io_channels]))
def masked_local_attention_1d(x, kv_channels, heads, window_size=128, master_dtype=tf.float32, slice_dtype=tf.float32, length_per_split=None, name=None): """Attention to the source position and a neighborhood to the left of it. Attention for a given query position p can only see memory positions in the range (p - window_size, p]. Args: x: a mtf.Tensor with shape batch_dims + [length, io_channels] kv_channels: a mtf.Dimension (the size of the key and value vectors) heads: a mtf.Dimension (the number of heads) window_size: an integer master_dtype: a tf.dtype slice_dtype: a tf.dtype length_per_split: an optional integer indicating the part of the length dimension per processor. You can omit if the length dimension is not split. name: an optional string. Returns: a Tensor with the same shape as x Raises: ValueError: if channels or depth don't match. """ with tf.variable_scope(name, default_name="masked_local_attention_1d", values=[x]): batch_dims = x.shape.dims[:-2] length, io_channels = x.shape.dims[-2:] q_var, k_var, v_var, o_var = multihead_attention_vars( x.mesh, heads, io_channels, kv_channels, master_dtype, slice_dtype, x.dtype) # Get query q, keys k and values v. qkv_shape = mtf.Shape(batch_dims + [heads, length, kv_channels]) q = mtf.einsum([x, q_var], qkv_shape) k = mtf.einsum([x, k_var], qkv_shape) v = mtf.einsum([x, v_var], qkv_shape) # Choose a suitable block size. # We choose the greatest divisor of length_per_split less than or equal # to max(window_size, 128) if length_per_split is None: length_per_split = length.size block_length = max(window_size, 128) while length_per_split % block_length != 0: block_length -= 1 query_block_length = mtf.Dimension("query_block_length", block_length) memory_block_length = mtf.Dimension("memory_block_length", block_length) # The num_blocks dimension gets the same name as the length dimension, # so it will be split in the same way. num_blocks = mtf.Dimension(length.name, length.size // block_length) q_shape = batch_dims + [ heads, num_blocks, query_block_length, kv_channels ] kv_shape = batch_dims + [ heads, num_blocks, memory_block_length, kv_channels ] q = mtf.reshape(q, q_shape) k = mtf.reshape(k, kv_shape) v = mtf.reshape(v, kv_shape) # augment the keys and values for each block with keys and values for # the previous window_size timesteps. k = mtf.left_halo_exchange(k, num_blocks, memory_block_length, window_size) v = mtf.left_halo_exchange(v, num_blocks, memory_block_length, window_size) padded_memory_block_length = mtf.Dimension("memory_block_length", window_size + block_length) mpos = mtf.range(x.mesh, padded_memory_block_length, tf.float32) qpos = mtf.range(x.mesh, query_block_length, tf.float32) + window_size # prevent looking forward mask = mtf.cast(mtf.greater(mpos, qpos), x.dtype) * -1e9 # prevent looking >=block_length timesteps backward mask += mtf.cast(mtf.less_equal(mpos, qpos - block_length), x.dtype) * -1e9 # Note: The first window_size-1 positions can see back into pre-time # where all the keys and values are zero. We could mask this out, but we # don't. o = dot_product_attention(q, k, v, mask=mask) o = mtf.reshape(o, batch_dims + [heads, length, kv_channels]) return mtf.einsum([o, o_var], mtf.Shape(batch_dims + [length, io_channels]))
def grow_topk(i, alive_seq, alive_log_probs, states=None): r"""Inner beam search loop. This function takes the current alive sequences, and grows them to topk sequences where k = 2*beam. We use 2*beam because, we could have beam_size number of sequences that might hit <EOS> and there will be no alive sequences to continue. With 2*beam_size, this will not happen. This relies on the assumption the vocab size is > beam size. If this is true, we'll have at least beam_size non <EOS> extensions if we extract the next top 2*beam words. Length penalty is given by = (5+len(decode)/6) ^ -\alpha. Pls refer to https://arxiv.org/abs/1609.08144. Args: i: loop index alive_seq: Topk sequences decoded so far [batch, beam, length] alive_log_probs: probabilities of these sequences. [batch, beam] states: optional list of mtf.Tensor Returns: Tuple of (Topk sequences extended by the next word, The log probs of these sequences, The scores with length penalty of these sequences, Flags indicating which of these sequences have finished decoding, list of transformed decoding states) """ logits, new_states = logits_fn(i, alive_seq, states) batch_dim, beam_dim, vocab_dim = logits.shape.dims # Convert logits to normalized log probs candidate_log_probs = mtf.log_softmax(logits, vocab_dim) # Multiply the probabilities by the current probabilities of the beam. # (batch_size, beam_size, vocab_size) + (batch_size, beam_size, 1) log_probs = candidate_log_probs + alive_log_probs length_penalty = mtf.pow(((5. + mtf.cast(i + 1, logits.dtype)) / 6.), alpha) curr_scores = log_probs / length_penalty # scores have shape [batch, beam, vocab] beam_and_vocab_dim = mtf.Dimension( "beam_and_vocab", beam_dim.size * vocab_dim.size) flat_shape = mtf.Shape([batch_dim, beam_and_vocab_dim]) double_beam = mtf.Dimension("double_beam", beam_dim.size * 2) # Flatten out (beam_size, vocab_size) probs in to a list of possibilities flat_curr_scores = mtf.reshape(curr_scores, flat_shape) top_ids, top_scores = mtf.top_k( flat_curr_scores, reduced_dim=beam_and_vocab_dim, new_dim=double_beam) # Recovering the log probs because we will need to send them back top_log_probs = top_scores * length_penalty # Work out what beam the top probs are in. top_beam_index = top_ids // vocab_dim.size top_ids %= vocab_dim.size # Unflatten the ids def my_gather(tensor): return mtf.gather( tensor, top_beam_index, beam_dim, output_shape=mtf.Shape( [double_beam if d == beam_dim else d for d in tensor.shape.dims])) # Gather up the most probable 2*beams both for the ids and finished_in_alive # bools top_seq = my_gather(alive_seq) if states: states = [my_gather(state) for state in new_states] # Append the most probable alive top_seq += top_ids * mtf.one_hot(i, length_dim, dtype=tf.int32) top_finished = mtf.equal(top_ids, eos_id) return top_seq, top_log_probs, top_scores, top_finished, states
def beam_search(logits_fn, initial_ids, alpha, states=None, eos_id=EOS_ID, stop_early=True, decode_length=None, use_tpu=True, dtype=tf.float32): """Beam search with length penalties. Requires a function that can take the currently decoded symbols and return the logits for the next symbol. The implementation is inspired by https://arxiv.org/abs/1609.08144. When running, the beam search steps can be visualized by using tfdbg to watch the operations generating the output ids for each beam step. These operations have the pattern: (alive|finished)_topk_(seq,scores) Operations marked `alive` represent the new beam sequences that will be processed in the next step. Operations marked `finished` represent the completed beam sequences, which may be padded with 0s if no beams finished. Operations marked `seq` store the full beam sequence for the time step. Operations marked `scores` store the sequence's final log scores. The beam search steps will be processed sequentially in order, so when capturing observed from these operations, tensors, clients can make assumptions about which step is being recorded. WARNING: Assumes 2nd dimension of tensors in `states` and not invariant, this means that the shape of the 2nd dimension of these tensors will not be available (i.e. set to None) inside logits_fn. Args: logits_fn: Interface to the model, to provide logits. Shoud take: step_num - mtf Scalar ids - mtf Tensor with shape [batch, beam, length] Should return: logits - [batch, beam, vocab_size], dtype=dtype initial_ids: a mtf.Tensor with shape [batch_dim, beam_dim, length_dim]) alpha: alpha for length penalty. states: list of mtf.Tensor eos_id: ID for end of sentence. stop_early: a boolean - stop once best sequence is provably determined. decode_length: a mtf Scalar of dtype tf.int32 - maximum length of decodes use_tpu: a boolean dtype: a tf.dtype Returns: Tuple of (decoded beams [batch, beam, length] decoding probabilities [batch, beam_size]) """ batch_dim, beam_dim, length_dim = initial_ids.shape.dims mesh = initial_ids.mesh batch_by_beam = mtf.Shape([batch_dim, beam_dim]) initial_log_probs = mtf.broadcast( mtf.one_hot( mtf.constant(mesh, 0, dtype=tf.int32), beam_dim, on_value=0.0, off_value=-INF, dtype=dtype), batch_by_beam) length_scalar = mtf.constant(mesh, length_dim.size, dtype=tf.int32) if decode_length is None: decode_length = length_scalar else: decode_length = mtf.minimum(decode_length, length_scalar) alive_log_probs = initial_log_probs alive_seq = initial_ids # Finished will keep track of all the sequences that have finished so far # Finished log probs will be negative infinity in the beginning # finished_flags will keep track of booleans finished_seq = initial_ids finished_scores = mtf.constant(mesh, -INF, batch_by_beam, dtype=dtype) # Setting the scores of the initial to negative infinity. finished_flags = mtf.constant(mesh, False, batch_by_beam, tf.bool) def grow_finished(finished_seq, finished_scores, finished_flags, curr_seq, curr_scores, curr_finished): """Given sequences and scores, will gather the top k=beam size sequences. Args: finished_seq: Current finished sequences. [batch, beam, length] finished_scores: scores for each of these sequences. [batch, beam] finished_flags: finished bools for each of these sequences. [batch, beam] curr_seq: current topk sequence that has been grown by one position. [batch, beam, length] curr_scores: scores for each of these sequences. [batch, beam] curr_finished: Finished flags for each of these sequences. [batch, beam] Returns: Tuple of (Topk sequences based on scores, log probs of these sequences, Finished flags of these sequences, None (no states)) """ # Set the scores of the unfinished seq in curr_seq to large negative # values curr_scores += (1. - mtf.cast(curr_finished, curr_scores.dtype)) * -INF unused_batch_dim, beam_dim, unused_length_dim = finished_seq.shape.dims # concatenating the sequences and scores along beam axis def _my_concat(a, b): a = mtf.rename_dimension(a, "beam", "triple_beam") b = mtf.rename_dimension(b, "double_beam", "triple_beam") return mtf.concat([a, b], "triple_beam") curr_finished_seq = _my_concat(finished_seq, curr_seq) curr_finished_scores = _my_concat(finished_scores, curr_scores) curr_finished_flags = _my_concat(finished_flags, curr_finished) return compute_topk_scores_and_seq( curr_finished_seq, curr_finished_scores, curr_finished_scores, curr_finished_flags, beam_dim, "grow_finished", states=None) def grow_alive(curr_seq, curr_scores, curr_log_probs, curr_finished, states): """Given sequences and scores, will gather the top k=beam size sequences. Args: curr_seq: current topk sequence that has been grown by one position. [batch, beam, length] curr_scores: scores for each of these sequences. [batch_size, beam_size] curr_log_probs: log probs for each of these sequences. [batch, beam] curr_finished: Finished flags for each of these sequences. [batch, beam] states: list of mtf.Tensor Returns: Tuple of (Topk sequences based on scores, log probs of these sequences, Finished flags of these sequences) """ # Set the scores of the finished seq in curr_seq to large negative # values curr_scores += mtf.cast(curr_finished, curr_scores.dtype) * -INF return compute_topk_scores_and_seq(curr_seq, curr_scores, curr_log_probs, curr_finished, beam_dim, "grow_alive", states) def grow_topk(i, alive_seq, alive_log_probs, states=None): r"""Inner beam search loop. This function takes the current alive sequences, and grows them to topk sequences where k = 2*beam. We use 2*beam because, we could have beam_size number of sequences that might hit <EOS> and there will be no alive sequences to continue. With 2*beam_size, this will not happen. This relies on the assumption the vocab size is > beam size. If this is true, we'll have at least beam_size non <EOS> extensions if we extract the next top 2*beam words. Length penalty is given by = (5+len(decode)/6) ^ -\alpha. Pls refer to https://arxiv.org/abs/1609.08144. Args: i: loop index alive_seq: Topk sequences decoded so far [batch, beam, length] alive_log_probs: probabilities of these sequences. [batch, beam] states: optional list of mtf.Tensor Returns: Tuple of (Topk sequences extended by the next word, The log probs of these sequences, The scores with length penalty of these sequences, Flags indicating which of these sequences have finished decoding, list of transformed decoding states) """ logits, new_states = logits_fn(i, alive_seq, states) batch_dim, beam_dim, vocab_dim = logits.shape.dims # Convert logits to normalized log probs candidate_log_probs = mtf.log_softmax(logits, vocab_dim) # Multiply the probabilities by the current probabilities of the beam. # (batch_size, beam_size, vocab_size) + (batch_size, beam_size, 1) log_probs = candidate_log_probs + alive_log_probs length_penalty = mtf.pow(((5. + mtf.cast(i + 1, logits.dtype)) / 6.), alpha) curr_scores = log_probs / length_penalty # scores have shape [batch, beam, vocab] beam_and_vocab_dim = mtf.Dimension( "beam_and_vocab", beam_dim.size * vocab_dim.size) flat_shape = mtf.Shape([batch_dim, beam_and_vocab_dim]) double_beam = mtf.Dimension("double_beam", beam_dim.size * 2) # Flatten out (beam_size, vocab_size) probs in to a list of possibilities flat_curr_scores = mtf.reshape(curr_scores, flat_shape) top_ids, top_scores = mtf.top_k( flat_curr_scores, reduced_dim=beam_and_vocab_dim, new_dim=double_beam) # Recovering the log probs because we will need to send them back top_log_probs = top_scores * length_penalty # Work out what beam the top probs are in. top_beam_index = top_ids // vocab_dim.size top_ids %= vocab_dim.size # Unflatten the ids def my_gather(tensor): return mtf.gather( tensor, top_beam_index, beam_dim, output_shape=mtf.Shape( [double_beam if d == beam_dim else d for d in tensor.shape.dims])) # Gather up the most probable 2*beams both for the ids and finished_in_alive # bools top_seq = my_gather(alive_seq) if states: states = [my_gather(state) for state in new_states] # Append the most probable alive top_seq += top_ids * mtf.one_hot(i, length_dim, dtype=tf.int32) top_finished = mtf.equal(top_ids, eos_id) return top_seq, top_log_probs, top_scores, top_finished, states def inner_loop(i, alive_seq, alive_log_probs, finished_seq, finished_scores, finished_flags, *states): """Inner beam search loop. There are three groups of tensors, alive, finished, and topk. The alive group contains information about the current alive sequences The topk group contains information about alive + topk current decoded words the finished group contains information about finished sentences, that is, the ones that have decoded to <EOS>. These are what we return. The general beam search algorithm is as follows: While we haven't terminated (pls look at termination condition) 1. Grow the current alive to get beam*2 topk sequences 2. Among the topk, keep the top beam_size ones that haven't reached EOS into alive 3. Among the topk, keep the top beam_size ones have reached EOS into finished Repeat To make things simple with using fixed size tensors, we will end up inserting unfinished sequences into finished in the beginning. To stop that we add -ve INF to the score of the unfinished sequence so that when a true finished sequence does appear, it will have a higher score than all the unfinished ones. Args: i: loop index alive_seq: Topk sequences decoded so far [batch_size, beam_size, i+1] alive_log_probs: probabilities of the beams. [batch_size, beam_size] finished_seq: Current finished sequences. [batch_size, beam_size, i+1] finished_scores: scores for each of these sequences. [batch_size, beam_size] finished_flags: finished bools for each of these sequences. [batch_size, beam_size] *states: mtf Tensors Returns: Tuple of (Incremented loop index New alive sequences, Log probs of the alive sequences, New finished sequences, Scores of the new finished sequences, Flags indicating which sequence in finished as reached EOS, dict of final decoding states) """ # Each inner loop, we carry out three steps: # 1. Get the current topk items. # 2. Extract the ones that have finished and haven't finished # 3. Recompute the contents of finished based on scores. (top2k_seq, top2k_log_probs, top2k_scores, top2k_finished, top2k_states) = grow_topk(i, alive_seq, alive_log_probs, states) alive_seq, alive_log_probs, _, states = grow_alive( top2k_seq, top2k_scores, top2k_log_probs, top2k_finished, top2k_states) finished_seq, finished_scores, finished_flags, _ = grow_finished( finished_seq, finished_scores, finished_flags, top2k_seq, top2k_scores, top2k_finished) return (i + 1, alive_seq, alive_log_probs, finished_seq, finished_scores, finished_flags) + tuple(states) def _is_finished(i, unused_alive_seq, alive_log_probs, unused_finished_seq, finished_scores, finished_in_finished, *unused_states): """Checking termination condition. We terminate when we decoded up to decode_length or the lowest scoring item in finished has a greater score that the highest prob item in alive divided by the max length penalty Args: i: loop index alive_log_probs: probabilities of the beams. [batch_size, beam_size] finished_scores: scores for each of these sequences. [batch_size, beam_size] finished_in_finished: finished bools for each of these sequences. [batch_size, beam_size] Returns: Bool. """ # TODO(noam): support a different decode length... # decode_length = mtf.constant(mesh, length_dim.size, dtype=tf.int32) # del alive_log_probs, finished_scores, finished_in_finished # return mtf.less(i, length_dim.size) if not stop_early: return mtf.less(i, decode_length) max_length_penalty = mtf.pow( ((5. + mtf.cast(decode_length, finished_scores.dtype)) / 6.), alpha) # The best possible score of the most likely alive sequence. lower_bound_alive_scores = mtf.gather( alive_log_probs, mtf.constant(mesh, 0, dtype=tf.int32), beam_dim) / max_length_penalty # Now to compute the lowest score of a finished sequence in finished # If the sequence isn't finished, we multiply it's score by 0. since # scores are all -ve, taking the min will give us the score of the lowest # finished item. lowest_score_of_finished_in_finished = mtf.reduce_min( finished_scores * mtf.cast(finished_in_finished, finished_scores.dtype), reduced_dim=beam_dim) # If none of the sequences have finished, then the min will be 0 and # we have to replace it by -ve INF if it is. The score of any seq in alive # will be much higher than -ve INF and the termination condition will not # be met. lowest_score_of_finished_in_finished += ( (1. - mtf.cast(mtf.reduce_any( finished_in_finished, reduced_dim=beam_dim), finished_scores.dtype)) * -INF) bound_is_met = mtf.reduce_all( mtf.greater(lowest_score_of_finished_in_finished, lower_bound_alive_scores)) return mtf.logical_and( mtf.less(i, decode_length), mtf.logical_not(bound_is_met)) initial_step_num = mtf.constant(mesh, 0, dtype=tf.int32) while_loop_inputs = [ initial_step_num, alive_seq, alive_log_probs, finished_seq, finished_scores, finished_flags] + states (_, alive_seq, alive_log_probs, finished_seq, finished_scores, finished_flags) = mtf.while_loop( _is_finished, inner_loop, while_loop_inputs, num_loop_vars=None if use_tpu else 6)[:6] # Accounting for corner case: It's possible that no sequence in alive for a # particular batch item ever reached EOS. In that case, we should just copy # the contents of alive for that batch item. tf.reduce_any(finished_flags, 1) # if 0, means that no sequence for that batch index had reached EOS. We need # to do the same for the scores as well. finished_seq = mtf.where( mtf.reduce_any(finished_flags, reduced_dim=beam_dim), finished_seq, alive_seq) finished_scores = mtf.where( mtf.reduce_any(finished_flags, reduced_dim=beam_dim), finished_scores, alive_log_probs) return finished_seq, finished_scores
def multihead_self_attention_memory_compressed(x, mask_right, compression_factor, kv_channels, heads, dropout=0.0, dropout_broadcast_dims=None, master_dtype=tf.float32, slice_dtype=tf.float32, name="multihead_attention"): """Memory-compressed self-attention. The memory is first average-pooled (strided) to make it shorter by a factor of compression_factor. Args: x: a mtf.Tensor with shape [<batch_dims>, query_length, io_channels] mask_right: a boolean compression_factor: an integer kv_channels: a mtf.Dimension (the size of the key and value vectors) heads: a mtf.Dimension (the number of heads) dropout: a floating point value dropout_broadcast_dims: an optional list of mtf.Dimension master_dtype: a tf.dtype slice_dtype: a tf.dtype name: an optional string. Returns: A mtf.Tensor with shape [batch, query_length, io_channels] Raises: ValueError: if the dimensions do not match. """ batch_dims = x.shape.dims[:-2] length, io_channels = x.shape.dims[-2:] with tf.variable_scope(name, default_name="compressed_attention", values=[x]): q_var, k_var, v_var, o_var = multihead_attention_vars( x.mesh, heads, io_channels, kv_channels, master_dtype, slice_dtype, x.dtype) memory_antecedent = compress_mean(x, length, compression_factor) memory_antecedent = rename_length_to_memory_length(memory_antecedent) memory_length = memory_antecedent.shape.dims[-2] q = mtf.einsum([x, q_var], mtf.Shape(batch_dims + [heads, length, kv_channels])) k = mtf.einsum([memory_antecedent, k_var], mtf.Shape(batch_dims + [heads, memory_length, kv_channels])) v = mtf.einsum([memory_antecedent, v_var], mtf.Shape(batch_dims + [heads, memory_length, kv_channels])) if mask_right: query_pos = mtf.range(x.mesh, length, dtype=tf.int32) memory_pos = (mtf.range(x.mesh, memory_length, dtype=tf.int32) * compression_factor + (compression_factor - 1)) mask = mtf.cast(mtf.greater(memory_pos, query_pos), x.dtype) * -1e9 else: mask = None o = dot_product_attention(q, k, v, mask, dropout, dropout_broadcast_dims, extra_logit=0.0) return mtf.einsum([o, o_var], mtf.Shape(batch_dims + [length, io_channels]))