Training Tips for the Transformer Model

In plain words

The authors trained a translation model that reads a sentence in one language and writes it in another. The model is a Transformer, a neural network whose main component is attention, which lets each word look at the other words in the sentence. The paper is a lab notebook of practical settings for training it: how many examples to process at once, how fast the model learns, how long to warm up, how many computers to use, and how much training data to feed it.

Most of the findings are about the trade-off between speed and quality. Bigger batches, more graphics cards and more data help, but each has a cost in memory or in the way the learning rate has to be set. The authors also show that when several cards are used, a setting that looks unchanged actually makes the learning rate larger, so it has to be adjusted.

The advice is aimed at people who will train a similar model on their own hardware. The authors do not claim their settings carry over to other languages or other tasks.

TL;DR

The paper tests the Transformer in Tensor2Tensor v1.2.9 on English-to-Czech translation, training on CzEng 1.7 (58M sentence pairs) and reporting full BLEU learning curves against wall-clock time instead of single final scores §1, §2.2, tab. 1. Larger batches, more GPUs and more data improve results, and the big model beats the base model after a few hours on the same batch size §4.3, §4.5, §4.7. With several GPUs, keeping learning_rate fixed raises the real learning rate by the square root of the GPU count, so warmup and learning-rate settings must be re-checked §4.8.3. Checkpoint averaging gives a small, cheap gain §4.10. The best 8-GPU model reaches BLEU 23.8 on WMT17 newstest2017, above the top-ranked WMT17 system in BLEU tab. 7.

Problem

Transformer results depend on many hyperparameters, and the usual practice of stopping training early or reporting a single score makes comparisons unreliable §2.1. Early stopping with a patience rule gave high variance in both training time and final BLEU, even across random seeds §2.1. The paper also notes that its framework's default BLEU proxy is not comparable with real BLEU §2.4. The authors target the English-to-Czech setting specifically and say that the Transformer's behaviour at large data and on multiple GPUs was not yet explored in the literature they cite §1.

Method

The paper is an experimental protocol, not a new model. All experiments use the Transformer as implemented in Tensor2Tensor v1.2.9 §1, §4.

1. Evaluation protocol §2. Report full learning curves: case-insensitive BLEU (sacreBLEU, tok.intl) against wall-clock hours, not a single best score §2.1, §2.2, §2.4 footnote. Training speed is steps per hour; throughput is subwords per hour, equal to speed times effective batch size §2.3. "Time Till Score" is the time after which BLEU never falls below a level §2.3. Use the paper's t2t-bleu script, not the built-in approximate BLEU §2.4.
2. Data §3. English-to-Czech. Main training set is CzEng 1.7 (57M pairs) plus Europarl, News Commentary and Common Crawl, 58M pairs in total tab. 1. Dev set is newstest2013; the final comparison uses newstest2017 §3. A shared English-Czech subword vocabulary of 32k is trained by the framework on a sample of the data §3.1.
3. Baseline configuration §4. The BIG model's defaults, except batch_size=1500, a very large step budget (6,000,000) so runs are stopped manually, checkpoints saved every 3600 s, and schedule=train with internal evaluation disabled §4.
4. Batch size and throughput §4.1, §4.5. Speed falls as batch size grows. Throughput grows sub-linearly tab. 2. For the base model, larger batches (up to 4500 in the tested range) converge faster in wall-clock time §4.5, fig. 5. For the big model, the batch size must be above a threshold near 1450 to converge at all, and 1400 diverged §4.5, fig. 6.
5. GPUs §4.7. batch_size is per GPU, so the effective batch is batch_size × gpus (1500 on 8 GPUs means 12k subwords per step) §2.3, §4.7. Training with batch 4000 on one GPU, batch 2000 on two and batch 1000 on four is reported to give identical loss and BLEU versus steps §4.7. The authors recommend more GPUs for faster convergence and running one multi-GPU experiment at a time rather than several smaller ones in parallel §4.7 tips.
6. Model size §4.3, tab. 4. BASE has hidden size 512, filter size 2048, 8 heads and Adam β₂ 0.980. BIG has 1024, 4096, 16 heads and β₂ 0.998. On a single GPU, BIG at batch 2000 overtakes BASE at batch 4500 in under a day §4.3, fig. 2. On 8 GPUs, BIG at batch 1500 overtakes BASE at batch 4500 after 18 hours §4.3, fig. 3.
7. Maximum sentence length §4.4, tab. 5. max_length drops training pairs longer than a threshold measured in subwords. It allows a larger batch within the same memory. Pairs longer than 70 subwords are 2.1% of training data; longer than 100 are 0.7% tab. 5. With batch 1500, max_length of 70 or lower produced a BLEU drop after one hour that the authors could not explain. With batch 2000, 70 and above matched the unrestricted curve §4.4, fig. 4.
8. Learning rate and warmup, single GPU §4.6. Default learning_rate is 0.20. Values from 0.05 to 0.25 gave almost the same curves. 0.01 converged slowly and 0.30 diverged §4.6. The schedule ramps linearly to the peak at learning_rate_warmup_steps (default 16,000), then decays with the inverse square root of the step §4.6, §4.8.3. Raising warmup to 32k allowed 0.30 and 0.50 without divergence. Learning rate 1.0 needed 60k warmup (40k diverged) §4.6.
9. Learning rate on multiple GPUs §4.8. Scaling the effective batch 8× with the same learning_rate gives an actual learning rate √8 times higher at the same step, because the decay uses the global step count and not epochs §4.8.3. The authors' reading: this matches the √k rule from Hoffer et al. §4.8.1, §4.8.3. Warmup must be re-tuned: 8k warmup on 8 GPUs matched 16k, 6k was slower and 2k diverged §4.8.2, §4.8.3. Gradient clipping at norm 1.0 let 0.40 train but 0.60 still diverged §4.8.2.
10. Resumed training §4.9. Resuming restores the checkpoint, including Adam momentum, but not the position in the data, so the run restarts from a random point. The authors suggest changing batch size or domain mid-run and adjusting the learning rate schedule to suit §4.9.
11. Checkpoint averaging §4.10, fig. 10. Average the last N checkpoints. The paper uses hourly checkpoints, while the original Transformer paper used 10-minute intervals over the last 20 §4.10. Averaging gives a typical gain of about 0.2 BLEU and lower variance §4.10. Averaging 8 checkpoints takes about 5 minutes §4.10 tips. Early in training, fewer checkpoints work better. Later, 16 checkpoints (16 h) were slightly better than 8, though not tested for significance §4.10.

