Karpathy's AutoResearch

Let an AI agent run ML experiments while you sleep

github.com/karpathy/autoresearch

47.2k stars · MIT License · Python

PART 1

ELI5

The simple version anyone can understand

The Big Idea

Imagine you have a science fair project due tomorrow.

The Old Way

  • You try one experiment
  • Wait for results
  • Think about what to try next
  • Try another experiment
  • Repeat... all day long

The AutoResearch Way

  • You give a robot your project
  • Tell it "try stuff and keep what works"
  • Go to sleep
  • Wake up to ~100 experiments done
  • The robot kept only the winners

Three Files, One Rule

The whole project is laughably simple:

program.md The instruction sheet for the robot
train.py The ONLY file the AI can change
prepare.py The test / scorecard (read-only)

The one rule: Each experiment gets exactly 5 minutes. No more. The AI tries to get the best score in that time.

What the AI Does All Night

1. Tweak the code

Change something in train.py — bigger model, different learning rate, new trick

2. Run for 5 minutes

Train the model, measure the score (val_bpb — lower is better)

3. Keep or discard

Score improved? Keep it. Score worse? Undo and try something else.

4. Never stop

Repeat forever until the human wakes up and says "enough"

Why This Matters

~100

experiments run overnight while you sleep

~12

experiments per hour, fully autonomous

The shift

Researchers used to babysit experiments all day. Now the AI does the tedious trial-and-error, and the human just reads the results in the morning.

Real results

The agent finds optimizations a human might take days or weeks to discover through manual experimentation.

PART 2

Technical Deep Dive

Architecture, code patterns, and implementation details

Repository Overview

FileRoleWho edits
train.pyModel + optimizer + training loopAI agent
prepare.pyData prep, tokenizer, eval harnessNobody (read-only)
program.mdAgent instructions & constraintsHuman
pyproject.tomlDependenciesNobody (locked)
analysis.ipynbResults analysis notebookHuman
results.tsvExperiment log (untracked)AI agent

Design Principle

Minimal surface area. One file to change, one metric to optimize, fixed time budget. The agent can't break the eval or install packages.

Python 83.5% · Jupyter 16.5%

Based on nanochat, Karpathy's single-GPU LLM training framework

The Experiment Loop

Defined in program.md — the agent's operating manual:

LOOP FOREVER:

  1. Look at git state (current branch/commit)
  2. Edit train.py with an experimental idea
  3. git commit
  4. Run: uv run train.py > run.log 2>&1
  5. Read results: grep "^val_bpb:" run.log
  6. If empty output → crash → check logs, attempt fix
  7. Log to results.tsv
  8. If val_bpb improved → keep commit (advance branch)
  9. If val_bpb worse    → git reset (revert)

NEVER STOP: The agent must not pause to ask the human. It runs indefinitely until manually interrupted. The human might be asleep.

Agent Constraints

CAN Do

  • Modify train.py — anything goes
  • Change model architecture
  • Swap optimizers or hyperparameters
  • Adjust batch size & model size
  • Restructure the training loop
  • Change activation functions

CANNOT Do

  • Modify prepare.py (read-only)
  • Install new packages
  • Change the evaluation harness
  • Modify the tokenizer or data pipeline
  • Change the time budget (fixed 5 min)
  • Blow up VRAM without justification

Simplicity Criterion

A 0.001 improvement adding 20 lines of hacky code? Not worth it. A 0.001 improvement from deleting code? Definitely keep.

GPT Model Architecture

A clean, single-file transformer in train.py:

@dataclass
class GPTConfig:
    sequence_len: int = 2048
    vocab_size:   int = 32768
    n_layer:      int = 12
    n_head:       int = 6
    n_kv_head:    int = 6
    n_embd:       int = 768
    window_pattern: str = "SSSL"

Model dim computed from: depth * ASPECT_RATIO (default 8 * 64 = 512, rounded to HEAD_DIM)

Key Components

  • RoPE — Rotary positional embeddings
  • RMSNorm — used everywhere (not LayerNorm)
  • Flash Attention 3 — via kernels library
  • Sliding Window — SSSL pattern (3 short, 1 long)
  • Value Embeddings — ResFormer-style
  • ReluSquared activation in MLP
  • Logit softcap at 15
  • Residual lambdas — learnable per-layer

Sliding Window Attention

Alternating local/global attention via a repeating pattern:

