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
What we are trying to find
Can a faithfully reproduced GPT-2 checkpoint support explanations that identify what actually controls a next-token prediction?
How we will do it
First reproduce and verify GPT-2 124M; then freeze that checkpoint and test explanation methods with controlled interventions.
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
What we are trying to find
Which engineering and scientific skills are required to move from building a Transformer to making defensible claims about its behavior?
How we will do it
Implement every major subsystem, inspect training signals, design causal controls, and turn the results into a reproducible study.
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
What we are trying to find
Can our implementation reproduce the architecture, numerical behavior, and training quality of the GPT-2 124M reference?
How we will do it
Build and test one causal attention head first, expand it into multi-head attention and a complete decoder block, stack 12 blocks, then validate the full training pipeline.
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.
Interactive build sequence
01 / 05Learn the architecture by constructing it
next tokenOutput probabilities
Vocabulary projectionLinear + Softmax50,257 logits
× 1 block
Pre-LayerNorm + residualFeed-forward network768 → 3,072 → 768
Pre-LayerNorm + residual1 causal attention headQ · K · V + causal mask
first working unit
H1
Scaled attention → outputStarting representationToken + position embeddingssequence × 768
The capital of France isshifted token sequence
Stage 01
Represent the input
- What we are trying to find
- Can token identity and position be combined into the vectors the model will process?
- How we will do it
- Tokenize the text, look up token embeddings, add learned positional embeddings, and verify every tensor shape.
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
What we are trying to find
For one predicted token, where do candidate causes appear across the input, attention heads, layers, and residual stream?
How we will do it
Trace the baseline prediction, ablate candidate input features, and patch internal activations while recording changes to the target logit.
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
What we are trying to find
Which interpretability methods reliably identify features or internal states that causally influence GPT-2's next-token prediction?
How we will do it
Compare five explanation methods on matched prompts, then test their claims against fluent counterfactuals, ablations, activation patches, and matched baselines.
Token ablation
What we are trying to findWhich input tokens matter?
How we will do it
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
What we are trying to findWhich inputs receive local attribution?
How we will do it
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
What we are trying to findWhere is information routed?
How we will do it
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
What we are trying to findWhere does a fact affect the computation?
How we will do it
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
What we are trying to findDoes the explanation generalize?
How we will do it
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
What we are trying to find
Did the reproduction succeed, and do the explanations predict real intervention effects rather than merely look plausible?
How we will do it
Apply fixed reproduction gates and score explanations with comprehensiveness, sufficiency, consistency, stability, patching effect, and matched baselines.
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
What we are trying to find
Can the full reproduction and explanation study fit within one A100, the available storage, and the project budget?
How we will do it
Estimate each workload separately, reserve debugging time, enforce health checks and a cost cap, and keep only verified data and artifacts.
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
What we are trying to find
Can the project reach a defensible answer within 16 weeks without letting training or interface work crowd out the experiments?
How we will do it
Sequence specification, data validation, training, instrumentation, benchmarking, and release, with an exit gate before each next phase.
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
What we are trying to find
What evidence must be available for someone else to inspect, reproduce, and challenge the project's conclusions?
How we will do it
Release the implementation, checkpoint and training record, XAI benchmark, interactive microscope, and a report that includes negative results.
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.