Set of hyperparameters.

def transformer_ae_small():
  """Set of hyperparameters."""
  hparams = transformer.transformer_small()
  hparams.batch_size = 2048
  hparams.learning_rate = 0.2
  hparams.learning_rate_warmup_steps = 4000
  hparams.num_hidden_layers = 3
  hparams.hidden_size = 384
  hparams.filter_size = 2048
  hparams.add_hparam("compress_filter_size", 2048 * 2)
  hparams.label_smoothing = 0.0
  hparams.optimizer = "adam"  # Can be unstable, maybe try Adam.
  hparams.optimizer_adam_epsilon = 1e-9
  hparams.optimizer_adam_beta1 = 0.9
  hparams.optimizer_adam_beta2 = 0.997  # Needs tuning, try 0.98 to 0.999.
  hparams.add_hparam("z_size", 14)
  hparams.add_hparam("noise_dev", 0.5)
  hparams.add_hparam("d_mix", 0.5)
  hparams.add_hparam("logit_normalization", True)
  hparams.add_hparam("word_dropout", 0.)
  # Bottleneck kinds supported: dense, vae, semhash, gumbel-softmax, dvq.
  hparams.add_hparam("bottleneck_kind", "semhash")
  hparams.add_hparam("num_blocks", 1)
  hparams.add_hparam("num_decode_blocks", 1)
  # Add an hparam for number of reiduals
  hparams.add_hparam("num_residuals", 1)
  # Reshape method for DVQ: slice, project
  hparams.add_hparam("word_shuffle", 0.5)
  hparams.add_hparam("causal", True)
  hparams.add_hparam("reshape_method", "slice")
  hparams.add_hparam("trainable_projections", False)
  hparams.add_hparam("unmasked_percentage", 0.1)
  hparams.add_hparam("do_ae", True)
  hparams.add_hparam("do_mask", True)
  hparams.add_hparam("use_predict_mask", True)
  hparams.add_hparam("do_refine", False)
  hparams.add_hparam("do_attend_compress", False)
  hparams.add_hparam("do_attend_decompress", True)
  hparams.add_hparam("do_residual_compress", False)
  hparams.add_hparam("drop_inputs", False)
  hparams.add_hparam("v_size", 1024*64)
  hparams.add_hparam("max_context_length", 64)
  hparams.add_hparam("num_compress_steps", 3)
  hparams.add_hparam("startup_steps", 10000)
  hparams.add_hparam("mask_startup_steps", 50000)
  hparams.add_hparam("z_dropout", 0.1)
  hparams.add_hparam("is_2d", 0)
  hparams.add_hparam("softmax_k", 0)
  hparams.add_hparam("decode_autoregressive", True)
  hparams.add_hparam("do_vae", True)
  hparams.add_hparam("bit_vae", True)
  hparams.add_hparam("beta", 0.25)
  hparams.add_hparam("epsilon", 1e-5)
  hparams.add_hparam("decay", 0.999)
  hparams.add_hparam("ema", True)
  hparams.add_hparam("random_top_k", 1)
  hparams.add_hparam("soft_em", False)
  hparams.add_hparam("num_samples", 10)
  hparams.add_hparam("inv_temp", 1.0)
  hparams.add_hparam("entropy_scale", 0.0)
  hparams.add_hparam("prior_scale", 1.0)
  hparams.add_hparam("do_hard_gumbel_softmax", False)
  hparams.add_hparam("num_flows", 0)
  hparams.add_hparam("approximate_gs_entropy", False)
  hparams.add_hparam("temperature_warmup_steps", 150000)
  hparams.add_hparam("sum_over_latents", False)
  hparams.force_full_predict = True

  # task params
  hparams.add_hparam("task", "translate")  # translate or image tasks supported
  return hparams