Card: training data, recipe, tokens and time; how the head was edited by hand. Add the fine-tune run record.

README.md

@@ -320,8 +320,57 @@ All splits below are out of distribution for the model: none of their items or i
 ## Provenance
 <hr style="margin: -0.7em 0 1.2em 0; border: 0; border-top: 1px solid rgba(127, 127, 127, 0.3);" />
 
-1. **Fine-tune.** d1-2B, full-parameter with the tied token embeddings / LM head frozen, two epochs on a family-deduplicated mix of agent-action security data (CLI command risk labels, agent traces with needs-review and injection labels, authorization labels), up to 16 questions per sequence. Checkpoint `aidendle94/d1-guard-ft-v0-20261028` (private).
-2. **Head edit.** Fine-tuning left several features intact in the last layer while the fixed head read them poorly: linear probes on last-layer activations separate LlamaFirewall injection at 98.9 AUROC (head: 87.4 on the dev split), ATBench injection at 89 (head: 74) and ATBench needs-review at 89 (head: 74). Three unit directions were fit on the final-norm output at the answer slot, each made orthogonal to what the other yes/no questions read and to the earlier directions, and added to the head rows of the "yes" tokens and subtracted from the "no" tokens: LlamaFirewall injection (weight 1.0), ATBench injection recorded under served packing (0.7), and needs-review pooled over ATBench and 900 in-distribution rows (0.5). The head was untied from the input embeddings so the input side is untouched. 10,240 weights changed. The directions and the full record, including the simulated effect on every probe, are in `provenance/`.
+Scantron Preview was made in two stages: a full fine-tune of d1-2B on one rented GPU, then a hand edit of five rows of the output head, planned and carried out by an AI agent on a laptop. Both are described below; the run record and the edit record are in `provenance/`.
+
+### Stage 1: fine-tune
+
+**Data.** 92,577 training rows from six public datasets, converted into one JSON state format (the user's request, the proposed tool call, recent actions, untrusted content, destination) and labeled with the 21 questions in `questions.json`. Rows were deduplicated by action template and split by family so that no command prefix or trace family appears on both sides; 10% of families were held out as the in-distribution eval. Every row that shared a key with the external benchmarks in the Performance section was excluded before training.
+
+| Source | Train rows | Labels contributed |
+|---|---|---|
+| ProCreations/auto-1b-data | 49,863 | authorization, injection |
+| karanxa/agent-action-safety-dataset | 21,147 | needs review, authorization, injection |
+| yannelli/cli-command-risk | 10,064 | risk level, risk categories, reversible, requires root |
+| tomngdev/tool-safety-v2 | 8,238 | needs review |
+| westenfelder/NL2SH-ALFA | 1,678 | risk level, categories, reversible, requires root |
+| b-mc2/cli-commands-explained | 1,587 | risk level, categories, reversible, requires root |
+
+Label coverage over the training rows: authorization 70,863 (38,024 authorized / 32,839 unauthorized), needs review 35,202, injection 31,295 (1,977 positive), risk level 14,831, risk categories 14,743, reversible 14,125, requires root 13,079. Risk labels on shell commands were adjudicated by a second model where the source and a labeling model disagreed; rows that stayed ambiguous were dropped.
+
+**Objective.** d1's native training view: each row's state is followed by its labeled questions in random order, packed up to 16 per sequence, and the loss is cross entropy over the option scores read at each question's answer slot (the best token form of each option, max-pooled). Nothing is generated. The 8,192-token state budget truncates untrusted content first, oldest first, and a row whose untrusted content was cut has its injection label masked.
+
+**Recipe and cost.**
+
+| | |
+|---|---|
+| Trainable parameters | 1.47B (the 24 backbone layers and the final norm); token embeddings / LM head (0.26B) and the vision tower frozen |
+| Optimizer | AdamW, lr 1e-5, betas 0.9 / 0.999, weight decay 0.01 on matrices only, gradient clip 1.0 |
+| Schedule | 3% linear warmup (98 steps), cosine to zero |
+| Batch | 8 micro-batches of 8,192 padded tokens per optimizer step; 3,293 steps over 2 epochs |
+| Precision | bf16 weights and compute, fp32 master weights, fp32 gradient accumulation and optimizer state |
+| Tokens processed | 171.9M (85.9M per epoch), 706,060 labeled decisions |
+| Hardware and time | one NVIDIA RTX 6000D (rented), 4 h 46 min wall clock, 10,014 tokens/s, 54.4 GB peak memory |
+| Checkpoint | epoch 2, `aidendle94/d1-guard-ft-v0-20261028` (private) |
+
+### Stage 2: editing the weights by hand
+
+The second stage used no gradient steps and no GPU rental. It was done on an Apple M3 Max laptop (128 GB, Metal) by [Claude Fable 5.1](https://www.anthropic.com/claude/fable) working as an autonomous agent in Claude Code: the agent designed the probes, ran them through the served engine, chose the edit, verified it, and wrote this card. A second model (DeepSeek V4.1 Flash, run locally) reviewed the labels and the code. No human labeled data or reviewed the edit.
+
+**What was found.** Fine-tuning moved the backbone while the output head stayed frozen, and the two drifted apart: the last layer still carried features that the head no longer read well. Linear probes trained on the last-layer activation at the answer slot separated LlamaFirewall injection at 98.9 AUROC where the head read 87.4 on the same dev rows, ATBench injection at 89 where the head read 74, and ATBench needs-review at 89 where the head read 74. The fix is to point the head at the feature, which is a change to a handful of rows of one matrix.
+
+**How the edit works.** Each yes/no question is answered by comparing the logits of the "yes" token forms against the "no" token forms at the answer slot. Those logits are dot products of the final-norm output `x` with the head rows of those tokens, so adding a unit direction `w` with weight `β` to the "yes" rows and subtracting it from the "no" rows shifts the served yes-vs-no log-odds by exactly `2β (w · x)`. Because the shift is linear in `x`, its effect on every recorded probe can be computed offline before anything is served, and a candidate edit can be rejected in seconds.
+
+**Procedure.**
+
+1. **Record.** For each question of interest, run the fine-tuned model over a dev split through the served engine (same prompt, same packing, same 16-question sets) and keep the hidden state of every layer at the question's answer slot. Probes were recorded under the exact question set the model is served with, because a d1 answer depends on which other questions share its set.
+2. **Locate.** Fit a logistic-regression probe at every layer and compare its AUROC with the head's. The gap at the last layer is the readout loss; it was large for injection and needs-review and near zero for the shell-command questions, which were left alone.
+3. **Fit the direction.** Logistic regression on the final-norm output at the answer slot (standardized features, C = 0.01), normalized to a unit vector. Three directions: LlamaFirewall injection (654 rows, 50 positive), ATBench injection recorded under served packing (80 rows, 18 positive), and needs-review pooled over ATBench and 900 in-distribution rows (1,077 rows, 652 positive).
+4. **Protect the other questions.** Because the "yes" and "no" rows are shared by all 20 yes/no questions, each direction was made orthogonal (QR) to the mean activation and top four principal components of every other question's probe, and to the directions fit before it. Cosine between the raw and the protected direction: 0.91, 0.93 and 0.94, so most of the signal survived.
+5. **Simulate.** Compute the shifted log-odds `s + 2 (δ · x)` with `δ = 1.0 w₁ + 0.7 w₂ + 0.5 w₃` on all 20 recorded probes: AUROC before and after, mean shift, and how many rows cross the decision point. Weights were chosen where the targets improved and the guard questions moved by at most one or two examples each.
+6. **Write the weights.** Untie the head from the input embeddings (so the input side is untouched), add `δ` (norm 1.32) to the two "yes" token rows and subtract it from the three "no" token rows. 5 rows × 2,048 = 10,240 weights changed out of 1.9B; everything else is byte-identical to the fine-tuned checkpoint.
+7. **Verify by serving.** Re-run the edited model on the dev splits and confirm the offline simulation (correlation 1.0 with the served shift), then score the frozen test splits once. The frozen results are the "Scantron Preview" column in the Performance section; the pre-edit column is the same checkpoint without the edit.
+
+The in-sample fit AUROC of each direction before and after the edit was 0.875 → 0.993 (LlamaFirewall injection), 0.741 → 0.928 (ATBench injection) and 0.978 → 0.981 (needs-review). `provenance/weight_edit.json` has the full record, including the simulated effect on every probe, and `provenance/head_delta.npy` is `δ` itself, so the edit can be undone or re-weighted.
 
 Modified files relative to d1-2B, as the license requires us to say: `model.safetensors` (all weights fine-tuned; `lm_head.weight` added), `config.json` (`tie_word_embeddings: false`), `calibration.json` (served temperature 1.0), `questions.json` (new). See `NOTICE`.
 

provenance/train_run.json

@@ -0,0 +1,185 @@
+{
+ "run_id": "ft-v0-20261028",
+ "status": "done",
+ "base_checkpoint": {
+  "repo": "LiquidAI/d1-2B",
+  "revision": "cfe5760ff5d5998f91d9ac56bd4b2dbd96996d41",
+  "local_path": "/workspace/d1-guard/models/d1-2B"
+ },
+ "dataset_version": {
+  "dataset_version": "ds-20261028-05072130",
+  "files": {
+   "clean/mix_v2/train.parquet": {
+    "sha256": "7f405a8c8b897461e502a6bbd488bd5a54d1b50761b38e43842444ce7d58d601"
+   },
+   "clean/mix_v2/eval.parquet": {
+    "sha256": "ff06a2dc9faac0eb51b1d4ca9615b631098905e832925f4f12b60d6bd4bbe76f"
+   }
+  },
+  "mix_manifest": "clean/mix_v2/manifest.json",
+  "note": "mix_v2 is not listed in DATASET_VERSION.json (it lists mix_v1); the mix files are pinned by the sha256 above"
+ },
+ "split_and_family_definitions": "scripts/40_build_mix.py (mix_v2); cli_capped / auth_aod_bench per DATASET_VERSION.json",
+ "taxonomy_version": "v0-draft-2026-10-27",
+ "question_set_version": "questions-v0-2026-10-28 (training/questions_v0.json)",
+ "trainable_parameter_set": {
+  "description": "language backbone: model.language_model.layers.* and model.language_model.embedding_norm",
+  "n_trainable": 1466657536,
+  "n_frozen": 451524352,
+  "tensors_trainable": 285,
+  "frozen": [
+   "vision_tower",
+   "multi_modal_projector",
+   "embeddings (embed_tokens, tied to lm_head)"
+  ]
+ },
+ "method": "full",
+ "objective": "d1 tree: cross entropy over each labeled question's option scores (max-pooled token forms at the answer slot), questions packed after their state in random order, groups of set_size",
+ "seed": 0,
+ "config": {
+  "lr": 1e-05,
+  "optimizer": "adamw (fused, betas 0.9/0.999, eps 1e-8)",
+  "weight_decay": 0.01,
+  "weight_decay_on": "matrices only (not norms)",
+  "schedule": "linear warmup then cosine to 0",
+  "warmup": 0.03,
+  "warmup_steps": 98,
+  "epochs": 2,
+  "precision": "bf16 weights and compute, fp32 master weights, grads accumulated in fp32, fp32 AdamW states",
+  "grad_clip": 1.0,
+  "set_size": 16,
+  "state_budget_tokens": 8192,
+  "micro_batch_padded_tokens": 8192,
+  "accum": 8,
+  "total_steps": 3293,
+  "grad_ckpt": false,
+  "loss_normalization": "mean over the step's decisions",
+  "truncation": "untrusted_context oldest-first, then recent_actions oldest-first; untrusted_origin masked on rows whose untrusted_context was cut"
+ },
+ "hardware": {
+  "worker": "vast.ai 58064827",
+  "gpu": "NVIDIA RTX 6000D"
+ },
+ "tokens_processed": 171862216,
+ "decisions_processed": 706060,
+ "wall_clock_s": 17176,
+ "throughput_tokens_per_s": 10014,
+ "peak_memory_gb": 54.43,
+ "correctness_checks": {
+  "see": "experiments/full_ft_v0/checks/checks.json"
+ },
+ "metrics": {
+  "A_public": null,
+  "B_private_ood": null,
+  "C_redteam": null,
+  "D_dev": null,
+  "paraphrase_gap": null,
+  "latency_ms_p50_p99": null
+ },
+ "epochs": [
+  {
+   "epoch": 1,
+   "step": 1647,
+   "tokens": 85931108,
+   "minutes": 146.2,
+   "quick_eval": {
+    "risk_level": {
+     "loss": 0.3628,
+     "acc": 0.8986,
+     "n": 552,
+     "hc_recall_at_0.5": 0.7941,
+     "safe_fpr_at_0.5": 0.0097
+    },
+    "category": {
+     "loss": 0.0252,
+     "acc": 0.993,
+     "n": 8265
+    },
+    "reversible": {
+     "loss": 0.1781,
+     "acc": 0.943,
+     "n": 544
+    },
+    "requires_root": {
+     "loss": 0.1824,
+     "acc": 0.9561,
+     "n": 524
+    },
+    "needs_review": {
+     "loss": 0.0835,
+     "acc": 0.9745,
+     "n": 432
+    },
+    "authorization": {
+     "loss": 0.0373,
+     "acc": 0.9908,
+     "n": 326
+    },
+    "untrusted_origin": {
+     "loss": 0.0213,
+     "acc": 0.9918,
+     "n": 122
+    },
+    "all": {
+     "loss": 0.0605,
+     "n": 10765
+    }
+   },
+   "checkpoint": "/workspace/d1-guard/checkpoints/ft-v0-20261028/epoch1"
+  },
+  {
+   "epoch": 2,
+   "step": 3293,
+   "tokens": 171862216,
+   "minutes": 285.5,
+   "quick_eval": {
+    "risk_level": {
+     "loss": 0.3868,
+     "acc": 0.9022,
+     "n": 552,
+     "hc_recall_at_0.5": 0.8627,
+     "safe_fpr_at_0.5": 0.0097
+    },
+    "category": {
+     "loss": 0.0248,
+     "acc": 0.9936,
+     "n": 8265
+    },
+    "reversible": {
+     "loss": 0.2041,
+     "acc": 0.9485,
+     "n": 544
+    },
+    "requires_root": {
+     "loss": 0.1836,
+     "acc": 0.9695,
+     "n": 524
+    },
+    "needs_review": {
+     "loss": 0.0612,
+     "acc": 0.9745,
+     "n": 432
+    },
+    "authorization": {
+     "loss": 0.0133,
+     "acc": 0.9939,
+     "n": 326
+    },
+    "untrusted_origin": {
+     "loss": 0.0081,
+     "acc": 0.9918,
+     "n": 122
+    },
+    "all": {
+     "loss": 0.061,
+     "n": 10765
+    }
+   },
+   "checkpoint": "/workspace/d1-guard/checkpoints/ft-v0-20261028/epoch2"
+  }
+ ],
+ "checkpoint_id": "/workspace/d1-guard/checkpoints/ft-v0-20261028/epoch2",
+ "log": "/workspace/d1-guard/logs/train_ft-v0-20261028.log",
+ "train_s": 17163,
+ "steps_done": 3293
+}