# Window pattern: "SSSL"
# S = short window (seq_len / 2 = 1024)
# L = long window (full seq_len = 2048)

def _compute_window_sizes(self, config):
    pattern = config.window_pattern
    long_window  = config.sequence_len
    short_window = long_window // 2

    # Map pattern chars to window sizes
    char_to_window = {
        "L": (long_window, 0),
        "S": (short_window, 0)
    }

    # Last layer ALWAYS uses full context
    window_sizes[-1] = (long_window, 0)

Why this pattern?

  • Short windows are cheaper — less memory, faster
  • Most layers only need local context
  • Periodic long windows let info flow globally
  • Final layer always full context for best output quality

With SSSL on 12 layers:

S S S L | S S S L | S S S L

9 layers use half-context, 3 use full. Massive compute savings.

Value Embeddings (ResFormer)

A key innovation: input-dependent gating mixes learned value embeddings into attention values.

# Alternating layers get value embeds
def has_ve(layer_idx, n_layer):
    return layer_idx % 2 == \
           (n_layer - 1) % 2

# In attention forward:
if ve is not None:
    ve = ve.view(B, T, n_kv_head, head_dim)
    # Input-dependent gate (32 channels)
    gate = 2 * torch.sigmoid(
        self.ve_gate(x[..., :32])
    )
    v = v + gate.unsqueeze(-1) * ve

How it works

  • Separate embedding table for values (per layer)
  • Gated by first 32 dims of input → input-dependent
  • sigmoid(0) = 0.5, scaled by 2 = 1.0 (neutral at init)
  • Only on alternating layers (saves params)

Why?

Standard attention recomputes values from scratch each layer. Value embeddings give a residual shortcut from token identity to attention output.

MLP: ReluSquared Activation

class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.c_fc = nn.Linear(
            config.n_embd,
            4 * config.n_embd, bias=False
        )
        self.c_proj = nn.Linear(
            4 * config.n_embd,
            config.n_embd, bias=False
        )

    def forward(self, x):
        x = self.c_fc(x)
        x = F.relu(x).square()  # ReluSquared!
        x = self.c_proj(x)
        return x

Design Choices

  • No bias on any linear layer
  • 4x expansion ratio (standard)
  • ReluSquared instead of GELU or SwiGLU

Why ReluSquared?

ReLU squared creates sharper sparsity than GELU. Values are either zero or grow quadratically, giving the network a stronger "on/off" signal while being fully differentiable.

No gating mechanism (unlike SwiGLU) = fewer params and simpler code.

Learnable Residual Scaling

Per-layer learned scalars control residual vs. skip connections:

# In GPT.__init__:
self.resid_lambdas = nn.Parameter(
    torch.ones(config.n_layer)    # start at 1.0
)
self.x0_lambdas = nn.Parameter(
    torch.zeros(config.n_layer)   # start at 0.0
)

# In forward pass:
x = norm(x)
x0 = x  # save initial embeddings
for i, block in enumerate(self.transformer.h):
    # Mix residual stream with x0
    x = self.resid_lambdas[i] * x \
      + self.x0_lambdas[i] * x0
    x = block(x, ve, cos_sin, window)

Two scalar families

  • resid_lambdas — scales the evolving residual stream (init 1.0)
  • x0_lambdas — adds back the initial token embedding (init 0.1 after init_weights)

Intuition

Early layers can learn to weight the "raw" token representation more heavily. Later layers can dial it down. This gives the network an identity shortcut all the way from input to deep layers.

MuonAdamW: Hybrid Optimizer

Different parameter types get different optimizers:

Parameter GroupOptimizerLRNotes
Embedding (wte)AdamW0.6High LR for embeddings
Value EmbeddingsAdamW0.6Same as token embeddings
Unembedding (lm_head)AdamW0.004Much lower LR
Residual scalarsAdamW0.005Gentle updates
x0 scalarsAdamW0.5Beta1=0.96 (more momentum)
Matrix params (2D)Muon0.04Grouped by shape

All LRs scaled by 1/sqrt(model_dim / 768) for dimension-independent tuning.

Muon: Orthogonalization via Polar Express

The core of Muon — approximate matrix polar decomposition using polynomial iteration:

# Pre-computed polynomial coefficients
polar_express_coeffs = [
    (8.156, -22.483, 15.878),
    (4.042, -2.808,  0.500),
    (3.891, -2.772,  0.506),
    ...
]

