Play with a new model with Transformer with a GAN'y part.

PiperOrigin-RevId: 174403736

Author: Lukasz Kaiser

tensor2tensor/layers/common_hparams.py

@@ -116,12 +116,15 @@ def basic_params1():
       # If set to True, drop sequences longer than max_length during eval.
       # This affects the validity of the evaluation metrics.
       eval_drop_long_sequences=int(False),
+      # TODO(lukaszkaiser): these parameters should probably be set elsewhere.
       # in SymbolModality, share the output embeddings and the softmax
       # variables.
       # You can also share the input embeddings with the output embeddings
       # by using a problem_hparams that uses the same modality object for
       # the input_modality and target_modality.
       shared_embedding_and_softmax_weights=int(False),
+      # In SymbolModality, skip the top layer, assume we're providing logits.
+      symbol_modality_skip_top=int(False),
       # For each feature for which you want to override the default input
       # modality, add an entry to this semicolon-separated string. Entries are
       # formatted "feature_name:modality_type:modality_name", e.g.

tensor2tensor/layers/modalities.py

@@ -115,6 +115,8 @@ def top(self, body_output, _):
     else:
       scope_name = "softmax"
       reuse = False
+    if self._model_hparams.symbol_modality_skip_top:
+      return tf.expand_dims(body_output, 3)
     with tf.variable_scope(scope_name, reuse=reuse):
       var = self._get_weights()
       if (self._model_hparams.factored_logits and

tensor2tensor/layers/modalities_test.py

@@ -40,6 +40,7 @@ def testSymbolModalityInputs(self):
         symbol_modality_num_shards=4,
         hidden_size=hidden_size,
         multiply_embedding_mode="sqrt_depth",
+        symbol_modality_skip_top=0,
         shared_embedding_and_softmax_weights=0)
     x = -1 + np.random.random_integers(
         vocab_size, size=(batch_size, length, 1, 1))
@@ -65,6 +66,7 @@ def testSymbolModalityTargets(self):
         symbol_modality_num_shards=4,
         hidden_size=hidden_size,
         label_smoothing=0.2,
+        symbol_modality_skip_top=0,
         shared_embedding_and_softmax_weights=0,
         factored_logits=0,
         mode=tf.estimator.ModeKeys.TRAIN)
@@ -99,6 +101,7 @@ def testSymbolModalityTargetsFactored(self):
         symbol_modality_num_shards=4,
         hidden_size=hidden_size,
         label_smoothing=0.2,
+        symbol_modality_skip_top=0,
         shared_embedding_and_softmax_weights=0,
         factored_logits=1,
         mode=tf.estimator.ModeKeys.TRAIN)

tensor2tensor/models/__init__.py

@@ -36,6 +36,7 @@
 from tensor2tensor.models import shake_shake
 from tensor2tensor.models import slicenet
 from tensor2tensor.models import transformer
+from tensor2tensor.models import transformer_adv
 from tensor2tensor.models import transformer_alternative
 from tensor2tensor.models import transformer_moe
 from tensor2tensor.models import transformer_revnet

tensor2tensor/models/transformer_adv.py

