Paper2code

paper2code — Orchestration

You are executing the paper2code skill. This file governs the high-level flow. Each stage dispatches to a detailed reasoning protocol in pipeline/. Do NOT skip stages. Do NOT combine stages. Execute them in order.

Parse arguments

Extract from the user's input:

• ARXIV_ID: the arxiv paper ID (e.g., 2106.09685). Strip any URL prefix.
• MODE: one of minimal (default), full, educational.
• FRAMEWORK: one of pytorch (default), jax, numpy.

If the user provided a full URL like https://arxiv.org/abs/2106.09685, extract the ID 2106.09685. If the user provided a versioned ID like 2106.09685v2, keep the version.

Set up working directory

Create a temporary working directory: .paper2code_work/{ARXIV_ID}/ This is where intermediate artifacts go. The final output goes in the current directory under {paper_slug}/.

Install dependencies

Run via Bash:

pip install pymupdf4llm pdfplumber requests pyyaml

Execute pipeline

Stage 1 — Paper Acquisition and Parsing

Read and follow: pipeline/01_paper_acquisition.md

Run the helper script to fetch and parse the paper:

python skills/paper2code/scripts/fetch_paper.py {ARXIV_ID} .paper2code_work/{ARXIV_ID}/

Then run structure extraction:

python skills/paper2code/scripts/extract_structure.py .paper2code_work/{ARXIV_ID}/paper_text.md .paper2code_work/{ARXIV_ID}/

Verify the outputs exist before proceeding. If extraction failed, follow the fallback protocol in pipeline/01_paper_acquisition.md.

The script also searches for official code repositories (in the paper text and on the arxiv page) and saves any found links to paper_metadata.json under the official_code key. Verify these links before relying on them — see Step 8 in pipeline/01_paper_acquisition.md.

Stage 2 — Contribution Identification

Read and follow: pipeline/02_contribution_identification.md

Read the parsed paper sections. Identify the single core contribution. Classify the paper type. Write the contribution statement. Save it to .paper2code_work/{ARXIV_ID}/contribution.md.

Stage 3 — Ambiguity Audit

Read and follow: pipeline/03_ambiguity_audit.md

Before reading this stage, also read: guardrails/hallucination_prevention.md

Go through every implementation-relevant detail. Classify each as SPECIFIED, PARTIALLY_SPECIFIED, or UNSPECIFIED. Save the audit to .paper2code_work/{ARXIV_ID}/ambiguity_audit.md.

Stage 4 — Code Generation

Read and follow: pipeline/04_code_generation.md

Before writing code, read:

• guardrails/scope_enforcement.md — to determine what's in and out of scope
• guardrails/badly_written_papers.md — if the paper is vague or inconsistent
• The relevant knowledge files in knowledge/ for the paper's domain
• The scaffold templates in scaffolds/ for the expected file structure

Determine the paper_slug from the paper title (lowercase, underscores, no special chars). Generate all files under {paper_slug}/ in the current working directory.

Stage 5 — Walkthrough Notebook

Read and follow: pipeline/05_walkthrough_notebook.md

Generate the walkthrough notebook that connects paper sections to code with runnable sanity checks. Save to {paper_slug}/notebooks/walkthrough.ipynb.

Cleanup

Remove the .paper2code_work/ directory after successful completion.

Final output

Print a summary:

✓ paper2code complete for: {paper_title}
  Output directory: {paper_slug}/
  Files generated: {list of files}
  Unspecified choices: {count} (see REPRODUCTION_NOTES.md)
  Mode: {MODE} | Framework: {FRAMEWORK}

Mode-specific behavior

• minimal (default): Core contribution only. Training loop only if contribution involves training. No data pipeline beyond Dataset skeleton.
• full: Core contribution + full training loop + data pipeline + evaluation pipeline. More code, same citation rigor.
• educational: Same as minimal but with extra inline comments explaining ML concepts, expanded walkthrough notebook with theory sections, and a PAPER_GUIDE.md that walks through the paper section by section.

Guardrails — always active

These apply at ALL stages. Read them if you haven't already:

• guardrails/hallucination_prevention.md — the most important file in this skill
• guardrails/scope_enforcement.md — what to implement and what to skip
• guardrails/badly_written_papers.md — what to do when the paper is unclear

Knowledge base — consult as needed

Before implementing any of these components, read the corresponding knowledge file:

• Transformer layers, attention, positional encoding → knowledge/transformer_components.md
• Optimizers, LR schedules, batch size semantics → knowledge/training_recipes.md
• Cross-entropy, contrastive loss, diffusion loss, ELBO → knowledge/loss_functions.md
• Framework-specific pitfalls, notation mismatches → knowledge/paper_to_code_mistakes.md

Guardrail: Handling Badly Written Papers

Reality check

Many papers are vague, inconsistent, or incomplete. This is not a criticism — writing a paper is hard, page limits are strict, and reviewers rarely check hyperparameter tables. But it means you will regularly encounter papers where the text alone is insufficient to produce a correct implementation.

This file tells you what to do when that happens. The answer is never "guess silently."

---

Decision tree for resolving ambiguity

Is the detail stated explicitly in the paper?
├── YES → Use it. Cite the section. Done.
├── PARTIALLY → Follow the partial specification protocol (below)
└── NO →
    Is there official code on GitHub?
    ├── YES → Use it as ground truth. Cite the file and line number.
    │         Mark as [FROM_OFFICIAL_CODE], not [SPECIFIED].
    └── NO →
        Is there a well-known reimplementation?
        ├── YES → Reference it in REPRODUCTION_NOTES.md.
        │         DO NOT blindly copy — they made their own choices.
        │         Mark as [UNSPECIFIED] and note the external reference.
        └── NO →
            Is there a "standard" choice in the field?
            ├── YES → Use it. Mark as [UNSPECIFIED] with alternatives.
            └── NO →
                Write a stub with a detailed docstring.
                Explain what should go here and what the options are.

---

Partial specification protocol

When the paper partially specifies something:

"We use Adam"

The paper says an optimizer but not all parameters.

# §4.1 — "We use the Adam optimizer"
# [PARTIALLY_SPECIFIED] Optimizer stated as Adam, but beta and epsilon values not specified
# Using: β₁=0.9, β₂=0.999, ε=1e-8 (PyTorch defaults)
# Alternatives: β₁=0.9, β₂=0.98, ε=1e-9 (Transformer default from Vaswani et al.)
optimizer = torch.optim.Adam(model.parameters(), lr=config.lr,
                              betas=(0.9, 0.999), eps=1e-8)

"We use dropout for regularization"

The paper mentions dropout but not the rate or placement.

# §3.3 — "We use dropout for regularization"
# [PARTIALLY_SPECIFIED] Dropout mentioned but rate not specified
# Using: 0.1 (common default for transformer models)
# Placement: after attention and FFN (not specified — this is our choice)
self.dropout = nn.Dropout(0.1)

"Similar to Transformer [Vaswani et al., 2017]"

The paper references another work for a component.

# §3.1 — "We use an architecture similar to the Transformer [Vaswani et al., 2017]"
# [PARTIALLY_SPECIFIED] "Similar" implies differences exist but none are specified
# We implement the standard Transformer from Vaswani et al. unless other sections
# describe modifications. Reader should check if the authors intended differences.

---

When the paper contradicts itself

Figure vs text

Figures are often created early in the paper-writing process and may show an earlier version of the architecture or method. Text is usually more up-to-date.

Rule: Trust the text. Flag the figure discrepancy.

# §3.2 — "We apply layer normalization before each sub-layer" (Pre-LN)
# NOTE: Figure 2 shows post-norm placement, inconsistent with this text.
# We implement pre-norm as stated in the text. If reproduction fails,
# try post-norm (change this to: x = self.norm(x + sublayer(x)))

Equation vs prose

Equations are peer-reviewed more carefully and are more precise.

Rule: Trust the equation. Flag the prose discrepancy.

Different numbers in different sections

The paper might say "learning rate of 1e-4" in one section and "learning rate of 3e-4" in another.

Rule: If one is in the appendix/hyperparameter table, trust that. If both are in prose, flag both and use the one from the main experimental section (usually the one paired with the best results).

---

When the paper is genuinely incomplete

Some papers are missing critical details that make reproduction impossible without additional information. Here's what to do:

Missing architecture details (for architecture papers)

If you cannot determine the model architecture from the paper: 1. Check for official code — this is the only reliable resolution 2. Check for well-known reimplementations (HuggingFace, lucidrains, etc.) 3. If neither exists, implement what you CAN determine and write stubs for the rest:

def mysterious_component(x: torch.Tensor) -> torch.Tensor:
    """§3.4 — The paper describes a 'novel aggregation mechanism' but does not
    provide sufficient detail to implement it.
    
    What we know:
      - Input: tensor of shape (batch, seq_len, d_model) — from §3.3
      - Output: tensor of shape (batch, d_model) — inferred from §3.5 which uses the output
      - It "aggregates information across the sequence" — from §3.4
    
    What we don't know:
      - The specific aggregation operation (mean pooling? attention? learned query?)
      - Whether it has learnable parameters
      - Whether it uses masking
    
    To complete this implementation, check:
      1. The authors' official code (if released after this generation)
      2. Follow-up papers that reproduce this work
      3. Contact the authors
    """
    raise NotImplementedError(
        "Cannot implement: paper §3.4 does not provide sufficient architectural detail. "
        "See docstring for what is known and what is missing."
    )

Missing training details (for training method papers)

If the training procedure is the contribution but is underspecified: 1. Implement the parts that ARE specified 2. For unspecified parts, use the most common baseline and flag it loudly 3. Add a section in REPRODUCTION_NOTES.md titled "Critical unspecified training details"

Missing evaluation details

Note which metric version you're using and that it might not match:

# [UNSPECIFIED] Paper reports "BLEU score" but does not specify which implementation
# Using: sacrebleu (the de facto standard for reproducible BLEU)
# NOTE: Different BLEU implementations can give significantly different numbers
# (sometimes 1-2 points) due to tokenization and smoothing differences

---

Using official code as ground truth

When official code exists, use it to resolve ambiguities. But be careful:

DO:

• Link to the exact file and line number
• Mark resolved items as [FROM_OFFICIAL_CODE]
• Note when the official code differs from the paper text
• Check the commit history — early commits are closer to the paper, later commits may reflect improvements

DON'T:

• Assume official code is always right — it may have bugs
• Copy code style or non-essential patterns
• Use code from Pull Requests or branches — stick to main/master
• Assume the code matches the paper — sometimes best results come from later improvements not in the paper

Format:

# [FROM_OFFICIAL_CODE] Weight initialization: normal with std=0.02
# Source: github.com/author/repo/blob/main/model.py#L42
# Paper does not specify initialization (§3.1)
nn.init.normal_(self.weight, std=0.02)

---

Using well-known reimplementations

HuggingFace, lucidrains, Andrej Karpathy, Phil Wang — some implementers are widely trusted. When using their implementations as reference:

1. Reference in REPRODUCTION_NOTES.md, not in inline comments (unless resolving a specific choice) 2. Note that they may have made their own unspecified choices — their choices are educated guesses, not specifications 3. Do not assume they are correct — well-known reimplementations sometimes have bugs or intentional deviations 4. Use them for sanity checking, not as ground truth

---

Known erratum handling

If the paper has a published erratum, correction, or the authors have acknowledged an error:

# §3.2, Eq. 4 — Original paper has a typo: denominator should be sqrt(d_k), not d_k
# ERRATUM: Acknowledged by authors in arxiv v2 (see changelog)
# We implement the corrected version
scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)

