code: def mtf_transformer_paper_lm(size):
    
    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
docs: 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