SKILL.md

---
name: paper-reproduction
description: "Skill for reproducing ML research papers from scratch when no official code exists. Use this whenever a user asks to implement, reproduce, or replicate a paper — especially papers involving novel loss functions, custom training loops, or non-standard architectures that aren't covered by existing HF trainers. Also use when the user mentions 'paper reproduction', 'implement this paper', 'no official code', or describes a method from a specific arxiv paper. Covers: reading papers systematically, extracting hyperparameters, building custom training pipelines, handling library-specific gotchas, VRAM estimation, checkpointing for multi-session training, and iterating on GPU results."
---

# Paper Reproduction Skill

Rules and procedures for reproducing ML research papers from scratch. All concrete mistakes, war stories, and examples live in [LEARNING.md](LEARNING.md). Next steps for this project live in [TODO.md](TODO.md).

---

## 1. Paper Reading

Read methodology sections (3, 4, 5) line by line. Read ALL appendices — they contain the actual recipe.

### Extraction checklist

```
□ Loss function — exact math, every symbol defined
□ Architecture — layers, dims, activations, normalization
□ Optimizer — type, lr, betas, weight decay, scheduler
□ Batch size — for each phase/component separately
□ Training iterations — for each phase/component separately
□ Dataset preprocessing — normalization range, image size, augmentation
□ Evaluation protocol — metrics, number of samples, special setup
□ Hyperparameters per experiment — papers often have different configs per dataset
□ Algorithm pseudocode — follow exactly before improvising
□ GPU hardware used — what the authors trained on (often buried in appendix)
□ Training time — how long did the authors' runs take?
```

---

## 2. Library API Verification

Before building ANY training loop that uses a third-party library (geomloss, POT, torchsde, torchdiffeq, etc.), write a 10-line test script that calls the library with the EXACT tensor shapes you'll use in every experiment. Not just the simplest one — all of them. If you have 2D points, MNIST images, and CIFAR images, test all three shapes.

---

## 3. VRAM Estimation

Estimate VRAM BEFORE running — not after OOM. Paper hyperparameters assume paper hardware.

**Formula for Sinkhorn (tensorized backend):** O(N² × D) per call. Pool building does ~10 calls per batch (2 potentials × 5 flow steps). Add model params × 4 bytes × 3 (params + grads + optimizer states).

**Rule:** If paper used A100 80GB and you have T4 16GB, re-derive batch sizes from VRAM constraints. Keep total samples seen (batch × iterations) constant by increasing iterations when you shrink batch.

Add CLI override flags (e.g. `--sinkhorn-batch`) so users can tune without editing config.

---

## 4. Architecture

- UNet skip connections: count pushes during downward pass, pops during upward pass. They must match exactly.
- Store config values (`num_res_blocks`, `num_levels`) as instance variables at init. Never infer them from module list lengths.
- `nn.GroupNorm(32, channels)` requires channels divisible by 32. Assert this at init for all levels.

---

## 5. Multi-Phase Training

Each phase gets its own trainer with its own optimizer. Previous phase's model goes to `eval()`.

### Shared state rules

- Never cache a DataLoader with a fixed batch size if different phases use different batch sizes. Track cached params and invalidate on change.
- `torch.cuda.empty_cache()` between phases. `del` large objects (pools, computation graphs) that won't be needed again.
- CLI overrides must touch ALL phases. If `--train-iters` should override 3 phases, grep the config for all 3 fields.

---

## 6. Checkpointing

### Phase-level (mandatory)

Save checkpoint after each phase completes. Include all model state dicts accumulated so far. Implement `--resume-phase N` that loads phase N-1 checkpoint and skips completed phases.

### Step-level (strongly recommended for phases > 10 min)

Save every N steps within a phase. Include model state, optimizer state, step number. Overwrite same file (keep latest only, unless you have disk space).

### Kaggle persistence

`/kaggle/working/` persists within a session but NOT across sessions. To carry checkpoints between sessions: commit notebook output, or copy checkpoints to a HF dataset, or download them before session ends.

---

## 7. Memory Management

- Trajectory pools / replay buffers live on CPU. Only the sampled minibatch goes to GPU via `.to(device)`.
- Pre-concatenate data structures after building: `finalize()` once → O(1) sampling per step. Never `torch.cat` the entire pool every step.
- Call `torch.cuda.empty_cache()` after pool building and between any phases with different GPU memory patterns.

---

## 8. Testing

### Before any GPU run:
1. Test EVERY experiment type with minimal configs — not just the simplest one
2. Test ALL training phases end-to-end — not just Phase 1
3. Test with `--train-iters 5 --pool-batches 2` — should complete in <60 seconds on CPU
4. Test `--resume-phase` actually works (save checkpoint → load → skip → continue)

### Before declaring code ready (pre-flight checklist):
```
□ All experiment types tested (2d, mnist, cifar10, etc.)
□ All training phases tested end-to-end
□ Library APIs tested with exact tensor shapes per experiment
□ Shared state across phases verified
□ CLI flags override ALL relevant config values
□ VRAM estimated for target hardware
□ Checkpointing works: save + resume + skip phases
□ No O(N) operations per training step where O(1) suffices
□ Expected runtimes documented per hardware tier
□ Multi-GPU limitations documented
□ Requirements.txt complete
```

---

## 9. Documentation for User

When the user runs on their own GPU (Kaggle, Colab, local):

1. Provide exact copy-paste commands
2. Document expected runtimes per hardware tier
3. Document GPU requirements and VRAM limits per experiment
4. Document what the code does NOT support (single-GPU only, no DDP, etc.)
5. If training exceeds one session, provide session-by-session commands with `--resume-phase`

---

## 10. Maintaining LEARNING.md

When a new mistake happens or a new principle is discovered:

1. Add the mistake to the **Mistake Catalog** in LEARNING.md with: What, Impact, Root cause, Prevention
2. If the mistake reveals a general principle, add it to the **Principles** section
3. If the mistake would have been caught by a pre-flight check, add that check to the checklist in section 8 above
4. Keep SKILL.md lean (rules only). LEARNING.md holds the stories and evidence.