Added options for configuring different types of processing on layer input and layer output (normalization, dropout, residuals). These settings are configured by common_hparams, and should work across many models. Normalization on layer input instead of after the residual seems to help in learning deep networks. This change breaks current model checkpoints.

PiperOrigin-RevId: 164630450

Author: nshazeer

tensor2tensor/layers/common_hparams.py

@@ -69,8 +69,23 @@ def basic_params1():
       sampling_method="argmax",  # "argmax" or "random"
       problem_choice="adaptive",  # "uniform", "adaptive", "distributed"
       multiply_embedding_mode="sqrt_depth",
+      # Sequences of operations to perform on layer input and layer output.
+      # Used by common_layers.layer_preprocess, common_layers.layer_postprocess
+      # Each character repsesnts an operation:
+      #   d: apply dropout
+      #   n: apply normalization (see norm_type and norm_epsilon)
+      #   a: add layer input (residual connection - only during postprocess)
+      # TODO(noam): The current settings ("", "dan") are the published version
+      # of the transformer.  ("n", "da") seems better for harder-to-learn
+      # models, so it should probably be the default.
+      layer_preprocess_sequence="",
+      layer_postprocess_sequence="dan",
+      # dropout rate to use during layer_preprocess and layer_postprocess
+      layer_prepostprocess_dropout=0.1,
+      # What type of normalization to use
       norm_type="none",  # "batch", layer", "noam", "none".
-      layer_norm_epsilon=1e-6,
+      # epsilon parameter to normalization function
+      norm_epsilon=1e-6,
       symbol_modality_num_shards=16,
       # setting the max length in a minibatch. 0 means default behavior,
       # max_length = hparams.batch_size * length_multiplier

tensor2tensor/layers/common_layers.py

@@ -462,64 +462,136 @@ def layer_norm(x, filters=None, epsilon=1e-6, name=None, reuse=None):
     return result
 
 
-def noam_norm(x, name=None):
+def noam_norm(x, epsilon=1.0, name=None):
   """One version of layer normalization."""
   with tf.name_scope(name, default_name="noam_norm", values=[x]):
     shape = x.get_shape()
     ndims = len(shape)
