def attention_bias_local_2d_block(mesh, h_dim, w_dim, memory_h_dim, memory_w_dim, dtype=tf.int32): """Bias for attention for local blocks where attention to right is disallowed. Create the bias matrix by using two separate masks, one for the memory part which doesn't overlap with the query and second which interacts with the query and should be disallowed to look to the right of the current query position. Args: mesh: a MeshTensorflow object h_dim: a mtf.Dimension w_dim: a mtf.Dimension memory_h_dim: a mtf.Dimension memory_w_dim: a mtf.Dimension dtype: a tf.dtype Returns: a mtf.Tensor with shape [block_length, memory_length] """ memory_height = mtf.Dimension(memory_h_dim.name, h_dim.size) memory_width = mtf.Dimension(memory_w_dim.name, w_dim.size) mask_top_visible = mtf.zeros(mesh, [h_dim, memory_height], dtype=dtype) mask_left_visible = mtf.zeros(mesh, [w_dim, memory_width], dtype=dtype) mask_query = mtf.greater(mtf.range(mesh, memory_height, dtype=tf.int32), mtf.range(mesh, memory_width, dtype=dtype)) width_mask = mtf.concat([mask_left_visible, mask_query], memory_width.name) mask = mtf.cast(mtf.concat([mask_top_visible, width_mask], memory_height.name), dtype=tf.float32) * -1e9 return mask
def attention_bias_local_block(mesh, block_length, memory_length, dtype=tf.int32): """Bias for attention for local blocks where attention to right is disallowed. Create the bias matrix by using two separate masks, one for the memory part which doesn't overlap with the query and second which interacts with the query and should be disallowed to look to the right of the current query position. Args: mesh: a MeshTensorflow object block_length: a mtf.Dimension memory_length: a mtf.Dimension dtype: a tf.dtype Returns: a mtf.Tensor with shape [block_length, memory_length] """ memory_length = mtf.Dimension(memory_length.name, block_length.size) memory_mask = mtf.zeros(mesh, [block_length, memory_length], dtype=dtype) mask = mtf.cast(mtf.less(mtf.range(mesh, block_length, dtype=dtype), mtf.range(mesh, memory_length, dtype=dtype)), dtype=dtype) mask = mtf.cast(mtf.concat([memory_mask, mask], memory_length.name), dtype=tf.float32) * -1e9 return mask
def compress_mean(x, dim, compression_factor): """Compress by taking group means. Args: x: a Tensor dim: a dimension in x.shape compression_factor: an integer Returns: a Tensor """ dims = x.shape.dims pos = dims.index(dim) compressed_dim = mtf.Dimension(dim.name, dim.size // compression_factor) compression_factor_dim = mtf.Dimension("compression_factor", compression_factor) new_shape = (dims[:pos] + [compressed_dim, compression_factor_dim] + dims[pos + 1:]) x = mtf.reshape(x, new_shape) x = mtf.reduce_mean(x, reduced_dim=compression_factor_dim) return x
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 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 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 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, new_states, first_selector) = grow_topk(i, alive_seq, alive_log_probs, states) alive_seq, alive_log_probs, _, second_selector = grow_alive( top2k_seq, top2k_scores, top2k_log_probs, top2k_finished) finished_seq, finished_scores, finished_flags, _ = grow_finished( finished_seq, finished_scores, finished_flags, top2k_seq, top2k_scores, top2k_finished) old_beam_dim = mtf.Dimension("old_beam", beam_dim.size) selector = mtf.einsum([ mtf.rename_dimension(first_selector, beam_dim.name, old_beam_dim.name), second_selector ], output_shape=[batch_dim, old_beam_dim, beam_dim]) new_states = [ mtf.einsum([ mtf.rename_dimension(state, beam_dim.name, old_beam_dim.name), mtf.cast(selector, state.dtype) ], reduced_dims=[old_beam_dim], output_shape=state.shape) for state in new_states ] return (i + 1, alive_seq, alive_log_probs, finished_seq, finished_scores, finished_flags) + tuple(new_states)
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]))