Check arxiv for paper versions — v2, v3 often contain corrections. The abstract page shows the version history.

---

When to give up

There is a point where a paper is so underspecified that a meaningful implementation is impossible. Signs:

• The core method relies on a component described only as "proprietary" or "internal"
• The paper describes results but not the method (sometimes in position papers)
• The paper is a survey or review with no novel method
• The paper's contribution is purely theoretical with no implementable component

In these cases: 1. Generate what you CAN (data loading, evaluation metrics, utility functions) 2. Write a clear REPRODUCTION_NOTES.md explaining why a full implementation is not possible 3. List exactly what information would be needed to complete it 4. Do NOT generate fake implementations to make the output look complete

Guardrail: Hallucination Prevention

The core problem

The single most dangerous failure mode of paper2code is confident invention — generating code that looks plausible but implements details the paper never specified, without flagging them. A researcher who trusts this output will waste days debugging differences that exist because the model guessed.

This file exists to prevent that. Read it before every code generation stage. Internalize it.

---

The bright line rule

If a detail is not explicitly stated in the paper text, the appendix, the paper's official GitHub repository, or a well-known replication paper — it is UNSPECIFIED.

Not "probably this." Not "standard practice." Not "everyone uses." UNSPECIFIED.

---

What counts as "stated in the paper"

Counts as SPECIFIED:

• Direct statement: "We use d_model = 512" → SPECIFIED
• Table entry: Table 3 shows "learning rate: 3e-4" → SPECIFIED
• Equation: Eq. 4 defines the loss function → SPECIFIED (for the equation, not for implementation details like numerical precision)
• Algorithm box: Algorithm 1 lines 3-5 describe the update rule → SPECIFIED
• Footnote: Footnote 3 says "we clip gradients at 1.0" → SPECIFIED
• Appendix: Appendix B Table 6 shows "warmup steps = 4000" → SPECIFIED

Does NOT count as SPECIFIED:

• "We use standard optimization" → UNSPECIFIED (standard according to whom? which optimizer? what hyperparameters?)
• "Following prior work [23]" → PARTIALLY_SPECIFIED (you must look up [23] and report what it says, or flag that the reader needs to)
• "We use Adam" → PARTIALLY_SPECIFIED (Adam has β₁, β₂, ε parameters — are they stated?)
• "Similar architecture to [X]" → PARTIALLY_SPECIFIED (similar is not identical — what differs?)
• "Standard hyperparameters" → UNSPECIFIED
• "Default settings" → UNSPECIFIED (whose defaults? PyTorch? TensorFlow? They differ.)
• Descriptions in related work of other people's methods → NOT A SPECIFICATION OF THIS PAPER'S METHOD
• Blog posts, tweets, or talks by the authors → NOT PEER-REVIEWED, note as supplementary only

---

The UNSPECIFIED comment protocol

When you make a choice for an UNSPECIFIED item, the code comment must have three parts:

# [UNSPECIFIED] {What the paper doesn't specify}
# Using: {your choice}
# Alternatives: {other reasonable choices}

Example:

# [UNSPECIFIED] Paper does not state activation function in the feed-forward network
# Using: GELU (most common in recent transformer implementations)
# Alternatives: ReLU (original transformer), SiLU/Swish (used in LLaMA, PaLM)
self.activation = nn.GELU()

You must NEVER write just:

self.activation = nn.GELU()  # standard choice

This hides the fact that the paper didn't specify it.

---

Equation ground truth rule

Equations are more precise than prose. If the prose description and the equation conflict:

1. Implement the equation 2. Flag the discrepancy explicitly:

# §3.2, Eq. 4 — loss = -log(exp(sim(z_i, z_j)/τ) / Σ_k exp(sim(z_i, z_k)/τ))
# NOTE: The prose in §3.2 says "we average over all positive pairs" but Eq. 4
# sums rather than averages. We implement Eq. 4 as written.

This applies even when the equation is clearly a typo. Implement the equation, flag the likely typo, and add a comment about what the authors probably meant.

---

The "standard" trap

Papers frequently use the word "standard" without definition. Here's what to do:

---

The number precision trap

Do not silently change numbers:

• If the paper says 512, use 512 — not 256 "for simplicity" without flagging it
• If the paper says 0.9, use 0.9 — not 0.99 because "that's what people usually use"
• If the paper says 100k steps, use 100000 — not 50000 because "it should converge faster"

Any deviation from stated numbers must be flagged:

# §5.2 states d_model = 512, but this walkthrough uses d_model = 64 for CPU execution
# Set d_model = 512 for actual reproduction

---

The initialization trap

Weight initialization is almost never specified in papers but matters enormously. If the paper does not specify initialization:

# [UNSPECIFIED] Paper does not describe weight initialization
# Using: PyTorch defaults (Kaiming uniform for Linear, uniform for Embedding)
# Alternatives: Xavier uniform (common for transformers), normal init with std=0.02
# NOTE: Initialization can significantly affect training stability and convergence

If the paper specifies initialization (rare but valuable), implement it exactly and cite the section.

---

The framework translation trap

Different frameworks have different defaults that papers don't always clarify:

Batch normalization momentum

• Paper says momentum = 0.1
• PyTorch BatchNorm uses momentum = 0.1 (but its definition is inverted: running_mean = (1 - momentum) * running_mean + momentum * batch_mean)
• TensorFlow BatchNorm uses momentum = 0.99 (with convention: running_mean = momentum * running_mean + (1 - momentum) * batch_mean)
• PyTorch momentum=0.1 ≈ TensorFlow momentum=0.9
• Always clarify which convention the paper uses

Dropout rate vs keep probability

• Paper says "dropout 0.1" — does this mean drop probability = 0.1 or keep probability = 0.1?
• Almost always means drop probability = 0.1 (keep = 0.9), but older papers sometimes use keep probability
• PyTorch nn.Dropout(p=0.1) means drop probability = 0.1

Layer normalization epsilon

• Papers almost never specify this
• PyTorch default: 1e-5
• Common in papers: 1e-6
• Some implementations: 1e-8
• Flag as UNSPECIFIED and note which you chose

---

The "we found" trap

When a paper says "we found that X works better," this is usually an empirical observation, not a derived result. Treat it as useful information but note that:

• "Better" often means "better for our specific setup/dataset/scale"
• The alternative (not-X) might work fine for different scenarios
• This is an [ASSUMPTION] not a [SPECIFIED] item when you implement it as the default

---

Prohibited phrases in generated code

Never write any of these in code comments:

