Can I rebuild the system faithfully?
Implement the architecture and training pipeline from scratch, reproduce the 124M configuration, and verify every stage with parity tests and logged metrics.
Output: a trained, reproducible checkpoint.Inside the prediction
InsideGPT
From-scratch reproduction and causal explanation of GPT-2 124M next-token predictions.
Buildone model from first principlesThentest what actually causes its predictions
124Mparameters
12 × 12layers × heads
10Btraining tokens
1 × A10040 GB GPU
01 / Scope
One project, two separate claims
Reproducing GPT-2 teaches how the model is built and optimized. It does not, by itself, make the project Explainable AI. The second phase freezes the trained checkpoint and tests explanations against controlled interventions.
Can I identify what controls a prediction?
Compare attribution, attention, and causal interventions on matched prompts, then quantify when each explanation succeeds or fails.
Output: an explanation benchmark, not just a visualization.02 / Learning
What I will know by the end
Transformer mechanics
Implement token and position embeddings, causal multi-head attention, MLP blocks, residual streams, layer normalization, weight tying, and next-token loss directly in PyTorch.
Training systems
Build sharded data loading, gradient accumulation, mixed precision, learning-rate scheduling, checkpointing, reproducible evaluation, and GPU-memory monitoring.
Training dynamics
Read loss curves, gradient norms, token throughput, validation gaps, and failure signals instead of treating training as a black-box command.
Causal interpretability
Separate an attractive visualization from evidence by testing whether highlighted tokens and internal states actually control a prediction.
Scientific evaluation
Define baselines, matched counterfactuals, quantitative metrics, uncertainty, and negative-result reporting before inspecting the final explanations.
Research engineering
Turn one reproducible model checkpoint into a documented experiment suite, an interactive inspection tool, and a defensible final report.
03 / Execution
How the model will be built and trained
The implementation follows the public build-nanogpt reference repository. I will pin its commit SHA before starting so “reproduction” refers to one exact code state.
01TokensBPE IDs
02Embeddingstoken + position
123456789101112
03Logitsnext token
Architecture
Rebuild GPT-2 124M
12 Transformer blocks · 12 attention heads · 768-dimensional residual stream · 1,024-token context · 50,257-token GPT-2 vocabulary.
Exit gateLoad the matching Hugging Face weights and verify logits on the same input within a small numerical tolerance.
Data
Prepare FineWeb-Edu
Tokenize the sample-10BT subset with the GPT-2 tokenizer and store deterministic uint16 train/validation shards.
Exit gateRecord dataset revision, shard hashes, token counts, split logic, and several decoded samples.
Pipeline
Prove the loop before renting hours
Overfit one batch, run a short training smoke test, save and reload a checkpoint, and resume without changing the loss trajectory.
Exit gateNo NaNs, stable memory, identical resumed step, and a visibly decreasing smoke-test loss.
Training
Run the 10B-token reproduction
Use BF16, Flash Attention, fused AdamW, gradient accumulation, gradient clipping, warmup, cosine decay, and periodic validation.
Exit gateComplete the planned token budget and retain the best plus final checkpoints with logs and configuration.
Evaluation
Measure language-model quality
Track validation loss and HellaSwag accuracy, then inspect fixed-seed text samples as a qualitative diagnostic.
Exit gateCompare results to the pinned reference run and explain any reproducible difference instead of hiding it.
04 / Instrument
A prediction microscope
This interactive preview shows the questions the final tool will answer. It is an experiment design, not a display of completed results.
Interactive experiment preview
illustrativeWhat caused this next token?
Original prompt
ThecapitalofFranceis
target tokenParis
Measurement path
123456789101112
Search all 12 blocksEstablish the model's next-token prediction and collect token, head, and residual-stream signals.
Measurement
Record the target logit before any intervention.
Interpretation rule
No causal claim yet: this is the measurement baseline.
05 / XAI study
Five methods, each with a causal check
Token ablation
Which input tokens matter?Remove or replace one token or phrase, rerun the model, and measure the change in the target token's logit.
Validity checkAblation can create unnatural text, so every ablation gets a fluent counterfactual control.
Integrated Gradients
Which inputs receive local attribution?Integrate gradients from a defined baseline embedding to the real embedding and attribute the target logit across tokens.
Validity checkAttribution is sensitivity along one path; it is compared with interventions rather than accepted as causal proof.
Attention analysis
Where is information routed?Inspect head-level attention patterns and test whether apparently relevant heads are consistent across prompt variants.
Validity checkAttention maps are diagnostics, not explanations by themselves.
Activation patching
Where does a fact affect the computation?Patch a source prompt's residual or head activation into a matched counterfactual and measure recovery of the source prediction.
Validity checkPatch cleanly matched prompts and compare against random-layer and random-token controls.
Counterfactual prompts
Does the explanation generalize?Change one semantic variable, such as France to Germany, while holding syntax and task structure constant.
Validity checkReport prediction changes across many pairs, not one hand-picked demonstration.
06 / Evidence
What counts as success
Reproduction gates
- Architecture and parameter-count tests pass.
- Imported reference weights produce matching logits.
- A smoke run overfits one batch and resumes exactly.
- The full run completes without NaNs or hidden data changes.
- Validation loss, HellaSwag, samples, throughput, and configuration are preserved.
Explanation metrics
- Comprehensiveness
- How much the target score falls after removing features marked important.
- Sufficiency
- How much of the target score remains when only the marked features are retained.
- Counterfactual consistency
- Whether the explanation changes when the causal fact changes.
- Explanation stability
- Whether meaning-preserving paraphrases receive similar explanations.
- Patching effect
- How strongly a patched internal state moves the target logit toward the source prompt.
Decision rule
I will call a method causally informative only when its highlighted features predict held-out intervention effects better than random, position-matched, and magnitude-matched baselines.
07 / Resources
Single-A100 run card
Target machineNVIDIA A100 PCIe
40 GB VRAM · 64 GB RAM minimum, 128 GB preferred · 120 GB disk
Data preparation2–5 h
Smoke + parity tests0.5–1 h
Full training5–10 h
XAI runs1–3 h
Debug reserve2–4 h
Expected total10–16 h
Quoted GPU rate$0.70 / hour
Expected compute$7.00–$11.20
Budget cap$15
Planning estimates, not guarantees. The run stops if loss, throughput, or memory fails a pre-set health check.
120 GB is workable with an explicit storage policy
Token shards require roughly 20–22 GB. The environment, checkpoints, logs, and XAI artifacts add approximately 20–35 GB. Raw dataset caches are cleared only after shard hashes and decoded samples are verified, keeping the run below the disk ceiling.
08 / Schedule
Six milestones with hard exit gates
Pin the reference commit, recreate the architecture, and write parity tests.
Build deterministic shards and validate training, checkpointing, and resume behavior.
Run the 10B-token job, monitor it, and preserve all metrics and checkpoints.
Implement attribution, ablation, counterfactual, attention, and activation-patching pipelines.
Run matched examples, baselines, controls, uncertainty estimates, and error analysis.
Build the prediction microscope, write the dissertation, and package reproducible artifacts.
09 / Output
What will be released
Model implementation
Readable PyTorch code, architecture tests, pinned dependencies, data manifest, and reproducible commands.
Training evidence
Checkpoint, configuration, train/validation curves, HellaSwag scores, throughput, and run-cost report.
XAI benchmark
Matched prompts, interventions, attribution outputs, activation patches, baselines, uncertainty, and failure taxonomy.
Prediction microscope
An interactive dashboard for tracing one next-token prediction across inputs, layers, heads, and counterfactuals.
Final report
A precise account of what replicated, which explanations survived causal tests, and which negative results changed the conclusion.