Key relations used in the card (stated in prose in the paper; no numbered equations):

g = \epsilon \left(\frac{N}{B} - 1\right) §4.8.1 (gradient noise scale; ε learning rate, N training set size, B effective batch)

\textit{actual\_lr}(\textit{steps}) = c \cdot \textit{steps}^{-0.5} §4.8.3 (inverse square root decay; c includes learning_rate)

The following pseudocode implements the arithmetic in the paper:

# Training-step arithmetic for multi-GPU runs  [§2.3] [§4.7]
# Input:  batch_size (subwords per GPU), gpus, steps, total_subwords
# Output: subwords per step, epochs completed

def effective_batch_size(batch_size, gpus):
    # batch_size is per GPU, so one step consumes batch_size x gpus  [§4.7]
    return batch_size * gpus

def epochs_completed(steps, batch_size, gpus, total_subwords):
    # Steps x subwords per step, divided by the subwords in the data  [§2.3]
    # total_subwords: 992M for CzEng 1.7, 327M for CzEng 1.0  [§3.1]
    per_step = effective_batch_size(batch_size, gpus)
    return steps * per_step / total_subwords

# Learning-rate schedule (linear warmup, then inverse-sqrt decay)  [§4.6]
# Input:  learning_rate, warmup_steps, step (global step, not epoch)
# Output: actual learning rate at this step

def actual_learning_rate(learning_rate, warmup_steps, step):
    # Ramp linearly up to the peak at warmup_steps  [§4.6]
    if step < warmup_steps:
        return learning_rate * step / warmup_steps
    # Decay proportional to step^-0.5; the exact constant is not stated
    # (the join at warmup_steps is inferred, not stated)  [§4.6] [§4.8.3]
    return learning_rate * (warmup_steps / step) ** 0.5

Hyperparameters

Defaults are those of Tensor2Tensor v1.2.9 unless stated. "Not stated" means the paper does not give the value.

Results

Author-reported. Setup: English-to-Czech; Tensor2Tensor v1.2.9; GTX 1080 Ti GPUs; BLEU from sacreBLEU §2.2, §2.4. The WMT17 comparison uses the authors' best run (BIG, 8 GPUs, 8 days, average of 8 checkpoints), trained on a subset of constrained data, without back-translation, reranking or ensembling §5.

Paper vs code

No official code.