• "standard practice" (without specifying what practice)
• "as usual" (usual for whom?)
• "obviously" (if it's obvious, you don't need to say it; and it's probably not obvious)
• "typically" (without a citation)
• "it's well known that" (without a citation)
• "for simplicity" (as justification for deviating from the paper)
• "should work" (either it's specified or it isn't)

---

The official code shortcut

If official code exists: 1. Note its URL in REPRODUCTION_NOTES.md 2. You may use it to resolve UNSPECIFIED items — but:

• Mark them as [FROM_OFFICIAL_CODE], not as SPECIFIED
• The official code may differ from the paper (bug fixes, improvements, errors)
• Link to the exact line: github.com/author/repo/blob/main/model.py#L42

3. Do NOT copy-paste code. Read the official code to understand the choice, then implement it yourself with a citation

---

Self-audit questions

Before finishing any code generation, ask yourself: 1. If I removed all my [UNSPECIFIED] comments, would a reader think the paper specified everything? If yes, I'm probably missing flags. 2. Did I add any implementation detail from my own ML knowledge without checking if the paper says it? Flag it. 3. Would the authors of this paper agree that my code matches their description? If I'm not sure, something needs a flag. 4. Is there a single magic number anywhere without a citation or [UNSPECIFIED] comment? Find it and fix it.

Guardrail: Scope Enforcement

Purpose

Determine what to implement and what to skip. Implementing too much adds noise. Implementing too little is useless. This decision tree makes scope decisions explicit and traceable.

---

The decision tree

For each component you're considering implementing, walk through this tree:

Is this component the paper's core contribution?
├── YES → ALWAYS IN SCOPE. Implement fully with citation anchoring.
│
└── NO →
    Does the paper describe this component in detail (equations, pseudocode, algorithm box)?
    ├── YES →
    │   Is this component necessary to understand the core contribution?
    │   ├── YES → IN SCOPE. Implement with citations.
    │   └── NO → OUT OF SCOPE. Note its existence, don't implement.
    │
    └── NO →
        Does the paper reference another paper for this component?
        ├── YES → REFERENCE ONLY. Import or note the dependency.
        │         Add a comment: "§X.Y — uses {component} from {reference}"
        │
        └── NO →
            Is this a standard ML component (optimizer, standard layer, etc.)?
            ├── YES → USE FRAMEWORK DEFAULT. Import from PyTorch/library.
            │         Don't reimplement. Document which version you're using.
            │
            └── NO →
                You've found something the paper assumes you know but doesn't explain.
                Flag it in REPRODUCTION_NOTES.md and use a reasonable default.

---

Scope by paper type

Type (a): New architecture paper

Type (b): New training method or loss function

Type (c): New inference technique

Type (d): New dataset or benchmark

Type (e): Theoretical with empirical validation

Type (f): System/engineering paper

---

Mode-specific scope overrides

minimal mode (default)

• Stick strictly to the decision tree above
• Training loop: only if the contribution IS training
• Data pipeline: skeleton only
• No distributed training, no checkpointing, no logging

full mode

• Same as minimal PLUS:
• Full training loop with all paper-specified hyperparameters
• Data pipeline with preprocessing (still needs user to provide data)
• Evaluation pipeline (not just metric functions)
• Checkpointing (minimal, for training)
• Learning rate schedule visualization

educational mode

• Same as minimal PLUS:
• Extra comments explaining ML concepts
• Extended walkthrough notebook with theory sections
• PAPER_GUIDE.md that walks through the paper section by section

---

Explicit exclusions (never in scope regardless of mode)

1. Baselines and comparison methods. The paper compares against X, Y, Z — we only implement the paper's contribution, not X, Y, Z.

2. Ablation variants. If the paper ablates components A, B, C — we implement the full model with A+B+C, not each variant separately.

3. Multi-GPU/distributed training infrastructure. Unless the paper's contribution IS the distributed strategy.

4. Experiment tracking and logging. No wandb, no tensorboard, no MLflow setup.

5. Docker/container configuration. Not relevant to understanding the paper.

6. CI/CD, linting, formatting configuration. This is a research scaffold, not a production project.

7. Unit tests. The walkthrough notebook serves as a verification layer. No pytest needed.

8. Dataset downloading. Provide instructions, never download automatically.

9. Pre-trained weight downloading. Reference where to get them, never download.

10. Web interfaces, APIs, or demos. Out of scope entirely.

---

How to handle scope creep

During code generation, you may be tempted to add "just one more thing" that would make the implementation more complete. Resist this unless:

1. The "one more thing" is directly specified in the paper AND 2. It is necessary for the core contribution to make sense AND 3. Leaving it out would make the implementation misleading

If it's "nice to have" — leave it out. A focused, correct implementation of the core contribution is infinitely better than a sprawling implementation with guessed peripherals.

---

How to document scope decisions

In REPRODUCTION_NOTES.md, include a section:

## Scope decisions

### Implemented
- {Component} — reason: {core contribution / necessary for X / etc.}

### Intentionally excluded
- {Component} — reason: {baseline method / standard component imported from X / not described in paper}

### Would need for full reproduction
- {Component} — what it is, where to get it, why we didn't include it

This makes scope decisions transparent. A researcher reading REPRODUCTION_NOTES.md should understand exactly what this implementation covers and what it doesn't.

Knowledge: Loss Functions

Purpose

Canonical implementations of loss functions commonly referenced in papers, with attention to numerical stability, framework-specific differences, and common implementation mistakes.

---

Cross-Entropy Loss

Standard cross-entropy (classification)

# PyTorch built-in:
loss_fn = nn.CrossEntropyLoss()
# Expects: predictions (batch, num_classes) as LOGITS (not probabilities)
#          targets (batch,) as class indices (not one-hot)

Common mistake: Passing probabilities (after softmax) instead of logits. PyTorch's CrossEntropyLoss applies log-softmax internally for numerical stability.

Cross-entropy with label smoothing

Papers describe label smoothing differently from what PyTorch does:

Paper formula (Szegedy et al., 2016): For C classes and smoothing ε:

• Target for correct class: 1 - ε
• Target for each incorrect class: ε / (C - 1)

PyTorch >= 1.10 built-in:

loss_fn = nn.CrossEntropyLoss(label_smoothing=0.1)

PyTorch's implementation distributes ε uniformly: correct class gets 1 - ε + ε/C, others get ε/C. This is slightly different from the paper formula (the paper distributes to incorrect classes only).

Custom implementation matching the paper exactly:

def label_smoothed_cross_entropy(logits: torch.Tensor, targets: torch.Tensor,
                                  n_classes: int, smoothing: float = 0.1) -> torch.Tensor:
    """Label smoothing exactly as described in Szegedy et al., 2016.
    
    Args:
        logits: (batch, n_classes) — raw logits, NOT probabilities
        targets: (batch,) — class indices
        n_classes: number of classes
        smoothing: label smoothing factor ε
    """
    log_probs = F.log_softmax(logits, dim=-1)  # (batch, n_classes)
    
    # NLL loss for the true class
    nll_loss = -log_probs.gather(dim=-1, index=targets.unsqueeze(-1)).squeeze(-1)
    
    # Smooth loss: uniform over all classes
    smooth_loss = -log_probs.mean(dim=-1)
    
    loss = (1.0 - smoothing) * nll_loss + smoothing * smooth_loss
    return loss.mean()

The difference matters: For large vocabularies (NLP), the difference between PyTorch's built-in and the paper formula is negligible. For small class counts (e.g., 10), it can be noticeable.

---

Contrastive Losses

NT-Xent / InfoNCE (SimCLR, Oord et al.)

def nt_xent_loss(z_i: torch.Tensor, z_j: torch.Tensor,
                 temperature: float = 0.5) -> torch.Tensor:
    """Normalized Temperature-scaled Cross Entropy loss.
    
    Used in SimCLR (Chen et al., 2020).
    
    Args:
        z_i: (batch, d) — representations of view 1
        z_j: (batch, d) — representations of view 2 (positive pairs)
        temperature: τ — scaling factor
        
    Returns:
        Scalar loss
    """
    batch_size = z_i.size(0)
    
    # Normalize representations
    z_i = F.normalize(z_i, dim=-1)  # (batch, d)
    z_j = F.normalize(z_j, dim=-1)  # (batch, d)
    
    # Concatenate representations: [z_i; z_j]
    z = torch.cat([z_i, z_j], dim=0)  # (2*batch, d)
    
    # Compute similarity matrix
    sim = torch.matmul(z, z.T) / temperature  # (2*batch, 2*batch)
    
    # Mask out self-similarity (diagonal)
    mask = ~torch.eye(2 * batch_size, dtype=torch.bool, device=z.device)
    sim = sim.masked_fill(~mask, float('-inf'))
    
    # Positive pairs: (i, i+batch) and (i+batch, i)
    labels = torch.cat([
        torch.arange(batch_size, 2 * batch_size),
        torch.arange(0, batch_size)
    ], dim=0).to(z.device)  # (2*batch,)
    
    return F.cross_entropy(sim, labels)

Critical: temperature matters enormously. SimCLR uses τ=0.5 in the paper, but the appendix shows results are sensitive to this value. CLIP uses τ as a learned parameter initialized to 0.07. Always check what temperature the paper uses.

Common mistake: Not normalizing the representations before computing similarity. Without normalization, the cosine similarity becomes a dot product, which is unbounded and destabilizes training.

Triplet loss

def triplet_loss(anchor: torch.Tensor, positive: torch.Tensor,
                 negative: torch.Tensor, margin: float = 1.0) -> torch.Tensor:
    """Standard triplet margin loss.
    
    Args:
        anchor: (batch, d)
        positive: (batch, d)
        negative: (batch, d)
        margin: minimum desired distance gap
    """
    pos_dist = F.pairwise_distance(anchor, positive)  # (batch,)
    neg_dist = F.pairwise_distance(anchor, negative)  # (batch,)
    loss = F.relu(pos_dist - neg_dist + margin)
    return loss.mean()

---

Diffusion Losses

DDPM loss (Ho et al., 2020)

The simplified loss from DDPM. This is L_simple from Eq. 14:

def ddpm_loss(model: nn.Module, x_0: torch.Tensor,
              noise_schedule: dict) -> torch.Tensor:
    """Denoising Diffusion Probabilistic Models simplified loss.
    
    L_simple = E_{t, x_0, ε}[||ε - ε_θ(x_t, t)||²]
    
    The model predicts the noise ε that was added, not the clean signal.
    
    Args:
        model: noise prediction network ε_θ
        x_0: (batch, C, H, W) — clean images
        noise_schedule: dict with 'betas', 'alphas_cumprod', etc.
    """
    batch_size = x_0.size(0)
    
    # Sample random timesteps
    t = torch.randint(0, len(noise_schedule['betas']), (batch_size,),
                      device=x_0.device)
    
    # Sample noise
    noise = torch.randn_like(x_0)
    
    # Get noisy image: x_t = sqrt(α̅_t) * x_0 + sqrt(1 - α̅_t) * ε
    alpha_cumprod_t = noise_schedule['alphas_cumprod'][t]  # (batch,)
    alpha_cumprod_t = alpha_cumprod_t.view(-1, 1, 1, 1)  # (batch, 1, 1, 1) for broadcasting
    
    x_t = torch.sqrt(alpha_cumprod_t) * x_0 + torch.sqrt(1 - alpha_cumprod_t) * noise
    
    # Predict noise
    predicted_noise = model(x_t, t)  # (batch, C, H, W)
    
    # MSE loss between true and predicted noise
    loss = F.mse_loss(predicted_noise, noise)
    return loss

Noise schedule:

def linear_noise_schedule(timesteps: int, beta_start: float = 1e-4,
                           beta_end: float = 0.02) -> dict:
    """Linear variance schedule from DDPM.
    
    β_t increases linearly from β_start to β_end over T timesteps.
    """
    betas = torch.linspace(beta_start, beta_end, timesteps)
    alphas = 1.0 - betas
    alphas_cumprod = torch.cumprod(alphas, dim=0)
    
    return {
        'betas': betas,
        'alphas': alphas,
        'alphas_cumprod': alphas_cumprod,
        'sqrt_alphas_cumprod': torch.sqrt(alphas_cumprod),
        'sqrt_one_minus_alphas_cumprod': torch.sqrt(1, - alphas_cumprod),
    }

Common mistakes with diffusion losses: 1. Indexing: alphas_cumprod[t] — make sure t is the right shape for broadcasting 2. The DDPM paper uses 1000 timesteps with linear schedule β₁=1e-4 to β_T=0.02 3. The model predicts noise ε, not x₀ (in the simplified loss; other parameterizations exist) 4. sqrt_one_minus_alphas_cumprod should NOT go to exactly 0 — check boundary conditions

---

VAE ELBO

Evidence Lower Bound

def vae_loss(recon_x: torch.Tensor, x: torch.Tensor,
             mu: torch.Tensor, log_var: torch.Tensor,
             beta: float = 1.0) -> tuple:
    """VAE loss = Reconstruction loss + β * KL divergence.
    
    ELBO = E_q[log p(x|z)] - β * KL(q(z|x) || p(z))
    
    Args:
        recon_x: (batch, ...) — reconstructed output
        x: (batch, ...) — original input
        mu: (batch, d_latent) — mean of q(z|x)
        log_var: (batch, d_latent) — log variance of q(z|x)
        beta: weight for KL term (β=1 gives standard VAE)
    """
    # Reconstruction loss (pixel-wise)
    recon_loss = F.mse_loss(recon_x, x, reduction='sum') / x.size(0)
    # Alternative: F.binary_cross_entropy for images in [0,1]
    
    # KL divergence: KL(N(μ, σ²) || N(0, 1))
    # = -0.5 * Σ(1 + log(σ²) - μ² - σ²)
    kl_loss = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp()) / x.size(0)
    
    total_loss = recon_loss + beta * kl_loss
    return total_loss, recon_loss, kl_loss

Common mistakes: 1. Reconstruction loss type: Binary cross-entropy for images normalized to [0,1], MSE for others. Papers often don't specify. 2. Reduction: sum vs mean — affects the relative weight of reconstruction vs KL. If using mean, the KL term effectively gets more weight relative to reconstruction for high-dimensional data. 3. β value: β=1 is standard VAE. β-VAE uses β>1 for more disentangled representations. If the paper says "VAE loss" without specifying β, use β=1 and flag it. 4. KL divergence sign: The formula gives a positive KL value. The loss ADDS KL (penalizes deviation from prior). If your KL is negative, the sign is wrong.

---

Numerical stability patterns

Log-sum-exp trick

When computing log(Σ exp(x_i)):

# WRONG (overflow for large values):
result = torch.log(torch.sum(torch.exp(x)))

# CORRECT:
result = torch.logsumexp(x, dim=-1)
# Internally: max_x + log(Σ exp(x_i - max_x))

Softmax stability

PyTorch's F.softmax and F.log_softmax are already numerically stable (subtract max internally). But if you implement softmax manually:

# WRONG:
weights = torch.exp(scores) / torch.exp(scores).sum(dim=-1, keepdim=True)

# CORRECT:
scores = scores - scores.max(dim=-1, keepdim=True).values
weights = torch.exp(scores) / torch.exp(scores).sum(dim=-1, keepdim=True)

Epsilon in denominators

When dividing by a value that could be zero (e.g., normalizing):

# Not just any epsilon — match the paper's convention or use a safe default
x_normalized = x / (x.norm(dim=-1, keepdim=True) + 1e-8)

Clamping log probabilities

# When computing log of probabilities that might be 0:
log_probs = torch.log(probs.clamp(min=1e-8))
# Or better: compute in log space from the start using log_softmax

Knowledge: Paper-to-Code Translation Mistakes

Purpose

A detailed catalog of patterns where naive paper-to-code translation produces wrong results. These are not bugs in the usual sense — the code runs fine but implements something different from what the paper describes. This file exists because these mistakes are systematic and predictable.

---

Notation mismatches

Batch normalization momentum

The problem: PyTorch and TensorFlow define momentum differently.

Paper says: running_mean = (1 - α) * running_mean + α * batch_mean where α (or momentum) = 0.1

• PyTorch nn.BatchNorm2d(momentum=0.1): uses running_mean = (1 - momentum) * running_mean + momentum * batch_mean → matches the formula above with momentum=0.1
• TensorFlow tf.keras.layers.BatchNormalization(momentum=0.99): uses running_mean = momentum * running_mean + (1 - momentum) * batch_mean → to get the same behavior, set momentum=0.9