# For each Nesterov-momentum gradient:
X = g.bfloat16()
X = X / (X.norm() * 1.02 + 1e-6)

# Iterative orthogonalization (ns_steps=5)
for a, b, c in polar_express_coeffs[:ns_steps]:
    A = X.mT @ X
    B = b * A + c * (A @ A)
    X = a * X + X @ B          # approaches orthogonal matrix

Key insight: Instead of SVD (expensive), use 5 cheap matrix multiply iterations to approximate the unitary part of the gradient. This makes weight updates direction-only, removing scale information.

NorMuon: Variance Reduction

After orthogonalization, Muon applies per-row/column variance normalization:

# After polar express orthogonalization...

# Compute per-row/col variance
v_mean = g.float().square().mean(dim=red_dim, keepdim=True)
v_norm = v_mean.sum(dim=(-2, -1)).sqrt()

# EMA of variance (like Adam's second moment)
second_momentum_buffer.lerp_(v_mean, 1 - beta2)
step_size = second_momentum_buffer.clamp_min(1e-10).rsqrt()

# Rescale to preserve overall gradient norm
final_scale = step_size * (v_norm / v_norm_new)
g = g * final_scale

# Cautious weight decay + update
mask = (g * params) >= 0   # only decay aligned dims
params.sub_(lr * g + lr * wd * params * mask)

Cautious weight decay: Only applies decay where the gradient and parameter point in the same direction. Prevents decay from fighting useful parameter values.

Data Pipeline

Tokenizer

  • rustbpe — Rust-based BPE training
  • GPT-4 style split pattern
  • vocab_size = 8192 (intentionally small)
  • 4 special tokens reserved
  • Wrapped in tiktoken for fast encoding

Dataset

  • climbmix-400b-shuffle from HuggingFace
  • 6,543 parquet shards
  • Default: download 10 shards for training
  • Last shard pinned as validation

Dataloader: Best-Fit Packing

# BOS-aligned, no padding
# 100% token utilization

while pos < row_capacity:
  # Find largest doc that fits
  best_idx = find_best_fit(
    doc_buffer, remaining
  )
  if best_idx >= 0:
    # Pack it in
    pack(doc_buffer[best_idx])
  else:
    # Crop shortest doc to fill
    crop_and_fill(shortest)

Every row starts with BOS. Documents are best-fit packed to minimize waste. When nothing fits, the shortest doc is cropped.

Evaluation: Bits Per Byte (BPB)

The single metric the agent optimizes. Lower is better.

@torch.no_grad()
def evaluate_bpb(model, tokenizer, bs):
    # Sum per-token cross-entropy (nats)
    # Sum target byte lengths
    # Convert nats/byte -> bits/byte

    total_nats = 0.0
    total_bytes = 0
    for _ in range(steps):
        loss_flat = model(x, y,
            reduction='none').view(-1)
        nbytes = token_bytes[y_flat]
        mask = nbytes > 0
        total_nats += (loss_flat * mask).sum()
        total_bytes += nbytes.sum()

    return total_nats / (log(2) * total_bytes)

Why BPB?

  • Vocab-size independent — comparing models with different tokenizers is fair
  • Measures compression in bits, normalized by actual byte count
  • Special tokens (0 bytes) excluded from both sums
  • Evaluated over 40 * 524,288 = ~21M tokens

This is in prepare.py and cannot be modified. It is the ground truth that keeps the agent honest.

Training Loop Mechanics

Time-Based Training

  • Fixed TIME_BUDGET = 300s (5 min)
  • First 10 steps excluded (JIT compilation warmup)
  • Training stops when wall-clock exceeds budget
  • Steps per run varies by hardware

Gradient Accumulation

# Total batch = 524K tokens
# Device batch = 128 * 2048
tokens_per_fwdbwd = 128 * 2048
grad_accum = 524288 // tokens_per_fwdbwd
# = 2 micro-steps per optimizer step

LR Schedule

  • WARMUP_RATIO = 0.0 — no warmup by default
  • WARMDOWN_RATIO = 0.5 — cosine decay for last half
  • FINAL_LR_FRAC = 0.0 — decays to zero

Safety Rails

  • NaN or loss > 100 → immediate exit(1)
  • GC disabled after step 0 (avoids 500ms stalls)
  • GC runs every 5000 steps
  • bf16 autocast for all forward/backward

