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.