Translation rule: PyTorch momentum=x ≈ TensorFlow momentum=1-x

If a paper reports momentum without specifying the framework convention → PARTIALLY_SPECIFIED. Check which framework the official code uses.

Dropout rate vs keep probability

The problem: Older papers (and TensorFlow v1) use "keep probability." Newer papers and PyTorch use "drop probability."

• Paper says "dropout 0.1" → almost certainly means drop probability = 0.1 (keep 90% of neurons)
• Paper says "dropout rate 0.9" → probably means keep probability = 0.9 (same thing)
• Paper says "we keep 90% of neurons" → keep probability = 0.9

# PyTorch: p = drop probability
nn.Dropout(p=0.1)  # drops 10%, keeps 90%

When in doubt: Check the dropout paper (Srivastava et al., 2014) which uses keep probability notation. Most post-2018 papers use drop probability.

Convolution padding

The problem: "same padding" means different things.

• Paper says "same convolution" → output size = input size
• PyTorch: nn.Conv2d(padding='same') or manually compute padding = kernel_size // 2
• TensorFlow: tf.keras.layers.Conv2D(padding='same') handles it automatically
• For even kernel sizes, "same" can't be symmetric → PyTorch requires explicit asymmetric padding via nn.functional.pad

Tensor dimension ordering

• PyTorch images: (batch, channels, height, width) — NCHW
• TensorFlow images: (batch, height, width, channels) — NHWC (default)
• If converting a TensorFlow paper to PyTorch, every convolution, pooling, normalization, and reshape must account for this

---

Activation function gotchas

GELU approximation

There are two common GELU implementations:

# Exact GELU (PyTorch default since 1.12):
nn.GELU()
# = x * Φ(x) where Φ is the standard normal CDF

# Tanh approximation (used in GPT-2, BERT):
nn.GELU(approximate='tanh')
# = 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x³)))

These give slightly different results. Some papers use one, some the other. BERT uses the tanh approximation. Modern papers usually use exact GELU but often don't specify.

Swish / SiLU

# These are the same thing:
nn.SiLU()  # PyTorch name
# = x * sigmoid(x)

# Some papers call it Swish with a trainable β:
# swish(x) = x * sigmoid(β * x)
# When β=1, this is SiLU. PyTorch's SiLU always uses β=1.

If the paper says "Swish" without specifying β, use β=1 (SiLU) and flag it.

Leaky ReLU slope

Papers rarely specify the negative slope for Leaky ReLU:

• Paper default: 0.01
• PyTorch default: 0.01 (matches)
• But some implementations use 0.1 or 0.2 (especially in GANs)

---

Weight initialization gotchas

Xavier/Glorot initialization

# Xavier uniform (Glorot & Bengio, 2010):
nn.init.xavier_uniform_(layer.weight)  # U(-a, a) where a = sqrt(6 / (fan_in + fan_out))

# Xavier normal:
nn.init.xavier_normal_(layer.weight)  # N(0, 2 / (fan_in + fan_out))

The problem: "Xavier initialization" doesn't specify uniform or normal. They're different. Papers rarely specify which one.

Kaiming/He initialization

# Kaiming uniform (He et al., 2015):
nn.init.kaiming_uniform_(layer.weight, mode='fan_in', nonlinearity='relu')

# Kaiming normal:
nn.init.kaiming_normal_(layer.weight, mode='fan_in', nonlinearity='relu')

The problem: mode='fan_in' vs mode='fan_out' matters. PyTorch defaults to fan_in, but some papers require fan_out for the decoder.

PyTorch default initialization

If the paper doesn't mention initialization, PyTorch applies:

• nn.Linear: Kaiming uniform (fan_in)
• nn.Conv2d: Kaiming uniform (fan_in)
• nn.Embedding: Normal(0, 1)
• nn.LayerNorm: weight=1, bias=0
• nn.BatchNorm: weight=1, bias=0

These defaults are reasonable but may not match what the paper used (and didn't specify).

---

Dimension and reshaping gotchas

view() vs reshape()

# view() requires contiguous memory — will error if not contiguous
x = x.view(batch, -1)

# reshape() handles non-contiguous tensors by copying if necessary
x = x.reshape(batch, -1)

After transpose() or permute(), tensors are NOT contiguous. Use .contiguous() before .view() or just use .reshape().

einsum notation

Papers increasingly use einsum notation. Common patterns:

# Matrix multiplication:
torch.einsum('bij,bjk->bik', A, B)  # same as torch.bmm(A, B)

# Attention scores:
torch.einsum('bhqd,bhkd->bhqk', q, k)  # (batch, heads, q_len, k_len)

# Bilinear:
torch.einsum('bi,ij,bj->b', x, W, y)

Gotcha: einsum is correct but can be slow for some patterns. For standard operations (matmul, dot product), use torch.matmul or @ operator — they're optimized.

---

Loss function gotchas

Cross-entropy: logits vs probabilities

If you see loss = -sum(y * log(p)) in a paper, this is cross-entropy applied to probabilities. PyTorch nn.CrossEntropyLoss expects LOGITS and applies log-softmax internally.

# WRONG: passing softmax output to CrossEntropyLoss
probs = F.softmax(logits, dim=-1)
loss = nn.CrossEntropyLoss()(probs, targets)  # WRONG — double softmax

# CORRECT:
loss = nn.CrossEntropyLoss()(logits, targets)  # applies log-softmax internally

MSE reduction

Papers often specify loss as a sum or average differently:

• PyTorch MSELoss(reduction='mean'): averages over ALL elements (batch × features)
• Some papers average only over features, then sum over batch
• Some papers sum everything, no averaging

Check the paper's equation: does it have $\frac{1}{N}$ (mean) or $\sum$ (sum)?

KL divergence direction

$KL(P || Q) \neq KL(Q || P)$

• Forward KL ($KL(P || Q)$): "I want Q to cover everywhere P has mass" — mode-covering
• Reverse KL ($KL(Q || P)$): "I want Q to be 0 wherever P is 0" — mode-seeking

PyTorch's nn.KLDivLoss computes KL(target || input) — note the argument order!

# PyTorch expects LOG probabilities as input, regular probabilities as target
loss = nn.KLDivLoss(reduction='batchmean')(
    F.log_softmax(predicted_logits, dim=-1),  # LOG probs
    target_probs  # regular probs
)

---

Training loop gotchas

optimizer.zero_grad() placement

# CORRECT: zero gradients BEFORE forward pass
optimizer.zero_grad()
output = model(input)
loss = loss_fn(output, target)
loss.backward()
optimizer.step()

# ALSO CORRECT (set_to_none is faster):
optimizer.zero_grad(set_to_none=True)

Learning rate scheduler step

Different schedulers step at different rates:

# Per-step schedulers (most warmup schedules):
for step in range(total_steps):
    train_step()
    scheduler.step()  # after every training step

# Per-epoch schedulers (ReduceLROnPlateau, StepLR):
for epoch in range(epochs):
    train_epoch()
    scheduler.step()  # after every epoch

Papers often don't specify when the scheduler steps. If it's a warmup schedule, it's per-step.

Gradient accumulation

# If effective_batch = micro_batch × accumulation_steps:
for i, batch in enumerate(dataloader):
    loss = model(batch) / accumulation_steps  # DIVIDE by accumulation steps
    loss.backward()  # accumulate gradients
    
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

Common mistake: Forgetting to divide the loss by accumulation_steps. Without this, the effective learning rate is scaled by the number of accumulation steps.

---

Evaluation gotchas

BLEU score implementations

Different BLEU implementations give different numbers:

• sacrebleu (recommended, reproducible): uses standard tokenization
• nltk.translate.bleu_score: requires pre-tokenized input
• torchtext.data.metrics.bleu_score: deprecated
• Papers before 2020 may use any of these

Difference can be 1-2 BLEU points — significant for reporting.

FID (Fréchet Inception Distance)

• Must use the SAME Inception v3 model weights (the default TF weights)
• PyTorch and TensorFlow Inception models give slightly different features
• Number of samples used for FID computation matters (more ≈ more stable)
• clean-fid package is the recommended implementation

Top-k accuracy

def topk_accuracy(output: torch.Tensor, target: torch.Tensor, k: int = 5) -> float:
    """Compute top-k accuracy.
    output: (batch, num_classes) — logits
    target: (batch,) — class indices
    """
    _, pred = output.topk(k, dim=-1)  # (batch, k)
    correct = pred.eq(target.unsqueeze(-1)).any(dim=-1)  # (batch,)
    return correct.float().mean().item()

---

Miscellaneous gotchas

model.eval() vs model.train()

model.eval()   # Disables dropout and uses running stats for batch norm
model.train()  # Enables dropout and computes batch stats for batch norm

# CRITICAL: Always call model.eval() before evaluation/inference
# Forgetting this is a very common source of poor eval performance

torch.no_grad() vs torch.inference_mode()

# For evaluation (no gradient computation):
with torch.no_grad():
    output = model(input)
    
# For inference (even faster, more restricted):
with torch.inference_mode():
    output = model(input)

Determinism

If the paper reports "we average over 3 random seeds":

def set_seed(seed: int):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

Note: cudnn.deterministic = True can significantly slow down training.

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

# 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]:

# [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

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

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)

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:

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)

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)

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

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.

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.

Knowledge: Transformer Components

Purpose

Canonical correct implementations of transformer building blocks that papers frequently reference but don't re-explain. When a paper says "standard transformer encoder," this file tells you what that means and what mistakes to avoid.

---

Multi-Head Attention

Canonical implementation

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.0,
                 bias: bool = True):
        super().__init__()
        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads  # head dimension

        self.W_q = nn.Linear(d_model, d_model, bias=bias)
        self.W_k = nn.Linear(d_model, d_model, bias=bias)
        self.W_v = nn.Linear(d_model, d_model, bias=bias)
        self.W_o = nn.Linear(d_model, d_model, bias=bias)
        self.dropout = nn.Dropout(dropout)

    def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,
                mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        batch_size = query.size(0)

        # Project and reshape: (batch, seq, d_model) -> (batch, n_heads, seq, d_k)
        q = self.W_q(query).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        k = self.W_k(key).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        v = self.W_v(value).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)

        # Scaled dot-product attention
        # (batch, n_heads, seq_q, d_k) @ (batch, n_heads, d_k, seq_k) -> (batch, n_heads, seq_q, seq_k)
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)

        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))

        attn_weights = F.softmax(scores, dim=-1)  # (batch, n_heads, seq_q, seq_k)
        attn_weights = self.dropout(attn_weights)

        # (batch, n_heads, seq_q, seq_k) @ (batch, n_heads, seq_k, d_k) -> (batch, n_heads, seq_q, d_k)
        context = torch.matmul(attn_weights, v)

        # Reshape back: (batch, n_heads, seq_q, d_k) -> (batch, seq_q, d_model)
        context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
        return self.W_o(context)  # (batch, seq_q, d_model)

Common mistakes

1. Scaling by sqrt(d_model) instead of sqrt(d_k)

• The scale factor is sqrt(d_k) where d_k = d_model / n_heads
• NOT sqrt(d_model). This is the single most common mistake.
• Vaswani et al. §3.2.1: "We suspect that for large values of d_k, the dot products grow large in magnitude"

2. Wrong mask convention

• Additive mask: add a large negative number (e.g., -1e9 or -inf) to scores BEFORE softmax
• Multiplicative mask: multiply attention weights by 0/1 AFTER softmax
• Both are valid, but papers rarely specify which. Additive is more common and numerically cleaner.
• Using -inf is cleaner than -1e9 (avoids non-zero attention for -1e9 with float16)