-    return (tf.nn.l2_normalize(x, ndims - 1, epsilon=1.0) *
+    return (tf.nn.l2_normalize(x, ndims - 1, epsilon=epsilon) *
             tf.sqrt(tf.to_float(shape[-1])))
 
 
-def get_norm(norm_type):
-  """Get the normalizer function."""
+def apply_norm(x, norm_type, depth, epsilon):
+  """Apply Normalization."""
   if norm_type == "layer":
-    return lambda x, name, filters=None, epsilon=1e-6: layer_norm(  # pylint: disable=g-long-lambda
-        x, filters=filters, epsilon=epsilon, name=name)
+    return layer_norm(x, filters=depth, epsilon=epsilon)
   if norm_type == "batch":
-    return tf.layers.batch_normalization
+    return tf.layers.batch_normalization(x, epsilon=epsilon)
   if norm_type == "noam":
-    return noam_norm
+    return noam_norm(x, epsilon)
   if norm_type == "none":
-    return lambda x, name: x
+    return x
   raise ValueError("Parameter normalizer_fn must be one of: 'layer', 'batch',"
                    "'noam', 'none'.")
 
 
-def residual_fn(x,
-                y,
-                norm_type,
-                residual_dropout,
-                filters=None,
-                epsilon=1e-16,
-                name=None,
-                reuse=None):
-  """Returns a function for combining layer input and layer output.
+def layer_prepostprocess(previous_value,
+                         x,
+                         sequence,
+                         dropout_rate,
+                         norm_type,
+                         depth,
+                         epsilon,
+                         name):
+  """Apply a sequence of functions to the input or output of a layer.
+
+  The sequence is specified as a string which may contain the following
+  characters:
+    a: add previous_value
+    n: apply normalization
+    d: apply dropout
 
-  The returned function on x (layer input) and y (layer output) computes:
-    norm_function(x + dropout(y))
+  For example, if sequence=="dna", then the output is
+    previous_value + normalize(dropout(x))
 
   Args:
-    x: tensor, input layer
-    y: tensor, output layer
-    norm_type: string, type of normalizer function
-    residual_dropout: integer, dropout value for residual connection
-    filters: integer, dimension for layer norm, optional
-    epsilon: integer, value of layer norm epsilon
-    name: string, name
-    reuse: bool, whether to reuse
+    previous_value: A Tensor, to be added as a residual connection ('a')
+    x: A Tensor to be transformed.
+    sequence: a string.
+    dropout_rate: a float
+    norm_type: a string (see apply_norm())
+    depth: an integer (size of last dimension of x).
+    epsilon: a float (parameter for normalization)
+    name: a string
 
   Returns:
-    residual layer output with applied norm_fn.
+    a Tensor
   """
-  with tf.variable_scope(
-      name, default_name="residual", values=[x, y], reuse=reuse):
-    norm_fn = get_norm(norm_type)
-    res = x + tf.nn.dropout(y, 1.0 - residual_dropout)
-    if norm_type == "layer":
-      return norm_fn(res, filters=filters, epsilon=epsilon, name=norm_type)
-    else:
-      return norm_fn(res, name=norm_type)
+  with tf.variable_scope(name):
+    for c in sequence:
+      if c == "a":
+        x += previous_value
+      elif c == "n":
+        x = apply_norm(x, norm_type, depth, epsilon)
+      else:
+        assert c == "d", ("Unknown sequence step %s" % c)
+        x = tf.nn.dropout(x, 1.0 - dropout_rate)
+    return x
+
+
+def layer_preprocess(layer_input, hparams):
+  """Apply layer preprocessing.
+
+  See layer_prepostprocess() for details.
+
+  A hyperparemeters object is passed for convenience.  The hyperparameters
+  that may be used are:
+
+    layer_preprocess_sequence
+    layer_prepostprocess_dropout
+    norm_type
+    hidden_size
+    norm_epsilon
+
+  Args:
+    layer_input: a Tensor
+    hparams: a hyperparameters object.
+
+  Returns:
+    a Tensor
+  """
+  assert "a" not in hparams.layer_preprocess_sequence, (
+      "No residual connections allowed in hparams.layer_preprocess_sequence")
+  return layer_prepostprocess(
+      None, layer_input,
+      sequence=hparams.layer_preprocess_sequence,
+      dropout_rate=hparams.layer_prepostprocess_dropout,
+      norm_type=hparams.norm_type,
+      depth=hparams.hidden_size,
+      epsilon=hparams.norm_epsilon,
+      name="layer_prepostprocess")
+
+
+def layer_postprocess(layer_input, layer_output, hparams):
+  """Apply layer postprocessing.
+
+  See layer_prepostprocess() for details.
+
+  A hyperparemeters object is passed for convenience.  The hyperparameters
+  that may be used are:
+
+    layer_postprocess_sequence
+    layer_prepostprocess_dropout
+    norm_type
+    hidden_size
+    norm_epsilon
+
+  Args:
+    layer_input: a Tensor
+    layer_output: a Tensor
+    hparams: a hyperparameters object.
+
+  Returns:
+    a Tensor
+  """
+  return layer_prepostprocess(
+      layer_input, layer_output,
+      sequence=hparams.layer_postprocess_sequence,
+      dropout_rate=hparams.layer_prepostprocess_dropout,
+      norm_type=hparams.norm_type,
+      depth=hparams.hidden_size,
+      epsilon=hparams.norm_epsilon,
+      name="layer_postprocess")
 
 
 def conv_block_internal(conv_fn,

tensor2tensor/layers/common_layers_test.py

@@ -303,74 +303,43 @@ def testDeconvStride2MultiStep(self):
       actual = session.run(a)
     self.assertEqual(actual.shape, (5, 32, 1, 16))
 
-  def testGetNormLayerFn(self):
-    norm_type = "layer"
+  def testApplyNormLayer(self):
     with self.test_session() as session:
-      a = common_layers.get_norm(norm_type)
       x1 = np.random.rand(5, 2, 1, 11)
-      x2 = a(tf.constant(x1, dtype=tf.float32), name="layer", filters=11)
+      x2 = common_layers.apply_norm(
+          tf.constant(x1, dtype=tf.float32), "layer", depth=11, epsilon=1e-6)
       session.run(tf.global_variables_initializer())
       actual = session.run(x2)
     self.assertEqual(actual.shape, (5, 2, 1, 11))
 
-  def testGetNormNoamFn(self):
-    norm_type = "noam"
+  def testApplyNormNoam(self):
     with self.test_session() as session:
-      a = common_layers.get_norm(norm_type)
       x1 = np.random.rand(5, 2, 1, 11)
-      x2 = a(tf.constant(x1, dtype=tf.float32), name="noam")
+      x2 = common_layers.apply_norm(
+          tf.constant(x1, dtype=tf.float32), "noam", depth=11, epsilon=1e-6)
       session.run(tf.global_variables_initializer())
       actual = session.run(x2)
     self.assertEqual(actual.shape, (5, 2, 1, 11))
 
-  def testGetNormBatchFn(self):
-    norm_type = "batch"
+  def testApplyNormBatch(self):
     with self.test_session() as session:
-      a = common_layers.get_norm(norm_type)
       x1 = np.random.rand(5, 2, 1, 11)
-      x2 = a(tf.constant(x1, dtype=tf.float32), name="batch")
+      x2 = common_layers.apply_norm(
+          tf.constant(x1, dtype=tf.float32), "batch", depth=11, epsilon=1e-6)
       session.run(tf.global_variables_initializer())
       actual = session.run(x2)
     self.assertEqual(actual.shape, (5, 2, 1, 11))
 
-  def testGetNormNoneFn(self):
-    norm_type = "none"
+  def testApplyNormNone(self):
     with self.test_session() as session:
-      a = common_layers.get_norm(norm_type)
       x1 = np.random.rand(5, 2, 1, 11)
-      x2 = a(tf.constant(x1, dtype=tf.float32), name="none")
+      x2 = common_layers.apply_norm(
+          tf.constant(x1, dtype=tf.float32), "none", depth=11, epsilon=1e-6)
       session.run(tf.global_variables_initializer())
       actual = session.run(x2)
     self.assertEqual(actual.shape, (5, 2, 1, 11))
     self.assertAllClose(actual, x1, atol=1e-03)
 
-  def testResidualFn(self):
-    norm_type = "batch"
-    with self.test_session() as session:
-      x1 = np.random.rand(5, 2, 1, 11)
-      x2 = np.random.rand(5, 2, 1, 11)
-      x3 = common_layers.residual_fn(
-          tf.constant(x1, dtype=tf.float32),
-          tf.constant(x2, dtype=tf.float32), norm_type, 0.1)
-      session.run(tf.global_variables_initializer())
-      actual = session.run(x3)
-    self.assertEqual(actual.shape, (5, 2, 1, 11))
-
-  def testResidualFnWithLayerNorm(self):
-    norm_type = "layer"
-    with self.test_session() as session:
-      x1 = np.random.rand(5, 2, 1, 11)
-      x2 = np.random.rand(5, 2, 1, 11)
-      x3 = common_layers.residual_fn(
-          tf.constant(x1, dtype=tf.float32),
-          tf.constant(x2, dtype=tf.float32),
-          norm_type,
-          0.1,
-          epsilon=0.1)
-      session.run(tf.global_variables_initializer())
-      actual = session.run(x3)
-    self.assertEqual(actual.shape, (5, 2, 1, 11))
-
   def testGlobalPool1d(self):
     x1 = np.random.rand(5, 4, 11)
     no_mask = np.ones((5, 4))

tensor2tensor/models/attention_lm.py

@@ -51,13 +51,10 @@ def model_fn_body(self, features):
     (decoder_input, decoder_self_attention_bias) = attention_lm_prepare_decoder(
         targets, hparams)
 
-    def residual_fn(x, y):
-      return common_layers.layer_norm(x + tf.nn.dropout(
-          y, 1.0 - hparams.residual_dropout))
-
-    decoder_input = tf.nn.dropout(decoder_input, 1.0 - hparams.residual_dropout)
-    decoder_output = attention_lm_decoder(decoder_input, residual_fn,
-                                          decoder_self_attention_bias, hparams)
+    decoder_input = tf.nn.dropout(
+        decoder_input, 1.0 - hparams.layer_prepostprocess_dropout)
+    decoder_output = attention_lm_decoder(
+        decoder_input, decoder_self_attention_bias, hparams)
     decoder_output = tf.expand_dims(decoder_output, 2)
 
     return decoder_output
@@ -84,15 +81,13 @@ def attention_lm_prepare_decoder(targets, hparams):
 
 
 def attention_lm_decoder(decoder_input,
-                         residual_fn,
                          decoder_self_attention_bias,
                          hparams,
                          name="decoder"):
   """A stack of attention_lm layers.
 
   Args:
     decoder_input: a Tensor
-    residual_fn: a function from (layer_input, layer_output) -> combined_output
     decoder_self_attention_bias: bias Tensor for self-attention
       (see common_attention.attention_bias())
     hparams: hyperparameters for model
@@ -105,25 +100,25 @@ def attention_lm_decoder(decoder_input,
   with tf.variable_scope(name):
     for layer in xrange(hparams.num_hidden_layers):
       with tf.variable_scope("layer_%d" % layer):
-        x = residual_fn(
-            x,
-            common_attention.multihead_attention(
-                x,
-                None,
-                decoder_self_attention_bias,
-                hparams.attention_key_channels or hparams.hidden_size,
-                hparams.attention_value_channels or hparams.hidden_size,
-                hparams.hidden_size,
-                hparams.num_heads,
-                hparams.attention_dropout,
-                name="decoder_self_attention"))
-        x = residual_fn(x,
-                        common_layers.conv_hidden_relu(
-                            x,
-                            hparams.filter_size,
-                            hparams.hidden_size,
-                            dropout=hparams.relu_dropout))
-  return x
+        with tf.variable_scope("self_attention"):
+          y = common_attention.multihead_attention(
+              common_layers.layer_preprocess(x, hparams),
+              None,
+              decoder_self_attention_bias,
+              hparams.attention_key_channels or hparams.hidden_size,
+              hparams.attention_value_channels or hparams.hidden_size,
+              hparams.hidden_size,
+              hparams.num_heads,
+              hparams.attention_dropout)
+          x = common_layers.layer_postprocess(x, y, hparams)
+        with tf.variable_scope("ffn"):
+          y = common_layers.conv_hidden_relu(
+              common_layers.layer_preprocess(x, hparams),
+              hparams.filter_size,
+              hparams.hidden_size,
+              dropout=hparams.relu_dropout)
+          x = common_layers.layer_postprocess(x, y, hparams)
+      return common_layers.layer_preprocess(x, hparams)
 
 
 @registry.register_hparams
@@ -158,7 +153,6 @@ def attention_lm_base():
   # when not in training mode.
   hparams.add_hparam("attention_dropout", 0.0)
   hparams.add_hparam("relu_dropout", 0.0)
-  hparams.add_hparam("residual_dropout", 0.1)
   hparams.add_hparam("pos", "timing")  # timing, none
   return hparams
 
@@ -178,5 +172,5 @@ def attention_lm_small():
   hparams.num_hidden_layers = 4
   hparams.hidden_size = 512
   hparams.filter_size = 2048
-  hparams.residual_dropout = 0.5
+  hparams.layer_prepostprocess_dropout = 0.5
   return hparams