def AttentionQKV(d_feature, n_heads=1, dropout=0.0, mode='train'): """Transformer-style multi-headed attention. Accepts inputs of the form q, k, v, mask. Args: d_feature: int: dimensionality of feature embedding n_heads: int: number of attention heads dropout: float: dropout rate mode: str: 'train' or 'eval' Returns: Multi-headed self-attention result and the mask. """ return [ cb.Parallel( core.Dense(d_feature), core.Dense(d_feature), core.Dense(d_feature), ), PureAttention( # pylint: disable=no-value-for-parameter n_heads=n_heads, dropout=dropout, mode=mode), core.Dense(d_feature), ]
def CausalAttention(d_feature, n_heads=1, dropout=0.0, mode='train'): """Transformer-style multi-headed causal attention. # TODO(jonni,lukaszkaiser): standardize and improve layer comments. Accepts inputs of the form x and constructs (q, k, v) and causal mask from x. Args: d_feature: int: dimensionality of feature embedding n_heads: int: number of attention heads dropout: float: dropout rate mode: str: 'train' or 'eval' Returns: Multi-headed self-attention result. """ return [ cb.Dup(), cb.Parallel([], CausalMask(axis=-2)), # pylint: disable=no-value-for-parameter Attention(d_feature, n_heads=n_heads, dropout=dropout, mode=mode), cb.Parallel([], cb.Drop()), # x ]
def GeneralGRUCell(candidate_transform, memory_transform_fn=None, gate_nonlinearity=core.Sigmoid, candidate_nonlinearity=core.Tanh, dropout_rate_c=0.1, sigmoid_bias=0.5): r"""Parametrized Gated Recurrent Unit (GRU) cell construction. GRU update equations: $$ Update gate: u_t = \sigmoid(U' * s_{t-1} + B') $$ $$ Reset gate: r_t = \sigmoid(U'' * s_{t-1} + B'') $$ $$ Candidate memory: c_t = \tanh(U * (r_t \odot s_{t-1}) + B) $$ $$ New State: s_t = u_t \odot s_{t-1} + (1 - u_t) \odot c_t $$ See combinators.Gate for details on the gating function. Args: candidate_transform: Transform to apply inside the Candidate branch. Applied before nonlinearities. memory_transform_fn: Optional transformation on the memory before gating. gate_nonlinearity: Function to use as gate activation. Allows trying alternatives to Sigmoid, such as HardSigmoid. candidate_nonlinearity: Nonlinearity to apply after candidate branch. Allows trying alternatives to traditional Tanh, such as HardTanh dropout_rate_c: Amount of dropout on the transform (c) gate. Dropout works best in a GRU when applied exclusively to this branch. sigmoid_bias: Constant to add before sigmoid gates. Generally want to start off with a positive bias. Returns: A model representing a GRU cell with specified transforms. """ gate_block = [ # u_t candidate_transform(), core.AddConstant(constant=sigmoid_bias), gate_nonlinearity(), ] reset_block = [ # r_t candidate_transform(), core.AddConstant( constant=sigmoid_bias), # Want bias to start positive. gate_nonlinearity(), ] candidate_block = [ cb.Dup(), reset_block, cb.Multiply(), # Gate S{t-1} with sigmoid(candidate_transform(S{t-1})) candidate_transform(), # Final projection + tanh to get Ct candidate_nonlinearity(), # Candidate gate # Only apply dropout on the C gate. Paper reports 0.1 as a good default. core.Dropout(rate=dropout_rate_c) ] memory_transform = memory_transform_fn() if memory_transform_fn else [] return cb.Model( cb.Dup(), cb.Dup(), cb.Parallel(memory_transform, gate_block, candidate_block), cb.Gate(), )
def test_parallel_no_ops(self): layer = cb.Parallel([], None) input_shape = ((3, 2), (4, 7)) expected_shape = ((3, 2), (4, 7)) output_shape = base.check_shape_agreement(layer, input_shape) self.assertEqual(output_shape, expected_shape)
def test_parallel_div_div(self): layer = cb.Parallel(core.Div(divisor=0.5), core.Div(divisor=3.0)) input_shape = ((3, 2), (4, 7)) expected_shape = ((3, 2), (4, 7)) output_shape = base.check_shape_agreement(layer, input_shape) self.assertEqual(output_shape, expected_shape)
def test_parallel_dup_dup(self): layer = cb.Parallel(cb.Dup(), cb.Dup()) input_shape = ((3, 2), (4, 7)) expected_shape = ((3, 2), (3, 2), (4, 7), (4, 7)) output_shape = base.check_shape_agreement(layer, input_shape) self.assertEqual(output_shape, expected_shape)