@@ -0,0 +1,229 @@
+# coding=utf-8
+# Copyright 2017 The Tensor2Tensor Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Adversarial Transformer."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+# Dependency imports
+
+from tensor2tensor.layers import common_layers
+from tensor2tensor.models import transformer
+from tensor2tensor.models import transformer_vae
+from tensor2tensor.utils import registry
+from tensor2tensor.utils import t2t_model
+
+import tensorflow as tf
+
+
+def encode(x, x_space, hparams, name):
+  """Transformer preparations and encoder."""
+  with tf.variable_scope(name):
+    (encoder_input, encoder_self_attention_bias,
+     ed) = transformer.transformer_prepare_encoder(x, x_space, hparams)
+    encoder_input = tf.nn.dropout(encoder_input, 1.0 - hparams.dropout)
+    return transformer.transformer_encoder(
+        encoder_input, encoder_self_attention_bias, hparams), ed
+
+
+def decode(encoder_output, encoder_decoder_attention_bias, targets,
+           hparams, name, reuse=False):
+  """Transformer decoder."""
+  with tf.variable_scope(name, reuse=reuse):
+    targets = common_layers.flatten4d3d(targets)
+
+    decoder_input, decoder_self_bias = transformer.transformer_prepare_decoder(
+        targets, hparams)
+
+    decoder_input = tf.nn.dropout(decoder_input,
+                                  1.0 - hparams.layer_prepostprocess_dropout)
+
+    decoder_output = transformer.transformer_decoder(
+        decoder_input,
+        encoder_output,
+        decoder_self_bias,
+        encoder_decoder_attention_bias,
+        hparams)
+
+    # Expand since t2t expects 4d tensors.
+    return tf.expand_dims(decoder_output, axis=2)
+
+
+def reverse_gradient(x, delta=1.0):
+  return tf.stop_gradient((1.0 + delta) * x) - delta * x
+
+
+def adversary(embedded, inputs, hparams, name, reuse=False):
+  with tf.variable_scope(name, reuse=reuse):
+    h0, i0 = common_layers.pad_to_same_length(
+        embedded, inputs, final_length_divisible_by=16)
+    h0 = tf.concat([h0, tf.expand_dims(i0, axis=2)], axis=-1)
+    h0 = tf.layers.dense(h0, hparams.hidden_size, name="io")
+    h1 = transformer_vae.compress(h0, None, False, hparams, "compress1")
+    h2 = transformer_vae.compress(h1, None, False, hparams, "compress2")
+    res_dense = tf.reduce_mean(h2, axis=[1, 2])
+    res_single = tf.squeeze(tf.layers.dense(res_dense, 1), axis=-1)
+    return tf.nn.sigmoid(res_single)
+
+
+def softmax_embed(x, embedding, batch_size, hparams):
+  """Softmax x and embed."""
+  x = tf.reshape(tf.nn.softmax(x), [-1, 34*1024])
+  x = tf.matmul(x, embedding)
+  return tf.reshape(x, [batch_size, -1, 1, hparams.hidden_size])
+
+
+def adv_transformer_internal(inputs, targets, target_space, hparams):
+  """Adversarial Transformer, main step used for training."""
+  with tf.variable_scope("adv_transformer"):
+    batch_size = tf.shape(targets)[0]
+    targets = tf.reshape(targets, [batch_size, -1, 1])
+    embedding = tf.get_variable("embedding", [34*1024, hparams.hidden_size])
+    targets_emb = tf.gather(embedding, targets)
+
+    # Noisy embedded targets.
+    targets_noisy = tf.one_hot(targets, 34*1024)
+    noise_val = hparams.noise_val
+    targets_noisy += tf.random_uniform(tf.shape(targets_noisy),
+                                       minval=-noise_val, maxval=noise_val)
+    targets_emb_noisy = softmax_embed(
+        targets_noisy, embedding, batch_size, hparams)
+
+    # Encoder.
+    if inputs is not None:
+      inputs_emb = common_layers.flatten4d3d(inputs)
+      inputs, ed = encode(inputs_emb, target_space, hparams, "input_enc")
+    else:
+      ed = None
+
+    # Masking.
+    masking = common_layers.inverse_lin_decay(60000)
+    masking *= common_layers.inverse_exp_decay(20000)  # Not much at start.
+    masking -= tf.random_uniform([]) * 0.4
+    mask = tf.less(masking, tf.random_uniform(tf.shape(targets)))
+    mask = tf.expand_dims(tf.to_float(mask), 3)
+    noise = tf.random_uniform(tf.shape(targets_emb))
+    targets_emb = mask * targets_emb + (1.0 - mask) * noise
+
+    # Decoder.
+    res_dec = decode(inputs, ed, targets_emb, hparams, "decoder")
+    res = tf.layers.dense(res_dec, 34*1024, name="res_sm")
+    res_emb = softmax_embed(res, embedding, batch_size, hparams)
+
+    # Extra steps.
+    extra_step_prob = masking * 0.6
+    if hparams.mode != tf.estimator.ModeKeys.TRAIN:
+      extra_step_prob = 1.0
+    for _ in xrange(hparams.extra_steps):
+      def another_step(emb):
+        res_dec = decode(inputs, ed, emb, hparams, "decoder", reuse=True)
+        res = tf.layers.dense(res_dec, 34*1024, name="res_sm", reuse=True)
+        return softmax_embed(res, embedding, batch_size, hparams), res
+      res_emb, res = tf.cond(tf.less(tf.random_uniform([]), extra_step_prob),
+                             lambda e=res_emb: another_step(e),
+                             lambda: (res_emb, res))
+
+    # Adversary.
+    delta = masking * hparams.delta_max
+    true_logit = adversary(tf.stop_gradient(targets_emb_noisy),
+                           tf.stop_gradient(inputs + inputs_emb),
+                           hparams, "adversary")
+    gen_logit = adversary(reverse_gradient(res_emb, delta),
+                          tf.stop_gradient(inputs + inputs_emb),
+                          hparams, "adversary", reuse=True)
+    losses = {"adv": gen_logit - true_logit}
+    res = tf.stop_gradient(masking * res) + (1.0 - masking) * res
+    return res, losses
+
+
+@registry.register_model
+class TransformerAdv(t2t_model.T2TModel):
+  """Adversarial Transformer."""
+
+  def model_fn_body(self, features):
+    inputs = features.get("inputs", None)
+    return adv_transformer_internal(
+        inputs, features["targets_raw"],
+        features["target_space_id"], self._hparams)
+
+  def infer(self, features=None, decode_length=50, beam_size=1, top_beams=1,
+            last_position_only=False, alpha=0.0):
+    """Produce predictions from the model."""
+    if not features:
+      features = {}
+    inputs_old = None
+    if "inputs" in features and len(features["inputs"].shape) < 4:
+      inputs_old = features["inputs"]
+      features["inputs"] = tf.expand_dims(features["inputs"], 2)
+
+    # Create an initial targets tensor.
+    if "partial_targets" in features:
+      initial_output = tf.convert_to_tensor(features["partial_targets"])
+    else:
+      batch_size = tf.shape(features["inputs"])[0]
+      length = tf.shape(features["inputs"])[1]
+      initial_output = tf.zeros((batch_size, 2 * length, 1, 1), dtype=tf.int64)
+
+    features["targets"] = initial_output
+    sharded_logits, _ = self.model_fn(
+        features, False, last_position_only=last_position_only)
+    sharded_samples = self._data_parallelism(tf.argmax, sharded_logits, 4)
+    samples = tf.concat(sharded_samples, 0)
+
+    # More steps.
+    how_many_more_steps = 5
+    for _ in xrange(how_many_more_steps):
+      with tf.variable_scope(tf.get_variable_scope(), reuse=True):
+        features["targets"] = samples
+        sharded_logits, _ = self.model_fn(
+            features, False, last_position_only=last_position_only)
+        sharded_samples = self._data_parallelism(tf.argmax, sharded_logits, 4)
+        samples = tf.concat(sharded_samples, 0)
+
+    if inputs_old is not None:  # Restore to not confuse Estimator.
+      features["inputs"] = inputs_old
+    return samples
+
+
+@registry.register_hparams
+def transformer_adv_small():
+  """Set of hyperparameters."""
+  hparams = transformer.transformer_small()
+  hparams.batch_size = 2048
+  hparams.learning_rate_warmup_steps = 4000
+  hparams.num_hidden_layers = 3
+  hparams.hidden_size = 384
+  hparams.filter_size = 2048
+  hparams.label_smoothing = 0.0
+  hparams.weight_decay = 0.1
+  hparams.symbol_modality_skip_top = int(True)
+  hparams.add_hparam("num_compress_steps", 2)
+  hparams.add_hparam("extra_steps", 0)
+  hparams.add_hparam("noise_val", 0.3)
+  hparams.add_hparam("delta_max", 2.0)
+  return hparams
+
+
+@registry.register_hparams
+def transformer_adv_base():
+  """Set of hyperparameters."""
+  hparams = transformer_adv_small()
+  hparams.batch_size = 1024
+  hparams.hidden_size = 512
+  hparams.filter_size = 4096
+  hparams.num_hidden_layers = 6
+  return hparams