Implementing the Transformer Architecture from Scratch: Training Acceleration and Performance Improvement in Practice

Why Implementing the Transformer from Scratch Is Worth Doing

A Reddit user implemented the original 2017 Transformer architecture from scratch using PyTorch and shared the code. He admitted that the model's predictions were poor, because he used only a 2-layer encoder and decoder and had only 1,000 words of training data, and training took too long. This is actually a common dilemma many developers face when reproducing classic papers: the architecture is built, but training efficiency and results get stuck. Based on this case, this article analyzes how to accelerate training and improve results without changing the core of the architecture.

The Root Cause of Slow Training: More Than Just Too Few Layers

The original post mentioned that the "2-layer encoder and decoder" results in small model capacity, but slow training comes more from implementation details. The original Transformer training ran for 3.5 days on 8 P100s, so reproducing it directly on personal devices is bound to be slow. Common bottlenecks include:

• Data loading and preprocessing: The corpus of 1,000 words is extremely small, but if tokenization or batch construction is redone on every iteration, the overhead can take up too high a proportion.
• Unoptimized attention computation: The manually implemented scaled dot-product attention does not use PyTorch's scaled_dot_product_attention (PyTorch 2.0+), which can automatically select efficient kernels such as Flash Attention on GPUs.
• No mixed-precision training: FP16/BF16 can significantly reduce memory usage and improve throughput, especially for Transformers, which are dense in matrix multiplications.
• Missing learning rate scheduling: The original paper's warmup strategy is crucial for convergence; without it, training oscillates and requires more steps.

Five Practical Improvements to Accelerate Training

1. Enable PyTorch 2.0 Compilation and Attention Optimization

Wrap the model with torch.compile(model) and replace the hand-written attention with F.scaled_dot_product_attention. On GPUs that support Tensor Cores, these two changes usually bring a 1.5–2x speedup without changing mathematical equivalence.

2. Use Mixed Precision and Gradient Accumulation

Enable mixed precision through torch.cuda.amp.autocast and GradScaler. If memory is insufficient, use gradient accumulation to simulate a larger batch size, which both stabilizes training and avoids frequent parameter updates.

3. Implement Standard Learning Rate Warmup and Decay

The original paper uses lr = d_model^{-0.5} * min(step^{-0.5}, step * warmup^{-1.5}). Even for small-scale training, warmup can prevent early gradient explosions, and decay helps with later convergence.

4. Optimize the Data Pipeline

For a small corpus of 1,000 words, it is recommended to encode it once into tensors, use DataLoader with num_workers=0 (with small data, multiprocessing actually increases overhead), and pre-generate all training sample pairs to avoid dynamic construction at every step.

5. Appropriately Increase Layers While Controlling Parameter Count

2 layers are indeed too shallow. You can try 4-layer encoders/decoders while reducing d_model from 512 to 256 and ffn_dim from 2048 to 1024; the total parameter count becomes smaller and training becomes faster. Deeper small models usually have stronger expressive power than shallow large models.

Strategies to Improve Results with Small Data

A training set of 1,000 words is almost insufficient for machine translation or language modeling. If the goal is to validate the architecture, you can adopt:

• Use synthetic tasks: Such as copy tasks, sequence reversal, and simple arithmetic; these tasks can quickly verify whether the attention mechanism is correct.
• Data augmentation: Perform back-translation on text, randomly replace synonyms, or use subword tokenization (BPE) to increase vocabulary combinations.
• Transfer learning: Load pre-trained small Transformer weights (such as Hugging Face's distilbert) and then fine-tune.
• Regularization: Add dropout (0.1–0.3) and label smoothing (0.1) to prevent overfitting on small data.

Quick Code-Level Checklist

• Confirm that positional encoding is correctly added to the embedding rather than concatenated.
• Check whether the mask is correctly applied in decoder self-attention to avoid information leakage.
• Verify whether the loss function ignores padding tokens (ignore_index).
• Use torch.no_grad() for inference and call model.eval().

Conclusion

Implementing the Transformer from scratch is an excellent path to understanding the attention mechanism. Slow training and poor results are not architecture errors but unoptimized engineering details. Prioritize enabling PyTorch 2.0's attention kernels and compilation, mixed precision, and standard learning rate scheduling, then adjust model depth and width; this usually achieves acceptable training speed on personal devices. For small data, switching to synthetic tasks or transfer learning is more effective than forcibly enlarging the model.

Background Supplement

scaled_dot_product_attention (SDPA) is an attention computation primitive introduced in PyTorch 2.0. It automatically selects the optimal backend based on input shape, data type, and hardware capabilities: on NVIDIA GPUs that support Tensor Cores, it may use Flash Attention (reducing memory reads and writes), Memory-Efficient Attention (block-wise computation), or a mathematically equivalent standard implementation. Compared with hand-written attention, SDPA avoids explicit materialization of the intermediate attention matrix, thereby reducing memory usage and improving speed. torch.compile, on the other hand, optimizes the entire model's forward and backward propagation through graph capture and kernel fusion, and is especially suitable for the repeated linear layers and normalization operations in Transformers. The combination of the two usually brings significant acceleration while fully preserving mathematical equivalence.