The paper uses Tensor2Tensor v1.2.9 §1, a Google project, and sacreBLEU §2.1 footnote. The authors' own scripts (t2t-bleu, t2t-translate-all, t2t-avg-all) were merged into Tensor2Tensor master, so no separate authors' repository is named §2.4 footnote. A GitHub repository search for the title "Training Tips for the Transformer Model" returned no results. No project page is linked from the arXiv record. The Tensor2Tensor repository was not cloned because it is not the authors' code.

Pitfalls

• Batch size is in subwords and is per GPU. batch_size=1500 on 8 GPUs means 12k subwords per step §2.3 §4.7. Sentence length for batching is the max of source and target subword counts, and padding means the batch only approximates the real subword count §2.3 footnote.
• Learning-rate scaling is a side effect of the schedule. The decay uses the global step, so keeping learning_rate fixed when you multiply the batch by k raises the real learning rate by √k §4.8.3. Lowering learning_rate by √k matches the single-GPU curve for the decay term only. The authors note this holds for the inverse-sqrt schedule, ignoring warmup §4.8.3.
• Warmup controls the peak. Cutting warmup steps, while keeping learning_rate, raises the peak real learning rate, because the peak is at warmup_steps §4.6 tips. On 8 GPUs, reducing warmup below 8k diverged §4.8.3. The authors found that reducing warmup does not improve final BLEU §4.8 tips.
• The inverse-sqrt formula is not given. The paper gives the shape (linear ramp, then decay ∝ step^-0.5) but no exact constants or the exact join point. The pseudocode above assumes continuity at warmup_steps; this is inferred, not stated §4.6 §4.8.3.
• Divergence is the main failure mode. It shows up in the first hours, when the real learning rate is highest §4.6. Remedies the authors used: more warmup, gradient clipping at 1.0, or a lower learning rate §4.6 §4.8.2. Divergence at batch below about 1450 for BIG is a single-setup observation, and the authors state the reason is unknown §4.5.
• max_length interacts with batch size. At batch 1500, max_length of 25, 50 and 70 gave BLEU drops after an hour; the authors have no explanation §4.4. At batch 2000, 70 and above matched the unrestricted curve §4.4. Models trained with a given max_length did not produce longer outputs, even after changing the decoding parameter alpha §4.4. The authors recommend a reasonably low max_length for memory and a reasonably high one for quality §4.4 tips.
• Batch limits depend on memory and optimizer. The maximum batch that fits in 11 GB varies with model size, optimizer and max_length tab. 5. Adafactor is cited as using less memory, but its experiments were left for future work §4.4 footnote.
• Approximate BLEU is not comparable. Tensor2Tensor's internal approx-bleu is 1.2–1.8 times the real BLEU and depends on the subword vocabulary, so do not report it §2.4.
• Non-determinism and seeds. The fixed random seed was not applied correctly in v1.2.9 §2.1 footnote. A change in hyperparameters also changes the initialization, so single-run differences can be noise §2.1 footnote. Learning curves "flicker" between checkpoints; Time Till Score depends on this and needs a larger dev set to reduce variance §2.3 footnote.
• Dataset size is inconsistent in the text. The smaller CzEng 1.0 set is 15M pairs in §3 and 16M in §4.2 and the caption of fig. 1. Use the number from the table or your own count.
• Data size changes epoch counts. Steps to one epoch: CzEng 1.0 with batch 1500 on 8 GPUs is 27k steps; CzEng 1.7 is 83k §4.2. Epochs are not reported by T2T itself §2.3.
• Vocabulary training needs enough data. In v1.2.9, file_byte_budget is a source-code constant. Logs showing a low min_count mean the subword vocabulary was built from words seen only once or twice §3.1 footnote.
• Resuming changes the data order. Resumed training starts from a random data position, so run comparisons across resumes are not exact §4.9.
• Compute cost is high. The whole study used about 4 years of GPU time, and the best model needed 8 GPUs for 8 days §6 §5. Hardware-specific speeds (GTX 1080 Ti, driver 375.66) do not transfer directly §2.2.
• Scope. All experiments are English-to-Czech on Transformer, with one framework version. The authors do not claim the recommendations carry over to other pairs or models §1.
• Statistics. Checkpoint-count comparisons were not tested for significance §4.10. Results with one run per curve should be read as indicative.

Lineage

• Builds on: arXiv:1706.03762 (Vaswani et al., Transformer), the model studied here §1. The paper also builds on the √k learning-rate rule of Hoffer et al. §4.8.1 and the gradient-noise-scale analysis of Smith and Le §4.8.1. Neither has a card on this wiki.
• Compared against: WMT17 English-to-Czech systems, including uedin-nmt, online-B and LIUM-NMT tab. 7. None has a card on this wiki.
• Superseded by: not known from content already in this wiki.