Config for language-model experiments.

  Train these on languagemodel_lm1b32k_packed for 136000 steps (10 epochs)

  The size parameter is an integer that controls the number of heads and the
  size of the size of the feedforward hidden layers.  Increasing size by 1
  doubles each of these.

  Results:
  size   params/10^9  log-ppl(per-token)
  -1     0.14         3.209
  0      0.22         3.119
  1      0.37         3.037
  2      0.67         2.969
  3      1.28         2.912
  4      2.48         2.874
  5      4.90         2.871

  (to get word-level log-ppl, multiply by 1.1078)

  Args:
    size: an integer
  Returns:
    a hparams object

def mtf_transformer_paper_lm(size):
  """Config for language-model experiments.

  Train these on languagemodel_lm1b32k_packed for 136000 steps (10 epochs)

  The size parameter is an integer that controls the number of heads and the
  size of the size of the feedforward hidden layers.  Increasing size by 1
  doubles each of these.

  Results:
  size   params/10^9  log-ppl(per-token)
  -1     0.14         3.209
  0      0.22         3.119
  1      0.37         3.037
  2      0.67         2.969
  3      1.28         2.912
  4      2.48         2.874
  5      4.90         2.871

  (to get word-level log-ppl, multiply by 1.1078)

  Args:
    size: an integer
  Returns:
    a hparams object
  """
  n = 2 ** size
  hparams = mtf_transformer_base_lm()
  hparams.batch_size = 256
  hparams.d_model = 1024
  hparams.d_ff = int(8192 * n)
  hparams.d_kv = 256
  hparams.num_heads = int(8 * n)
  hparams.shared_embedding_and_softmax_weights = False
  # one epoch for languagemodel_lm1b32k_packed = 13600 steps
  hparams.learning_rate_decay_steps = 13600
  return hparams