3. Forgetting `.contiguous()` after transpose

• After transpose(1, 2), the tensor may not be contiguous
• .view() requires a contiguous tensor
• This will crash, not silently fail — but it's a common "why doesn't my code run" bug

4. Causal mask shape

• Should be (1, 1, seq_len, seq_len) for broadcasting with (batch, n_heads, seq_len, seq_len) scores
• Mask where mask[i][j] = 1 if position j is allowed for position i
• Upper triangular = disallowed, not lower triangular (common mistake)

Causal masking

def create_causal_mask(seq_len: int, device: torch.device) -> torch.Tensor:
    """Creates a causal (autoregressive) attention mask.
    Returns a (1, 1, seq_len, seq_len) boolean tensor where True = attend, False = mask.
    """
    mask = torch.tril(torch.ones(seq_len, seq_len, device=device, dtype=torch.bool))
    return mask.unsqueeze(0).unsqueeze(0)  # (1, 1, seq_len, seq_len)

---

Positional Encodings

Sinusoidal (Vaswani et al., 2017)

class SinusoidalPositionalEncoding(nn.Module):
    """Fixed sinusoidal positional encoding from 'Attention Is All You Need'.

    PE(pos, 2i) = sin(pos / 10000^(2i/d_model))
    PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
    """
    def __init__(self, d_model: int, max_len: int = 5000, dropout: float = 0.0):
        super().__init__()
        self.dropout = nn.Dropout(dropout)

        pe = torch.zeros(max_len, d_model)  # (max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1).float()  # (max_len, 1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
        )  # (d_model/2,)

        pe[:, 0::2] = torch.sin(position * div_term)  # even indices
        pe[:, 1::2] = torch.cos(position * div_term)  # odd indices
        pe = pe.unsqueeze(0)  # (1, max_len, d_model)
        self.register_buffer('pe', pe)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (batch, seq_len, d_model)
        x = x + self.pe[:, :x.size(1)]
        return self.dropout(x)

Common mistakes with sinusoidal PE:

• Using arange(0, d_model) instead of arange(0, d_model, 2) for div_term
• Off-by-one in position indexing (should start at 0)
• Forgetting to register_buffer (so it's not a parameter but moves with the model to GPU)

Learned positional embeddings

class LearnedPositionalEmbedding(nn.Module):
    def __init__(self, max_len: int, d_model: int):
        super().__init__()
        self.embedding = nn.Embedding(max_len, d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (batch, seq_len, d_model)
        positions = torch.arange(x.size(1), device=x.device)  # (seq_len,)
        return x + self.embedding(positions)  # broadcast over batch

Rotary Position Embedding (RoPE) — Su et al., 2021

class RotaryPositionalEmbedding(nn.Module):
    """RoPE: Enhanced Transformer with Rotary Position Embedding.

    Applied to each head individually within the attention computation,
    AFTER the Q and K projections but BEFORE the dot product.
    """
    def __init__(self, d_head: int, max_len: int = 8192, base: float = 10000.0):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, d_head, 2).float() / d_head))
        self.register_buffer('inv_freq', inv_freq)
        self.max_len = max_len

    def forward(self, x: torch.Tensor, seq_len: int) -> Tuple[torch.Tensor, torch.Tensor]:
        t = torch.arange(seq_len, device=x.device).float()
        freqs = torch.outer(t, self.inv_freq)  # (seq_len, d_head/2)
        emb = torch.cat([freqs, freqs], dim=-1)  # (seq_len, d_head)
        return emb.cos(), emb.sin()


def apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """Apply RoPE to query or key tensor.
    x: (batch, n_heads, seq_len, d_head)
    """
    d_half = x.shape[-1] // 2
    x1, x2 = x[..., :d_half], x[..., d_half:]
    return torch.cat([
        x1 * cos[..., :d_half] - x2 * sin[..., :d_half],
        x2 * cos[..., d_half:] + x1 * sin[..., d_half:]
    ], dim=-1)

Key difference: RoPE is applied to Q and K individually, NOT summed onto embeddings like sinusoidal PE.

ALiBi (Press et al., 2022)

ALiBi doesn't use positional embeddings at all. Instead, it adds a linear bias to the attention scores:

def get_alibi_slopes(n_heads: int) -> torch.Tensor:
    """Compute ALiBi slopes for each head.
    Head i gets slope 2^(-8i/n_heads) for i = 1, ..., n_heads
    """
    ratio = 2 ** (-8 / n_heads)
    slopes = torch.tensor([ratio ** i for i in range(1, n_heads + 1)])
    return slopes  # (n_heads,)

def apply_alibi(scores: torch.Tensor, slopes: torch.Tensor) -> torch.Tensor:
    """Apply ALiBi bias to attention scores.
    scores: (batch, n_heads, seq_q, seq_k)
    slopes: (n_heads,)
    """
    seq_q, seq_k = scores.size(-2), scores.size(-1)
    # Position difference: relative distance between query and key positions
    positions = torch.arange(seq_k, device=scores.device).unsqueeze(0) - \
                torch.arange(seq_q, device=scores.device).unsqueeze(1)  # (seq_q, seq_k)
    bias = slopes.unsqueeze(-1).unsqueeze(-1) * positions.unsqueeze(0)  # (n_heads, seq_q, seq_k)
    return scores + bias.unsqueeze(0)  # broadcast over batch

---

Layer Normalization

Pre-norm vs Post-norm — THIS MATTERS ENORMOUSLY

Post-norm (original Transformer):

# Post-norm: normalize AFTER the residual addition
x = self.norm(x + self.sublayer(x))

Pre-norm (GPT-2, most modern transformers):

# Pre-norm: normalize BEFORE the sublayer, residual OUTSIDE the norm
x = x + self.sublayer(self.norm(x))

Why it matters:

• Post-norm requires learning rate warmup and careful initialization
• Pre-norm is much more stable to train at scale
• They produce different quality models — not interchangeable
• Many papers show post-norm in figures but use pre-norm in experiments — ALWAYS CHECK

RMSNorm (Zhang & Sennrich, 2019)

class RMSNorm(nn.Module):
    """Root Mean Square Layer Normalization.
    Used in LLaMA, T5. Simpler than LayerNorm (no centering, no bias).
    """
    def __init__(self, d_model: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(d_model))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps)
        return x / rms * self.weight

---

Feed-Forward Network

Standard (Vaswani et al.)

class FeedForward(nn.Module):
    """Two-layer feed-forward network with expansion factor.
    FFN(x) = W_2 * activation(W_1 * x + b_1) + b_2
    """
    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.0,
                 activation: str = "relu"):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

        if activation == "relu":
            self.activation = nn.ReLU()
        elif activation == "gelu":
            self.activation = nn.GELU()
        elif activation == "silu":
            self.activation = nn.SiLU()
        else:
            raise ValueError(f"Unknown activation: {activation}")

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # (batch, seq, d_model) -> (batch, seq, d_ff) -> (batch, seq, d_model)
        return self.linear2(self.dropout(self.activation(self.linear1(x))))

SwiGLU (Shazeer, 2020) — used in LLaMA, PaLM