Weight Initialization Strategy

Careful init ensures stable training from step 0 — no warmup needed:

ComponentInit StrategyRationale
Token embeddings (wte)Normal(0, 1)Standard for embeddings
LM headNormal(0, 0.001)Near-zero to start uniform
Q, K, V projectionsUniform(-s, s)s = 3^0.5 * n_embd^-0.5
Output projectionsZerosBlocks are identity at init
MLP up-projectionUniform(-s, s)Same scale as attention
MLP down-projectionZerosBlock starts as no-op
VE gatesZerossigmoid(0)*2 = 1.0 (neutral)
resid_lambdas1.0Full residual pass-through
x0_lambdas0.1Small initial skip connection

Key pattern: Output projections initialized to zero make each block an identity function at init. The network starts as a simple embedding lookup and gradually learns transformations.

Flash Attention 3 Integration

from kernels import get_kernel

cap = torch.cuda.get_device_capability()

# Hopper (H100) gets the fast path
if cap == (9, 0):
    repo = "varunneal/flash-attention-3"
else:
    repo = "kernels-community/flash-attn3"

fa3 = get_kernel(repo).flash_attn_interface

# Used in attention:
y = fa3.flash_attn_func(q, k, v, causal=True, window_size=window_size)

Why FA3?

  • IO-aware: minimizes HBM reads/writes
  • Fused kernel: no materialized attention matrix
  • Native sliding window support
  • Loaded dynamically via kernels package

QK-Norm Before Attention

# After RoPE, before FA3:
q, k = norm(q), norm(k)
# RMSNorm on head dimension
# Stabilizes attention logits

Git Workflow & Results Tracking

Branch Strategy

# Each run gets a dated branch
git checkout -b autoresearch/mar5

# Good result? Branch advances
git commit -m "increase LR to 0.04"

# Bad result? Revert
git reset --hard HEAD~1

The branch becomes a linear history of improvements. Every commit on the branch represents a positive delta.

results.tsv Format

commit  val_bpb  memory_gb  status  description
a1b2c3d  0.9979  44.0  keep     baseline
b2c3d4e  0.9932  44.2  keep     increase LR
c3d4e5f  1.0050  44.0  discard  switch to GeLU
d4e5f6g  0.0000  0.0   crash    double width

TSV is untracked by git — it is a local lab notebook. Tab-separated because commas break in descriptions.

Dependencies & Hardware

Minimal Dependency Stack

PackagePurpose
torch 2.9.1Core framework (CUDA 12.8)
kernelsFlash Attention 3 kernels
rustbpeRust-based BPE tokenizer training
tiktokenFast token encoding at runtime
pyarrowParquet shard reading
numpy / pandasAnalysis utilities
matplotlibVisualization in notebook
requestsData downloading

Hardware

  • Designed for single NVIDIA GPU
  • Tested on H100
  • Peak FLOPS reference: 989.5 TF (bf16)
  • ~45 GB VRAM at default config

Low-Resource Tips

Use TinyStories dataset, reduce vocab to ~4096, lower seq_len, decrease depth. Community forks exist for MacOS, Windows, and AMD.

Design Philosophy

Fixed Time Budget

5 minutes per experiment regardless of hardware. Faster GPUs run more steps, slower GPUs run fewer — but the wall-clock is always the same. This makes ~12 experiments/hour predictable.

Single File Scope

The agent can only edit train.py. This prevents it from gaming the eval, breaking the data pipeline, or installing dependencies. The attack surface is one file.

Simplicity Over Cleverness

The simplicity criterion explicitly penalizes complexity. Hacky improvements are rejected. Code deletion that maintains performance is celebrated.

Autonomous by Default

The agent must never ask "should I continue?". It runs until stopped. This mirrors real overnight batch processing — the human is not available.

Key Takeaways

  • Constraint-driven design — one file, one metric, fixed time makes the search space tractable
  • Modern architecture — RoPE, QK-norm, FA3, sliding windows, value embeddings, ReluSquared
  • Hybrid optimizer — Muon for 2D matrices (orthogonal updates), AdamW for everything else
  • Git as experiment log — every good result is a commit, every failure is a revert
  • 100 experiments overnight — AI agents as tireless research assistants

github.com/karpathy/autoresearch · 47.2k stars · MIT License