File size: 9,304 Bytes
256c9c2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 | # Knowledge: Training Recipes
## Purpose
Common optimizer configurations, learning rate schedules, and training details that papers frequently assume you know but don't re-explain. When a paper says "we use standard training settings," this file helps you determine what they probably mean (and helps you flag what's still ambiguous).
---
## Optimizers
### Adam (Kingma & Ba, 2014)
Default parameters in the paper: Ξ²β=0.9, Ξ²β=0.999, Ξ΅=1e-8
**PyTorch defaults match the Adam paper.**
But in practice:
- Transformer papers often use Ξ²β=0.98 (from Vaswani et al.)
- Some NLP papers use Ξ²β=0.95
- Ξ΅ values vary: 1e-8 (default), 1e-6 (BERT), 1e-9 (Vaswani et al.)
- "We use Adam" does NOT mean these defaults β you must check
### AdamW (Loshchilov & Hutter, 2019)
AdamW is NOT the same as Adam with L2 regularization:
- Adam + L2: weight decay is applied to the gradient, which scales with the adaptive learning rate
- AdamW: weight decay is applied directly to the weights, decoupled from the gradient update
- This difference matters for large learning rates
```python
# Adam with L2 (WRONG if paper says "AdamW" or "decoupled weight decay"):
optimizer = torch.optim.Adam(params, lr=lr, weight_decay=0.01)
# AdamW (correct decoupled weight decay):
optimizer = torch.optim.AdamW(params, lr=lr, weight_decay=0.01)
```
Critical detail: **Weight decay is usually NOT applied to bias terms and LayerNorm parameters.** If the paper doesn't state this, implement it anyway and flag as [ASSUMPTION]:
```python
# [ASSUMPTION] Not applying weight decay to biases and normalization layers
# (standard practice, but paper does not specify)
no_decay = ['bias', 'LayerNorm.weight', 'LayerNorm.bias']
param_groups = [
{'params': [p for n, p in model.named_parameters()
if not any(nd in n for nd in no_decay)],
'weight_decay': config.weight_decay},
{'params': [p for n, p in model.named_parameters()
if any(nd in n for nd in no_decay)],
'weight_decay': 0.0}
]
optimizer = torch.optim.AdamW(param_groups, lr=config.lr)
```
### SGD with momentum
Common in vision papers (ResNet, EfficientNet):
- Momentum: 0.9 (almost always)
- Weight decay: 1e-4 (vision standard, not universal)
- Nesterov: sometimes on, sometimes off β papers often don't specify
---
## Learning rate schedules
### Linear warmup then constant
```python
def get_warmup_schedule(optimizer, warmup_steps):
def lr_lambda(step):
if step < warmup_steps:
return step / warmup_steps
return 1.0
return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
```
### Linear warmup then linear decay
```python
def get_linear_schedule(optimizer, warmup_steps, total_steps):
def lr_lambda(step):
if step < warmup_steps:
return step / warmup_steps
return max(0.0, (total_steps - step) / (total_steps - warmup_steps))
return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
```
### Linear warmup then cosine decay (very common)
```python
def get_cosine_schedule(optimizer, warmup_steps, total_steps, min_lr_ratio=0.0):
def lr_lambda(step):
if step < warmup_steps:
return step / warmup_steps
progress = (step - warmup_steps) / (total_steps - warmup_steps)
return min_lr_ratio + (1.0 - min_lr_ratio) * 0.5 * (1.0 + math.cos(math.pi * progress))
return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
```
### Transformer schedule (Vaswani et al., 2017)
The original transformer uses a unique schedule that most papers reference:
```python
def get_transformer_schedule(optimizer, d_model, warmup_steps):
"""lr = d_model^(-0.5) * min(step^(-0.5), step * warmup_steps^(-1.5))"""
def lr_lambda(step):
step = max(step, 1) # avoid division by zero
return d_model ** (-0.5) * min(step ** (-0.5), step * warmup_steps ** (-1.5))
return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
```
Note: This schedule doesn't need a base learning rate β it computes the LR from scratch.
When papers say "following Vaswani et al. for learning rate," this is what they mean.
---
## Batch size semantics
### The nomenclature problem
When a paper says "batch size 256," it could mean:
1. **Per-GPU batch size**: each GPU processes 256 samples
2. **Global batch size**: 256 total across all GPUs
3. **Effective batch size after gradient accumulation**: might be 32 per-GPU Γ 8 accumulation steps
This is CRITICAL for reproduction because learning rate scaling depends on batch size.
### How to determine which
Check for:
- "per-GPU batch size" or "micro-batch size" β clear
- "total batch size" or "effective batch size" β clear
- Just "batch size" β ambiguous. Look for:
- Number of GPUs mentioned β if N GPUs and batch size B, probably B total
- "gradient accumulation steps = K" β if batch B and K steps, effective batch = B Γ K
- If they mention both batch size and number of GPUs but don't clarify β flag as PARTIALLY_SPECIFIED
### Learning rate linear scaling rule
Many papers implicitly use the linear scaling rule (Goyal et al., 2017):
- If base LR is `lr` for batch size `B`, then LR for batch size `B'` is `lr * B' / B`
- This applies to SGD. For Adam, the scaling is usually sqrt: `lr * sqrt(B' / B)`
- If the paper trains on 8 GPUs with batch 256 and you use 1 GPU with batch 32, adjust the LR
- **Flag this if you change the batch size from what the paper specifies**
---
## Gradient clipping
### Max gradient norm (most common)
```python
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
```
Apply AFTER `loss.backward()` and BEFORE `optimizer.step()`.
Common values:
- 1.0 (most common)
- 0.5 (some NLP papers)
- 5.0 (some RL papers)
- If paper says "gradient clipping" without a value β UNSPECIFIED
### Max gradient value (less common)
```python
torch.nn.utils.clip_grad_value_(model.parameters(), clip_value=0.5)
```
Different from norm clipping β clips each gradient element independently.
---
## Mixed precision training
### What it means
Train with FP16 (or BF16) for speed, but keep master weights in FP32 for stability.
### PyTorch implementation
```python
scaler = torch.amp.GradScaler()
with torch.amp.autocast(device_type='cuda', dtype=torch.float16):
output = model(input)
loss = loss_fn(output, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
```
### What papers don't tell you
- BF16 doesn't need loss scaling (GradScaler) β it has the same dynamic range as FP32
- FP16 DOES need loss scaling to prevent gradient underflow
- Some operations must stay in FP32: softmax, layer norm, loss computation
- PyTorch autocast handles most of this automatically, but custom operations may not be covered
- "We use mixed precision" without specifying FP16 vs BF16 is **PARTIALLY_SPECIFIED**
---
## Exponential Moving Average (EMA)
Common in diffusion models, GANs, and some vision models. Papers often mention EMA without specifying the decay rate.
```python
class EMA:
"""Maintains exponential moving average of model parameters."""
def __init__(self, model: nn.Module, decay: float = 0.9999):
self.decay = decay
self.shadow = {name: param.clone().detach()
for name, param in model.named_parameters()}
@torch.no_grad()
def update(self, model: nn.Module):
for name, param in model.named_parameters():
self.shadow[name].mul_(self.decay).add_(param.data, alpha=1 - self.decay)
def apply(self, model: nn.Module):
"""Load EMA weights into model for evaluation."""
for name, param in model.named_parameters():
param.data.copy_(self.shadow[name])
```
Common decay values:
- 0.9999 (diffusion models β DDPM, DDIM)
- 0.999 (some GAN papers)
- 0.99 (faster averaging)
- If paper says "EMA" but not the decay rate β UNSPECIFIED
---
## Common training recipes by domain
### Language models (GPT-style)
- Optimizer: AdamW with Ξ²β=0.9, Ξ²β=0.95, Ξ΅=1e-8
- Weight decay: 0.1 (not on biases/norms)
- LR schedule: cosine decay with linear warmup
- Gradient clipping: 1.0 (max norm)
- Batch size: typically reported as total tokens/batch (e.g., "batch of 0.5M tokens")
### Vision transformers (ViT-style)
- Optimizer: AdamW with Ξ²β=0.9, Ξ²β=0.999
- Weight decay: 0.05-0.3
- LR schedule: cosine decay with linear warmup (5-10 epochs warmup)
- Data augmentation: RandAugment, Mixup, CutMix, random erasing (but specific combination varies)
- Label smoothing: 0.1
### Diffusion models (DDPM-style)
- Optimizer: Adam with Ξ²β=0.9, Ξ²β=0.999 (or 0.9999)
- LR: constant (often 2e-4 or 1e-4)
- No LR warmup (sometimes)
- EMA: 0.9999
- Gradient clipping: sometimes but not always
### Contrastive learning (SimCLR-style)
- Optimizer: SGD with momentum 0.9 (or LARS)
- LR schedule: cosine decay after linear warmup (10 epochs)
- Weight decay: 1e-6 (very small)
- Batch size: LARGE (4096+) β crucial for performance
- Temperature: 0.1 or 0.5
**WARNING:** These are guidelines, not specifications. If a paper says "we use standard training," it MIGHT mean the recipe above β but it's still UNSPECIFIED unless the paper states specifics. Use this knowledge to choose reasonable defaults, not to skip flagging.
|