class SwiGLU(nn.Module):
    """Gated feed-forward with SiLU activation.
    SwiGLU(x) = (SiLU(W_1 * x) ⊙ W_3 * x) * W_2
    Note: uses 3 weight matrices, not 2. This changes parameter count.
    """
    def __init__(self, d_model: int, d_ff: int):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff, bias=False)
        self.w2 = nn.Linear(d_ff, d_model, bias=False)
        self.w3 = nn.Linear(d_model, d_ff, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

---

Embedding with Weight Tying

class TransformerEmbedding(nn.Module):
    """Token + positional embedding with optional weight tying to output projection.

    Weight tying (Press & Wolf, 2017): The embedding matrix and the output
    projection matrix are the SAME tensor. This reduces parameters and often
    improves performance. Many papers do this without mentioning it explicitly.
    """
    def __init__(self, vocab_size: int, d_model: int, max_len: int,
                 dropout: float = 0.0, scale: bool = True):
        super().__init__()
        self.token_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb = LearnedPositionalEmbedding(max_len, d_model)
        self.dropout = nn.Dropout(dropout)
        self.scale = math.sqrt(d_model) if scale else 1.0
        # Vaswani et al. §3.4: "we multiply those weights by sqrt(d_model)"

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (batch, seq_len) of token IDs
        tok = self.token_emb(x) * self.scale  # (batch, seq_len, d_model)
        return self.dropout(self.pos_emb(tok))

Weight tying note: If the paper ties embedding and output weights, the output projection is F.linear(x, model.embedding.token_emb.weight) — not a separate nn.Linear. Many papers do this without stating it. Check the parameter count in the paper against your model — if yours is higher, weight tying might be missing.

---

Complete Transformer Block

Post-norm variant (original)

class TransformerBlockPostNorm(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, dropout):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.ff = FeedForward(d_model, d_ff, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        x = self.norm1(x + self.dropout1(self.attn(x, x, x, mask)))
        x = self.norm2(x + self.dropout2(self.ff(x)))
        return x

Pre-norm variant (modern standard)

class TransformerBlockPreNorm(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, dropout):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.ff = FeedForward(d_model, d_ff, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        x = x + self.dropout1(self.attn(self.norm1(x), self.norm1(x), self.norm1(x), mask))
        x = x + self.dropout2(self.ff(self.norm2(x)))
        return x

Key difference: In pre-norm, LayerNorm is applied BEFORE each sublayer. The residual connection adds the UN-normalized input. This is more stable for training deep models.

Stage 1: Paper Acquisition and Parsing

Purpose

Fetch the arxiv paper, extract its full text with mathematical notation preserved, and produce a structured markdown representation that downstream stages can consume section by section.

Input

• ARXIV_ID: e.g., 2106.09685 or 2106.09685v2

Output

• .paper2code_work/{ARXIV_ID}/paper_text.md — full paper text in markdown
• .paper2code_work/{ARXIV_ID}/paper_metadata.json — title, authors, abstract, categories
• .paper2code_work/{ARXIV_ID}/sections/ — individual section files
• .paper2code_work/{ARXIV_ID}/algorithms/ — extracted algorithm boxes
• .paper2code_work/{ARXIV_ID}/equations/ — extracted numbered equations
• .paper2code_work/{ARXIV_ID}/tables/ — extracted tables (especially hyperparameter tables)
• .paper2code_work/{ARXIV_ID}/footnotes.md — all footnotes collected

---

Reasoning protocol

Step 1: Normalize the input

Ask yourself:

• Is this a full URL or a bare ID?
• Does it have a version suffix (v1, v2)?
• Is the ID format valid? (should be YYMM.NNNNN or older format like arch-ive/NNNNNNN)

Strip to just the ID. Keep version suffix if present.

Step 2: Fetch the paper

Run scripts/fetch_paper.py. The script handles: 1. Downloading from https://arxiv.org/pdf/{id}.pdf 2. Extraction via pymupdf4llm (preferred — preserves math notation as LaTeX) 3. Fallback to pdfplumber if pymupdf4llm fails 4. Fallback to HTML from https://ar5iv.labs.arxiv.org/html/{id} if PDF extraction produces garbled output

How to detect garbled output: After extraction, scan the first 500 characters. If more than 20% are non-ASCII, non-LaTeX special characters, or if the text has no recognizable English words, the extraction is garbled. Fall back.

Step 3: Verify extraction quality

Read the extracted paper_text.md and check:

• [ ] Can you identify the paper title?
• [ ] Is the abstract present and readable?
• [ ] Are section headings identifiable?
• [ ] Are equations present (even if in LaTeX notation)?
• [ ] Is the references section present at the end?

If any of these fail, attempt the ar5iv HTML fallback. If that also fails, inform the user that automatic extraction failed and ask them to paste the paper text directly.

Step 4: Run structure extraction

Run scripts/extract_structure.py on the extracted text. This script:

• Identifies section boundaries using heading patterns (#, numbered headings, ALL CAPS headings)
• Extracts algorithm boxes (text between "Algorithm N" and the next section)
• Extracts numbered equations
• Extracts tables (especially those containing hyperparameters, learning rates, dimensions)
• Extracts footnotes

Step 5: Verify all critical sections exist

You MUST find these sections (they may be named differently):

• Abstract
• Introduction (or "1 Introduction" or similar)
• Method/Model/Approach section (this is the core — it may be named anything)
• Experiments/Results
• Conclusion

You MUST actively look for these (authors hide crucial details here):

• Appendix — check for content AFTER the references. Many papers have appendices with implementation details, hyperparameter tables, ablation studies, prompts, and proofs that are essential for reproduction
• Supplementary material references — if the paper mentions "see supplementary" or "see appendix," note what is referenced
• Footnotes — often contain critical caveats about implementation choices

Step 6: Special handling for appendices

This deserves its own step because it's that important:

1. After the References section, look for any additional content (Appendix A, B, C, etc.) 2. If appendices exist, extract them as separate section files 3. Pay special attention to:

• Hyperparameter tables (often in Appendix A or B)
• System Prompts
• Architecture diagrams described in text
• Training details (often Appendix C or D)
• Ablation studies (contain information about what matters and what doesn't)

4. If the paper references a supplementary PDF, note this — you may need to fetch it separately from the arxiv page

Step 7: Extract metadata

From the paper text or arxiv page, extract:

• Title
• Authors
• Year
• Arxiv categories (e.g., cs.LG, cs.CV)
• Abstract (first 500 words)

Save to paper_metadata.json.

Step 8: Search for official code repositories

The fetch_paper.py script automatically searches for official code in two places:

1. Inside the paper text — scans for GitHub/GitLab/Bitbucket URLs and phrases like "code available at," "our implementation is released at," etc. 2. The arxiv abstract page — checks for code repository links in the page HTML.

Results are saved to paper_metadata.json under the official_code key. Each entry has:

• url — the repository URL
• source — where it was found (paper_text or arxiv_page)
• context — surrounding text that confirms it's the authors' code

After the script runs, verify the links:

• Open the repository URL. Is it actually the authors' official code for THIS paper, or an unrelated repo?
• Does the repo contain a working implementation? Some repos are empty placeholders or "coming soon."
• Note the primary language/framework — this may inform your implementation choices.

If official code is found, it becomes a critical resource for Stage 3 (Ambiguity Audit). Every [UNSPECIFIED] item should be checked against the official code before choosing a default. Choices resolved this way get the [FROM_OFFICIAL_CODE] tag instead.

---

Fallback protocol

If PDF download fails (403, 404, network error):

1. Try the ar5iv HTML version: https://ar5iv.labs.arxiv.org/html/{id} 2. If that also fails, try the abstract page: https://arxiv.org/abs/{id} to verify the paper exists 3. If the paper exists but can't be fetched, ask the user to download it manually and provide the path

If extraction produces garbled text:

1. Try pdfplumber instead of pymupdf4llm 2. If still garbled, fetch ar5iv HTML 3. If HTML is also bad, try extracting just the text without math preservation 4. Last resort: ask the user to paste the paper text

If the paper is very long (>50 pages):

1. Still extract everything — don't truncate 2. The section-level files in sections/ allow reading parts individually 3. Focus on Method and Appendix sections for implementation details

---

Quality checklist before proceeding to Stage 2

• [ ] paper_text.md exists and is readable
• [ ] paper_metadata.json has title and authors
• [ ] At least one section file exists in sections/
• [ ] You've checked for appendices
• [ ] You've checked for algorithm boxes
• [ ] Equations are present (even if in LaTeX form)
• [ ] The Method/Model section is identified and readable
• [ ] You've checked for official code repositories (results in paper_metadata.json under official_code)

If the Method section is garbled but other sections are fine, attempt to re-extract just that section from the HTML version. The Method section is the most critical — you cannot proceed without a readable version of it.

Stage 2: Contribution Identification

Purpose

Before writing any code, you must understand exactly what this paper contributes and what "implementing this paper" means. This stage forces clarity.

Input

• Parsed paper sections from Stage 1
• paper_metadata.json (includes official_code links if any were found)

Output

• .paper2code_work/{ARXIV_ID}/contribution.md — contains the contribution statement, paper classification, and implementation scope

---

Reasoning protocol

Step 1: Read the abstract and introduction

First read: understand what the paper claims to do. Second read: identify the single sentence that states the core contribution. This is usually in the abstract ("We propose...," "We introduce...," "In this work, we...") or the last paragraph of the introduction ("Our contributions are...").

Tasks:

• What is the ONE thing this paper does that didn't exist before?
• If I had to explain this paper in 10 seconds to an ML engineer, what would I say?
• What would be missing from ML if this paper didn't exist?

Write this down as a one-sentence contribution statement.

Step 2: Classify the paper type

Determine which category best fits. This directly affects what you implement:

(a) New architecture

Examples: Transformer, ResNet, Vision Transformer, U-Net What "implementing the core contribution" means: The model architecture — layer definitions, forward pass, how components connect. The architecture IS the contribution. Priority sections: Model/Architecture section, Figure showing the architecture, any Algorithm boxes What's secondary: Training procedure (use standard), specific dataset preprocessing

(b) New training method or loss function

Examples: DDPM, SimCLR, RLHF, DPO What "implementing the core contribution" means: The training loop, the loss function, the optimization procedure. The model architecture is usually borrowed from existing work. Priority sections: Training/Objective section, loss function equations, Algorithm boxes describing the training procedure What's secondary: Model architecture (often standard), dataset specifics

(c) New inference or generation technique

Examples: Beam search variants, nucleus sampling, classifier-free guidance What "implementing the core contribution" means: The inference algorithm. There may be no training at all. Priority sections: Inference/Generation section, Algorithm boxes, any sampling or decoding procedures What's secondary: Training (none or standard), model architecture (usually pre-existing)

(d) New dataset or benchmark

Examples: ImageNet, GLUE, MMLU What "implementing the core contribution" means: The evaluation framework, metrics, data loading pipeline. There may be baseline models but they're not the contribution. Priority sections: Dataset description, evaluation methodology, metric definitions What's secondary: Baseline model architectures (reference only)

(e) Theoretical analysis with empirical validation

Examples: Lottery Ticket Hypothesis, Scaling Laws What "implementing the core contribution" means: The experimental setup that validates the theory. Often involves running specific experiments with specific measurements. Priority sections: Experimental setup, what is measured and how, specific procedures What's secondary: The theory proofs (important to understand but not to implement)

(f) System or engineering paper

Examples: Megatron-LM, FlashAttention, vLLM What "implementing the core contribution" means: The system design — often involves low-level optimizations, custom kernels, or infrastructure. Priority sections: System design section, performance benchmarks, implementation details Note: These papers often cannot be reproduced in a minimal implementation because the contribution IS the engineering. Flag this honestly.

Step 3: Find the Algorithm box

Search the paper for formal algorithm descriptions (usually labeled "Algorithm 1," "Algorithm 2," etc.). These are gold:

• They are the most precise specification of the procedure
• They resolve ambiguities in prose descriptions
• They define variable names and control flow explicitly
• If the prose says one thing and the Algorithm box says another, implement the Algorithm box and flag the discrepancy

If there is no Algorithm box, note this — it means the implementation will rely more on equations and prose, which are more ambiguous.

Step 4: Check for official code

Check paper_metadata.json for the official_code field. If official code repositories were found in Stage 1:

1. Verify the link — open the URL. Confirm it's the authors' code for this paper, not an unrelated project or empty placeholder. 2. Note the framework — is it PyTorch, JAX, TensorFlow? This may inform your implementation. 3. Note the repo structure — what files exist? This helps you understand the implementation scope. 4. Do NOT read the code in detail yet — that happens in Stage 3 when resolving ambiguities. At this stage, you just need to know it exists and what it covers.

Record the official code status in the contribution statement (Step 5).

Step 5: Identify what to implement vs. reference

Make three lists:

IMPLEMENT (we will write this code):

• The core contribution (as identified in Step 2)
• Any component described in enough detail to implement from the paper alone

REFERENCE (we will import or note as dependency):

• Standard components the paper uses but didn't invent (e.g., "standard transformer encoder")
• Pre-trained models used as backbones
• Existing techniques referenced with "following [X]" or "similar to [X]"

OUT OF SCOPE (we will not touch):

• Baselines and comparison methods
• Ablation variants (unless the user specifically requests full mode)
• Dataset collection or annotation procedures
• Deployment-specific optimizations

Step 6: Write the contribution statement

Write a structured contribution statement and save it to contribution.md:

# Contribution Analysis

## Paper
{title} ({year})
{authors}
arxiv: {id}

## One-sentence summary
{what this paper does in one sentence}

## Paper type
{(a)-(f) classification with justification}

## Core contribution to implement
{One paragraph describing exactly what we will implement, referencing specific paper sections}

## Algorithm specification
{Whether an Algorithm box exists, and if so, which one is the primary specification}

## Official code
{URL if found, or "None found"}
{Framework used in official code, if applicable}
{Brief note on what the official repo covers}

## Implementation scope
### Will implement:
- {item 1 — with paper section reference}
- {item 2 — with paper section reference}

### Will reference (not reimplement):
- {item 1 — what it is, where to get it}

### Out of scope:
- {item 1 — why}

## Key sections for implementation
1. {Section X.Y — what it specifies}
2. {Section X.Y — what it specifies}
3. {Appendix A — what it specifies}

Step 7: Self-check

Before proceeding:

• Could a competent ML engineer read my contribution statement and know exactly what to build? If not, it's too vague.
• Have I identified the paper type correctly? Read the abstract one more time.
• Have I checked the appendices for additional specification of the core contribution?
• Is there anything in my "IMPLEMENT" list that the paper doesn't actually describe in enough detail? Move it to "REFERENCE" or flag it.

---

Common mistakes at this stage

1. Trying to implement everything. A paper about a new loss function doesn't need you to reimplement the backbone architecture. Import it.

2. Misidentifying the contribution. A paper might describe a new architecture but the actual contribution is the training recipe. Read what the authors claim, not what takes up the most pages.

3. Ignoring paper type effects on scope. The same checklist item (e.g., "training loop") can be in-scope or out-of-scope depending on the paper type. A training method paper MUST include the training loop. An architecture paper doesn't need more than a minimal example.

4. Skipping the Algorithm box. If Algorithm 1 exists and you don't read it carefully, you will implement based on noisy prose instead of the precise specification.

5. Ignoring official code. Stage 1 automatically searches for code links, but you must verify them in Step 4. If official code exists and you skip it, you'll waste time on ambiguities that the authors already resolved.

Stage 3: The Ambiguity Audit

Purpose

This is the most important anti-hallucination stage. Before generating a single line of code, systematically audit every implementation-relevant detail and classify it as SPECIFIED, PARTIALLY_SPECIFIED, or UNSPECIFIED. This stage produces the raw material for REPRODUCTION_NOTES.md.

Input

• Parsed paper sections from Stage 1
• Contribution analysis from Stage 2 (includes official code status)
• Official code repository (if found in Stage 1 — URL in paper_metadata.json under official_code)

Output

• .paper2code_work/{ARXIV_ID}/ambiguity_audit.md — complete audit with classification for every item

---

Critical mindset

Assume nothing. Verify everything.

When you read "we use standard settings," ask: standard according to whom? When you read "following prior work," ask: which prior work, and which specific detail from that work? When something seems obvious, ask: is it obvious because it's specified, or because you're making an assumption?

The goal of this stage is to produce a document where a researcher can look at any implementation choice and immediately know whether the paper specified it or whether you chose it.

---

Reasoning protocol

The complete checklist

Go through every item below. For each one, search the paper (method section, experiments section, appendix, footnotes — all of them) and classify:

• SPECIFIED: The paper states this explicitly. Record the exact quote and section.
• PARTIALLY_SPECIFIED: The paper mentions this but is ambiguous. Record the quote and what's unclear.
• UNSPECIFIED: The paper does not state this at all. Record common choices and which one you'll use as default.

---

Architecture details

Training details

Data details

Evaluation details

---

How to search the paper

For each item in the checklist:

1. Search the Method section — the obvious place 2. Search the Experiments section — often contains "Implementation Details" or "Training Details" subsection 3. Search the Appendix — THE most common hiding spot for training details 4. Search figure captions — sometimes specify architectural details not in text 5. Search footnotes — authors sometimes put caveats and clarifications in footnotes 6. Search the table captions — hyperparameter tables often have footnotes with additional info 7. Search for the specific term — Ctrl+F the paper text for "learning rate," "dropout," "optimizer," etc.

---

How to handle "following X" references

When the paper says "following [X]" or "similar to [X]" or "as in [X]":

1. Note the reference — what is paper X? 2. If X is a well-known paper (e.g., "following Vaswani et al."), describe what X specifies 3. If X is an obscure paper, flag this as requiring the reader to look up that paper 4. Record: [PARTIALLY_SPECIFIED] Paper says "following [X]" for {component} — X specifies {detail} but reader should verify

---

How to resolve ambiguities using official code

If official code was found in Stage 1, use it as a primary resource for resolving UNSPECIFIED and PARTIALLY_SPECIFIED items. This is the single most effective way to reduce ambiguity.

Process: 1. For each UNSPECIFIED item, search the official repo for the relevant implementation detail (e.g., grep for "eps", "dropout", "lr" in config files and model code). 2. If you find the answer, change the classification:

• UNSPECIFIED → SPECIFIED with tag [FROM_OFFICIAL_CODE]
• PARTIALLY_SPECIFIED → SPECIFIED with tag [FROM_OFFICIAL_CODE]

3. Record the exact file and line from the official repo: github.com/author/repo/blob/main/model.py#L42

What official code can resolve:

• Activation functions, normalization epsilon, dropout placement
• Initialization schemes (often the biggest silent assumption)
• Learning rate schedule details, optimizer parameters
• Data preprocessing steps
• Any "standard settings" or "following prior work" references

What official code cannot resolve:

• Whether the code matches the paper's intent (bugs exist in official code too)
• Whether a detail was chosen for the paper's experiments vs. for engineering convenience

Important: Official code is a reference, not ground truth. If the official code contradicts the paper, note both and flag the discrepancy. See the errata section below.

---

How to handle errata and contradictions

1. Figure vs text disagreement: Flag it explicitly. Implement what the text/equations say (they're usually more carefully reviewed than figures) but note the figure shows something different.

2. Abstract vs method section disagreement: Method section wins. The abstract is a summary and may oversimplify.

3. Equation vs prose disagreement: Equation wins. It's more precise. Flag the discrepancy.

4. Paper vs official code disagreement: Note both. The official code may fix bugs or reflect a later understanding, but the paper is what was peer-reviewed. Record: Paper says X in §Y.Z, but official code does W — we implement {your choice} because {reason}

---

Output format

Save to .paper2code_work/{ARXIV_ID}/ambiguity_audit.md:

# Ambiguity Audit: {paper_title}

## Architecture

| Component | Status | Paper Quote | Section | Our Choice | Alternatives |
|-----------|--------|-------------|---------|------------|-------------- |
| Hidden dim | SPECIFIED | "d_model = 512" | §3.1 | 512 | — |
| Activation | UNSPECIFIED | — | — | GELU | ReLU, SiLU |
| Pre/post norm | PARTIALLY_SPECIFIED | "we use layer norm" | §3.2 | Pre-norm | Post-norm (figure suggests post-norm) |

## Training
{same table format}

## Data
{same table format}

## Evaluation
{same table format}

## Contradictions found
- {description of any inconsistencies}

## References to check
- {any "following X" references that need verification}

## Official code
- URL: {if found}
- Used to resolve: {list of items resolved by reading official code}

---

Self-check before proceeding

• [ ] Every item in the checklist has been classified (not just the obvious ones)
• [ ] Every SPECIFIED item has an exact quote and section reference
• [ ] Every UNSPECIFIED item has a default choice AND alternatives listed
• [ ] You checked the appendix (not just the method section)
• [ ] You checked footnotes
• [ ] You looked for an official code repository
• [ ] Any "following [X]" references are flagged
• [ ] Contradictions between text/figures/equations are noted

Stage 4: Code Generation

Purpose

Generate the implementation repository. Every line of code must be anchored to the paper. Every unspecified choice must be flagged. The output should look like it was written by someone who read the paper three times and implemented it carefully.

Input

• Contribution analysis from Stage 2
• Ambiguity audit from Stage 3
• Paper sections from Stage 1
• Official code repository (if found — URL in paper_metadata.json under official_code)

Output

• {paper_slug}/ directory with all generated files

---

Pre-generation checklist

Before writing any code, verify:

• [ ] You have the contribution statement (what exactly are you implementing?)
• [ ] You have the ambiguity audit (what is specified vs. unspecified?)
• [ ] You know the paper type (architecture / training method / inference technique / etc.)
• [ ] You've read guardrails/scope_enforcement.md to know what's in/out of scope
• [ ] You've read the relevant knowledge/ files for domain-specific gotchas
• [ ] You've checked knowledge/paper_to_code_mistakes.md for relevant pitfalls
• [ ] You've checked if official code exists (in paper_metadata.json) and reviewed any [FROM_OFFICIAL_CODE] items in the ambiguity audit

---

File generation order

Generate files in this order (dependencies flow downward):

1. configs/base.yaml — all hyperparameters first, so code can reference them 2. src/utils.py — shared utilities (masking, positional encoding, helpers) 3. src/model.py — architecture 4. src/loss.py — loss functions 5. src/data.py — dataset and dataloader 6. src/train.py — training loop (if in scope) 7. src/evaluate.py — evaluation 8. requirements.txt — dependencies 9. REPRODUCTION_NOTES.md — from the ambiguity audit 10. README.md — project readme

---

Citation anchoring — mandatory for all code

Every class definition, every non-trivial function, every implementation choice must have a citation comment. Use this format consistently:

# §3.2 — "We apply layer normalization before each sub-layer"
# This is the Pre-LN variant (Ba et al., 2016)

# §3.2, Eq. 2 — attention_weights = softmax(QK^T / sqrt(d_k))

# §4.1, Table 2 — d_model = 512, h = 8, d_ff = 2048

# [UNSPECIFIED] Paper does not state epsilon for LayerNorm — using 1e-6
# Alternatives: 1e-5 (PyTorch default), 1e-8 (some older implementations)

# [PARTIALLY_SPECIFIED] "We use dropout for regularization" (§3.3)
# Rate of 0.1 stated in §4.1, but placement not specified — applying after attention and FFN

# [ASSUMPTION] Using pre-norm based on "we found pre-norm more stable" (§4.1)
# Figure 1 shows post-norm, creating ambiguity. We follow the text.

# [FROM_OFFICIAL_CODE] Using learned positional embeddings
# github.com/author/repo/blob/main/model.py#L42

The § symbol is non-negotiable. Use it consistently for section references.

---

Shape comments — mandatory for all tensor operations

Every tensor operation must have a comment showing the shape transformation:

x = self.embedding(input_ids)  # (batch, seq_len) -> (batch, seq_len, d_model)
q = self.W_q(x)  # (batch, seq_len, d_model) -> (batch, seq_len, d_model)
q = q.view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)  # (batch, n_heads, seq_len, d_k)
scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)  # (batch, n_heads, seq_len, seq_len)

---

configs/base.yaml

Every hyperparameter in one place. No magic numbers in code. Format:

# Configuration for {paper_title}
# Paper: {arxiv_url}
# All values cited to paper section or flagged as [UNSPECIFIED]

model:
  d_model: 512          # §3.1 — "d_model = 512"
  n_heads: 8            # §3.1 — "h = 8"
  n_layers: 6           # §3.1 — "N = 6"
  d_ff: 2048            # §3.1 — "d_ff = 2048"
  dropout: 0.1          # §4.1 — "P_drop = 0.1"
  activation: "relu"    # [UNSPECIFIED] — Paper does not specify activation in FFN
  norm_eps: 1.0e-6      # [UNSPECIFIED] — LayerNorm epsilon not stated

training:
  optimizer: "adam"      # §4.1 — "We use the Adam optimizer"
  lr: 0.0001            # §4.1 — from the LR schedule formula
  betas: [0.9, 0.98]    # §4.1 — "β1 = 0.9, β2 = 0.98"
  eps: 1.0e-9           # §4.1 — "ε = 10^-9"
  warmup_steps: 4000    # §4.1 — "warmup_steps = 4000"
  # ... etc

---

src/model.py

Structure

"""
{Paper title} — Model Architecture

Paper: {arxiv_url}
Implements: {one-line description of what this file implements}

Section references:
  §X.Y — {what it describes}
  §X.Z — {what it describes}
"""

import math
from dataclasses import dataclass
from typing import Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F


@dataclass
class ModelConfig:
    """Configuration for {model name}.
    
    All defaults from {paper reference} unless marked [UNSPECIFIED].
    """
    # §X.Y — "description from paper"
    param_name: type = default_value


class ComponentA(nn.Module):
    """§X.Y — Implements {component description from paper}.
    
    "{Exact quote from paper describing this component}"
    """
    # ...


class ComponentB(nn.Module):
    """§X.Z — Implements {component description from paper}."""
    # ...


class MainModel(nn.Module):
    """§X — {Paper's name for the model}.
    
    Composed of:
      - ComponentA (§X.Y)
      - ComponentB (§X.Z)
    """
    
    def __init__(self, config: ModelConfig):
        super().__init__()
        # Build components
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass following §X.Y description.
        
        Args:
            x: {description} — shape: (batch, ...)
            
        Returns:
            {description} — shape: (batch, ...)
        """
        # Implementation with shape comments on every operation

Rules for model.py

1. Match paper notation in variable names. If the paper uses Q, K, V — use q, k, v. If it uses h for heads — use n_heads (slightly more readable but add comment # h in paper notation).

2. One class per conceptual component. If the paper describes attention, feed-forward, and the overall block as separate concepts — make them separate classes. Do not put everything in one monolithic class.

3. Config dataclass, not scattered parameters. All hyperparameters come from a config dataclass. No 512 literals anywhere.

4. Forward method mirrors paper flow. The operations in forward() should follow the same order as the paper describes them. A reader should be able to follow along in the paper.

5. No dead code. Don't add methods "for convenience" that aren't needed. No save/load, no from_pretrained, no num_parameters — unless the paper contribution involves them.

---

src/loss.py

Rules

1. Implement the exact equation from the paper. Reference the equation number. 2. Note numerical stability considerations. Log-sum-exp instead of naive softmax, epsilon in denominators, etc. 3. If the paper uses a standard loss with modifications, implement only the modifications and wrap the standard loss. 4. Comment where the paper's definition differs from PyTorch's built-in. This is a common source of bugs.

---

src/data.py

Rules

1. Provide a Dataset class skeleton, not a complete implementation. The class should show:

• What the expected data format is
• What preprocessing is needed
• What __getitem__ returns (shapes and types)

2. Include clear TODOs for dataset-specific logic (download path, preprocessing steps) 3. Do not hardcode dataset paths. Use config or arguments. 4. Do not download datasets automatically. Add instructions in docstring.

---

src/train.py

Include only if training is in scope (see contribution analysis)

1. If the contribution IS the training method (type b paper): This is a primary deliverable. The training loop must precisely follow the paper's algorithm. 2. If the contribution is architectural (type a paper): Provide a minimal training example that instantiates the model and runs one forward/backward pass. Not a full training pipeline. 3. Always include the optimizer setup with all hyperparameters from the config. 4. Always include the learning rate schedule if specified in the paper.

---

src/evaluate.py

1. Implement the exact metrics the paper reports in its main results table. 2. Note which metrics package/computation is used (e.g., sacrebleu vs nltk BLEU — they give different numbers). 3. This is metric computation code, not a full eval pipeline. No data loading, no model loading, just functions that take predictions and targets and return numbers.

---

src/utils.py

1. Only include utilities that are shared across multiple files (e.g., masking functions, positional encoding). 2. Do not create a "utils" dumping ground. If a utility is only used in one file, put it in that file.

---

requirements.txt

torch>=2.0.0
pyyaml>=6.0
numpy>=1.24.0
# Add paper-specific dependencies below

Pin major versions only. Add comments for why each dependency is needed.

---

REPRODUCTION_NOTES.md

Transform the ambiguity audit from Stage 3 into a researcher-facing document. Use the template from scaffolds/reproduction_notes_template.md. This is a first-class deliverable — not an afterthought. See the scaffold template for the exact structure.

---

README.md (for the generated project)

Use the template from scaffolds/readme_template.md. Should include:

• Paper title, authors, link
• What this implementation covers (from contribution statement)
• Quick-start: how to run the model
• File structure with one-line descriptions
• Citation

---

Code quality final check

Before declaring this stage complete:

• [ ] No magic numbers — every literal is from the config
• [ ] No missing citations — every class and non-trivial function has a § reference
• [ ] No missing shape comments — every tensor operation has a shape annotation
• [ ] No silent assumptions — every UNSPECIFIED item has a [UNSPECIFIED] comment
• [ ] No dead code — everything that exists is used
• [ ] No relative imports — all imports are absolute
• [ ] Type hints on all function signatures
• [ ] The model's forward pass mirrors the paper's description order
• [ ] The config file has a citation comment on every parameter
• [ ] REPRODUCTION_NOTES.md is complete and covers every ambiguity from Stage 3

---

What NOT to generate

Read this list before writing any code. If you find yourself about to write any of these, stop:

• Do NOT reimplement standard components the paper references but doesn't describe. Import or note the dependency.
• Do NOT implement distributed training, checkpointing, experiment tracking, or logging infrastructure (unless it IS the paper's contribution).
• Do NOT implement baseline methods or comparison approaches.
• Do NOT download or embed datasets.
• Do NOT write unit tests (this is a reproduction scaffold, not a production library).
• Do NOT add CLI argument parsing beyond loading a config file.
• Do NOT add visualization code unless visualization IS the contribution.
• Do NOT add if __name__ == "__main__" blocks with complex logic — keep scripts minimal.

Stage 5: Walkthrough Notebook

Purpose

Generate a Jupyter notebook that connects paper sections to code, making the implementation pedagogically useful. This is not for training — it's for understanding. A researcher should be able to read this notebook and understand exactly how every paper section maps to code.

Input

• All generated code from Stage 4
• Paper sections from Stage 1
• Contribution analysis from Stage 2

Output

• {paper_slug}/notebooks/walkthrough.ipynb

---

Notebook structure

The notebook follows a strict section pattern. For each major component of the implementation:

Cell pattern (repeat for each component)

Cell 1: Paper excerpt (Markdown)

A quoted block of the relevant paper passage that describes this component. Include the section number and, if applicable, the equation.

## Multi-Head Attention (§3.2)

> "An attention function can be described as mapping a query and a set of key-value
> pairs to an output, where the query, keys, values, and output are all vectors.
> The output is computed as a weighted sum of the values, where the weight assigned
> to each value is computed by a compatibility function of the query with the
> corresponding key."
>
> — §3.2, Attention Is All You Need

**Equation 1:**
$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

Cell 2: Code (Python)

The corresponding code from src/model.py or wherever this component is defined. Don't just import — inline the relevant class/function so the reader can see it without switching files.

# From src/model.py — Multi-Head Attention

class MultiHeadAttention(nn.Module):
    """§3.2 — Multi-Head Attention mechanism."""
    # ... full implementation

Cell 3: Sanity check (Python)

A runnable check that verifies the component works correctly with toy dimensions:

# Sanity check: verify output shapes match paper specification
config = ModelConfig(d_model=64, n_heads=4)  # Toy dimensions for CPU
attn = MultiHeadAttention(config)

batch, seq_len = 2, 10
x = torch.randn(batch, seq_len, config.d_model)

output = attn(x)
assert output.shape == (batch, seq_len, config.d_model), \
    f"Expected (2, 10, 64), got {output.shape}"
print(f"✓ MultiHeadAttention: input {x.shape} → output {output.shape}")

---

Required sections

1. Setup and imports

import torch
import torch.nn as nn
import math
# ... minimal imports

# Use tiny dimensions so everything runs on CPU quickly
BATCH_SIZE = 2
SEQ_LEN = 16
D_MODEL = 64  # Paper uses 512, we use 64 for walkthrough
# ... etc

Explain why dimensions are small: "We use reduced dimensions so the notebook runs instantly on CPU. The architecture is identical — only the numbers change."

2. Configuration

Show the config dataclass with paper citations on each field. Instantiate with toy dimensions.

3. One section per architectural component

Follow the paper's presentation order. If the paper describes components A, B, C used to build the model, the notebook should have sections for A, B, C in that order, then the composed model.

For each: paper excerpt → code → sanity check.

4. Full model assembly

Show how the components compose into the full model. Run a full forward pass with random input. Verify output shape.

# Full model forward pass
model = TheModel(config)
x = torch.randn(BATCH_SIZE, SEQ_LEN)  # or appropriate input
output = model(x)
print(f"✓ Full model: input {x.shape} → output {output.shape}")
print(f"  Parameters: {sum(p.numel() for p in model.parameters()):,}")

5. Loss function (if in scope)

Paper excerpt for the loss function → code → compute loss on random predictions.

6. Training step (if in scope)

Show one forward + backward pass. Verify gradients flow:

# One training step
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
loss = loss_fn(output, target)
loss.backward()
optimizer.step()

# Verify gradients exist
for name, param in model.named_parameters():
    if param.grad is None:
        print(f"⚠ No gradient for: {name}")
    else:
        print(f"✓ {name}: grad norm = {param.grad.norm():.4f}")

7. Common pitfalls (Markdown)

A cell listing what typically goes wrong when implementing this paper:

## Common Pitfalls

1. **Pre-norm vs post-norm**: The paper's figure shows post-norm but experiments
   use pre-norm. This significantly affects training stability. [§3.2]

2. **Attention scale factor**: Must divide by √d_k, not √d_model. If using
   multi-head attention, d_k = d_model / n_heads. [§3.2, Eq. 1]

3. **Positional encoding**: The sinusoidal encoding uses specific frequencies.
   Off-by-one errors in the position index are common. [§3.5]

4. **Label smoothing**: The paper uses ε=0.1 label smoothing. PyTorch's
   CrossEntropyLoss does not support this directly — you need a custom
   implementation. [§5.4]

5. **Learning rate schedule**: The warmup schedule has a specific formula.
   Using a generic linear warmup gives different results. [§5.3]

---

Notebook generation rules

1. Everything must run on CPU. No GPU required. Use tiny dimensions. 2. No external data needed. All inputs are random tensors. 3. No training to convergence. Just verify shapes, gradients, and basic numerics. 4. Each cell should be self-contained enough to understand. Imports can be at the top, but each section should be readable on its own. 5. Markdown cells should quote the paper directly. Don't paraphrase — use exact quotes with section references. 6. Sanity checks should be assertions, not just prints. Assert the expected shapes so failures are loud. 7. Use the paper's variable names in code. If the paper uses Q, K, V — the notebook code should too.

---

Mode-specific behavior

minimal mode

Standard notebook as described above.

educational mode

Add additional cells:

• Before each component: a "Background" markdown cell explaining the ML concept (what is attention? what is a residual connection?) for readers who are learning
• After each sanity check: a "What would happen if..." cell exploring edge cases (what if you remove the scale factor? what if you use post-norm instead of pre-norm?)
• A "Further Reading" section at the end pointing to tutorials and explanations

full mode

Same as minimal but includes:

• Full data loading example (with mock data)
• Complete training loop (3-5 steps, not to convergence)
• Evaluation metric computation on random predictions

---

Self-check

• [ ] Notebook runs end-to-end without errors (all cells execute)
• [ ] All assertion-based sanity checks pass
• [ ] Every paper section relevant to the implementation is quoted
• [ ] Toy dimensions are used throughout (runs on CPU in seconds)
• [ ] Common pitfalls section exists and is genuinely useful
• [ ] No external data or GPU required

{{PAPER_TITLE}}

Implementation of {{PAPER_TITLE}} ({{AUTHORS}}, {{YEAR}}).

What this implements

{{ONE_PARAGRAPH_CONTRIBUTION_STATEMENT}}

Quick start

pip install -r requirements.txt

from src.model import {{MODEL_CLASS}}, ModelConfig

config = ModelConfig()
model = {{MODEL_CLASS}}(config)

# Example forward pass
import torch
x = torch.randn({{EXAMPLE_INPUT_SHAPE}})
output = model(x)
print(output.shape)  # {{EXPECTED_OUTPUT_SHAPE}}

File structure

{{PAPER_SLUG}}/
├── README.md                 # This file
├── REPRODUCTION_NOTES.md     # Ambiguity audit — what's specified vs. assumed
├── requirements.txt          # Dependencies
├── src/
│   ├── model.py              # {{MODEL_DESCRIPTION}}
│   ├── loss.py               # {{LOSS_DESCRIPTION}}
│   ├── data.py               # {{DATA_DESCRIPTION}}
│   ├── train.py              # {{TRAIN_DESCRIPTION}}
│   ├── evaluate.py           # {{EVAL_DESCRIPTION}}
│   └── utils.py              # {{UTILS_DESCRIPTION}}
├── configs/
│   └── base.yaml             # All hyperparameters — each cited or flagged
└── notebooks/
    └── walkthrough.ipynb     # Paper sections → code → sanity checks

Important: Read REPRODUCTION_NOTES.md

This implementation flags every choice that the paper does not specify. Before using this code for research, read REPRODUCTION_NOTES.md to understand which implementation details are from the paper and which are our choices.

Citation

{{BIBTEX_ENTRY}}

---

Generated by [paper2code](https://github.com/PrathamLearnsToCode/paper2code) — citation-anchored paper implementation.