def universal_transformer_highway(layer_inputs,
                                  step, hparams,
                                  ffn_unit,
                                  attention_unit,
                                  pad_remover=None):
  """Universal Transformer with highway connection.


  It transforms the state using a block contaaining sel-attention and transition
  function  and wrap the whole block with a highway connection.
  (the new state is a combination of the state and the transformed-state
  based on cary/transform gates.)

  Interesting observation:
    Controlling the cary/transform gate with the original inputs works usually
    better (i.e. hparams.gates_inputs="i")

  Args:
    layer_inputs:
      - state: state
      - inputs: the original embedded inputs (= inputs to the first step)
    step: indicates number of steps taken so far
    hparams: model hyper-parameters.
    ffn_unit: feed-forward unit
    attention_unit: multi-head attention unit
    pad_remover: to mask out padding in convolutional layers (efficiency).

  Returns:
    layer_output:
        new_state: new state
        inputs: the original embedded inputs (= inputs to the first step)

  """

  state, inputs, memory = layer_inputs
  new_state = step_preprocess(state, step, hparams)

  for i in range(hparams.num_inrecurrence_layers):
    with tf.variable_scope("rec_layer_%d" % i):
      new_state = ffn_unit(attention_unit(new_state))

  transformed_state = new_state

  gate_inputs = []
  if "s" in hparams.gates_inputs:
    gate_inputs.append(state)

  if "t" in hparams.gates_inputs:
    gate_inputs.append(transformed_state)

  if "i" in hparams.gates_inputs:
    gate_inputs.append(inputs)

  gate_ffn_layer = hparams.gate_ffn_layer

  transform_gate = _ffn_layer_multi_inputs(
      gate_inputs,
      hparams,
      ffn_layer_type=gate_ffn_layer,
      name="transform",
      bias_initializer=tf.constant_initializer(hparams.transform_bias_init),
      activation=tf.sigmoid,
      pad_remover=pad_remover,
      preprocess=True,
      postprocess=True)

  if hparams.couple_carry_transform_gates:
    carry_gate = tf.subtract(1.0, transform_gate, name="carry")

  else:
    carry_gate = _ffn_layer_multi_inputs(
        gate_inputs,
        hparams,
        ffn_layer_type=gate_ffn_layer,
        name="carry",
        bias_initializer=tf.constant_initializer(-hparams.transform_bias_init),
        activation=tf.sigmoid,
        pad_remover=pad_remover,
        preprocess=True,
        postprocess=True)

  new_state = state * carry_gate + transformed_state * transform_gate

  tf.contrib.summary.scalar("highway_transform_gate_layer",
                            tf.reduce_mean(transform_gate))

  tf.contrib.summary.scalar("highway_carry_gate_layer",
                            tf.reduce_mean(carry_gate))

  return new_state, inputs, memory
def universal_transformer_highway ( layer_inputs , step , hparams , ffn_unit , attention_unit , pad_remover = None ) : state , inputs , memory = layer_inputs new_state = step_preprocess ( state , step , hparams ) for i in range ( hparams . num_inrecurrence_layers ) : with tf . variable_scope ( "rec_layer_%d" % i ) : new_state = ffn_unit ( attention_unit ( new_state ) ) transformed_state = new_state gate_inputs = [ ] if "s" in hparams . gates_inputs : gate_inputs . append ( state ) if "t" in hparams . gates_inputs : gate_inputs . append ( transformed_state ) if "i" in hparams . gates_inputs : gate_inputs . append ( inputs ) gate_ffn_layer = hparams . gate_ffn_layer transform_gate = _ffn_layer_multi_inputs ( gate_inputs , hparams , ffn_layer_type = gate_ffn_layer , name = "transform" , bias_initializer = tf . constant_initializer ( hparams . transform_bias_init ) , activation = tf . sigmoid , pad_remover = pad_remover , preprocess = True , postprocess = True ) if hparams . couple_carry_transform_gates : carry_gate = tf . subtract ( 1.0 , transform_gate , name = "carry" ) else : carry_gate = _ffn_layer_multi_inputs ( gate_inputs , hparams , ffn_layer_type = gate_ffn_layer , name = "carry" , bias_initializer = tf . constant_initializer ( - hparams . transform_bias_init ) , activation = tf . sigmoid , pad_remover = pad_remover , preprocess = True , postprocess = True ) new_state = state * carry_gate + transformed_state * transform_gate tf . contrib . summary . scalar ( "highway_transform_gate_layer" , tf . reduce_mean ( transform_gate ) ) tf . contrib . summary . scalar ( "highway_carry_gate_layer" , tf . reduce_mean ( carry_gate ) ) return new_state , inputs , memory
Universal Transformer with highway connection.


  It transforms the state using a block contaaining sel-attention and transition
  function  and wrap the whole block with a highway connection.
  (the new state is a combination of the state and the transformed-state
  based on cary/transform gates.)

  Interesting observation:
    Controlling the cary/transform gate with the original inputs works usually
    better (i.e. hparams.gates_inputs="i")

  Args:
    layer_inputs:
      - state: state
      - inputs: the original embedded inputs (= inputs to the first step)
    step: indicates number of steps taken so far
    hparams: model hyper-parameters.
    ffn_unit: feed-forward unit
    attention_unit: multi-head attention unit
    pad_remover: to mask out padding in convolutional layers (efficiency).

  Returns:
    layer_output:
        new_state: new state
        inputs: the original embedded inputs (= inputs to the first step)