Heuristic:Lucidrains X transformers Gradient Clipping And Accumulation
| Knowledge Sources | |
|---|---|
| Domains | Optimization, Deep_Learning |
| Last Updated | 2026-02-08 18:00 GMT |
Overview
Training stability techniques using gradient clipping at 0.5 norm and gradient accumulation to simulate larger effective batch sizes.
Description
The x-transformers training examples demonstrate two complementary techniques for stable training: gradient norm clipping to prevent exploding gradients, and gradient accumulation to achieve larger effective batch sizes without increasing memory usage. These patterns are consistent across the example training scripts and represent the author's recommended approach.
Usage
Use this heuristic when setting up a training loop for any x-transformers model, especially character-level language models or tasks requiring large effective batch sizes on limited GPU memory.
The Insight (Rule of Thumb)
- Action: Apply gradient clipping with max norm of 0.5 after accumulation, before optimizer step.
- Value:
- Gradient clip norm: `0.5`
- Gradient accumulation steps: `4` (effective batch = `BATCH_SIZE * 4`)
- Learning rate: `1e-4` (character-level) to `3e-4` (synthetic tasks)
- Trade-off: Gradient accumulation increases training time per effective batch but reduces memory. Clipping at 0.5 is conservative and may slow convergence slightly but prevents training instability.
Training Pattern
| Parameter | Character-Level (enwik8) | Synthetic Task (copy) |
|---|---|---|
| Batch size | 4 | 32 |
| Gradient accumulation | 4 (effective: 16) | 1 (none) |
| Learning rate | 1e-4 | 3e-4 |
| Gradient clipping | 0.5 | Not used |
| Optimizer | Adam | Adam |
| Total iterations | 100,000 | 100,000 |
Reasoning
Character-level language modeling with transformers can produce sudden gradient spikes, especially early in training. Clipping at 0.5 is more aggressive than the common 1.0 default, reflecting the author's experience with smaller transformer models (dim=512, depth=6). The absence of gradient clipping in the copy task (train_copy.py) suggests that simpler synthetic tasks with smaller models (dim=128, depth=3) are inherently more stable.
Gradient accumulation with loss division (`loss / GRADIENT_ACCUMULATE_EVERY`) ensures correct gradient scaling when simulating larger batches.
Code Evidence
Gradient accumulation pattern from `train_enwik8.py:102-114`:
for i in tqdm.tqdm(range(NUM_BATCHES), mininterval=10., desc='training'):
model.train()
for __ in range(GRADIENT_ACCUMULATE_EVERY):
loss = model(next(train_loader))
(loss / GRADIENT_ACCUMULATE_EVERY).backward()
print(f'training loss: {loss.item()}')
torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
optim.step()
optim.zero_grad()
Simpler training loop without clipping from `train_copy.py:51-61`:
for i in tqdm.tqdm(range(NUM_BATCHES), mininterval=10., desc='training'):
model.train()
src, tgt, src_mask = next(cycle())
loss = model(src, tgt, mask=src_mask)
loss.backward()
optim.step()
optim.zero_grad()