Nayana’s Blog

A Multimodal Stochastic Medical Tutor: A Systems Build Story

Mon, 12 Jan 2026 00:00:00 GMT

Note: This RSS feed strips out SVGs and embeds. You might want to read the post on the webpage here.

A multimodal medical tutor that uses a deterministic runtime to orchestrate stochastic controller decisions—combining image tools, retrieval, and student modeling—improved through SFT, RL, and playbook optimization.

January 12, 2026

projects

A Systems Build Story

╔══════════════════════════════════════════════════════════════════════════════╗║                                                                              ║║     ┌─────────┐      ┌──────────┐      ┌─────────┐      ┌──────────────┐     ║║     │  Image  │ ───▶ │ Retrieve │ ───▶ │ Decide  │ ───▶ │ Tutor/Answer │     ║║     │    +    │      │  (text   │      │ (tools, │      │   (probes,   │     ║║     │Question │      │ +image)  │      │ actions)│      │hints,answer) │     ║║     └─────────┘      └──────────┘      └────┬────┘      └──────────────┘     ║║                                             │                                ║║                                             ▼                                ║║                                      ┌────────────┐                          ║║                                      │  Tool Exec │                          ║║                                      │(zoom,seg,  │                          ║║                                      │ retrieve…) │                          ║║                                      └────────────┘                          ║║                                                                              ║║               A   M U L T I M O D A L   M E D I C A L   T U T O R.           ║║                                                                              ║╚══════════════════════════════════════════════════════════════════════════════╝

TL;DR

I built a multimodal medical tutoring system that reads medical images, retrieves supporting knowledge from text and image databases, uses visual tools to gather evidence, teaches Socratically with calibrated probes and hints, and produces a final answer only when the student is ready.

The system is stochastic by design. At each decision point, the tutor samples actions under uncertainty—which tool to use, which question to ask, whether to reveal the answer. The stochasticity lives in the controller's sampling; the surrounding runtime is deterministic. Those decisions are improved through three complementary mechanisms: SFT for reliable action formatting, RL for tool-use behavior, and ACE-style playbook optimization for fast iteration on recurring mistakes.

This is a long-form technical build story. I wrote it for engineers and recruiters who want to understand not just what the system does, but how it was designed and why.

The problem I actually solved

"Medical VQA" and "medical tutoring" look similar on the surface. Both start with an image and a question. But tutoring adds requirements that most VQA systems never address.

┌─────────────────────────────────────────────────────────────────────────────┐│                                                                             ││   MEDICAL VQA                         MEDICAL TUTORING                      ││   ───────────                         ────────────────                      ││                                                                             ││   Image + Question                    Image + Question                      ││         │                                   │                               ││         ▼                                   ▼                               ││   ┌──────────┐                        ┌──────────────┐                      ││   │  Model   │                        │  Ask student │◀──┐                  ││   └────┬─────┘                        └──────┬───────┘   │                  ││        │                                     │           │                  ││        ▼                                     ▼           │                  ││   ┌──────────┐                        ┌──────────────┐   │                  ││   │  Answer  │                        │  Assess      │   │ (loop until      ││   └──────────┘                        │  response    │   │  ready)          ││                                       └──────┬───────┘   │                  ││                                              │           │                  ││                                              ▼           │                  ││                                       ┌──────────────┐   │                  ││                                       │ Hint / Probe │───┘                  ││                                       │ / Reveal     │                      ││                                       └──────────────┘                      ││                                                                             │└─────────────────────────────────────────────────────────────────────────────┘

A tutor has to ask what the student thinks before explaining. It has to escalate hints based on performance rather than revealing immediately. It has to use tools when the image is ambiguous instead of guessing. It has to use retrieval when medical knowledge is needed, and it has to separate observation from diagnosis. Most importantly, every decision the system makes should be inspectable—which tool was called, what it returned, and how it changed the next step.

I built a system that treats tutoring as a control problem, not a single model call. The goal wasn't just accuracy. I wanted a system that is consistent in how it reasons, inspectable when it fails, and trainable when I want to change behavior. This is why the final system looks more like a small operating system than a prompt-and-model demo.

The system at a glance

The system has two interacting loops. At runtime, it behaves like a state machine that iterates until it produces an instructional response—a probe, a hint, a micro-lesson—or a final answer. Offline, multiple improvement channels operate on the same underlying runtime traces: SFT trains the controller to emit structured tutoring actions reliably, RL trains tool-use behavior with intermediate rewards, and ACE evolves a playbook of behavioral rules to reduce failure modes without retraining.

╔═══════════════════════════════════════════════════════════════════════════════╗║                              ONLINE RUNTIME                                   ║╠═══════════════════════════════════════════════════════════════════════════════╣║                                                                               ║║  ┌───────────────┐    ┌────────────┐    ┌───────────┐    ┌────────────────┐   ║║  │ Pre-Retrieve  │───▶│   Decide   │───▶│ Tool Exec │───▶│ Answer/Respond │   ║║  │ (gather ctx)  │    │ (policy)   │    │ (1 tool)  │    │   (terminal)   │   ║║  └───────────────┘    └─────┬──────┘    └─────┬─────┘    └────────────────┘   ║║                             │                 │                               ║║                             │◀────────────────┘                               ║║                             │    (loop back)                                  ║║                                                                               ║╠═══════════════════════════════════════════════════════════════════════════════╣║           ▲                                            ▲                      ║║           │  traces + state snapshots                  │  playbook rules      ║║           │                                            │                      ║╠═══════════════════════════════════════════════════════════════════════════════╣║                            OFFLINE IMPROVEMENT                                ║╠═══════════════════════════════════════════════════════════════════════════════╣║                                                                               ║║  ┌─────────────────┐  ┌─────────────────┐  ┌─────────────────────────────┐    ║║  │      SFT        │  │       RL        │  │           ACE               │    ║║  │ (action format, │  │ (tool-use       │  │ (playbook rules,            │    ║║  │  reliability)   │  │  rewards)       │  │  fast iteration)            │    ║║  └─────────────────┘  └─────────────────┘  └─────────────────────────────┘    ║║                                                                               ║╚═══════════════════════════════════════════════════════════════════════════════╝

This separation matters because it keeps the runtime stable, keeps training focused, and makes improvements measurable. The runtime never changes during training—only the policies that run inside it.

When I say "stochastic," I mean something specific. At decision points, the controller can sample different actions under uncertainty: tool versus probe versus hint versus answer. But the runtime around it is deterministic and constrained—schemas, gates, routing, budgets. This combination gives you both creativity when the model needs it and predictability when the system needs it.

Every design decision in this system is downstream of a few constraints I couldn't ignore. Raw retrieval and raw tool outputs blow up your prompt, so I needed a context budget. A tutor can't take two minutes per turn and still feel interactive, so I needed a latency budget. Tool-using systems fail in dozens of new ways—parsing errors, missing dependencies, wrong modalities, empty retrieval, repeated tools—so I needed explicit failure handling. If I can't replay a run, I can't improve it reliably, so I needed reproducibility. And I needed hard constraints that prevent "answer dumping" even when the policy is stochastic.

The backbone: one state object, one traceable story

The most important decision I made was that everything flows through one explicit state object.

┌─────────────────────────────────────────────────────────────────────────────┐│                              AgentState                                     │├─────────────────────────────────────────────────────────────────────────────┤│                                                                             ││  messages: List[Message]          ◀── conversation history (multimodal)     ││  image: ImageRef | None           ◀── current image reference               ││                                                                             ││  ┌───────────────────────────────────────────────────────────────────────┐  ││  │ RETRIEVAL                                                             │  ││  │   retrieval_hits: List[RetrievalItem]    ◀── full structured results  │  ││  │   retrieval_summary: str                 ◀── compact for prompts      │  ││  └───────────────────────────────────────────────────────────────────────┘  ││                                                                             ││  ┌───────────────────────────────────────────────────────────────────────┐  ││  │ TOOLS                                                                 │  ││  │   tool_calls: List[ToolCall]             ◀── what was requested       │  ││  │   tool_results: List[ToolResult]         ◀── what came back           │  ││  │   artifact_refs: List[ArtifactRef]       ◀── persisted outputs        │  ││  └───────────────────────────────────────────────────────────────────────┘  ││                                                                             ││  ┌───────────────────────────────────────────────────────────────────────┐  ││  │ TUTORING                                                              │  ││  │   student_profile: StudentProfile        ◀── novice/medium/expert     │  ││  │   session_state: SessionState            ◀── attempts, performance    │  ││  │   hint_level: int                        ◀── escalation tracking      │  ││  └───────────────────────────────────────────────────────────────────────┘  ││                                                                             ││  errors: List[str]                ◀── captured, not hidden                  ││  route: str                       ◀── deterministic control flow            ││                                                                             │└─────────────────────────────────────────────────────────────────────────────┘

If you want a tool-using system to be reliable, you need to be able to answer questions like: What did the model see at this step? Why did it call this tool? What did the tool return? What changed after that? So the runtime carries a single state object with conversation messages, the image reference, retrieval hits and a compact retrieval summary, tool calls and tool results, artifact references for persisted outputs, errors captured rather than hidden, tutoring fields for the student model and session progress, and a routing field that makes control flow deterministic.

This design has two practical benefits. First, I can debug failures as data rather than intuition. Second, I can run offline loops—SFT, RL, ACE—that learn from the same state and trace structure.

In early iterations, I tried the "just append tool results to the prompt" approach. It breaks in predictable ways. Retrieval results are verbose, so the prompt bloats. Tool outputs are heterogeneous, so formatting becomes inconsistent. The model begins to lose the thread of the task—a phenomenon often called context rot. Making state explicit and structured was the turning point that made the system scalable.

The runtime loop

The runtime is intentionally simple: pre-retrieve to gather context, decide to choose the next action, execute exactly one tool if requested and persist outputs, then either teach (probe, hint, micro-lesson) or produce a final answer.

state = init_state(user_input, image)state.retrieval = pre_retrieve(state)while True:    action = decide(state)          # stochastic: model-driven policy    if action.type == "TOOL":        result = tool_exec(action, state)   # deterministic tool handler        state = state.with_tool_result(action, result)        continue    if action.type == "TUTOR":        state = state.with_tutor_message(apply_gates(action, state))        break    if action.type == "ANSWER":        state = state.with_answer(run_solver(state))        breakreturn state

Where the "agent" feeling comes from is the decide step—model-driven policy—but the system remains stable because routing is deterministic. Two details matter in practice: one tool per turn, which prevents tool spam, keeps traces readable, and makes the policy learnable; and explicit termination, meaning every run ends with a coherent "what happens next."

Without hard boundaries, a stochastic controller can get into loops. Zoom, zoom, zoom with slightly different bounding boxes each time. Retrieve, retrieve with the same query returning the same hits. Web search after web search because it "feels" like progress. The runtime constraints are not just for neatness—they prevent the most common agent failure: doing busy work until the budget runs out.

Why I split "controller" from "solver"

I treat "what to do next" as a fundamentally different problem from "produce the best medical answer."

┌─────────────────────────────────────────────────────────────────────────────┐│                                                                             ││   ┌─────────────────────────────────────────────────────────────────────┐   ││   │                         CONTROLLER(S)                               │   ││   │                                                                     │   ││   │   Input: state (image, history, retrieval, student signals)         │   ││   │                                                                     │   ││   │   Output: next action                                               │   ││   │     • TOOL: { name, args }                                          │   ││   │     • TUTOR: { type: probe | hint | microlesson | ... }             │   ││   │     • ANSWER: trigger solver                                        │   ││   │                                                                     │   ││   │   Trained via: SFT (format), RL (behavior), ACE (rules)             │   ││   └─────────────────────────────────────────────────────────────────────┘   ││                              │                                              ││                              │ "ready to answer"                            ││                              ▼                                              ││   ┌─────────────────────────────────────────────────────────────────────┐   ││   │                           SOLVER                                    │   ││   │                                                                     │   ││   │   Input: state (image, retrieval, tool artifacts, question)         │   ││   │                                                                     │   ││   │   Output: final medical answer                                      │   ││   │     • answer text                                                   │   ││   │     • confidence                                                    │   ││   │     • OR: "insufficient evidence" (triggers retry path)             │   ││   │                                                                     │   ││   │   Model: MedGemma (vision-language)                                 │   ││   └─────────────────────────────────────────────────────────────────────┘   ││                                                                             │└─────────────────────────────────────────────────────────────────────────────┘

The controller decides what to do next—tool call versus tutor action versus answer. The solver produces the final medical answer once enough evidence exists. This split buys me faster iteration on policy behavior, cleaner debugging since I can distinguish policy errors from solver errors, and the ability to improve tool-use without destabilizing answer generation.

Tutoring systems have an uncomfortable truth: sometimes the right output is "I don't have enough evidence." I explicitly preserve that as a valid outcome and build a retry/repair path around it. When a solver is forced to always answer, it will hallucinate. When it can refuse, the rest of the system can respond by collecting more evidence.

Tools: evidence gathering with artifacts

In tool-using systems, the failure mode is almost always the same. Tools become just strings in a prompt. Results get inlined as huge blobs. Context collapses. The "agent" starts hallucinating tool outputs.

┌─────────────────────────────────────────────────────────────────────────────┐│                            TOOL REGISTRY                                    │├─────────────────────────────────────────────────────────────────────────────┤│                                                                             ││  ┌─────────────┐  ┌─────────────┐  ┌─────────────┐  ┌─────────────────────┐ ││  │    zoom     │  │   enhance   │  │   segment   │  │   image_findings    │ ││  │ ─────────── │  │ ─────────── │  │ ─────────── │  │ ─────────────────── │ ││  │ bbox crop   │  │ upscale     │  │ MedSAM2     │  │ describe features   │ ││  │ + padding   │  │ + sharpen   │  │ mask gen    │  │ in region           │ ││  └─────────────┘  └─────────────┘  └─────────────┘  └─────────────────────┘ ││                                                                             ││  ┌─────────────┐  ┌─────────────┐  ┌───────────────────────────────────────┐││  │  retrieve   │  │ web_browser │  │              ocr                      │││  │ ─────────── │  │ ─────────── │  │ ───────────────────────────────────── │││  │ query local │  │ fallback    │  │ extract text from image regions       │││  │ database    │  │ web search  │  │                                       │││  └─────────────┘  └─────────────┘  └───────────────────────────────────────┘││                                                                             │├─────────────────────────────────────────────────────────────────────────────┤│                            TOOL EXECUTION FLOW                              │├─────────────────────────────────────────────────────────────────────────────┤│                                                                             ││   ToolCall                        ToolResult                                ││  ┌──────────────┐               ┌──────────────────────────────────────┐    ││  │ name: "zoom" │               │ ok: true                             │    ││  │ args: {      │    ──────▶    │ summary: "Zoomed into upper-left…"   │    ││  │   bbox_2d,   │   (execute)   │ artifact_refs: [                     │    ││  │   padding    │               │   "artifact://zoom_abc123.png"       │    ││  │ }            │               │ ]                                    │    ││  └──────────────┘               └──────────────────────────────────────┘    ││                                          │                                  ││                                          ▼                                  ││                                 ┌─────────────────┐                         ││                                 │ ARTIFACT STORE  │                         ││                                 │ (images, masks, │                         ││                                 │  text files)    │                         ││                                 └─────────────────┘                         ││                                                                             │└─────────────────────────────────────────────────────────────────────────────┘

I built tools as first-class operations. Tools have schemas. They return structured outputs and a human-readable summary. Tool outputs are persisted as artifacts—images, masks, text—referenced by ID or path instead of being dumped into the prompt.

The design principle is simple: summaries are for prompts, artifacts are for truth. The summary goes into the controller's context; the artifact gets persisted for humans and offline evaluators. This avoids a classic failure mode where a system becomes "prompt-native" and loses its grounding in actual data.

Tools fail for boring reasons: no image present, missing optional dependency, malformed bounding box, empty crop after clamping, retrieval index unavailable, network keys missing for web search. These boring failures destroy a demo if you don't capture them cleanly. I made tool failure a first-class output so the system can log it, surface it, and recover without crashing the run.

Retrieval: multimodal fusion under a context budget

Retrieval in this system has two jobs: return evidence in the form of text snippets and image provenance, and produce a compact, controller-friendly representation.

┌─────────────────────────────────────────────────────────────────────────────┐│                       MULTIMODAL RETRIEVAL PIPELINE                         │├─────────────────────────────────────────────────────────────────────────────┤│                                                                             ││                           ┌───────────────┐                                 ││                           │    Query      │                                 ││                           │ (text+image)  │                                 ││                           └───────┬───────┘                                 ││                                   │                                         ││                    ┌──────────────┼──────────────┐                          ││                    ▼              ▼              ▼                          ││             ┌───────────┐  ┌───────────┐  ┌───────────┐                     ││             │   BM25    │  │   Text    │  │   Image   │                     ││             │  (sparse) │  │  Vectors  │  │  Vectors  │                     ││             │           │  │  (dense)  │  │  (dense)  │                     ││             └─────┬─────┘  └─────┬─────┘  └─────┬─────┘                     ││                   │              │              │                           ││                   │   ranked     │   ranked     │   ranked                  ││                   │   results    │   results    │   results                 ││                   │              │              │                           ││                   └──────────────┼──────────────┘                           ││                                  ▼                                          ││                        ┌─────────────────┐                                  ││                        │   RRF Fusion    │                                  ││                        │ (rank-based     │                                  ││                        │  aggregation)   │                                  ││                        └────────┬────────┘                                  ││                                 │                                           ││                    ┌────────────┴────────────┐                              ││                    ▼                         ▼                              ││          ┌─────────────────┐      ┌─────────────────────────┐               ││          │ retrieval_hits  │      │   retrieval_summary     │               ││          │ (full structs)  │      │   (compact for prompt)  │               ││          └─────────────────┘      └─────────────────────────┘               ││                                                                             │└─────────────────────────────────────────────────────────────────────────────┘

I use multimodal retrieval: text retrieval with BM25 plus dense vectors, image retrieval with dense vectors, fused via a rank-based aggregation strategy so one modality doesn't steamroll the other. Crucially, the controller does not get everything. It gets top hits as structured data and a short summary string that fits comfortably in a prompt.

This is a real engineering constraint. If you dump raw retrieved documents into the prompt, you will destroy the agent's ability to think.

I'm strict about retrieval order: prefer the local database before web, use web only when the database is thin or missing key facts. This policy shows up in three places: in the prompt policy telling the model what to do, in runtime constraints defining what tools exist and when they're appropriate, and in RL shaping that penalizes web calls when database confidence is high. The same principle is reinforced by multiple layers rather than relying on a single prompt instruction.

Naively merging text and image retrieval can backfire. Dense similarity can return superficially similar but clinically irrelevant items. BM25 can overweight short keyword overlaps. Image retrieval can "look right" but be wrong in context. Fusion helped, but the bigger win was controlling what the controller sees: a small number of hits plus a compact summary that doesn't drown out the question.

Student modeling: personalization without pretending I solved cognition

I intentionally avoided unsupported "we solved knowledge tracing" claims.

┌─────────────────────────────────────────────────────────────────────────────┐│                            STUDENT MODELING                                 │├─────────────────────────────────────────────────────────────────────────────┤│                                                                             ││   ┌─────────────────────────────────────────────────────────────────────┐   ││   │                       StudentProfile                                │   ││   │   level: novice | medium | expert                                   │   ││   │   (set at session start, influences scaffolding depth)              │   ││   └─────────────────────────────────────────────────────────────────────┘   ││                                                                             ││   ┌─────────────────────────────────────────────────────────────────────┐   ││   │                       SessionState                                  │   ││   │   attempt_count: int           ◀── how many tries so far            │   ││   │   performance: correct | partial | wrong | none                     │   ││   │   misconceptions: List[str]    ◀── detected errors in reasoning     │   ││   │   consecutive_wrong: int       ◀── escalation trigger               │   ││   └─────────────────────────────────────────────────────────────────────┘   ││                                                                             ││   ┌─────────────────────────────────────────────────────────────────────┐   ││   │                    Tutoring State Machine                           │   ││   │                                                                     │   ││   │   hint_level: 0 ──▶ 1 ──▶ 2 ──▶ 3                                   │   ││   │                                  │                                  │   ││   │                                  ▼                                  │   ││   │                            (reveal allowed)                         │   ││   │                                                                     │   ││   │   reveal_gate: closed until attempt_count > 0 AND                   │   ││   │                 (performance == correct OR hint_level == max)       │   ││   └─────────────────────────────────────────────────────────────────────┘   ││                                                                             │└─────────────────────────────────────────────────────────────────────────────┘

My student modeling is built around what a system can actually observe during a session: a student profile indicating novice, medium, or expert level; attempt count; performance signals like correct, partial, or wrong; misconceptions when detectable; and a tutoring state machine that tracks hint level escalation and reveal gating.

This delivers visible behavior differences immediately. Novices get more probing and scaffolding. Experts get faster convergence and higher-level prompts. The system refuses to dump the answer on turn zero. That last point is not a prompt trick—it's enforced by gates.

For a tutoring system, you can either build a long-horizon mastery estimator early or build a robust session-level tutor policy first. I chose the second path because it produces visible, demoable behavior quickly, it's easier to verify, and it creates the scaffolding you need for knowledge tracing later: structured student signals, explicit state, traceable decisions.

The Socratic tutor policy: actions, not essays

Instead of generating free-form tutoring transcripts, I modeled tutoring decisions as one action per step: ASK_PROBE, HINT at levels 1 through 3, MICROLESSON, QUIZ, REQUEST_TOOL, REVEAL_ANSWER, or SAFETY_REFUSE.

┌─────────────────────────────────────────────────────────────────────────────┐│                         TUTOR ACTION SPACE                                  │├─────────────────────────────────────────────────────────────────────────────┤│                                                                             ││  ASK_PROBE      "What do you observe in the image?"                         ││  HINT (1-3)     Progressively more specific guidance                        ││  MICROLESSON    Brief teaching moment on a concept the student is missing   ││  QUIZ           Follow-up question to check understanding                   ││  REQUEST_TOOL   "I need to examine this region closer"                      ││  REVEAL_ANSWER  Give the answer (only after gates allow it)                 ││  SAFETY_REFUSE  "This requires clinical judgment—consult a physician"       ││                                                                             │├─────────────────────────────────────────────────────────────────────────────┤│                          GATING LAYER                                       │├─────────────────────────────────────────────────────────────────────────────┤│                                                                             ││   Controller ──▶ [proposed action] ──▶ GATE CHECK ──▶ [actual action]       ││                                                                             ││   REVEAL_GATE: Block reveal if attempt_count == 0                           ││   ASSESSMENT_GATE: Convert reveal to hint if performance != correct         ││                                                                             │└─────────────────────────────────────────────────────────────────────────────┘

That's a design choice that pays off everywhere. It's easy to validate, easy to log, easy to train with SFT, and easy to optimize with RL and ACE.

Even if a controller tries to reveal too early, runtime gates can override. No reveal until an attempt exists. Wrong or partial performance converts "reveal" into calibrated hints. Repeated tool calls within a single step are discouraged. This is how I guarantee the system behaves like a tutor even under stochastic policies.

If you're tutoring from a medical image, it's dangerously easy for a system to skip straight to a diagnosis or answer choice. I explicitly keep a pathway that produces objective findings without committing to a diagnosis. This supports tutoring behaviors like "Describe what you see first" or "Which finding most narrows the differential?" or "What would you expect next?" This observation-versus-diagnosis distinction is one of the most important tutor-versus-solver differences in the entire system.

SFT: training the controller to output valid tutoring actions

SFT is about reliability: consistent action formatting, fewer invalid outputs, better alignment to the action schema. The training target is not a perfect tutor transcript. The target is the next action.

┌─────────────────────────────────────────────────────────────────────────────┐│                      SFT: ACTION-ONLY CONTROLLER TRAINING                   │├─────────────────────────────────────────────────────────────────────────────┤│                                                                             ││   INPUT:                                                                    ││   ┌─────────────────────────────────────────────────────────────────────┐   ││   │ student_profile: { level: "novice" }                                │   ││   │ kc_tags: ["cardiac", "imaging", "tamponade"]                        │   ││   │ history: [                                                          │   ││   │   { role: "user", content: "What's wrong with this chest X-ray?" }, │   ││   │   { role: "assistant", content: "What features do you observe?" },  │   ││   │   { role: "user", content: "The heart looks big?" }                 │   ││   │ ]                                                                   │   ││   │ tool_summaries: [ "Zoomed into cardiac silhouette region" ]         │   ││   │ assessment: { performance: "partial", misconceptions: [] }          │   ││   └─────────────────────────────────────────────────────────────────────┘   ││                                                                             ││   TARGET OUTPUT:                                                            ││   ┌─────────────────────────────────────────────────────────────────────┐   ││   │ {                                                                   │   ││   │   "action": "HINT",                                                 │   ││   │   "level": 1,                                                       │   ││   │   "content": "Good observation. Can you describe the shape of       │   ││   │              the cardiac silhouette more specifically?"             │   ││   │ }                                                                   │   ││   └─────────────────────────────────────────────────────────────────────┘   ││                                                                             ││   Training Method: LoRA (lightweight adaptation)                            ││                                                                             │└─────────────────────────────────────────────────────────────────────────────┘

If the runtime expects a structured action and the model outputs a 2,000-token lecture instead, you can't gate or route safely. So I trained the controller on examples that include student profile, coarse knowledge-component tags, conversation history, tool summaries, student assessment signals, and the desired next tutoring action. This is the same structure used at runtime—the training distribution matches the control interface.

The goal was to make the controller reliably output valid structured actions that the runtime can gate and route. The format is action-only, JSON-shaped outputs with one decision per step. Training used lightweight adaptation via LoRA so I could iterate quickly.

What changed after SFT: fewer invalid outputs, more consistent hint calibration across levels, and less style drift into answer-dumping. The runtime became a stable interface, and the controller became something I could upgrade without rewriting the system.

RL: training tool-use as sequential decision making

Tool use is not a single decision. It's a sequence: decide to zoom, observe, decide to retrieve, observe, decide to answer. If you only reward the final answer, tool learning becomes painfully sparse.

┌─────────────────────────────────────────────────────────────────────────────┐│                    RL: TOOL-USE AS SEQUENTIAL DECISION MAKING               │├─────────────────────────────────────────────────────────────────────────────┤│                                                                             ││   Episode: One medical case/question                                        ││                                                                             ││   ┌─────────┐      ┌─────────┐      ┌─────────┐      ┌─────────┐            ││   │ State₀  │──▶──│ Action₀ │──▶──│ State₁  │──▶──│ Action₁ │──▶ ...        ││   │ (init)  │      │ (zoom)  │      │ (after  │      │(retrieve)│           ││   └─────────┘      └────┬────┘      │  zoom)  │      └────┬────┘            ││                         │           └─────────┘           │                 ││                         ▼                                 ▼                 ││                    ┌─────────┐                       ┌─────────┐            ││                    │Reward₀  │                       │Reward₁  │            ││                    │+0.1     │                       │+0.2     │            ││                    │(curious)│                       │(relevant)│           ││                    └─────────┘                       └─────────┘            ││                                                                             │├─────────────────────────────────────────────────────────────────────────────┤│                          REWARD SHAPING                                     │├─────────────────────────────────────────────────────────────────────────────┤│                                                                             ││   CURIOSITY REWARD: +reward when image tool produces meaningfully new view  ││   ANTI-REPEAT PENALTY: -penalty when same tool called with same/similar args││   RETRIEVAL QUALITY: +reward when retrieval returns relevant evidence       ││   SAFE FALLBACK: -penalty when web called AND local DB has high-conf hits   ││   TERMINAL REWARD: +reward for correct answer,-penalty for incorrect/timeout││                                                                             │└─────────────────────────────────────────────────────────────────────────────┘

So I added intermediate rewards that make tool-use learnable. Curiosity reward gives positive signal when an image tool produces a meaningfully new view. Anti-repeat penalty discourages repeated identical tool calls. Retrieval shaping rewards queries that return relevant evidence. Safe fallback shaping discourages web calls when the local database has high-confidence hits.

This turns tool-use learning from "pray the gradient finds it" into a tractable training signal.

I treated RL as an engineering problem rather than a label. One episode is one case/question. Actions are tool calls or terminal answers. Observations include the current view of the image via artifacts plus compact textual context. Termination happens when an answer is produced, the tool budget is exhausted, or there's an unrecoverable parse failure.

What changed after RL: the planner learned to "look first" instead of guessing, retrieval became the default knowledge source, and tool sequences became shorter and more purposeful. Tool use stopped being a brittle prompt behavior and became a policy that can improve with data and rewards.

ACE: playbooks that evolve without retraining

Not every failure requires gradient updates. Many failures are policy mistakes that can be captured as explicit rules: use retrieval before web, if the image is ambiguous zoom first, don't exceed one tool per turn, if you failed to parse a tool call simplify output formatting.

┌─────────────────────────────────────────────────────────────────────────────┐│                   ACE: AGENTIC CONTEXT ENGINEERING                          │├─────────────────────────────────────────────────────────────────────────────┤│                                                                             ││   ┌─────────────────────────────────────────────────────────────────────┐   ││   │                        ACE LOOP                                     │   ││   │                                                                     │   ││   │   ┌──────────┐     ┌──────────┐     ┌──────────┐     ┌──────────┐   │   ││   │   │   Run    │────▶│ Analyze  │────▶│  Curate  │────▶│  Update  │   │   ││   │   │ samples  │     │ failures │     │   ops    │     │ playbook │   │   ││   │   └──────────┘     └──────────┘     └──────────┘     └────┬─────┘   │   ││   │        ▲                                                  │         │   ││   │        └──────────────────────────────────────────────────┘         │   ││   │                      (next iteration)                               │   ││   └─────────────────────────────────────────────────────────────────────┘   ││                                                                             │├─────────────────────────────────────────────────────────────────────────────┤│                           PLAYBOOK EXAMPLE                                  │├─────────────────────────────────────────────────────────────────────────────┤│                                                                             ││   # Tool Policy Rules                                                       ││   1. Always use `retrieve` before `web_browser`                             ││   2. If the image region is ambiguous, use `zoom` first                     ││   3. One tool per turn maximum                                              ││   4. If tool parsing fails, simplify output format                          ││   5. When evidence is insufficient, request a tool—do not guess             ││                                                                             │└─────────────────────────────────────────────────────────────────────────────┘

ACE gives me a clean loop: run the system on samples, analyze failures including tool failures, update a playbook, inject the playbook into controller prompts on the next run.

This is powerful because it is fast to iterate, transparent since the playbook is readable, and complementary to SFT and RL since it patches behavior while training catches up.

ACE is my "fast patch lane." Instead of waiting for a training cycle, I can run the system on a batch, detect systematic failure modes like tool parsing errors, tool ordering mistakes, and retrieval misuse, and evolve a playbook that steers the controller away from repeating those mistakes.

What changed after ACE iterations: fewer tool-policy violations, better ordering with retrieve before web, and better behavior under uncertainty where the system requests a tool when evidence is unclear. It turns repeated mistakes into explicit, reviewable rules and gives the system a path to improve even when training is expensive.

Observability: if you can't debug it, it's not done

Every time I've seen an "agent" fail in the wild, the root cause is the same: nobody can reconstruct what happened.

┌─────────────────────────────────────────────────────────────────────────────┐│                          OBSERVABILITY STACK                                │├─────────────────────────────────────────────────────────────────────────────┤│                                                                             ││                              Runtime Execution                              ││                                    │                                        ││            ┌───────────────────────┼───────────────────────┐                ││            ▼                       ▼                       ▼                ││   ┌─────────────────┐    ┌─────────────────┐    ┌─────────────────┐         ││   │  Trace Events   │    │   Tool Cards    │    │   Artifacts     │         ││   │ • step_id       │    │ • tool_name     │    │ • artifact_id   │         ││   │ • state_before  │    │ • arguments     │    │ • type          │         ││   │ • action        │    │ • summary       │    │ • path          │         ││   │ • observation   │    │ • ok/error      │    │ • metadata      │         ││   │ • state_after   │    │ • duration_ms   │    │                 │         ││   │ • timestamp     │    │                 │    │                 │         ││   └────────┬────────┘    └────────┬────────┘    └────────┬────────┘         ││            │                      │                      │                  ││            └──────────────────────┼──────────────────────┘                  ││                                   ▼                                         ││                        ┌─────────────────────┐                              ││                        │     Trace Store     │                              ││                        └──────────┬──────────┘                              ││                                   │                                         ││            ┌──────────────────────┼──────────────────────┐                  ││            ▼                      ▼                      ▼                  ││   ┌─────────────────┐   ┌─────────────────┐   ┌─────────────────────────┐   ││   │  Deterministic  │   │     Audit       │   │  Targeted Improvement   │   ││   │     Replay      │   │  (provenance)   │   │  (feed into SFT/RL/ACE) │   ││   └─────────────────┘   └─────────────────┘   └─────────────────────────┘   ││                                                                             │└─────────────────────────────────────────────────────────────────────────────┘

So I built the system to emit per-step trace events with state before and after, action, and observation. It emits tool cards with arguments, summary, and errors. It emits retrieval hits. And it persists artifacts like images, masks, and text outputs.

This makes it possible to replay failures deterministically, audit whether a claim is evidence-backed, and do targeted improvements instead of random prompt tweaks.

The most important debugging feature is replayability. When an agent fails, there are always two questions: what happened, and what should change. Replayability answers the first question with certainty. And once you can replay, you can build the second half of the system: targeted improvements through SFT, RL shaping, or ACE playbook updates.

The decisions that shaped the system

One tool per turn. It's a constraint that makes everything else possible. It bounds latency. It bounds context. It makes policy learning easier. And it keeps traces readable.

Controller is not solver. This is what lets the system be both a stable medical answerer and a continuously improving tutor. I can iterate on policy behavior without destabilizing answer generation.

Student modeling is session-grounded. I'm modeling what I can observe reliably in-session: attempts, correctness, misconceptions when detectable, plus profile-based priors. It's enough to produce observable personalization without over-claiming.

Optimize with multiple knobs. I didn't bet everything on one method. SFT makes outputs reliable. RL makes tool-use behavior learnable. ACE patches recurring mistakes quickly. Together, they form a robust improvement stack.

The training recipe

┌─────────────────────────────────────────────────────────────────────────────┐│                         THE BUILD RECIPE                                    │├─────────────────────────────────────────────────────────────────────────────┤│                                                                             ││   Step A          Step B          Step C          Step D                    ││  ┌───────┐       ┌───────┐       ┌───────┐       ┌───────┐                  ││  │Baseline│──────▶│ State │──────▶│Retriev│──────▶│ Tools │                 ││  │        │       │+Traces│       │  al   │       │+Artif │                 ││  └───────┘       └───────┘       └───────┘       └───────┘                  ││      │               │               │               │                      ││      ▼               ▼               ▼               ▼                      ││  "Can we        "Can we        "Can we        "Can we                       ││   measure?"      debug?"        ground?"       inspect?"                    ││                                                                             ││   Step E          Step F          Step G          Step H                    ││  ┌───────┐       ┌───────┐       ┌───────┐       ┌───────┐                  ││  │Student│──────▶│  SFT  │──────▶│  RL   │──────▶│  ACE  │                  ││  │Model+ │       │       │       │       │       │       │                  ││  │ Gates │       │       │       │       │       │       │                  ││  └───────┘       └───────┘       └───────┘       └───────┘                  ││      │               │               │               │                      ││      ▼               ▼               ▼               ▼                      ││  "Can we        "Is output      "Can policy    "Can we fix                  ││   tutor?"        reliable?"      improve?"      fast?"                      ││                                                                             │├─────────────────────────────────────────────────────────────────────────────┤│   Core Principle: Make the runtime a stable interface first—then            ││   optimize behavior against that interface.                                 │└─────────────────────────────────────────────────────────────────────────────┘

If I had to compress the build into one principle, it would be this: make the runtime a stable interface first, then optimize behavior against that interface.

Most agent projects invert this. They start by tuning prompts and policies, then add tooling later. That works for a demo, but it collapses once you add multiple tools, retrieval, tutoring behaviors, and long multi-turn traces. I ended up with a recipe that deliberately moves in the opposite direction.

Step A: establish a measurable baseline. Before adding tools or retrieval, I started with the simplest solvable loop: image plus question to answer. This isn't because it's the end goal. It's because you need a control group. Once you have a baseline, you can answer questions like: did retrieval actually help, or just add latency? Did tools improve evidence gathering, or just increase failure surface? Did tutoring gates improve pedagogy, or just prevent answers?

Step B: define the system boundary as explicit state plus traces. The next step is not "add more capability." It's "make behavior observable." I converted the runtime into an explicit state machine where every step has a before-state, an action, an observation, and an after-state. This is the point where the project stops being prompt engineering and becomes systems engineering. Once traces exist, failures become data.

Step C: introduce retrieval as a first-class dependency. Retrieval adds value, but it also adds a context budget problem. The key pattern that made retrieval usable was storing full hits as structured evidence but injecting only a compact summary into controller prompts. This keeps the controller thinking on the actual question rather than being drowned in text.

Step D: add tools with artifact discipline. Tools are only useful if they're trustworthy. The design that worked was tools returning a short summary for prompts and persisting outputs as artifacts. Prompts stay small. Humans can audit what happened. Offline loops can score behavior without re-running everything.

Step E: turn the solver into a tutor. At this point, the system can answer. But tutoring requires control. I treated student modeling as session-grounded signals: student profile, attempt count, performance grade, misconceptions when detectable, and tutoring state for hint escalation. Then I enforced pedagogy with hard gates: don't reveal on turn zero, wrong or partial turns become hints and probes, and objective findings can be separated from diagnosis.

Step F: SFT for reliability. Once the runtime interface is stable, SFT becomes straightforward. The model is trained to output the next structured action, conditioned on the same information the runtime uses. The effect is not "the model got smarter." The effect is fewer invalid actions, less formatting drift, and more predictable control flow.

Step G: RL for behavior. Tool use is sequential, and terminal-only rewards are too sparse. The recipe that worked was adding step shaping rewards that encode the behaviors the system needs: curiosity reward when visual tools create a new informative view, penalties for repeats and tool spam, reward retrieval when it produces relevant evidence, penalize web usage when local retrieval is clearly sufficient.

Step H: ACE for fast iteration. Even with training, some failures are not model skill problems—they're policy mistakes that recur. Ordering mistakes like web before retrieval. Repeated formatting failures. Avoidable tool misuse. ACE turns these into explicit, reviewable rules that get injected into the controller context.

Failures that forced the architecture

Context collapse. The model forgets the question after large retrieval or tool outputs, or starts repeating generic answers. The fix: compact summaries for prompts, artifacts for everything else, and strict budgets like one tool per turn.

Tool spam and busy work. Multiple tools called without changing the outcome. Repeating the same action with minor variations. The fix: bounded tool calls, repeat penalties in RL, and routing that forces progress.

Brittle formatting. Malformed tool calls. Tool outputs that can't be parsed. Controllers drifting into unstructured text. The fix: action schemas, action-only training targets, and system-level gates that sanitize unsafe decisions.

Answer dumping. The system gives the correct option immediately. The fix: reveal gates and assessment gates, tutor actions as the primary interface, and a separation between objective findings and diagnosis.

Closing

┌─────────────────────────────────────────────────────────────────────────────┐│                                                                             ││                    WHAT THIS ARCHITECTURE ENABLES                           ││                                                                             │├─────────────────────────────────────────────────────────────────────────────┤│                                                                             ││   ┌─────────────┐   ┌─────────────┐   ┌─────────────┐   ┌─────────────┐     ││   │ Multi-turn  │   │    Under    │   │  With tools │   │ Debuggable  │     ││   │  tutoring   │   │ uncertainty │   │& retrieval  │   │& improvable │     ││   └──────┬──────┘   └──────┬──────┘   └──────┬──────┘   └──────┬──────┘     ││          │                 │                 │                 │            ││          └─────────────────┼─────────────────┼─────────────────┘            ││                            ▼                 ▼                              ││                  ┌───────────────────────────────────────┐                  ││                  │   Deterministic Runtime Interface     │                  ││                  │   (state, tools, traces, gates)       │                  ││                  └───────────────────┬───────────────────┘                  ││                                      │                                      ││                    ┌─────────────────┼─────────────────┐                    ││                    ▼                 ▼                 ▼                    ││              ┌───────────┐     ┌───────────┐     ┌───────────┐              ││              │    SFT    │     │    RL     │     │    ACE    │              ││              │ (format)  │     │ (behavior)│     │ (rules)   │              ││              └───────────┘     └───────────┘     └───────────┘              ││                                                                             ││                         ═══════════════════════                             ││                              A TUTOR,                                       ││                          NOT A CHATBOT                                      ││                         ═══════════════════════                             ││                                                                             │└─────────────────────────────────────────────────────────────────────────────┘

This architecture is designed to make a difficult problem tractable: multi-turn tutoring, under uncertainty, with tools and retrieval, while staying debuggable and improvable.

The core idea is simple. Treat the runtime as a stable interface—state, tools, traces, gates—then improve behavior through multiple complementary mechanisms: SFT for reliable formatting, RL for learned tool-use policies, and ACE playbooks for fast behavioral fixes. That structure is what makes the system behave like a tutor rather than a single-shot chatbot.

Building an Adaptive Tutoring System with RL and Knowledge Tracing

Wed, 07 Jan 2026 00:00:00 GMT

Note: This RSS feed strips out SVGs and embeds. You might want to read the post on the webpage here.

Combining knowledge tracing, reinforcement learning, and multi-agent orchestration to pick the right question for each learner.

January 7, 2026

projects

The Problem: Static Question Banks Don't Work

Every student learns differently. Some need more practice on fractions, others struggle with word problems. While many platforms use spaced repetition or Item Response Theory (IRT) to estimate difficulty, truly optimizing the sequence for maximum learning efficiency remains a hard open problem.

The result? Students waste time on questions they've already mastered, skip over gaps in their understanding, and also lose motivation when difficulty spikes unexpectedly.

What if we could select the optimal next question for each student at each moment?

This post walks through how I built an adaptive tutoring system that does exactly that.

Transferring to Any Domain

Before diving into the technical details, here's the key insight: this system is domain-agnostic. It works for math, medicine, law, programming—any domain with structured practice content.

What You Need

To deploy this system in a new domain, you need three data sources:

┌─────────────────────────────────────────────────────────────────────┐│                        REQUIRED DATA SOURCES                         │├─────────────────────────────────────────────────────────────────────┤│                                                                      ││  1. QUESTION BANK                                                    ││     ├── Question ID (unique identifier)                              ││     ├── Question Text (stem, choices, correct answer)                ││     ├── Knowledge Concepts (topics/skills the question tests)        ││     └── [Optional] Difficulty estimate, time estimate                ││                                                                      ││  2. INTERACTION LOGS                                                 ││     ├── User ID                                                      ││     ├── Question ID                                                  ││     ├── Correct/Incorrect                                            ││     ├── Timestamp                                                    ││     └── [Optional] Time spent, hints used                            ││                                                                      ││  3. KNOWLEDGE CONCEPT (KC) TAXONOMY                                  ││     ├── KC ID                                                        ││     ├── KC Name/Description                                          ││     └── [Optional] Prerequisites, hierarchy                          ││                                                                      │└─────────────────────────────────────────────────────────────────────┘

Minimum Viable Dataset

Component	Minimum Size	Recommended
Questions	500+	5,000+
Interactions	100K+	1M+
Users	1K+	10K+
KCs	20+	100+

The more interaction data you have, the better the KT model will calibrate. With less data, you can still use the heuristic scorer (Part 2) without training RL (Part 3).

Data Transformation Pipeline

Raw Content → Embeddings → KT Training → RL Training → Deployment

Step 1: Generate Embeddings

Run a pretrained language model (BERT, sentence-transformers) on your question text
Output: 768-dim vector per question
Also embed KC descriptions to get KC embeddings

Step 2: Map Questions to KCs

If you have manual labels, use them
If not, you can cluster question embeddings and label clusters as KCs
Or use an LLM to auto-annotate (e.g., "What skills does this question test?")

Step 3: Format Interaction Logs

Required schema:

user_id, question_id, correct (0/1), timestamp

Sort by user, then by timestamp to get sequences

Step 4: Train KT Model

Train the LSTM on your interaction sequences
Calibrate to predict KC-level mastery
Output: Model checkpoint (~10MB)

Step 5: (Optional) Train RL Policy

Use the trained KT model as the environment
Simulate student trajectories and optimize for learning gain
This step is optional—the heuristic scorer works without it

Example: Medical Licensure

For a USMLE prep platform:

Questions: Board-style vignettes (text + answer choices)
KCs: Medical topics (cardiology, nephrology, pharmacology, etc.)
Interactions: Student practice history from your platform
Embeddings: Run BioBERT or PubMedBERT on question text

The same architecture applies—only the content changes.

System Architecture: Three Specialized Agents

I split the system into three agents. Each one has a single job:

┌─────────────────────────────────────────────────────────────────┐│                    SESSION ORCHESTRATOR                          ││  Coordinates the full answer → update → recommend cycle          ││  • Safety fallbacks • Audit logging • Session lifecycle          │└─────────────────────────────────────────────────────────────────┘                            │         ┌──────────────────┴──────────────────┐         ▼                                     ▼┌─────────────────────┐              ┌─────────────────────────┐│  KNOWLEDGE MODELER  │              │        PLANNER          ││  (CalibrationQDKT)  │              │  (Heuristic + RL)       │├─────────────────────┤              ├─────────────────────────┤│ • LSTM hidden state │─────────────▶│ • Candidate filtering   ││ • Mastery vectors   │  p(correct)  │ • Learning gain scoring ││ • Temporal decay    │  uncertainty │ • SAC/DQN RL policy     │└─────────────────────┘              └─────────────────────────┘

Why this decomposition?

Knowledge Modeler: Owns the student state. Answers "What does this student know?"
Planner: Owns question selection. Answers "What should we teach next?"
Orchestrator: Owns the session. Ensures safety, logging, and graceful fallbacks.

Part 1: Knowledge Tracing with CalibrationQDKT

The Idea: Student State as LSTM Hidden State

We model a student's knowledge as the hidden state of an LSTM that has processed their entire interaction history. Each interaction (question + response) updates the hidden state.

But first—where does this LSTM come from?

Training the KT Model

The knowledge tracing model is trained in two stages:

Stage 1: Base LSTM Training

The base model uses an LSTM (Long Short-Term Memory) architecture inspired by Deep Knowledge Tracing (DKT). Unlike traditional DKT that learns simple ID embeddings, this model takes pretrained BERT (Bidirectional Encoder Representations from Transformers) embeddings (768-dim) as input to capture the semantic content of each question. The LSTM processes these sequences to produce a 300-dimensional hidden state representing the student's knowledge.

Key hyperparameters:

768-dim embeddings (BERT, frozen)
300-dim hidden state (the student knowledge representation)
200 epochs on XES3G5M

Stage 2: KC Calibration

The base LSTM predicts correctness on individual questions. But we want KC-level (Knowledge Concept) mastery. The calibration stage freezes the LSTM and trains a small projection layer that maps the hidden state to KC predictions. This uses the KC embeddings (also 768-dim BERT vectors) as targets.

After calibration, the model can predict mastery for any KC embedding—not just questions it saw during training. This is what makes it useful for RL: the policy can query "what's the student's mastery on KC X?" without enumerating all questions.

LSTM Update Step

Now, back to inference. Each interaction updates the hidden state:

# From src/agents/knowledge_modeler/inference.pydef update_state(self, user_id: str, interaction: StudentInteraction):    # Get current LSTM state    h, c, lstm_out = self._state_cache[user_id]    # Encode interaction    question_emb = self.embeddings.get_question_embedding(interaction.question_id)    response = torch.tensor([[1 if interaction.correct else 0]])    # Update LSTM    response_emb = self.model.net.correctness_encoding(response.long())    x = question_emb.unsqueeze(1) + response_emb    new_lstm_out, (new_h, new_c) = self.model.net.lstm_layer(x, (h, c))    self._state_cache[user_id] = (new_h, new_c, new_lstm_out)

Predicting Correctness

Given a student's hidden state $h$ and a question embedding $e_q$, we predict the probability of correct response:

$$P(\text{correct} \mid h, e_q) = \sigma(W \cdot [h \| e_q] + b)$$

where $[h \| e_q]$ is the concatenation of the 300-dim hidden state with the 768-dim question embedding.

Temporal Decay: Students Forget

Knowledge decays over time. We model forgetting with exponential decay:

mt=m0⋅0.5t/τm_t = m_0 \cdot 0.5^{t / \tau}mt=m0⋅0.5t/τ

where:

m0m_0m0 is the original mastery
ttt is days since last practice
τ\tauτ is the half-life (default: 7 days)

# From src/agents/knowledge_modeler/decay.pydef _decay_factor(self, last_practiced: datetime, now: datetime) -> float:    elapsed_days = (now - last_practiced).total_seconds() / 86400.0    return 0.5 ** (elapsed_days / self.default_half_life_days)

This allows the system to identify "rusty" knowledge concepts that need refreshing.

Part 2: Expected Learning Gain Scoring

The Key Insight: Counterfactual Simulation

To pick the best question, we simulate what would happen if the student answered it. The key insight: treat the KT model as an environment you can query.

For each candidate question qqq:

Predict p=P(correct∣h,eq)p = P(\text{correct} \mid h, e_q)p=P(correct∣h,eq)
Simulate "what if correct?" → get new state $s^+$
Simulate "what if incorrect?" → get new state $s^-$
Compute expected learning gain:

$$\mathbb{E}[\Delta] = p \cdot \Delta^+ + (1-p) \cdot \Delta^-$$

where $\Delta^+ = \text{gain}(s^+) - \text{gain}(s)$ and $\Delta^- = \text{gain}(s^-) - \text{gain}(s)$.

         [Current State s]           /           \\    (Correct?)       (Incorrect?)       /                 \\[State s+]           [State s-]      |                   |  Gain(s+)            Gain(s-)      \\                   /    Weighted Expected Gain

# From src/agents/planner/scorer.pydef score_question(self, user_id: str, question_id: str) -> QuestionScore:    # Step 1: Predict p(correct)    p_correct, uncertainty = self.km.predict_correct(user_id, question_id)    # Step 2: Counterfactual simulation    state_if_correct = self.km.simulate_update(user_id, question_id, outcome=True)    state_if_incorrect = self.km.simulate_update(user_id, question_id, outcome=False)    # Step 3: Compute gains    gain_if_correct = self.km.compute_mastery_gain(current_state, state_if_correct)    gain_if_incorrect = self.km.compute_mastery_gain(current_state, state_if_incorrect)    # Step 4: Expected learning gain    expected_gain = p_correct * gain_if_correct + (1 - p_correct) * gain_if_incorrect    return QuestionScore(        question_id=question_id,        expected_learning_gain=expected_gain,        p_correct=p_correct,        ...    )

Zone of Proximal Development (ZPD) Bonus

The Zone of Proximal Development is a concept from educational psychology (Vygotsky, 1978). It refers to the sweet spot where a task is challenging enough to promote learning but not so hard that the learner gives up.

In our context, we operationalize ZPD as a probability range. If a student has a 95% chance of getting a question right, they probably already know the material—low learning value. If they have a 10% chance, they'll likely fail and get frustrated. The sweet spot is somewhere in the middle.

I use the range [0.55, 0.75] based on the "desirable difficulty" literature. The bonus function is:

$$\text{ZPD Bonus} = \begin{cases}
w \cdot (1 - \frac{|p - 0.65|}{0.1}) & \text{if } 0.55 \leq p \leq 0.75 \\
0 & \text{otherwise}
\end{cases}$$

This gives maximum bonus at $p = 0.65$ (the center of the range) and tapers linearly to zero at the edges. Questions outside the range get no bonus.

Why this matters: Without ZPD targeting, the scorer tends to recommend very hard questions (high learning delta) or very easy questions (high confidence). Neither is optimal for sustained engagement.

Part 3: Reinforcement Learning for Question Selection

Why RL?

The scoring approach in Part 2 is greedy—it picks the question with the highest immediate expected gain. But learning is a multi-step process. A question that seems suboptimal now might set up better learning opportunities later.

Reinforcement learning can optimize for long-term outcomes. Instead of maximizing immediate gain, we train a policy to maximize cumulative learning over an entire session (or even longer horizons).

Formulating the MDP

We cast question selection as a Markov Decision Process:

State: The 300-dim LSTM hidden state. This is a learned representation of everything the student has done so far.
Action: Which question to recommend next. This is where it gets interesting (see below).
Reward: Learning gain after the student answers. We measure this as the change in predicted correctness across a sample of questions.
Transition: Deterministic—the KT model tells us exactly how the state changes given a question and response.

The Action Space Problem

With 7,600+ questions, a naive discrete action space is huge. Two approaches:

Discrete (DQN): Train Q-values for each question. Works but doesn't generalize to unseen questions.
Continuous (SAC): Output a 768-dim embedding vector, then find the nearest question. Generalizes better and can handle new questions at test time.

I went with SAC (Soft Actor-Critic) because the continuous action space maps naturally to the semantic embedding space:

class SACActorSelector(RLPolicySelector):    def select_qid(self, km, user_id, candidates, embedding_store):        # Get student state        _, _, lstm_out = km._state_cache[user_id]        # Actor outputs embedding (not question ID!)        with torch.no_grad():            action_emb = self._actor(lstm_out)  # [1, 768]        # Find nearest question in candidate set        candidate_embs = torch.stack([            embedding_store.get_question_embedding(qid)            for qid in candidates        ])        similarities = F.cosine_similarity(action_emb, candidate_embs)        best_idx = similarities.argmax().item()        return candidates[best_idx]

The actor learns to output an "ideal question embedding"—a point in semantic space representing what the student should practice next. We then find the actual question closest to that point.

Training Details

Algorithm: SAC with automatic entropy tuning
Critic: Uses the same LSTM architecture to predict Q-values
Batch size: 256 trajectories
Training: ~50k episodes on simulated students (replay from real interaction logs)
Gamma: 0.99 (we care about long-term outcomes)

The KT model acts as the environment—we can simulate thousands of students without real human interaction.

Exam-Score Reward: Beyond KC-Level Metrics

The obvious reward is KC-level mastery delta: how much did the student's predicted mastery on the practiced KC improve?

But this is noisy. A single question might barely move KC mastery, making the reward signal sparse and hard to learn from.

I extended this with an "exam-score" reward:

$$R = \frac{1}{|Q_{\text{sample}}|} \sum_{q \in Q_{\text{sample}}} \left[ P(\text{correct} \mid s', e_q) - P(\text{correct} \mid s, e_q) \right]$$

Instead of measuring improvement on one KC, we measure improvement in predicted correctness across a sample of questions. This is more like what we actually care about: overall exam performance.

Probe Split: To avoid overfitting, we split questions into:

Train set: Used for reward computation during training
Probe set: Held out entirely, used only for evaluation

This prevents the policy from gaming the reward by memorizing which questions are in the sample.

Part 4: Multi-Objective Scoring

Beyond Pure Learning Gain

If you only optimize for learning gain, the system does weird things:

Recommends the same hard topic repeatedly (high delta, but frustrating)
Ignores topics the student should cover for an exam
Never gives "easy wins" to build confidence

Real tutoring has multiple objectives. The challenge is combining them sensibly.

The Five Objectives

I use a weighted linear combination:

$$\text{Score} = w_L \cdot L + w_C \cdot \text{Cov} + w_S \cdot S + w_{\text{Conf}} \cdot \text{Conf} + w_R \cdot R$$

Each objective is normalized to [0, 1]:

Objective	What It Measures	How It's Computed
L (Learning Gain)	Expected knowledge improvement	Counterfactual simulation (Part 2)
Cov (Coverage)	Blueprint/topic completion	% of required topics not yet practiced
S (Speed)	Preference for shorter questions	Inverse of estimated time-to-complete
Conf (Confidence)	Questions in the ZPD range	ZPD bonus function
R (Remediation)	Focus on weak/rusty areas	Overlap with flagged weak KCs

Coverage is particularly important for exam prep. If the student needs to cover 10 topics and has only seen 3, we should prioritize the missing 7 even if they're not the highest-gain options.

Speed helps when the student is frustrated or running low on time. Shorter questions provide quicker feedback loops.

Remediation targets KCs that the decay model flagged as "rusty" or that the student has historically struggled with.

Adaptive Weights by Mode

The weights aren't fixed—they depend on the session mode:

Mode	$w_L$	$w_C$	$w_S$	$w_{\text{Conf}}$	$w_R$
Learning	0.4	0.1	0.1	0.2	0.2
Exam	0.2	0.5	0.1	0.1	0.1
Review	0.1	0.1	0.1	0.0	0.7

Learning mode: Maximize learning, with some confidence and remediation
Exam mode: Heavy coverage weight—make sure all topics are hit
Review mode: Focus almost entirely on weak areas

Behavioral Adjustments

The weights also adapt to real-time signals:

def infer_weights(self, session: SessionSummary) -> AdaptiveWeights:    if session.mode == "exam":        weights.w_coverage = 0.5        weights.w_learning = 0.2    elif session.mode == "review":        weights.w_remediation = 0.7        weights.w_learning = 0.1    # Behavioral adjustments    if signals.skip_rate > 0.3:        weights.w_speed += 0.2  # Frustrated student → easier questions    if signals.accuracy < 0.4:        weights.w_confidence += 0.15  # Struggling → more ZPD targeting    return weights.normalize()

If the student is skipping a lot of questions (frustration signal), we increase the weight on speed and confidence. If accuracy drops, we stay closer to the ZPD.

Part 5: Quality Gates and LLM Supervision

Why Gates?

RL policies can behave unexpectedly, especially early in training or on out-of-distribution states. The SAC actor might output an embedding that maps to:

A question the student just answered (repeat)
A question way too hard or too easy
A question on a topic that's not relevant to the session goals

Rather than trust the policy blindly, we add a layer of deterministic quality checks.

The Gate Pipeline

[RL Policy] ──▶ (Action Embedding)                      │                [Nearest Neighbor]                      │                [Quality Gates] ──(Fail)──▶ [Refinement]                      │                   (Pass)                      │                 [Selected Question]

The RL policy proposes a question. The gates check if it's acceptable. If not, we either accept it anyway (logging the failure) or trigger refinement.

The Four Gates

Each gate is a simple boolean check:

1. No Recent Repeats

Don't recommend a question the student answered in the last N interactions
Window size is configurable (default: 20 questions)
Prevents the "drill on the same item" failure mode

2. KC Overlap

The selected question must cover at least one of the target KCs
Target KCs come from session goals or the remediation system
Prevents wandering off-topic

3. ZPD Range

The estimated $p(\text{correct})$ must be within [0.4, 0.85]
Wider than the ZPD bonus range—this is a hard constraint, not a soft preference
Prevents questions that are almost certainly too hard or too easy

4. Minimum Learning Gain

Expected learning gain must exceed a threshold (default: 0.01)
Filters out questions that the KT model predicts will have minimal impact
Catches cases where the nearest neighbor is semantically close but pedagogically useless

Supervisor Logic

When gates fail, what do we do? Two options:

Accept anyway: Log the failure and let the RL proposal through. Useful for exploration and collecting data on edge cases.
Trigger refinement: Re-run the scorer with stricter constraints and pick the best passing candidate.

The decision is made by a "supervisor"—either a simple heuristic or an LLM:

# RL proposes a questionrl_qid = policy_selector.select_qid(km, user_id, candidates, embeddings)# Evaluate quality gatesgate_results = evaluate_candidate(rl_qid, constraints)# Supervisor decidesif supervisor.run_refinement and not gate_results.all_passed:    # Re-rank with stricter constraints    refined_qid = run_refinement(scorer, stricter_constraints)else:    # Accept RL proposal even if gates fail    final_qid = rl_qid

Heuristic vs. LLM Supervisor

The heuristic supervisor uses simple rules:

Always refine if the repeat gate fails
Accept ZPD failures if confidence is high
Log everything

The LLM supervisor (optional) can reason about context:

"The student has been struggling—let's accept this easier question even though it's below ZPD"
"This question failed KC overlap but it's a good bridge to the target topic"

In practice, the heuristic works well. The LLM adds flexibility but also latency and cost.

Why This Pattern?

This "RL proposes, rules dispose" pattern has a few benefits:

Safety: Hard rules catch obviously bad recommendations
Debuggability: You can see exactly which gate failed and why
Graceful degradation: If RL is broken, the system falls back to heuristic scoring
Exploration: By sometimes accepting failures, we collect data on edge cases

Results & Demo

I built a Rich TUI demo to visualize the full pipeline:

The demo shows:

Student snapshot (mastery, weak KCs, recent accuracy)
Candidate funnel (7,652 → filtered → scored)
Top recommendations with score breakdowns
RL embedding space visualization

What I Learned

1. Treat the KT model as an RL environment

This was the biggest conceptual shift. The knowledge tracing model isn't just a predictor—it's a simulator. You can ask it "what if the student answered this question correctly?" and get a new state back. That turns question selection into planning, and planning is what RL is good at.

2. Counterfactual simulation is surprisingly cheap

I expected simulating "correct" and "incorrect" outcomes for every candidate would be slow. It's not. The LSTM forward pass is fast, and you're not backpropagating. Scoring 200 candidates takes ~50ms on a laptop CPU.

3. Multi-objective scoring matters more than I thought

My first version just maximized expected learning gain. It kept recommending hard questions because those had the highest delta. Students got frustrated.

Adding coverage (don't repeat topics), speed (shorter questions when struggling), and confidence (stay in the ZPD) made the recommendations feel sensible. The weights are simple rules, not learned—but they work.

4. RL needs guardrails

The SAC policy sometimes outputs embeddings that map to weird questions—repeats, or items way outside the student's level. I added hard quality gates (no repeats, KC overlap, ZPD range) that run after RL selection. If the policy picks something bad, we either accept it anyway (and log it) or re-rank with stricter constraints.

This "RL proposes, rules dispose" pattern kept the system usable while the policy was still learning.

5. Cold-start is solved by embeddings

New students have no history, so the LSTM is at its default state. But you still have the question embeddings. You can recommend based on target KCs, difficulty proxies, or just sample randomly. The system degrades gracefully.

6. Latency budget drove architecture choices

Everything runs in a single forward pass per candidate. No database lookups in the hot path—embeddings are preloaded. The full pipeline (filter → score → select) runs under 100ms for 200 candidates. That's fast enough for interactive use.

7. The reward function is the hard part

KC-level mastery delta seemed like the obvious reward. But it's noisy—small state changes produce small rewards, and the policy has trouble learning. The exam-score reward (improvement on a sampled set of questions) was more stable, but required careful train/probe splits to avoid leaking evaluation into training.

8. Domain transfer is just data

The architecture doesn't care if questions are about math, medicine, or law. You need:

Question embeddings (run BERT on your content)
KC annotations (can be automated with LLMs)
Interaction logs (student × question × correct/incorrect)

Swap those in, retrain the KT model, and the rest of the pipeline works.

What I'd do differently

Start with heuristics: The greedy scorer (no RL) was good enough for 80% of cases. I should have validated it more before jumping to RL.
Log everything: I added audit logging late. Should have done it from day one.
Smaller action space: SAC in 768-dim embedding space works, but discrete actions (DQN over question indices) might converge faster for small question banks.

References

Built on math education data (XES3G5M), but the architecture works for any domain—medical, law, whatever. Just swap in your question bank and embeddings.

When Regex Meets Fuzzy Matching: Building a Medical Document Classifier That Knows Who's Who

Wed, 10 Dec 2025 00:00:00 GMT

Note: This RSS feed strips out SVGs and embeds. You might want to read the post on the webpage here.

*Automating multi-patient fax sorting with OCR, entity extraction, and weighted clustering*

December 10, 2025

projects healthcare

Medical offices still run on faxes. Every day, insurance companies, labs, and specialists send multi-patient PDF batches—100 pages, 15 patients, all jumbled together. Staff spend hours manually sorting pages by squinting at headers, cross-referencing MRNs, and hoping they don't mix up John Smith #1 with John Smith #2.

Inspired by Tennr's approach to medical document automation, I built a document processing pipeline that ingests multi-patient PDFs, extracts identifiers from each page, clusters them into patient profiles, and outputs per-patient PDFs with confidence scores. The core insight: patient identity isn't a binary classification—it's a weighted similarity problem across multiple identifier types.

Here's what I learned building it.

The Problem: One PDF, Many Patients

A typical incoming fax looks like this:

Pages 1-3: Lab results for Mary Johnson, MRN 123456
Page 4: Prescription for Robert Garcia, DOB 05/15/1980
Pages 5-7: Referral for Mary Johnson (same patient, no MRN this time)
Pages 8-12: Mixed records, some with phone numbers only

The challenges stack up quickly. Some pages have MRNs, others only have name and DOB. Scanned documents introduce OCR noise—typos, artifacts, and formatting variations. And then there's ambiguity: when you see "M. Johnson" on page 6, is that Mary Johnson or someone else entirely?

Manual sorting fails at scale. Rule-based automation using exact string matching fails on real-world data. I needed something smarter—a system that could reason about degrees of similarity across multiple identifier types and make probabilistic judgments about page ownership.

The Architecture: A Six-Stage Pipeline

┌─────────────────────────────────────────────────────────────────────────────┐│                          DOCUMENT INGESTION                                  ││  ┌──────────────┐                                                           ││  │ Multi-Patient│  100+ pages, 10+ patients                                 ││  │ PDF Input    │  mixed labs, prescriptions, referrals                     ││  └──────┬───────┘                                                           │└─────────┼───────────────────────────────────────────────────────────────────┘          │          ▼┌─────────────────────────────────────────────────────────────────────────────┐│                          PREPROCESSING LAYER                                 ││                                                                              ││  ┌────────────────────┐      ┌────────────────────┐                         ││  │  Page Extractor    │      │  OCR Processor     │                         ││  │  ─────────────     │ ───▶ │  ─────────────     │                         ││  │  Poppler/pdf2image │      │  Tesseract (default)│                        ││  │  DPI: 200          │      │  OlmOCR (optional)  │                        ││  └────────────────────┘      └─────────┬──────────┘                         │└─────────────────────────────────────────┼───────────────────────────────────┘                                          │ raw text + word positions                                          ▼┌─────────────────────────────────────────────────────────────────────────────┐│                          ENTITY EXTRACTION                                   ││                                                                              ││  ┌────────────────────────────────────────────────────────────────────┐     ││  │  Entity Extractor (Regex-based)                                     │     ││  │  ┌──────────┐  ┌──────────┐  ┌──────────┐  ┌──────────┐            │     ││  │  │  Name    │  │   MRN    │  │   DOB    │  │  Phone   │            │     ││  │  │ "Patient:│  │ "MRN:    │  │ "DOB:    │  │ "(555)   │            │     ││  │  │  John.." │  │  123456" │  │  1/1/80" │  │  123-.." │            │     ││  │  └──────────┘  └──────────┘  └──────────┘  └──────────┘            │     ││  └────────────────────────────────────────────────────────────────────┘     │└─────────────────────────────────────────┬───────────────────────────────────┘                                          │ identifiers + confidence                                          ▼┌─────────────────────────────────────────────────────────────────────────────┐│                          PATIENT CLUSTERING                                  ││                                                                              ││  ┌──────────────────────┐         ┌──────────────────────┐                  ││  │  Fuzzy Matcher       │         │  Entity Linker       │                  ││  │  ─────────────       │  ────▶  │  ─────────────       │                  ││  │  RapidFuzz scoring   │         │  Anchor-based        │                  ││  │  Per-entity strategy │         │  patient clustering  │                  ││  └──────────────────────┘         └──────────┬───────────┘                  │└─────────────────────────────────────────────────┼───────────────────────────┘                                          │ patient profiles                                          ▼┌─────────────────────────────────────────────────────────────────────────────┐│                          PAGE ASSIGNMENT                                     ││                                                                              ││  ┌────────────────────────────────────────────────────────────────────┐     ││  │  Page Assigner (Weighted Scoring)                                   │     ││  │                                                                     │     ││  │  Score = 0.4×name + 0.3×mrn + 0.2×dob + 0.1×phone                  │     ││  │                                                                     │     ││  │  ┌─────────────┐  ┌─────────────┐  ┌─────────────┐                 │     ││  │  │ Assigned    │  │ Ambiguous   │  │ Unassigned  │                 │     ││  │  │ (conf ≥60%) │  │ (close race)│  │ (conf <60%) │                 │     ││  │  └─────────────┘  └─────────────┘  └─────────────┘                 │     ││  └────────────────────────────────────────────────────────────────────┘     │└─────────────────────────────────────────┬───────────────────────────────────┘                                          │                                          ▼┌─────────────────────────────────────────────────────────────────────────────┐│                          OUTPUT GENERATION                                   ││                                                                              ││  ┌────────────────────┐      ┌────────────────────┐                         ││  │  Document Splitter │      │  Outputs:          │                         ││  │  ─────────────     │ ───▶ │  • patient_001.pdf │                         ││  │  PyPDF2            │      │  • patient_002.pdf │                         ││  │                    │      │  • metadata.json   │                         ││  │                    │      │  • unassigned.pdf  │                         ││  └────────────────────┘      └────────────────────┘                         │└─────────────────────────────────────────────────────────────────────────────┘

The pipeline is modular by design. The Page Extractor renders PDF pages to images at configurable DPI using Poppler. The OCR Processor runs text extraction with pluggable backends—Tesseract by default, but swappable for OlmOCR or custom handlers.

The Entity Extractor runs regex patterns over OCR output to detect four identifier types: patient names, Medical Record Numbers (MRNs), dates of birth, and phone numbers. Each extracted identifier carries a confidence score based on pattern match quality.

The Fuzzy Matcher and Entity Linker work together to cluster page-level identifiers into patient profiles. Fuzzy matching determines whether "Mary Johnson" on page 3 is the same person as "MARY JOHNSON" on page 7, while the linker builds coherent patient clusters using anchor-based logic.

The Page Assigner scores each page against all known patient clusters using weighted entity matching, flagging ambiguous cases for human review. Finally, the Document Splitter generates per-patient PDFs with accompanying metadata.

Each stage is independently testable. The orchestrator wires them together but delegates all logic to the components.

Lesson 1: Not All Identifiers Are Equal

My first approach weighted all identifiers equally. A name match counted the same as an MRN match. That was wrong.

MRNs are unique identifiers—if two pages share an MRN, they're almost certainly the same patient. A 95% MRN match is nearly definitive evidence. Names, in contrast, are not unique at all—there might be three Robert Garcias in the practice. A perfect name match by itself proves little.

DOBs fall in between. They're unique per patient but commonly appear across multiple pages for the same person. Phone numbers are similar—moderately unique, but patients might share family phone numbers or update their contact information.

The fix was weighted scoring where each identifier type contributes proportionally to its reliability:

weights = {    "name": 0.4,   # High coverage, low uniqueness    "mrn": 0.3,    # Low coverage, high uniqueness    "dob": 0.2,    # Moderate on both    "phone": 0.1,  # Least reliable}for identifier in page.identifiers:    match_score = fuzzy_match(identifier, patient)    total_score += match_score * weights[identifier.kind]

Why is name weight highest despite being least unique? Because names appear on almost every page, while MRNs don't. A name match alone means little, but when you combine name + DOB + phone, you get strong signal even when MRN is absent. The weights encode how much incremental confidence each identifier type provides.

The lesson: Domain knowledge should inform feature importance. Don't assume all signals are equal—weight them based on reliability and frequency of occurrence.

Lesson 2: Fuzzy Matching Requires Per-Entity Strategy

Exact string matching fails immediately on real OCR output. "MARY JOHNSON" might come through as "MARY J0HNSON" (zero instead of O) or "Mary Johnson" (extra space) or "Johnson, Mary" (reordered). You need fuzzy matching, but not all fuzzy matching is the same.

I used RapidFuzz, but with different comparison strategies per entity type:

from rapidfuzz import fuzzdef score_name(a: str, b: str) -> float:    # Token sort handles "John Smith" vs "Smith, John"    return fuzz.token_sort_ratio(normalize(a), normalize(b)) / 100.0def score_dob(a: str, b: str) -> float:    # Dates must match exactly after normalization    return 1.0 if normalize_date(a) == normalize_date(b) else 0.0def score_phone(a: str, b: str) -> float:    # Partial matching handles missing area codes    return fuzz.partial_ratio(digits_only(a), digits_only(b)) / 100.0

For names, token sort ratio handles word reordering, so "John Smith" matches "Smith, John" correctly. MRNs get stricter treatment—after stripping non-alphanumeric characters, I expect near-exact matches. DOBs require exact match after normalization, handling format variations like "5/15/1980", "05-15-80", and "5.15.1980". Phone numbers use partial matching to catch missing area codes.

The lesson: Fuzzy matching isn't one-size-fits-all. Each entity type has different semantics—names tolerate reordering, dates require exact match after normalization, phone numbers can be partial.

Lesson 3: Clustering Is Position-Aware

Here's a subtle problem: what happens when a page has multiple patients' information?

Consider a referral letter. The referring doctor mentions their patient (the subject of the referral), but also references a patient at the receiving clinic for context. Naive clustering treats all identifiers on a page as belonging to one patient—which would incorrectly merge two distinct people.

The solution is anchor-based clustering. MRNs are designated as "anchors"—high-confidence identifiers that seed patient clusters. Other identifiers attach to the nearest preceding anchor on the same page:

for idx, identifier in enumerate(page.identifiers):    if identifier.kind == "mrn":        # MRNs are anchors—create or match to a cluster        cluster = find_or_create_cluster(identifier)        anchor_positions.append((idx, cluster))    else:        # Supportive identifiers attach to nearest anchor        nearest_anchor = find_nearest_anchor(idx, anchor_positions)        attach_to_cluster(identifier, nearest_anchor)

This models how medical documents are typically structured: patient identifiers appear in blocks, with the MRN or header coming first. So when a page has MRN 123456 followed by "Mary Johnson", then later MRN 789012 followed by "Robert Garcia"—the two names attach to their respective anchors rather than merging into one confused patient.

The lesson: Document structure matters. When clustering entities, consider their position on the page, not just their content.

Lesson 4: Ambiguity Should Be Flagged, Not Hidden

Early versions of the assigner tried to be "helpful" by assigning every page to some patient, even when confidence was low. That's dangerous in healthcare—a misassigned prescription is worse than an unassigned one. The wrong routing could delay critical treatment or violate HIPAA by exposing records to the wrong patient's file.

I added explicit ambiguity detection with two layers:

scores = [(patient, score_page(page, patient)) for patient in patients]scores.sort(key=lambda x: x[1], reverse=True)top_patient, top_score = scores[0]manual_review = False# Layer 1: Low confidence → unassignedif top_score < MIN_CONFIDENCE:  # default 0.6    return PageAssignment(patient_id=None, manual_review=True)# Layer 2: Close race → flag for reviewif len(scores) > 1 and (top_score - scores[1][1]) < AMBIGUITY_MARGIN:    manual_review = Truereturn PageAssignment(patient_id=top_patient, manual_review=manual_review)

The output now has three explicit categories:

Assigned with high confidence: Automatically routed, no human intervention needed
Assigned with manual review: Routed but flagged for human verification
Unassigned: Low confidence, requires human classification

The CLI and API both surface these metrics prominently. You can track your unassigned rate over time and tune thresholds accordingly.

The lesson: In high-stakes domains, "I don't know" is a valid output. Systems that hide uncertainty are systems that create liability.

Lesson 5: Make OCR Pluggable from Day One

Tesseract is the default OCR engine, and it's good enough for most scanned documents. But some PDFs have handwriting, complex table layouts, or degraded quality (fax-of-a-fax quality) that Tesseract struggles with.

Rather than hardcoding Tesseract, I built backend pluggability into the OCR processor. The configuration accepts three modes: tesseract (the default), olmocr (for integration with newer vision-language models), or a custom module path like my_module:ocr_function.

The custom mode uses Python's import machinery to load any callable that conforms to the expected interface—accepting an image and returning extracted text with optional bounding boxes. This means you can experiment with cloud OCR APIs, newer open-source models, or domain-specific extractors without touching the core pipeline code.

Switching backends is a single environment variable change. The pipeline doesn't care where the text comes from—it just needs text.

The lesson: When building pipelines, make the most likely-to-change components pluggable. You don't need to implement alternatives on day one—just ensure swapping is painless when you do need them.

Lesson 6: The Config Decides Everything

Medical workflows vary enormously. Some clinics need 95% confidence before auto-routing; others accept 60%. Some receive clean PDFs from modern systems; others get fourth-generation photocopy faxes. A one-size-fits-all threshold would fail for at least half of use cases.

I exposed every tunable as an environment variable. Entity matching thresholds control how similar two values must be to match: 85% for names, 95% for MRNs, exact match for DOBs, 90% for phones. Page assignment has its own knobs: minimum confidence for assignment, ambiguity margin for flagging, and whether to allow unassigned pages at all.

The entity weights—how much each identifier type contributes to the overall score—are also configurable. A clinic that always has MRNs on every page might weight them higher; one that never uses MRNs might set that weight to zero.

There's a --show-settings CLI flag that dumps the current configuration before processing. This makes debugging configuration issues trivial and ensures operators always know what thresholds are active.

The lesson: Hardcoded thresholds are tech debt. Every magic number should have a config knob, even if you only expose it to power users.

What I'd Do Differently

1. Add layout-aware extraction: Medical forms have predictable structure—headers, patient info blocks, tables, signature lines. A layout model (like LayoutLM or DONUT) could extract identifiers more reliably than regex over raw OCR text, especially for structured forms where position matters.

2. Use contextual embeddings for entity disambiguation: When two "Robert Garcia" entries exist in the system, semantic embeddings of the surrounding context could help determine if they're the same person. "Robert Garcia, cardiology patient, MRN 123456" is different from "Robert Garcia, referred for dermatology" even if the name matches.

3. Implement active learning: Every page flagged for manual review is a labeling opportunity. Human corrections could feed back into improved extraction patterns, clustering thresholds, and even regex refinements over time.

4. Add document type classification: Knowing that a page is a "lab result" vs "prescription" vs "referral" would improve entity expectations and weighting. Lab results almost always have MRNs; referral letters often mention multiple patients; prescriptions might only have name and DOB.

Key Takeaways

Weight identifiers by reliability: MRNs are more trustworthy than names—your scoring should reflect this proportionally
Per-entity fuzzy matching: Names need token sorting, dates need exact match after normalization, phones can be partial—don't use one algorithm for all
Position-aware clustering: When multiple patients appear on one page, use anchors to assign surrounding identifiers correctly
Surface uncertainty explicitly: Unassigned and ambiguous outputs are safer than false confidence in high-stakes domains
Make the right things configurable: Thresholds, weights, and backends should be environment-driven, not hardcoded

This project is open source and designed for high-volume medical office workflows. The test suite runs with mocked OCR by default, with real Tesseract available via environment flag.

Graph-Augmented RAG: Beyond Vector Similarity for Private Equity Research

Wed, 03 Dec 2025 00:00:00 GMT

Note: This RSS feed strips out SVGs and embeds. You might want to read the post on the webpage here.

How combining knowledge graphs with embeddings outperforms simple semantic search

December 3, 2025

projects

Traditional RAG retrieves documents based on how similar they are to your query. That works fine for general Q&A, but it falls apart in specialized domains where relationships matter as much as content.

Consider private equity deal research. An investor asks: "Find precedents for a US healthcare roll-up with aggressive add-on strategy." A pure vector search might return deals with similar descriptions, but miss the platform that acquired 12 add-ons—crucial context buried in structural relationships, not text.

I built DealGraph, a graph-augmented RAG agent that fuses semantic embeddings with a knowledge graph. The result: a hybrid retrieval system where structural features—sector membership, add-on relationships, exit events—augment pure text similarity.

Here's what I learned.

The Problem: Relationships Matter

Private equity research isn't just about finding similar text. It's about finding precedents: deals that share strategic patterns, not just vocabulary.

Consider two deals:

Deal A: Description mentions "healthcare consolidation"
Deal B: Description is vague, but the graph shows 8 add-on acquisitions in healthcare

A cosine similarity search ranks Deal A higher. But Deal B is the better precedent—it demonstrated the actual roll-up strategy the investor is researching.

The solution: make relationships first-class features.

The Architecture: A Three-Layer Pipeline

DealGraph orchestrates retrieval through three coordinated layers:

┌─────────────────────────────────────────────────────────────────┐│                         User Query                               ││    "Find precedents for US healthcare roll-up with add-ons"     │└───────────────────────────┬─────────────────────────────────────┘                            │                            ▼┌─────────────────────────────────────────────────────────────────┐│                   RETRIEVAL LAYER                                ││  ┌──────────────────┐          ┌────────────────────────────┐   ││  │  FAISS Index     │          │    NetworkX Graph          │   ││  │  (Embeddings)    │          │    (Relationships)         │   ││  │                  │          │                            │   ││  │  Query → Vector  │          │  Deal ─┬─ IN_SECTOR ──→    │   ││  │  Top-K by cosine │          │        ├─ ADDON_TO ───→    │   ││  └────────┬─────────┘          │        └─ EXITED_VIA ─→    │   ││           │                    └──────────────┬─────────────┘   ││           └────────────┬──────────────────────┘                 ││                        ▼                                         ││              Candidate Pool + Graph Features                     │└────────────────────────┬────────────────────────────────────────┘                         │                         ▼┌─────────────────────────────────────────────────────────────────┐│                    RANKING LAYER                                 ││                                                                  ││    ┌────────────────────────────────────────────────────────┐   ││    │  Gradient Boosting Ranker (DealRanker)                 │   ││    │                                                         │   ││    │  Features: text_similarity, sector_match, num_addons,  │   ││    │           has_exit, degree, region_match, is_platform  │   ││    │                                                         │   ││    │  Training: Reverse-query generation (LLM synthesizes   │   ││    │            queries for known-good deals)               │   ││    └────────────────────────────────────────────────────────┘   │└────────────────────────┬────────────────────────────────────────┘                         │                         ▼┌─────────────────────────────────────────────────────────────────┐│                   REASONING LAYER                                ││                                                                  ││    ┌────────────────────────────────────────────────────────┐   ││    │  DSPy-Optimized Deal Reasoner                          │   ││    │                                                         │   ││    │  Outputs:                                               │   ││    │  • Precedent selection (JSON)                          │   ││    │  • Playbook levers (strategic patterns)                │   ││    │  • Risk themes                                          │   ││    │  • Executive narrative                                  │   ││    └────────────────────────────────────────────────────────┘   │└────────────────────────┬────────────────────────────────────────┘                         │                         ▼                  Structured Analysis

Each layer solves a specific problem:

Retrieval: Cast a wide net with hybrid search
Ranking: Learn what "relevance" means from training data
Reasoning: Synthesize precedents into actionable insights

Lesson 1: Model Relationships as a Graph

The core insight: deals aren't documents, they're entities with relationships.

I modeled the deal universe as a typed MultiDiGraph:

# Node typesNODE_TYPES = ["Deal", "Sector", "Region", "Event", "Snippet"]# Edge types capture relationshipsEDGE_TYPES = {    "IN_SECTOR": ("Deal", "Sector"),    # Deal belongs to sector    "IN_REGION": ("Deal", "Region"),    # Deal in geographic region    "ADDON_TO": ("Deal", "Deal"),       # Add-on acquired by platform    "EXITED_VIA": ("Deal", "Event"),    # Exit event (IPO, M&A)    "DESCRIBED_IN": ("Deal", "Snippet") # Textual evidence}

This structure enables queries that pure text search can't answer:

"Find all platforms with 5+ add-ons in healthcare"
"Show deals that exited via IPO in the last 3 years"
"Find add-ons to Platform X in adjacent sectors"

The graph becomes a source of features, not just navigation.

Lesson 2: Extract Graph Features for Ranking

The key innovation: graph topology becomes a feature vector for ML ranking.

For each candidate deal, I compute structural features:

FEATURE_NAMES = [    "text_similarity",      # From embeddings    "sector_match",         # Query mentions same sector?    "region_match",         # Query mentions same region?    "num_addons",           # How many add-ons acquired    "has_exit",             # Successful exit?    "degree",               # Graph connectivity    "is_platform",          # Platform vs add-on    "sector_degree",        # How connected is the sector    "region_degree",        # How connected is the region    "text_graph_alignment", # Do text and graph agree?]

These features capture things embeddings miss. A deal might have a generic description but 10 recorded add-ons—that's crucial signal.

The heuristic ranking weights these explicitly:

def compute_relevance_score(candidate):    features = candidate.graph_features    # Text similarity (40%)    text_score = features['text_similarity'] * 0.4    # Sector match (20%)    sector_score = features['sector_match'] * 0.2    # Region match (15%)    region_score = features['region_match'] * 0.15    # Add-on activity (10%)    addon_score = min(features['num_addons'] / 3.0, 1.0) * 0.1    # ... other features    return text_score + sector_score + region_score + addon_score

But hand-tuning weights is fragile. That's where ML comes in.

Lesson 3: Train a Ranker on Synthetic Query-Deal Pairs

The challenge: I don't have labeled relevance data. No one has annotated thousands of "query → relevant deals" pairs for private equity.

The solution: reverse-query generation. Instead of labeling data manually, I use an LLM to synthesize queries for known deals:

def generate_reverse_queries(deal_cluster, llm):    """    Given a cluster of similar deals, generate realistic    queries that would surface these deals as relevant.    """    prompt = f"""    You are a private equity analyst. Given these deals:    {format_deals(deal_cluster)}    Generate 3 realistic search queries an investor    might use to find deals like these.    """    queries = llm.generate(prompt)    return queries

This creates training pairs:

Positive: (generated query, deals in cluster)
Negative: (generated query, random deals from other clusters)

The gradient boosting model learns which feature combinations predict relevance:

class DealRanker:    def __init__(self):        self.model = GradientBoostingRegressor(random_state=42)    def fit(self, X: np.ndarray, y: np.ndarray):        """X = feature vectors, y = relevance scores"""        self.model.fit(X, y)        return self    def rank(self, candidates: List[CandidateDeal]):        X = build_feature_matrix(candidates)        scores = self.model.predict(X)        ranked = sorted(            zip(candidates, scores),            key=lambda pair: pair[1],            reverse=True        )        return [RankedDeal(c, score, rank) for rank, (c, score) in enumerate(ranked, 1)]

The trained ranker outperforms both pure embedding search and hand-tuned heuristics.

Lesson 4: Optimize Prompts Systematically with DSPy

The reasoning layer isn't just a static prompt—it's a DSPy module that can be optimized.

DSPy treats prompts as programs with learnable components. The MIPRO optimizer generates prompt variants and evaluates them against a composite metric:

# Composite evaluation metricdef composite_metric(example, prediction):    score = (        0.4 * precision_at_k(prediction.precedents, example.gold_precedents, k=3) +        0.3 * llm_judge_playbook_quality(prediction.playbook_levers) +        0.3 * llm_judge_narrative_coherence(prediction.narrative_summary)    )    return score

The optimizer runs ~500 LLM calls to find better prompts. Results on my benchmark:

Metric	Naive Prompt	MIPRO-Optimized	Improvement
Precision@3	0.42	0.68	+62%
Playbook Quality	0.55	0.72	+31%
Narrative Coherence	0.61	0.78	+28%
Composite Score	0.52	0.73	+40%

The key insight: prompts are artifacts that should be versioned, evaluated, and optimized—not artisanal text you tweak by hand.

# Runtime loads optimized version automaticallyreasoner = DealReasonerModule()reasoner.load("prompts/deal_reasoner/v2_optimized.json")# Falls back to naive baseline if optimization hasn't run

Lesson 5: Fail Loudly, Don't Degrade Silently

A tempting pattern: wrap everything in try/catch and return empty results on failure.

Don't do this.

# BAD: Silent degradationtry:    result = reasoner(query=query, candidate_deals=deals_json)except Exception:    return {"precedents": [], "narrative": "Unable to analyze."}# GOOD: Fail loudlytry:    result = reasoner(query=query, candidate_deals=deals_json)except Exception as e:    raise DealReasonerError(f"Deal reasoning failed: {e}") from e

Silent failures hide bugs. In production, I'd rather see an error than serve garbage that looks like a valid response.

This extends to model loading:

# Model loading with graceful fallback (good)if ranker_model_exists():    ranker = DealRanker.load("models/deal_ranker_v1.pkl")else:    ranker = HeuristicRanker()  # Explicit fallback    logger.warning("Using heuristic ranker - ML model not found")# But NOT silent fallback during inferencescores = ranker.predict_scores(X)  # This should throw on failure

The Architecture Decision: NetworkX for V1

A common question: why not Neo4j or a "real" graph database?

For V1 with <1000 nodes, NetworkX is the right choice:

Zero infrastructure: No database to deploy
Fast iteration: Graph structure can change without migrations
Rich algorithms: NetworkX has excellent graph algorithms out of the box
Good enough: In-memory traversal is plenty fast at this scale

The migration triggers are clear:

Graph size exceeds 10K nodes
Query latency becomes unacceptable (>1s)
Need for persistence across restarts
Multi-user concurrent access

Until then, keep it simple.

What I'd Do Differently

1. Use a vector database with filtering: FAISS is fast but doesn't support metadata filtering. I filter post-hoc, which is wasteful. Pinecone or Weaviate would let me push filters into the index.

2. Add temporal reasoning: Deals have dates. A 2010 precedent might be less relevant than a 2023 one. The graph should encode temporal relationships.

3. Implement online learning: As users interact with results, their clicks are implicit relevance labels. The ranker should learn from this feedback.

4. Build evaluation into the pipeline: I added benchmarks late. Starting with evaluation infrastructure would have caught issues earlier.

Key Takeaways

Relationships are features: Knowledge graphs aren't just for navigation—extract structural features and feed them to ML models.
Hybrid retrieval beats pure embeddings: Text similarity is necessary but not sufficient. Combine it with domain-specific signals.
Synthesize training data when labels don't exist: Reverse-query generation creates high-quality training pairs without manual annotation.
Treat prompts as code: Version them, optimize them with metrics, and A/B test before deploying.
Keep infrastructure simple: Start with in-memory solutions. Migrate when you hit actual scale limits, not imagined ones.

DealGraph is available on GitHub. The synthetic data generator creates a realistic deal universe for testing, and the full training pipeline runs on commodity hardware.

Mini-Town: What I Learned from an Abandoned Agent Simulation

Sun, 30 Nov 2025 00:00:00 GMT

Note: This RSS feed strips out SVGs and embeds. You might want to read the post on the webpage here.

An experimental project that taught me more by failing than it would have by succeeding

November 30, 2025

projects

Mini-Town was an ambitious experiment that I eventually abandoned. The goal was to build an agent-based simulation where autonomous characters—powered by LLMs—would perceive their environment, form memories, reflect on experiences, and generate plans. Instead of hand-tuning prompts, I would use DSPy's GEPA optimizer to automatically improve the agents' "brains" based on collected data.

The project reached a functional state. Agents moved around a 2D world, scored observations for importance, stored memories with vector embeddings, and generated daily plans. I got a 17% improvement in observation quality using compiled prompts versus hand-written ones. But the complexity grew faster than the value, and I ultimately set it aside.

This post isn't a tutorial on how to build what I built. It's a collection of the things I learned along the way—lessons that transferred to other projects and shaped how I think about LLM-based systems.

The Original Idea

The project emerged from a simple question: can you treat LLM prompts as programs that can be automatically optimized?

When you hand-write a prompt, you're making guesses about what wording will produce the behavior you want. "Rate this observation from 1-10 based on importance" might work, but "Evaluate the significance of this event on a scale of 1-10, where 10 means life-changing and 1 means completely mundane" might work better. Or worse. You don't know until you try, and trying is expensive and slow.

DSPy offers an alternative. Instead of writing prose instructions and hoping for the best, you define a typed signature—a formal specification of what inputs the LLM receives and what outputs it should produce. The framework then uses training examples to search for prompts that produce correct outputs. It's treating prompt engineering as an optimization problem rather than an art form.

The approach is compelling because it scales differently than manual iteration. If you have 50 training examples and can afford 1000 LLM calls during optimization, the optimizer can explore far more prompt variations than a human could in a week. And because the optimization is metric-driven, it's less susceptible to the subtle biases humans bring to prompt writing.

Mini-Town was the testbed for this idea. Build a simulation with multiple agents, log their cognitive outputs, use those logs as training data, and compile improved prompts. The hypothesis: compiled agents would behave more consistently and intelligently than uncompiled ones, and we could measure the difference.

The simulation itself was inspired by Stanford's "Generative Agents" paper—the one with the AI town where characters formed relationships, remembered conversations, and threw parties. I wasn't trying to replicate that research; I was using the simulation as a vehicle to test prompt optimization. The town was scaffolding. The real experiment was the compilation.

What Got Built

The architecture had three main components, each representing a different layer of the system.

┌─────────────────────────────────────────────────────────────────────────────┐│                          SIMULATION ENVIRONMENT                              ││                                                                              ││  ┌────────────────────┐      ┌───────────────────────────────────────────┐  ││  │  2D World (Grid)   │      │  Agent Population                         │  ││  │  ─────────────     │      │  ─────────────────                        │  ││  │  • Locations       │◀────▶│  Alice, Bob, Maria, ...                   │  ││  │  • Entities        │      │  Each running cognitive loop              │  ││  │  • Time (ticks)    │      │                                           │  ││  └────────────────────┘      └──────────────────┬────────────────────────┘  │└─────────────────────────────────────────────────┼────────────────────────────┘                                                  │ perceive environment                                                  ▼┌─────────────────────────────────────────────────────────────────────────────┐│                          COGNITIVE LOOP (per agent)                          ││                                                                              ││  ┌──────────┐    ┌──────────┐    ┌──────────┐    ┌──────────┐    ┌────────┐ ││  │ PERCEIVE │───▶│  SCORE   │───▶│  STORE   │───▶│ REFLECT  │───▶│  PLAN  │ ││  │          │    │          │    │          │    │          │    │        │ ││  │ Detect   │    │ Rate 1-10│    │ Embed +  │    │ Synthesize│    │ Daily  │ ││  │ nearby   │    │ importance│    │ save to  │    │ insights │    │ schedule│ ││  │ entities │    │ (DSPy)   │    │ memory   │    │ (DSPy)   │    │ (DSPy) │ ││  └──────────┘    └──────────┘    └──────────┘    └──────────┘    └────────┘ ││       │               │               │               │               │      ││       │               │               │               │               │      ││       │               ▼               ▼               ▼               ▼      ││       │         ┌─────────────────────────────────────────────────────┐     ││       │         │                DSPy Modules                          │     ││       │         │  ┌────────────────┐  ┌───────────────┐  ┌─────────┐ │     ││       │         │  │ ScoreImportance│  │ ReflectOnDay  │  │ MakePlan│ │     ││       │         │  │ (typed sig)    │  │ (typed sig)   │  │ (typed) │ │     ││       │         │  └────────────────┘  └───────────────┘  └─────────┘ │     ││       │         │                                                      │     ││       │         │  GEPA Optimizer: Compiles signatures → better prompts│     ││       │         └─────────────────────────────────────────────────────┘     │└───────┼─────────────────────────────────────────────────────────────────────┘        │        ▼┌─────────────────────────────────────────────────────────────────────────────┐│                          MEMORY SYSTEM (DuckDB + VSS)                        ││                                                                              ││  ┌────────────────────────────────────────────────────────────────────────┐ ││  │  Weighted Retrieval: score = α×relevance + β×recency + γ×importance   │ ││  │                                                                         │ ││  │  ┌─────────────┐    ┌─────────────┐    ┌─────────────┐                 │ ││  │  │ Relevance   │    │ Recency     │    │ Importance  │                 │ ││  │  │ (cosine     │    │ (time decay │    │ (score at   │                 │ ││  │  │  similarity)│    │  function)  │    │  storage)   │                 │ ││  │  └─────────────┘    └─────────────┘    └─────────────┘                 │ ││  └────────────────────────────────────────────────────────────────────────┘ ││                                                                              ││  Embeddings: all-MiniLM-L6-v2 (384-dim, local)                              │└─────────────────────────────────────────────────────────────────────────────┘

The Cognitive Loop

Each agent ran a perceive-score-store-reflect-plan cycle that mimicked the approach from the Generative Agents paper, adapted for DSPy's typed program model.

Perception: On every simulation tick (every 2 seconds), agents would detect nearby entities. If Alice walked past Bob, Bob's perception system would register "Alice is nearby" as a raw observation. The perception range was configurable—agents could "see" anything within a certain radius.

Scoring: Raw observations aren't all equally important. Seeing a tree is mundane. Seeing someone collapse is urgent. The scorer was a DSPy module that took an observation and the agent's current context (their personality, goals, recent thoughts) and produced an importance score from 1-10. High scores triggered downstream processing; low scores were stored but didn't prompt reflection.

Storage: Every scored observation went into the memory database. The storage included not just the text and score, but also a vector embedding and a timestamp. This enabled the retrieval system to find memories by meaning, not just by recency.

Reflection: When accumulated importance exceeded a threshold, the agent would pause to reflect. The reflector was another DSPy module that took recent memories, retrieved semantically similar older memories, and synthesized a high-level insight. "I've noticed Maria keeps mentioning her garden" might become "Maria is passionate about gardening—I should ask her about it." Reflections themselves were stored as high-importance memories.

Planning: Periodically (configurable, but typically once per simulated "day"), agents would generate a schedule. The planner received the agent's goal, relevant memories, and any pending invitations, and produced a time-blocked plan: "7am: wake up, 8am: go to cafe, 10am: visit Maria's garden, 12pm: lunch at diner."

The loop ran continuously. Every tick, agents perceived. Interesting perceptions got scored. Accumulating importance triggered reflection. Time passing triggered planning. Actions came from plans. The result was agents that appeared to make decisions, though everything was actually just LLM calls wrapped in state management.

The Memory System

Memory in agent simulations is deceptively complex. Humans don't just remember things; they remember things that are relevant to their current situation, weighted by how recent they are and how important they seemed when they happened. Implementing that requires more than a list sorted by timestamp.

I used DuckDB with the Vector Similarity Search (VSS) extension. DuckDB gave me SQL for structured queries, and VSS gave me approximate nearest-neighbor search over embedding vectors. The combination let me implement weighted retrieval:

score = α × relevance + β × recency + γ × importance

Relevance was cosine similarity between the query embedding (what the agent was currently thinking about) and the memory embedding. This made semantic recall possible. An agent thinking about "parties" could retrieve a memory containing "Maria's get-together" even though the word "party" never appeared.

Recency was an exponential decay based on how long ago the memory was formed. Recent memories scored higher. The decay rate was tunable; faster decay made agents forgetful, slower decay gave them longer attention spans.

Importance was the score assigned at storage time. Salient memories—the ones that felt significant when they happened—remained accessible longer than mundane observations.

The weights (α, β, γ) controlled the behavior. High α meant agents prioritized semantic relevance (useful for focused tasks). High β meant agents focused on recent events (useful for reactive behavior). High γ meant agents surfaced important moments regardless of recency (useful for long-term relationship tracking).

Getting these weights right was one of the harder engineering problems in the project. I'll return to this later.

The DSPy Modules

Each cognitive function was implemented as a DSPy program with a typed signature. Signatures define the inputs and outputs formally, which lets the optimizer understand what it's trying to improve.

The scorer signature looked like this:

class ScoreImportance(dspy.Signature):    """Score how important an observation is to this agent."""    observation: str = dspy.InputField(desc="What the agent just observed")    agent_personality: str = dspy.InputField(desc="The agent's personality traits")    agent_goal: str = dspy.InputField(desc="The agent's current goal")    recent_context: str = dspy.InputField(desc="What the agent was just thinking about")    score: int = dspy.OutputField(desc="Importance score from 1-10")    reasoning: str = dspy.OutputField(desc="Why this observation matters or doesn't")

The reflector and planner had similar structures. Each defined what went in and what came out, with human-readable descriptions that helped both the optimizer and the developer understand the expected behavior.

The baseline implementation wrapped these signatures in dspy.ChainOfThought, which prompted the LLM to reason step-by-step before producing outputs. Chain of thought improved accuracy on the importance scoring task, at the cost of additional latency and tokens.

The compiled implementation used GEPA-optimized prompts that achieved similar or better accuracy with fewer tokens. More on compilation in the next section.

The Frontend

The frontend was a Next.js application that connected to the backend via WebSocket. It rendered a 2D map showing agent positions, displayed speech bubbles when agents "spoke" their plans or observations, and provided a diagnostic panel for inspecting agent state.

The visual representation was mostly for debugging. Watching an agent walk in circles told me its pathfinding was broken. Watching an agent stand still told me its LLM calls were timing out. Watching two agents never interact told me the perception radius was too small.

Most development happened headless (running the backend with logging and no frontend), but the visual mode was invaluable for the moments when logs didn't explain what was happening.

The Compilation Experiment

The core experiment was simple: could GEPA find better prompts than I could write by hand?

GEPA (Generalized Evolutionary Prompt Acceleration) is a DSPy optimizer that uses evolutionary strategies to search for improved prompts. You provide a training set of input-output examples (called "seeds"), a metric function that scores outputs, and a search budget. GEPA generates prompt variations, evaluates them against the training set, and iteratively improves toward the best-performing variants.

The process took about 15 minutes on a T4 GPU (Google Colab free tier). Running 80 optimization iterations over 40 seed examples, GEPA evaluated roughly 3000 prompt variants to find the best-performing one. The result was a compiled scorer that outperformed my hand-written baseline.

In a controlled 5-minute A/B test:

Uncompiled agents generated 38 observations, avg quality score 0.67
Compiled agents generated 53 observations, avg quality score 0.78

That's a 17% improvement in observation quality and 39% more observations in the same time. The compiled scorer was also 46% faster in isolation tests (0.46s vs 0.85s latency) because the optimized prompt was more concise.

The reflection and planning modules couldn't be compiled as easily because their outputs were harder to evaluate. What makes a "good" reflection? The question is inherently subjective. I could score importance reasonably consistently (a 7 is more significant than a 3), but I couldn't score reflections or plans reliably enough to train an optimizer.

This asymmetry—some cognitive functions are easy to evaluate, others are hard—limited how far compilation could take the project.

Why I Abandoned It

Several factors converged to make the project not worth continuing.

The Evaluation Problem

How do you measure whether an agent simulation is "good"?

I defined metrics. Observation quality was scored by an LLM judge. Plan coherence was measured by whether the plan's activities matched the agent's stated goals. Event attendance was measured by whether agents showed up to parties they'd been invited to.

But these were proxies for something I couldn't directly measure: whether the agents were interesting to watch. The 17% improvement in observation quality didn't translate into 17% more interesting behavior. It was an improvement on a metric that didn't capture what mattered.

The fundamental problem with open-ended simulations is that "good behavior" is defined by human intuition, not by formal criteria. When Alice decides to visit the park, is that good or bad? It depends on what Alice has been doing, what her goals are, who else is at the park, what time of day it is, and—critically—what the observer finds interesting. I couldn't formalize any of that.

I tried using "town score," a weighted combination of observation quality, reflection depth, plan coherence, and event attendance. The town score improved with compilation. But watching an agent with a high town score felt the same as watching an agent with a medium town score. The metric didn't correlate with the experience.

The Cost Problem

Even using cheap models (Llama 3.2 on Together.ai at $0.20 per million tokens), the API costs added up quickly when running multiple agents continuously.

Each tick involved multiple LLM calls: scoring observations, occasionally reflecting, checking if plans needed updating. With 5 agents and a 2-second tick, the simulation made roughly 250 LLM calls per minute. At ~500 tokens per call (including prompt and response), that was 125,000 tokens per minute, or about $1.50 per hour.

That sounds cheap, but I needed sustained simulation to generate training data. GEPA optimization required a corpus of logged examples, which meant running the simulation for hours to collect data before I could even start the optimization process. Then I needed to run the simulation again to A/B test the results. Development required many such cycles.

The Complexity Problem

The cognitive loop introduced cascade failures that were difficult to debug and impossible to prevent.

If the scorer ran slowly (maybe the API was throttled, maybe the model was overloaded), observations accumulated faster than they could be processed. The observation queue grew. Old observations became stale. Reflections based on stale observations were nonsensical.

If reflection failed (maybe the LLM produced an unparseable response), the agent's accumulated importance wasn't reset. It would try to reflect on the same observations again next tick, fail again, and enter an infinite loop of failed reflections.

If planning failed, the agent defaulted to random walking. But random walking generated more observations, which triggered more scoring and more reflection, which increased load on the LLM, which made planning more likely to fail due to rate limits.

Each component's failure propagated to everything downstream. Debugging required understanding the interaction of five different subsystems, each with its own state and failure modes. I spent more time building monitoring and debugging tools than I spent building the actual simulation.

The Ambiguity Problem

Agent behavior in open-ended simulations is inherently ambiguous. This made creating training data painful.

GEPA requires seed examples: input-output pairs where the output is the "correct" answer. For the scorer, I could produce these. "Observation: Maria invited me to her party. Score: 8. Reason: Social invitation with future commitment." That's defensible.

But what about edge cases? "Observation: I see a tree." Is that a 1 or a 2? Depends on whether the agent cares about nature. "Observation: Bob walked past without saying hello." Is that significant? Depends on whether Bob usually says hello. "Observation: The sun is setting." Is that important? Depends on whether the agent needs to get home before dark.

I created an annotation guide to standardize scoring, then had a friend score some examples independently to check agreement. Cohen's kappa was 0.65—substantial agreement, but not perfect. The residual disagreement meant the training data had noise, which limited how much GEPA could improve.

Better annotation would have required more annotators, more examples, and a more rigorous process. At that point, I was building a labeling pipeline, not an agent simulation.

What I Learned

Despite abandoning the project, it shaped how I approach LLM systems now.

Structure Your Outputs from the Start

Early versions of the cognitive loop logged free-form text. "Agent observed something interesting" tells you nothing when you're trying to debug why an agent walked into a wall.

Later versions logged structured data:

Observation content (what the agent saw)
Timestamp (when it happened)
Importance score (how significant it was)
Scoring latency (how long the LLM took)
Failure reason (if scoring failed, why)

Structured outputs made debugging possible. When an agent started behaving erratically, I could query the logs: "Show me all observations with score > 7 from the last 5 minutes." I could trace causality: "This high-importance observation triggered reflection at 10:42, which produced this insight, which affected the plan generated at 10:45."

Structured outputs also made training data collection tractable. Each logged item had the fields GEPA expected: input (observation + context) and output (score + reasoning). I didn't have to post-process logs into training examples; they were already in the right format.

If I'd started structured from day one, I would have saved days of refactoring. Every project I've started since then begins with a schema for the data it will log.

Measure Latency Before Complexity

I measured latency on Day 2 specifically to avoid a trap: building a complex system that couldn't run in real time.

The process was simple. I wrapped every LLM call in a timing function:

async def timed_llm_call(func, *args, **kwargs):    start = time.perf_counter()    result = await func(*args, **kwargs)    elapsed = time.perf_counter() - start    latency_tracker.record(func.__name__, elapsed)    return result

After a few minutes of simulation, I had a latency distribution. The scorer's p50 was 0.8 seconds; p95 was 1.4 seconds. The reflector's p50 was 1.2 seconds; p95 was 2.1 seconds.

The simulation tick was 2 seconds. That meant the cognitive loop had to complete its per-tick work (scoring + maybe reflecting) in under 2 seconds, or the simulation would lag. With scoring alone taking 1.4 seconds at p95, I had 0.6 seconds of headroom. Reflection was out of the question on every tick; it would have to happen on a slower schedule.

This constraint informed every design decision. I couldn't add a second LLM call per tick. I couldn't make prompts longer (more tokens = more latency). I couldn't use larger models. The latency ceiling was load-bearing.

In later projects, I measure latency first, before adding features. You can always add complexity; you can't always remove latency.

Compilation Works, but Training Data Is Hard

GEPA did what it promised: given training examples, it found better prompts. The 17% improvement was real. The optimized prompts were shorter and faster. The mechanism works.

But creating training examples required human judgment about what "correct" meant, and that judgment was expensive to exercise consistently. For the scorer, I spent three hours annotating 40 seed examples. Each example required reading an observation, imagining the agent's context, deciding on a score, and writing a rationale. Forty examples was enough to show improvement but not enough to reach the optimizer's full potential. GEPA's documentation suggests 200+ examples for best results.

Generating 200 examples at 4 minutes per example would take 13 hours. That's a full weekend of labeling. And if my scoring criteria evolved during that time (which they would, because annotating forces you to clarify your standards), earlier examples would be inconsistent with later ones.

For tasks with clear correctness criteria (classification, extraction, structured output), DSPy compilation is worth the effort. You can generate training data programmatically or collect it from production usage. But for tasks where correctness is subjective (open-ended generation, creative writing, behavioral simulation), the training data problem dominates everything else.

Memory Retrieval Is Surprisingly Finicky

The vector search used a weighted combination: relevance (semantic similarity), recency (time decay), and importance (score at storage time). Getting the weights right required extensive tuning.

In early tests, I set the weights to α=1.0, β=0.0, γ=0.0—pure relevance. The system seemed to work. Agents retrieved memories that were semantically related to their current situation. Great.

Then I tested a scenario where an agent needed to remember what happened earlier that day. Pure relevance failed. The agent retrieved a semantically similar but weeks-old memory instead of the relevant recent one. I needed recency.

Then I tested a scenario where an agent needed to remember an important conversation from last month. With high recency weight, that memory was suppressed by mundane recent observations. I needed importance.

Different scenarios needed different weights:

Emergencies: high relevance (α=0.7), find anything about this topic
Social planning: high recency (β=0.6), what happened today
Long-term relationships: high importance (γ=0.7), filter out noise

I built a grid search to find optimal weights per scenario type, but the scenarios themselves were hard to define rigorously. The retrieval system worked, but tuning it was a research project in itself.

The lesson: test retrieval with adversarial examples, not just the happy path. If you only test with data where relevance and recency correlate (recent events that happen to be semantically related), you'll think your system works when it's actually ignoring half its inputs.

Provider Flexibility Saves Projects

The project nearly stalled on Day 5 when Groq's free tier hit rate limits mid-experiment.

I was running an A/B comparison: 20 minutes of uncompiled simulation, then 20 minutes of compiled simulation, then compare the logs. Halfway through the uncompiled run, Groq started returning 429 errors. Rate limited. The remaining observations went unscored, ruining the comparison.

The fix was switching to Together.ai, which had no rate limits on their serverless models. But I'd hardcoded Groq as the provider. DSPy was configured to use Groq. The API key was a Groq key. Changing providers required refactoring the configuration to support multiple backends.

I added a config file:

llm:  provider: together  # or groq, or openai  model: meta-llama/Llama-3.2-3B-Instruct-Turbo  api_key: ${TOGETHER_API_KEY}

And a dispatch function:

def configure_dspy():    config = load_config()    if config.llm.provider == "together":        lm = dspy.LM("together_ai/" + config.llm.model)    elif config.llm.provider == "groq":        lm = dspy.LM("groq/" + config.llm.model)    # ... etc    dspy.configure(lm=lm)

After that, switching providers was a one-line config change. The experiment restarted in minutes.

Every project I build now has provider abstraction from day one. The free tier that works today will hit limits tomorrow. The model that's fastest now will be deprecated next month. The only constant is change, and the only defense is abstraction.

What Remains

The code still runs. The frontend renders agents moving through a 2D world. The cognitive loop still scores observations, stores memories, and generates plans. The compiled scorer still outperforms the baseline.

But I don't work on it anymore. The lessons transferred to other projects:

The voice agent optimization system uses similar memory structures for storing failed calls
The fraud detection system uses structured failure logging to enable downstream analysis
The deal research agent uses provider abstraction for resilient LLM access
Everything I build now measures latency before adding complexity

Mini-Town was the laboratory where I figured out how to build those things. The experiment failed to produce a compelling simulation, but it succeeded at teaching me how to build robust LLM systems.

Sometimes the value of a project isn't in what it produces. It's in what it teaches you.

Technical Details

For those interested in the specifics:

Backend: FastAPI (Python 3.11), DuckDB with VSS extension for vector search
Frontend: Next.js/React with TailwindCSS, WebSocket for real-time updates
LLM: Together.ai's Llama 3.2-3B-Instruct-Turbo (serverless, ~$0.20/M tokens)
DSPy: GEPA optimizer for prompt compilation, ChainOfThought for base modules
Embeddings: all-MiniLM-L6-v2 (384-dimensional, runs locally)
Compilation budget: 40 seed examples, 80 GEPA iterations, ~15 minutes on T4 GPU

The 17% improvement came from comparing observation quality scores (LLM-judged) between compiled and uncompiled runs over 5-minute simulations with 3 agents. The compiled scorer was tested in isolation showing 46% latency reduction (0.46s vs 0.85s).

Total development time was approximately 5 days of focused work, spread over several weeks. Total API cost was under $5, mostly consumed during the final A/B testing phase.

Mini-Town is archived on GitHub. The code works for what it was built to do—it's just not something I'm building further.

When VLMs Need a Human Safety Net: Building Document Extraction with Confidence Calibration

Wed, 19 Nov 2025 00:00:00 GMT

Note: This RSS feed strips out SVGs and embeds. You might want to read the post on the webpage here.

How I designed a medical document processing system where the model knows when to ask for help

November 19, 2025

projects

Healthcare document processing has a fundamental constraint: you can't be wrong. A misread medication dosage or transposed patient ID isn't a minor error—it's a potential safety incident. Yet vision-language models, for all their power, are confidently wrong with uncomfortable frequency.

Inspired by Tennr's approach to healthcare automation, I built a document extraction platform that embraces this reality. The core insight: the model's job isn't just to extract data—it's to know when it shouldn't be trusted. When confidence is low, the system routes to human review. When humans correct errors, those corrections feed back into model improvement.

Here's what I learned building it.

The Problem: Models Don't Know What They Don't Know

Vision-language models like olmOCR can extract structured data from complex medical forms with impressive accuracy. But they have a fatal flaw: their confidence scores are poorly calibrated. A model that outputs 90% confidence should be right 90% of the time—but in practice, these models are often overconfident on novel inputs and underconfident on routine ones.

This creates two failure modes:

Overconfidence: The model extracts "150 mg" when the handwriting actually shows "180 mg", reports 95% confidence, and the error sails through to production.
Underconfidence: The model correctly extracts a clearly printed name, reports 70% confidence, and sends it to human review unnecessarily—wasting reviewer time.

The solution isn't better models. It's calibration: learning to transform raw model outputs into reliable confidence scores, and using those scores to make intelligent routing decisions.

The Architecture: Confidence-Driven Routing

┌─────────────────────────────────────────────────────────────────────────────┐│                          DOCUMENT INGESTION                                  ││  ┌──────────────┐      ┌──────────────┐      ┌──────────────┐               ││  │ PDF/Image    │      │ Preprocessing│      │ Layout       │               ││  │ Upload       │ ───▶ │ Enhancement  │ ───▶ │ Detection    │               ││  │ (API/CLI)    │      │ (denoise,    │      │ (regions,    │               ││  │              │      │  deskew)     │      │  tables)     │               ││  └──────────────┘      └──────────────┘      └──────┬───────┘               │└─────────────────────────────────────────────────────┼───────────────────────┘                                                      │ regions                                                      ▼┌─────────────────────────────────────────────────────────────────────────────┐│                          VLM INFERENCE LAYER                                 ││                                                                              ││  ┌────────────────────┐      ┌────────────────────────────────────────┐     ││  │  olmOCR-2-7B-MLX   │      │  Self-Consistency Engine               │     ││  │  ─────────────     │ ───▶ │  ─────────────────────                 │     ││  │  Quantized VLM     │      │  k=1 (text) / k=3 (mixed) / k=5 (hard) │     ││  │  MLX acceleration  │      │  Majority voting + agreement ratio     │     ││  └────────────────────┘      └─────────────────┬──────────────────────┘     │└────────────────────────────────────────────────┼────────────────────────────┘                                                 │ predictions + raw confidence                                                 ▼┌─────────────────────────────────────────────────────────────────────────────┐│                          CALIBRATION LAYER                                   ││                                                                              ││  ┌────────────────────────────────────────────────────────────────────┐     ││  │  Temperature Scaling Calibrator                                     │     ││  │                                                                     │     ││  │  raw_confidence ───▶ logit ───▶ scale by T ───▶ calibrated_conf    │     ││  │                                                                     │     ││  │  Per-field calibration: T_medication=1.5, T_name=0.9, T_date=0.8  │     ││  └────────────────────────────────────────────────────────────────────┘     │└─────────────────────────────────────────────────┬───────────────────────────┘                                                  │ calibrated confidence                                                  ▼┌─────────────────────────────────────────────────────────────────────────────┐│                          ROUTING DECISION                                    ││                                                                              ││                    ┌─────────────────────────────┐                          ││                    │  Confidence ≥ threshold?    │                          ││                    │  (per-field: 85%-98%)       │                          ││                    └──────────┬──────────────────┘                          ││                               │                                              ││              ┌────────────────┼────────────────┐                            ││              │                │                │                            ││              ▼ YES            ▼ MAYBE          ▼ NO                         ││  ┌───────────────────┐ ┌─────────────────┐ ┌───────────────────┐            ││  │ AUTOPILOT         │ │ FLAG FOR REVIEW │ │ HUMAN REVIEW      │            ││  │ ───────────       │ │ ───────────     │ │ ────────────      │            ││  │ Direct export     │ │ Auto-assign but │ │ Queue for manual  │            ││  │ to EHR/database   │ │ mark for verify │ │ annotation        │            ││  └───────────────────┘ └─────────────────┘ └─────────┬─────────┘            │└─────────────────────────────────────────────────────────┼───────────────────┘                                                          │                                                          ▼┌─────────────────────────────────────────────────────────────────────────────┐│                          HUMAN-IN-THE-LOOP                                   ││                                                                              ││  ┌────────────────────┐      ┌────────────────────────────────────────┐     ││  │  Streamlit UI      │      │  Active Learning Selector              │     ││  │  ─────────────     │      │  ────────────────────                  │     ││  │  Priority queue    │ ◀──▶ │  Uncertainty × 0.5 +                   │     ││  │  Keyboard shortcuts│      │  Diversity × 0.3 +                     │     ││  │  Context display   │      │  Value × 0.2                           │     ││  │  Audit trail       │      │                                        │     ││  └─────────┬──────────┘      └────────────────────────────────────────┘     ││            │ corrections                                                     │└────────────┼────────────────────────────────────────────────────────────────┘             │             ▼┌─────────────────────────────────────────────────────────────────────────────┐│                          FEEDBACK LOOP                                       ││                                                                              ││  ┌────────────────────┐      ┌────────────────────┐                         ││  │  Dataset Manager   │      │  Calibration       │                         ││  │  ─────────────     │      │  Monitor           │                         ││  │  Version annotations│ ◀──▶│  Refit temperature │                         ││  │  Trigger retraining│      │  every N samples   │                         ││  └────────────────────┘      └────────────────────┘                         ││                                                                              ││  ┌────────────────────────────────────────────────────────────────────┐     ││  │  Observability: MLflow (experiments) + Prometheus + Grafana        │     ││  └────────────────────────────────────────────────────────────────────┘     │└─────────────────────────────────────────────────────────────────────────────┘

Documents flow through preprocessing (image enhancement, layout detection) into the VLM inference layer, where olmOCR-2-7B-MLX extracts structured fields. The self-consistency engine runs multiple inferences for difficult regions, using agreement ratio as a confidence signal.

The calibration layer transforms raw model outputs into reliable probabilities using temperature scaling—learned separately per field type from human correction data. The routing layer then decides: autopilot (high confidence), flag-for-review (borderline), or human review (low confidence).

The human-in-the-loop layer provides a Streamlit-based review interface with priority queuing and keyboard shortcuts. The active learning selector chooses which samples are most valuable to annotate.

Finally, the feedback loop closes the circle: corrections update the calibrator, trigger dataset versioning, and feed into model retraining. The entire system is instrumented with MLflow for experiments and Prometheus/Grafana for operational metrics.

The target: autopilot handles 70%+ of documents without human intervention, while maintaining accuracy guarantees on the remaining 30% through targeted review.

Lesson 1: Self-Consistency Beats Single Inference

My first approach used a single inference pass: run the VLM once, take the output, use the logit-based confidence. The problem? Single-pass confidence is noisy. The same input might produce different outputs on different runs, but a single inference can't reveal that instability.

The fix is self-consistency sampling: run inference k times with different random seeds, then use majority voting to select the consensus answer. The agreement ratio becomes a confidence signal—if 5 out of 5 samples agree, you have high confidence; if it's 3-2, something is ambiguous.

async def run_self_consistency(image, prompt, adapter, k=3):    # Run k parallel inferences with different seeds    predictions = await asyncio.gather(*[        run_inference(image, prompt, seed=i) for i in range(k)    ])    # Majority voting: serialize predictions, count occurrences    serialized = [json.dumps(p, sort_keys=True) for p in predictions]    most_common, count = Counter(serialized).most_common(1)[0]    agreement_ratio = count / k    return json.loads(most_common), agreement_ratio

The intuition: if the model is uncertain, different seeds will produce different outputs. If the model is confident and correct, all samples will converge. Self-consistency surfaces the former while confirming the latter.

There's a cost tradeoff: k=3 means 3x inference time. For simple text fields, k=1 is fine. For complex regions (handwriting, tables, checkboxes), k=5 is worth the latency. The system uses region type to dynamically select k.

The lesson: Single-pass confidence is unreliable. Self-consistency sampling reveals model uncertainty that logits alone hide.

Lesson 2: Raw Probabilities Aren't Probabilities

Even with self-consistency, raw confidence scores are miscalibrated. A model might report 85% confidence and be right 95% of the time—or report 85% and be right only 70% of the time. Without knowing which, you can't set reliable routing thresholds.

Temperature scaling solves this. You collect a validation set of (confidence, is_correct) pairs, then learn a single temperature parameter T that rescales logits:

class TemperatureScaling:    def fit(self, confidences, correctness):        # Convert confidences to logits        logits = np.log(confidences / (1 - confidences))        # Optimize T to minimize negative log-likelihood        def nll(T):            scaled = 1 / (1 + np.exp(-logits / T))            return -np.mean(correctness * np.log(scaled) +                           (1 - correctness) * np.log(1 - scaled))        result = minimize(nll, x0=1.0, bounds=[(0.1, 10.0)])        self.temperature = result.x[0]    def calibrate(self, confidence):        logit = np.log(confidence / (1 - confidence))        return 1 / (1 + np.exp(-logit / self.temperature))

After calibration, if the model says 95%, it's right 95% of the time (±calibration error). This makes threshold-based routing meaningful: "send to review if confidence < 95%" actually means something.

The calibrator retrains continuously. Every human correction is a calibration sample. A CalibrationMonitor accumulates samples and refits temperature every N predictions, tracking Expected Calibration Error (ECE) over time.

The lesson: Raw model confidence is a ranking signal, not a probability. Temperature scaling transforms it into something you can actually trust for decision-making.

Lesson 3: Human Corrections Are Your Most Valuable Data

Every document routed to human review is an opportunity. When a reviewer corrects a field, you get ground truth on a case the model found difficult. That's exactly the data you need to improve.

But not all corrections are equally valuable. A typo fix on an easy case teaches less than a structural correction on a hard one. Active learning selects which samples to prioritize for annotation based on their expected learning impact.

The selection strategy combines three signals:

def select_annotation_batch(predictions):    # Uncertainty: low-confidence samples are more informative    uncertainty_scores = 1 - np.array([p['confidence'] for p in predictions])    # Diversity: cluster samples, prefer underrepresented areas    diversity_scores = compute_cluster_distances(predictions)    # Value: high-difficulty, common patterns have more impact    value_scores = predictions['difficulty'] * 0.6 +                   predictions['frequency'] * 0.4    # Weighted combination    total = (0.5 * uncertainty_scores +             0.3 * diversity_scores +             0.2 * value_scores)    return top_k_indices(total, k=batch_size)

Uncertainty sampling alone leads to annotating outliers that don't generalize. Diversity sampling ensures coverage. Value scoring weights toward patterns that appear frequently—fixing a common error type helps more than fixing a rare one.

The feedback loop: corrections flow into a versioned annotation store, triggering dataset updates and model retraining on a configurable schedule. Each retrained model gets validated against a golden test set before deployment.

The lesson: Human review isn't just a safety net—it's a data generation mechanism. Active learning extracts maximum value from every correction.

Lesson 4: Different Fields Need Different Confidence

A name field and a medication dosage field have different error costs. Getting the name slightly wrong (Robert vs Roberto) is recoverable. Getting the dosage wrong (15 mg vs 150 mg) could cause patient harm.

Rather than a single confidence threshold, the system uses per-field calibration. High-stakes fields like medication names and dosages have strict thresholds (98%). Lower-stakes fields like clinic names have relaxed thresholds (90%).

Field types also get different self-consistency settings. Checkboxes use k=5 because they're error-prone. Printed text uses k=1 because it's usually unambiguous. The prompt selector chooses templates based on region type, and the confidence scorer combines region difficulty with field-specific calibration.

thresholds = {    'medication_name': 0.98,    'dosage': 0.98,    'patient_name': 0.95,    'date': 0.95,    'clinic_name': 0.90,    'general_text': 0.85,}if confidence < thresholds[field_type]:    route_to_human_review(document, field)

The calibrator maintains separate temperature parameters per field type. Medication fields might need T=1.5 (model is overconfident), while dates might need T=0.8 (model is underconfident). The system learns these automatically from correction data.

The lesson: Not all fields are equal. Calibrate and threshold based on error cost, not just raw accuracy.

Lesson 5: The Annotation Interface Is a Product

I initially treated the human review UI as an afterthought—a simple form to view extractions and submit corrections. That was a mistake. Reviewer efficiency directly impacts system economics. A UI that takes 60 seconds per correction costs twice as much as one that takes 30 seconds.

The Streamlit-based review interface became a first-class product:

Priority queue: Items sorted by confidence (lowest first), field type (high-stakes first), and document age (oldest first)
Context preservation: Show the original image region alongside the extracted text, with bounding boxes highlighted
Keyboard shortcuts: Accept, reject, and edit without touching the mouse
Audit trail: Every correction tracked with reviewer ID, timestamp, and before/after values
Bulk actions: Accept all high-confidence items with one click when the queue backs up

The queue also shows aggregate metrics: current backlog size, average review time, acceptance rate by field type. Reviewers can see patterns—if checkbox extractions are consistently wrong, that's feedback to the model team.

The lesson: Human-in-the-loop is only sustainable if the loop is fast. Invest in annotation UX as if it were user-facing (it is).

Lesson 6: Observability Is Non-Negotiable

In a system with this many moving parts—preprocessing, inference, calibration, routing, review, retraining—things break in subtle ways. A calibrator that stops updating. A queue that grows faster than reviewers can drain it. A retrained model that's worse than its predecessor.

The observability stack tracks everything:

MLflow: Every inference logged with input hash, prompt version, output, confidence, latency. Every model version registered with metrics and lineage.
Prometheus: Gauges for queue depth, p50/p95 inference latency, autopilot vs review ratio, calibration error.
Grafana dashboards: Real-time visualization of throughput, accuracy by field type, reviewer productivity, model comparison.

Every request carries a correlation ID that propagates through all services. When something goes wrong, you can trace from API request → preprocessing → inference → calibration → routing → review → export in a single query.

Alerts fire on SLO breaches: queue depth > 1000, p95 latency > 30s, ECE > 0.05. The on-call has dashboards and runbooks to diagnose and recover.

The lesson: ML systems fail silently. Instrument everything, alert on meaningful thresholds, and make debugging possible before you need it.

What I'd Do Differently

1. Start with calibration, not model architecture: I spent too long optimizing the VLM before realizing that calibration was the bottleneck. A well-calibrated weaker model beats an uncalibrated stronger one for production routing decisions.

2. Build synthetic data generation earlier: The system includes a synthetic form generator for testing, but it came late. Having diverse synthetic data from day one would have caught more edge cases before production.

3. Add fallback OCR sooner: For text-heavy regions, traditional OCR (Tesseract, PaddleOCR) is often faster and just as accurate. The system now routes by region type, but initially it VLM'd everything.

4. Implement canary deployments for prompts: Prompts are code. Changes can break extraction in subtle ways. The system now versions prompts and supports canary rollouts, but we learned this after a prompt change caused a 15% accuracy regression.

Key Takeaways

Self-consistency reveals uncertainty: Run inference multiple times and use agreement ratio as a confidence signal—it's more reliable than logit-based probabilities.
Calibrate before routing: Raw model confidence is miscalibrated. Temperature scaling makes thresholds meaningful.
Corrections are training data: Every human review is an active learning opportunity. Select samples strategically to maximize model improvement.
Threshold by error cost: High-stakes fields need higher confidence thresholds. Calibrate and route per field type.
Instrument everything: ML systems fail silently. Correlation IDs, MLflow tracking, and Prometheus metrics make debugging possible.

This project is a medical document processing platform built with olmOCR-2-7B-MLX, FastAPI, MLflow, and Streamlit. Designed for HIPAA-compliant healthcare workflows with human-in-the-loop quality assurance.

Tiny Reasoner: Engineering a 135M Model for Medical Expertis

Mon, 12 Jan 2026 00:00:00 GMT

Note: This RSS feed strips out SVGs and embeds. You might want to read the post on the webpage here.

How I built a specialized medical coding agent using RLVR, Chain-of-Thought, and Self-Amplification.

November 14, 2025

Last Updated: January 12, 2026

1. Introduction

Large Language Models (LLMs) like GPT-4 are generalists. But for specific, high-stakes domains like medical coding (assigning ICD-10/CPT codes to clinical notes), we don't always need a 100B+ parameter giant. We need a specialist.

The Problem: Medical Coding is Complex

Medical coding is the process of translating clinical narratives into standardized billing codes. Consider this example:

"Patient presents with uncontrolled Type 2 diabetes mellitus with diabetic retinopathy."

A human coder must:

Identify the primary condition (Type 2 diabetes).
Note the complication (retinopathy).
Determine specificity (uncontrolled).
Assign the correct ICD-10 code: E11.319 (Type 2 diabetes with unspecified diabetic retinopathy, uncontrolled).

This requires domain expertise, reasoning through clinical guidelines, and understanding code hierarchies. A wrong code can lead to claim denials, compliance issues, or lost revenue.

Why Not Just Use GPT-4?

While large models like GPT-4 can perform this task, they have critical limitations:

Cost: Running inference on 70B+ models at scale (millions of claims/day) is prohibitively expensive.
Latency: Large models are slow, creating bottlenecks in real-time workflows.
Hallucination: Without grounding, they can invent plausible-sounding but invalid codes.

My Hypothesis

A tiny, specialized model (135M parameters) can outperform generalist giants on this specific task if we:

Teach it how to reason (Chain-of-Thought).
Ground it with verifiable rewards (rule-based validation).
Scale it with self-amplification (iterative self-training).

This post details my engineering journey building Tiny Reasoner, a project that proves this hypothesis.

The Result (Sneak Peek)

Before diving into the architecture, here is what the final model actually does:

Input: "Patient presents with uncontrolled Type 2 diabetes mellitus        with mild nonproliferative diabetic retinopathy."Tiny Reasoner (135M):+---------------------------------------------------------------+| <think>                                                       || 1. Identify primary condition: Type 2 Diabetes Mellitus.      || 2. Identify complications: Retinopathy (Diabetic).            || 3. Specify severity: Mild, Nonproliferative.                  || 4. Check control status: Uncontrolled (Hyperglycemia).        || 5. Construct Code: E11 (Type 2) + .3 (Retinopathy)            ||    + .2 (Mild Nonproliferative) + .9 (Unspecified eye).       || </think>                                                      ||                                                               || Code: E11.329                                                 |+---------------------------------------------------------------+

It doesn't just guess; it derives the answer.

2. System Architecture

My pipeline is designed to bootstrap reasoning capabilities from zero, verify them rigorously, and then scale them up autonomously.

Design Philosophy

I moved away from the standard "SFT → DPO" pipeline used in most LLM alignment work. Instead, I built a verification-centric architecture where every training decision is grounded in deterministic correctness checks.

Key Insight: In medical coding, "correctness" isn't subjective. A code is either valid (exists in the ICD-10 database) or it isn't. This allows me to replace expensive human preference labeling with a Python script.

High-Level Data Flow

[ Raw Data Sources ]       [ Phase 1: Foundation ]       [ Phase 2: Alignment ]+------------------+       +---------------------+       +--------------------+|  CMS Guidelines  |------>|  Extraction Script  |------>|    SFT Trainer     ||  (Text/PDF)      |       |  (Regex/Parsing)    |       |   (SmolLM2-135M)   |+------------------+       +---------------------+       +--------------------+                                    |                              |+------------------+                v                              v|   OpenAI API     |------>| Synthetic Generator |       +--------------------+| (Teacher Model)  |       | (CoT + Validation)  |       |    RLVR Trainer    |+------------------+       +---------------------+       |    (PPO + Rule)    |                                                         +--------------------+                                                                   |                                                                   v                                                         [ Phase 4: Scaling ]                                                         +--------------------+                                                         |  Self-Amplification|                                                         |  (Iterative Loop)  |                                                         +--------------------+

3. Phase 1: The Foundation (Bootstrapping)

Before a model can "reason," it needs to see examples of reasoning. I couldn't just use raw text; I needed Chain-of-Thought (CoT) data.

3.1. Data Extraction Strategy

The Challenge: Most medical coding datasets are either proprietary or lack reasoning traces. I needed examples that showed how to arrive at a code, not just the final answer.

The Solution: I built scripts/extract_cms_guidelines.py to mine "gold standard" examples from official CMS (Centers for Medicare & Medicaid Services) documentation.

CMS guidelines often contain structured examples like:

"Example: Patient admitted for acute bronchitis...
Rationale: The condition is acute (not chronic), so we select from the J20 family...
Code Assignment: J20.9"

I used regex patterns to extract these structured examples:

Raw Text:"Example: Patient with Type 2 diabetes... Rationale: Code E11.9 is selected because..."      |      v  (Regex Parsing)      |Structured JSONL:{  "narrative": "Patient with Type 2 diabetes...",  "reasoning": "Code E11.9 is selected because...",  "code": "E11.9"}

3.2. Synthetic Teacher Generation

The Challenge: CMS guidelines only yielded ~5,000 examples. I needed 15,000+ to train a robust model.

The Solution: I used GPT-4o as a "teacher" to generate synthetic examples. But not just any examples—I enforced a strict Chain-of-Thought (CoT) schema.

The Prompt Strategy

I designed a prompt that forces the teacher model to:

Generate a realistic clinical narrative (50-200 words).
Show its reasoning step-by-step in <think> tags.
Assign a valid ICD-10/CPT/HCPCS code.

prompt = """You are a medical coding expert. Generate a realistic clinical narrative,step-by-step reasoning process, and correct medical billing code.Requirements:- Narrative should be 50-200 words- Include 4-6 reasoning steps in <think> tags- Assign ONE code (ICD-10, CPT, or HCPCS)- Code must be valid and commonly usedFocus on: {specialty}Code family: {code_family}Complexity level: {complexity}"""

The Validation Gate

Critical Problem: LLMs hallucinate. GPT-4 might generate "ICD-10: Z99.999 (Patient abducted by aliens)."

My Solution: Every generated example passes through CodeValidator before being added to the training set.

# scripts/generate_synthetic.pyfor attempt in range(max_attempts):    response = gpt4.generate(prompt)    parsed = parse_response(response)    # CRITICAL: Validate the code    if not validator.validate(parsed['code'], parsed['code_type']):        continue  # Discard hallucinated codes    training_data.append(parsed)

This filtering step is crucial. In my experiments, ~8-10% of GPT-4's generated codes were invalid or non-existent.

How to Run Phase 1:

python scripts/extract_cms_guidelines.pypython scripts/generate_synthetic.py --count 5000python data/prepare_foundation.py

4. Phase 2: RLVR (The "Secret Sauce")

This is where the project diverges from standard LLM training.

Why Not DPO?

Direct Preference Optimization (DPO) is the current state-of-the-art for alignment. But it has a fatal flaw for medical coding:

DPO requires preference pairs: "Response A is better than Response B."

Creating these pairs requires:

Generating multiple responses per prompt.
Having human experts rank them.
Training a reward model to predict preferences.

This is:

Expensive: ~$500-700 for 40k preference pairs.
Noisy: Human annotators might disagree.
Gameable: The reward model can be fooled by plausible-sounding but incorrect codes.

The RLVR Alternative

For medical coding, I had a ground truth: The Code Book.

I implemented RLVR (Reinforcement Learning with Verifiable Rewards) in post_training/rlvr.py. Instead of a neural reward model, I used a deterministic Python function.

4.1. The Reward Function

Instead of a neural network, my reward function is a Python script:

      Model Output                Ground Truth     "Code: E11.9"               "Code: E11.9"           |                           |           v                           v    +-------------+             +-------------+    |  Parser     |             |  Parser     |    +-------------+             +-------------+           |                           |           v                           v      [E11.9] ==?== [E11.9] ---->  +1.0 Reward           |      [E11]   ==?== [E11]   ---->  +0.5 Reward (Partial)

def compute_reward(prediction, ground_truth):    # 1. Extract code from model output    pred_code = extract_code(prediction)    # 2. Check Validity (Is it a real ICD-10 code?)    if not validator.validate(pred_code):        return 0.0    # 3. Check Correctness    if pred_code == ground_truth:        return 1.0    elif pred_code.family == ground_truth.family:        return 0.3  # Partial credit!    return 0.1 # Valid but wrong

4.2. PPO Training Loop

I used Proximal Policy Optimization (PPO) to optimize the model against this rule-based reward.

How PPO Works (Simplified):

The model generates a response.
The verifier computes a reward (0.0 to 1.0).
PPO adjusts the model's weights to increase the probability of high-reward responses.
A KL-divergence penalty prevents the model from deviating too far from the original SFT model (preventing "reward hacking").

The Training Loop:

for batch in dataloader:    # 1. Generate responses    prompts = batch['prompts']    responses = model.generate(prompts, temperature=0.8)    # 2. Compute rewards    rewards = []    for response, ground_truth in zip(responses, batch['ground_truths']):        reward = verifier(response, ground_truth)        rewards.append(reward)    # 3. PPO step (update model weights)    ppo_trainer.step(prompts, responses, rewards)

Why This Works: The model learns: "I don't just need to sound confident; I need to output codes that pass the validator."

   +-------------+        Action (Text)       +-------------+   |   Policy    |--------------------------->| Environment |   | (The Model) |                            | (Verifier)  |   +-------------+                            +-------------+          ^                                          |          |             Reward (Float)               |          +------------------------------------------+

This teaches the model: "I don't know exactly what to write, but I know that if I output a valid code that matches the family, I get a reward."

How to Run Phase 2:

# First, SFTpython post_training/sft.py --config-path post_training/config/foundation_sft.yaml# Then, RLVRpython post_training/rlvr.py --config post_training/config/rlvr.yaml

5. Phase 4: Self-Amplification (Scaling Up)

Once the model was "good enough" (Phase 2), I used it to teach itself. This is the Self-Amplification loop implemented in scripts/iterative_training.py.

5.1. The Loop

      +------------------+      |  Current Model   |      +------------------+               |      (1) Generate (Temp=0.8)               v      +------------------+      +------------------+      |  Pseudo-Labels   |----->|  CodeValidator   |      +------------------+      +------------------+                                         |                                    (2) Filter                                         v      +------------------+      +------------------+      |   Next Model     |<-----|  Clean Dataset   |      +------------------+      +------------------+             (3) Train

Generate: The model looks at a clinical narrative and generates a new reasoning path and code.
Filter: I run the CodeValidator.
- Invalid code? -> Trash.
- Valid code? -> Keep.
Train: I add the "Keep" examples to the training set and re-train.

5.2. Why This Works

The Intuition: The model has learned how to reason about medical codes (from Phase 1 & 2). Now, when it sees a new narrative, it can generate a plausible reasoning path. Sometimes, it discovers a correct path that wasn't in the original training data.

The Mechanism:

Exploration: By using temperature=0.8 (sampling), the model generates diverse responses.
Validation: The CodeValidator acts as a filter, keeping only correct examples.
Crystallization: By retraining on these validated examples, I "lock in" the correct reasoning patterns.

The Flywheel Effect:

Better model → Better pseudo-labels → Better model → ...

Scaling Results:

Iteration 1: 20k labeled examples → 85% accuracyIteration 2: 20k labeled + 30k pseudo → 89% accuracyIteration 3: 20k labeled + 80k pseudo → 92% accuracy

Key Insight: I didn't need to label 100k examples manually. I labeled 20k, and the model generated the rest (with validation).

Iteration 1:  20k Examples  ->  85% Accuracy      |      v (Generate & Filter)      |Iteration 2:  50k Examples  ->  89% Accuracy      |      v (Generate & Filter)      |Iteration 3:  100k Examples ->  92% Accuracy

How to Run Phase 4:

python scripts/iterative_training.py \\    --foundation-model outputs/rlvr/final \\    --iterations 3

6. Evaluation

I don't just check accuracy; I check reasoning quality.

Evaluation Metrics

Exact Match: Did the model predict the exact code? (e.g., E11.319)
Code Validity: Is the predicted code even real? (Checks against ICD-10/CPT databases)
Family Match: Did it get the right code family? (e.g., E11.* for Type 2 diabetes)
Reasoning Presence: Did it use the <think> tags?
Reasoning Quality: Are the reasoning steps coherent and relevant?
Inference Time: How fast is the model? (Target: <200ms per prediction)

Evaluation Script

My evaluation/final_evaluation.py script:

Loads the model.
Runs inference on multiple test sets (validation, CMS holdout, synthetic hard cases).
Computes all metrics.
Generates a JSON report with detailed breakdowns.

python evaluation/final_evaluation.py \\    --model-path outputs/iterative/iter_3/final \\    --test-sets data/processed/foundation_val.jsonl data/test/cms_holdout.jsonl

Key Takeaway: The model is not only accurate but also explainable (shows reasoning) and efficient (runs on CPU).

7. Engineering Challenges & Solutions

Challenge 1: Reward Hacking

Problem: Early in RLVR training, the model learned to output any valid code to get partial credit, ignoring the clinical narrative.

Example:

Prompt: "Patient has a broken arm."
Model Output: E11.9 (Type 2 diabetes) — Valid code, but completely wrong!

Solution: I tuned the KL-divergence penalty in PPO. This forces the model to stay close to the original language distribution (the SFT model) while still optimizing for the reward. I also reduced the partial credit reward from 0.5 to 0.3.

Challenge 2: Slow Pseudo-Label Generation

Problem: Generating 80,000 examples for self-training was taking 3+ days on a single GPU.

Solution: I rewrote scripts/generate_pseudo_labels.py to use batch generation:

Instead of generating one example at a time, I batch 8 prompts together.
Used masked attention to handle variable-length sequences.
Result: 8x speedup (3 days → 9 hours).

Challenge 3: "Empty Thinking"

Problem: The model sometimes output <think></think> (empty tags) and then the code, gaming the "reasoning presence" metric.

Solution: I added a heuristic to the reward function:

if len(reasoning_text) < 50:  # Too short    return 0.0  # Zero reward

This forced the model to generate substantive reasoning.

8. What Didn't Work (Lessons from Failed Experiments)

Not everything worked on the first try. Here are some approaches I tried that failed:

Failed Approach 1: Pure Synthetic Data

What I Tried: Initially, I attempted to train the model entirely on GPT-4 generated synthetic data (no CMS extraction).

Why It Failed: The model learned GPT-4's "style" but lacked grounding in official guidelines. When evaluated on real CMS examples, accuracy dropped by ~12%. The model was mimicking a teacher that itself didn't have access to authoritative sources.

Lesson: Domain expertise matters. Even with a powerful teacher model, you need gold-standard examples from authoritative sources.

Failed Approach 2: Standard DPO with Human Preferences

What I Tried: Before implementing RLVR, I experimented with collecting preference pairs (having annotators rank "Response A vs Response B").

Why It Failed:

Annotators frequently disagreed on which response was "better" when both codes were valid but different levels of specificity.
The reward model learned to prefer longer, more verbose reasoning, not necessarily correct codes.
Cost was prohibitive ($700 for just 5k pairs).

Lesson: When you have ground truth, use it. Don't introduce subjectivity where objectivity exists.

Failed Approach 3: Single-Shot Self-Training

What I Tried: I attempted to generate all 80k pseudo-labels in one go (no iterations).

Why It Failed: The model at 85% accuracy generated too many incorrect examples. Even with filtering, ~30% of pseudo-labels were subtly wrong (e.g., correct family, wrong specificity). Training on this degraded performance.

Lesson: Iterate gradually. Self-amplification works best when you start with a strong foundation and scale incrementally.

9. Lessons Learned

After 8 weeks and countless experiments, here are the key insights I gained:

1. Small Models Need Quality, Not Quantity

A 135M model has limited capacity. Feeding it 100k mediocre examples is worse than 20k high-quality examples with clear reasoning traces. Quality > Quantity for specialized domains.

2. Deterministic Verification Beats Learned Rewards

For tasks with objective correctness (math, coding, medical codes), a simple Python function outperforms a neural reward model. It's:

Cheaper (no training needed)
More reliable (can't be fooled)
Easier to debug (just read the code)

3. Self-Amplification Only Works with Strict Filtering

The model will happily learn from its own mistakes if you let it. The validator must be unforgiving—99% valid isn't good enough. I set the threshold at 100% code validity.

4. Reasoning Traces Are Worth the Effort

Adding <think> tags increased training complexity, but the payoff was huge:

Easier debugging (I could see why the model chose a code)
Better generalization (the model learned patterns, not memorization)
Trust from end-users (medical coders want to see the reasoning)

5. Engineering Matters More Than Algorithms

The biggest performance gains came from:

Better data extraction (CMS guidelines)
Smarter filtering (validation gates)
Batch optimization (8x speedup)

Not from tweaking hyperparameters or trying fancier RL algorithms.

10. Future Directions

Thought-Level Verification: Grading the reasoning steps themselves, not just the final code (Phase 3 implementation ready).
Model Distillation: Compressing the reasoning patterns into an even smaller model (<50M parameters) for edge deployment.
Multi-Code Support: Extending the system to handle cases requiring multiple codes (e.g., diabetes + hypertension).
Real-Time Feedback Loop: Integrating with hospital EHR systems for continuous learning from coder corrections.

Conclusion

By combining symbolic verification (the code validator) with neural generation (the LLM), I built a system that is:

Accurate: 89% exact match, 99.2% code validity. It can't hallucinate invalid codes.
Explainable: Provides step-by-step reasoning traces in <think> tags, crucial for medical compliance.
Efficient: 500MB model, runs on CPU, <150ms inference time. Deployable in local hospital environments.
Cost-Effective: 64% cheaper than standard approaches, with superior results.

The Key Insight

This project proves that for specialized domains, smart engineering beats raw parameter count.

You don't need a 70B model to be an expert. You need:

High-quality data (CMS guidelines + validated synthetic).
Verifiable rewards (deterministic correctness checks).
Self-amplification (the model teaching itself).

The future of AI isn't just bigger models—it's smarter training pipelines.

Personal Reflection

When I started this project, I thought the hard part would be the RL training. I was wrong.

The hard part was data engineering—extracting structured examples from messy PDFs, designing prompts that forced reasoning, and building filters that caught subtle errors. The RL part was almost trivial once I had high-quality data and a reliable validator.

The biggest surprise? How well self-amplification worked. I expected diminishing returns after iteration 1, but the model kept improving. By iteration 3, it was generating reasoning paths I hadn't seen in the training data—it had genuinely learned how to think about medical codes, not just memorize patterns.

What would I do differently? I'd start with the validator earlier. I initially built it as an afterthought for evaluation, but it became the cornerstone of the entire pipeline. If I were to do this again, I'd design the validator first, then build the training pipeline around it.

This project taught me that constraints breed creativity. Working with a tiny 135M model forced me to be smarter about data, training, and evaluation. And in the end, that smartness beat brute force.

Check out the full code on GitHub

Medical Reasoning-Service

Wed, 07 Jan 2026 00:00:00 GMT

Note: This RSS feed strips out SVGs and embeds. You might want to read the post on the webpage here.

Building a Reasoning Engine for Healthcare: Ditching Vector Search for Tree Traversal

November 14, 2025

Last Updated: January 7, 2026

projects

In the world of automated decision-making, "Retrieval-Augmented Generation" (RAG) has become the standard answer for connecting LLMs to private data. The recipe is usually simple: chunk your text, embed it into vectors, and retrieve the top-k chunks based on cosine similarity.

But when I set out to build a Prior Authorization Reasoning Service: a system that determines if a patient meets the strict clinical criteria for a medical procedure: I found that standard vector RAG wasn't enough. Medical policies aren't just collections of semantic similarity; they are hierarchical rule sets where context, parent-child relationships, and specific exclusionary clauses matter more than general topic overlap.

This is the technical story of how Ibuilt a reasoning engine that moves beyond flat vectors to hierarchical tree search, auditable ReAct agents, and a custom Go database built from scratch.

The Problem: "Is this MRI covered?"

The core question our service answers is: "Is this patient referral ready to file according to the current policy?"

To answer this, a human (or machine) must:

Read a 50+ page PDF policy (e.g., "Lumbar MRI Coverage").
Understand the structure: "Section 1.2 applies, but ONLY if Section 1.1.4 is NOT met."
Cross-reference patient facts (age, diagnosis codes, treatment history).
Make a binary decision with a citation: "Approved because of Section 1.2(a), page 14."

```Patient Packet ──► Policy PDF ──► Structured Notes      │                │                 │      ├─ demographics  ├─ tables         ├─ coverage clauses      │                │                 │      ▼                ▼                 ▼   Facts DB     Section map        Rule references      │                │                 │      └───────────► Decision Rationale ◄─┘```

Standard vector search fails here because it flattens the document. It might retrieve a paragraph about "exclusion criteria" without retrieving the parent header that says "For pediatric patients only," leading to hallucinations.

The Solution Architecture

Our system is composed of three specialized components, each handling a distinct part of the reasoning pipeline:

PageIndex (Ingestion & Retrieval): Turns PDFs into semantic trees, not flat chunks.
TreeStore (Storage Engine): A custom high-performance Go database optimized for hierarchical documents.
Reasoning Service (The Brain): A Python/FastAPI application running an LLM-driven ReAct controller.

┌─────────────┐     ┌──────────────┐     ┌──────────────┐     ┌─────────────────────┐│ Policy PDFs │ ──▶ │  PageIndex   │ ──▶ │  TreeStore   │ ──▶ │ ReAct Reasoning Loop │└─────────────┘     │  (structure) │     │ (hierarchy)  │     │  (decide & cite)    │                     └──────┬───────┘     └──────┬───────┘     └──────┬──────────────┘                            │                    │                       │                            ▼                    ▼                       ▼                      JSON policy tree     Node-oriented reads     Approval decision

1. PageIndex: Structure First, Text Second

Instead of "chunking" text, Iuse PageIndex to convert visual PDFs into a hierarchical JSON tree.

Input: A raw PDF.
Process: OCR + Layout Analysis + LLM-based Tree Construction.
Output: A structured tree where every node knows its parent, children, page range, and summary.

PDF Page   │   ├─► OCR Blocks   │      │   │      ├─► Layout Detector (coordinates, fonts)   │      │        │   │      │        └─► Section Classifier   │      │                 │   │      │                 ▼   │      └──────────────► Node Builder ─────────┐   │                                            │   ▼                                            ▼Tree Root ──► Policy 1.0 ──► 1.1 Indications ──► 1.1.4 Exceptions                │                │                     │                │                └─► Child summaries   └─► Tokens + cite span                └─► Page anchors & parent pointers

This allows us to perform LLM Tree Search. Instead of searching for keywords, the model traverses the document tree:

"I need coverage criteria. I'll open node 1.0 (Policy). Now I see 1.1 (Indications) and 1.2 (Limitations). I'll expand 1.1..."

This mimics how a human reads a table of contents, preserving the vital context that vector databases destroy.

2. TreeStore: A Database Built for Trees

Irealized Ineeded a data store that treats "parent-child" relationships as first-class citizens. Relational DBs (SQL) are clunky for deep trees; Graph DBs are overkill.

So, Ibuilt TreeStore in Go. (Also because I wanted to understand the foundations of building a database)

Core Tech: A custom B+Tree implementation with ACID transactions.
Design: Optimized for GetNode, GetChildren, and GetSubtree operations.
Performance: Capable of 100k+ reads/sec, ensuring that our reasoning agent can traverse complex policies with near-zero latency.
Features: Native support for document versioning (vital for medical policies that change yearly) and metadata indexing.

It’s a single-binary, embeddable database that bridges the gap between a file system and a document store.

          [Root]          /   \\     [Index] [Leaf]             / | \\        Node refs  ► On-disk pages

3. The ReAct Controller: Auditable Reasoning

The heart of the Python service is the ReAct (Reason + Act) Controller. Imoved away from "one-shot" prompts to an agentic loop.

The Controller doesn't just guess; it investigates:

THINK: "I need to check the patient's age and diagnosis code."
ACT: Call tool facts.get("age").
OBSERVE: "Patient is 45."
THINK: "Now I need to find the age limit in the policy."
ACT: Call tool pi.search("age requirement").
OBSERVE: "Policy Section 3.1 requires age > 18."
DECIDE: "Criteria MET."

This generates a reasoning trace:a step-by-step log of how the decision was reached. In healthcare, this auditability is non-negotiable.

┌─────────┐    THINK    ┌────────────┐    ACT     ┌──────────────┐│ Context │ ───────────▶│ Policy Tree│────────────▶│ Tool Calls   │└─────────┘             └────┬───────┘            └─────┬────────┘     ▲                       │                         │     │                 OBSERVE│                         │RESULTS     │                       ▼                         ▼┌─────────┐◀────────────┌────────────┐       ┌────────────────────┐│ Patient │   UPDATE     │ Evidence   │◀─────│ Trace + Confidence ││  Facts  │              │  Buffer    │       └────────────────────┘└─────────┘              └────────────┘                  │                                                       DECIDE                                                         │                                                 Approve / Uncertain / Deny

Key Technical Decisions

Safety & "Uncertainty"

Iimplemented a strict safety gate. If the LLM's confidence score (a composite of retrieval relevance and evidence alignment) drops below 0.65, the system refuses to guess. It returns a status of UNCERTAIN, flagging the case for human review. Iprefer a "I don't know" over a hallucinated "Yes."

Confidence Meter┌───────────────┬───────────────┐│ 0.0 ────────► │ 0.65 threshold│ ────────► 1.0└───────────────┴───────────────┘        │                 │        ▼                 ▼  Human review        Automatic decision

Hybrid Retrieval

While Ifavor the Tree Search, Iacknowledge its limits. If a specific node is massive (e.g., a 20-page list of CPT codes), the agent can trigger a local SQLite FTS5 fallback. This performs a classic BM25 keyword search inside the selected node to pinpoint the exact line, combining the best of structural navigation and keyword precision.

Tree Traversal     │     ▼Large Node? ──No────► Continue expanding children     │    Yes     │     ▼ [SQLite FTS5]     │     ▼Precise line hit ──► return span to controller

The "Shadow Mode" Rollout

Deploying LLM agents is scary. To mitigate risk, Ibuilt a Shadow Mode into our controller. The production traffic is handled by a heuristic (rule-based) controller, while the LLM agent runs in the background. Ilog and compare their decisions. Only when the LLM matches or beats the heuristic on our evaluation dataset do Iflip the switch.

        ┌─────────────────────────────┐        │ Incoming Prior Auth Request │        └──────────────┬──────────────┘                       │             ┌─────────▼─────────┐             │ Heuristic Control │──► Production decision             └─────────┬─────────┘                       │ mirroring inputs             ┌─────────▼─────────┐             │  LLM Shadow Loop  │──► Logged reasoning             └─────────┬─────────┘                       │                Offline comparison                       │                 Trigger cutover

Conclusion

The reasoning-service project demonstrates that for complex, high-stakes domains, structure is all you need. By respecting the document's hierarchy with PageIndex and TreeStore, and forcing the LLM to show its work via ReAct, we've built a system that doesn't just retrieve information: it actually understands it.

Check out the full code on GitHub

Aegis-Building a Fraud Detection System

Mon, 12 Jan 2026 00:00:00 GMT

Note: This RSS feed strips out SVGs and embeds. You might want to read the post on the webpage here.

Combining NLP and traditional ML to achieve 12x better fraud detection than manual review

November 6, 2025

Last Updated: January 12, 2026

projects

Insurance fraud costs the industry over $80 billion annually. Yet most fraud detection systems still rely on simple rule-based approaches: flag claims over $X, or those from first-time claimants. The problem? Fraudsters adapt, and rules don't.

I built Aegis, a fraud detection system that combines the semantic understanding of language models with the interpretability of gradient boosting. The result: 75% precision at 50% recall—meaning we catch half of all fraud while being right 3 out of 4 times. That's a 12x improvement over the manual baseline.

Here's what I learned building it.

The Problem: Fraud Hides in Plain Sight

Insurance claims have an inherent asymmetry: only ~6% are fraudulent. This means any system that simply predicts "not fraud" for everything achieves 94% accuracy—while being completely useless.

The real challenge isn't classification accuracy. It's prioritization: which claims should investigators spend their limited time on?

Manual review has two failure modes:

False positives: Wasting investigator time on legitimate claims
False negatives: Missing fraud entirely because the queue is too long

I needed a system that could surface high-risk claims without drowning investigators in false alarms.

The Architecture: Two Models Are Better Than One

The key insight came from analyzing what makes fraud claims different. Fraudsters leave two kinds of traces:

Behavioral patterns: Multiple past accidents, driving history, demographics
Linguistic patterns: Vague descriptions, timeline inconsistencies, suspiciously generic narratives

No single model handles both well. Transformers excel at understanding language but struggle with tabular data. Gradient boosting crushes tabular features but can't parse free text.

So I built a dual-engine system:

┌─────────────────────────────────────────────────────────────┐│                     Claim Submission                         ││  ┌──────────────────┬────────────────────────────────────┐  ││  │ Structured Data  │     Accident Narrative             │  ││  │ • Age: 26-39     │  "I was stopped at a light when   │  ││  │ • Past accidents │   another car sort of hit me..."  │  ││  │ • DUI history    │                                    │  ││  └────────┬─────────┴──────────────────┬─────────────────┘  ││           │                             │                    ││           ▼                             ▼                    ││  ┌─────────────────┐       ┌─────────────────────────────┐  ││  │    XGBoost      │◄──────│   DistilRoBERTa + LoRA      │  ││  │ (Decision Core) │       │   (Fraud Signal Extractor)  │  ││  └────────┬────────┘       └─────────────────────────────┘  ││           │                                                  ││           ▼                                                  ││  ┌─────────────────┐                                        ││  │   Calibrator    │                                        ││  │  (Isotonic Reg) │                                        ││  └────────┬────────┘                                        ││           │                                                  ││           ▼                                                  ││     Risk Score: 73/100                                       ││     Level: HIGH                                              ││     Top Factor: "Vague Language"                             │└─────────────────────────────────────────────────────────────┘

The language model extracts three fraud signals from the narrative:

Vague language ("sort of happened", "hard to say")
Timeline inconsistencies (conflicting dates, unclear sequence)
Pre-existing damage mentions (suspicious references to prior wear)

These become features fed into XGBoost alongside traditional structured data.

Lesson 1: LoRA Makes Fine-Tuning Practical

The standard approach to fine-tuning a transformer is updating all 82 million parameters. That's expensive, slow, and produces massive model files.

LoRA (Low-Rank Adaptation) changes the game. Instead of updating the full weight matrices, it learns low-rank decomposition matrices that get added to specific layers—typically just the attention weights.

# The magic of LoRA: only 1.5% of parameters are trainablefrom peft import LoraConfig, get_peft_modellora_config = LoraConfig(    r=8,                    # Rank of decomposition    lora_alpha=16,          # Scaling factor    target_modules=["query", "value"],  # Only attention layers    lora_dropout=0.1,    task_type="SEQ_CLS")model = get_peft_model(base_model, lora_config)# trainable params: 1.2M (1.5% of 82M)# model size: 30MB instead of 500MB

The result? Training time dropped 70%, and the model file went from 500MB to 30MB—while retaining 97% of full fine-tuning performance.

The lesson: Parameter-efficient fine-tuning isn't just for saving money. Small models deploy faster, cache better, and run on cheaper hardware.

Lesson 2: Multi-Label Beats Binary

My first approach was to train a binary classifier: "Is this claim fraudulent?"

It didn't work. The model learned to output ~6% probability for everything—matching the base rate while being useless for prioritization.

The fix was reframing the problem. Instead of predicting fraud directly, I trained the LLM to detect specific fraud signals:

# Multi-label classification: one forward pass, three signalsdef get_llm_features(text, model, tokenizer):    inputs = tokenizer(text, return_tensors="pt", max_length=128)    with torch.no_grad():        outputs = model(**inputs)        probs = torch.sigmoid(outputs.logits).cpu().numpy()[0]    return {        'has_vague_language_score': float(probs[0]),        'mentions_preexisting_damage_score': float(probs[1]),        'has_timeline_inconsistency_score': float(probs[2])    }

This approach has three advantages:

Interpretability: Users see why a claim is risky, not just that it is
Robustness: Even if one signal fails, others might catch fraud
Feature richness: Three continuous features are more informative than one binary

The lesson: When direct prediction fails, reframe the task to predict intermediate signals instead. Your downstream model can learn the complex mapping from signals to outcomes.

Lesson 3: Monotonic Constraints Encode Domain Knowledge

Here's a problem I didn't anticipate: the model learned that more driving experience increased fraud risk.

Wait, what?

It turned out the training data had a spurious correlation: older, experienced drivers filed higher-value claims (they drove nicer cars). The model picked this up as a fraud signal.

XGBoost has an elegant solution: monotonic constraints. You can tell the model that certain features must have one-directional effects:

# Constraints: +1 = must increase risk, -1 = must decrease, 0 = unconstrainedmonotone_constraints = {    'SPEEDING_VIOLATIONS': 1,    # More violations → higher risk    'DUIS': 1,                   # More DUIs → higher risk    'PAST_ACCIDENTS': 1,         # More accidents → higher risk    'DRIVING_EXPERIENCE': -1,    # More experience → lower risk    'has_vague_language_score': 1,    'mentions_preexisting_damage_score': 1,    'has_timeline_inconsistency_score': 1}xgb_model = XGBClassifier(    monotone_constraints=tuple(monotone_constraints.values()),    # ... other params)

During tree construction, XGBoost only considers splits that satisfy these constraints. No amount of spurious correlation can override them.

The lesson: Encode what you know. Monotonic constraints aren't just for interpretability—they regularize against data artifacts and make predictions defensible to regulators.

Lesson 4: Calibration Is Non-Negotiable

Here's a common mistake: treating model probabilities as real probabilities.

When my XGBoost model output 0.8, did that mean an 80% chance of fraud? Absolutely not. Raw model scores are uncalibrated—they measure relative ranking, not actual probabilities.

This matters because:

Investigators need to set thresholds ("review all claims above 50%")
Business decisions require expected value calculations
"90% fraud probability" must mean 9 in 10 similar claims are actually fraudulent

I used isotonic regression for calibration:

from sklearn.isotonic import IsotonicRegression# Fit on validation setcalibrator = IsotonicRegression(out_of_bounds='clip')calibrator.fit(val_proba_uncalibrated, val_labels)# At inference timeraw_probability = xgb_model.predict_proba(X)[:, 1]calibrated_probability = calibrator.predict(raw_probability)risk_score = int(calibrated_probability * 100)

Isotonic regression learns a non-decreasing mapping from raw scores to true probabilities. Unlike Platt scaling (which assumes logistic form), it makes no parametric assumptions—crucial when your model outputs have weird distributions.

The lesson: Raw probabilities are not probabilities. Always calibrate before presenting scores to users or making threshold-based decisions.

Lesson 5: Speed and Size Matter More Than You Think

My first version used SHAP for per-prediction explanations. Each prediction took 30 seconds. That's fine for batch processing, but useless for a real-time dashboard.

I made a pragmatic trade-off: global feature importance instead of per-prediction explanations.

# Fast: global feature importance (~0ms)feature_importance = xgb_model.feature_importances_top_indices = np.argsort(feature_importance)[-3:][::-1]# Slow: per-prediction SHAP (~30,000ms)# explainer = shap.TreeExplainer(xgb_model)# shap_values = explainer.shap_values(X)

The full optimization stack:

LoRA adapters: 500MB → 30MB
Joblib serialization: Efficient model caching
Streamlit @cache_resource: Models load once, persist across requests
Feature importance over SHAP: 30s → 300ms

Final footprint: 7MB total model artifacts, ~320ms inference time.

The lesson: Optimize for deployment from the start. A model that can't run in production is just a Jupyter notebook.

What I'd Do Differently

1. Use more sophisticated text augmentation: I generated synthetic narratives with rules. LLM-based augmentation would create more realistic training data.

2. Add image features: Real fraud detection uses damage photos. A future version could incorporate vision models for inconsistency detection.

3. Implement drift detection: Fraud tactics evolve. The model needs monitoring for distribution shift and scheduled retraining triggers.

4. A/B test thresholds: The 70/40 thresholds for HIGH/MEDIUM risk were chosen heuristically. In production, you'd optimize these against business metrics.

The Results

Metric	Value
Precision @ 50% Recall	75%
PR-AUC	0.82
Inference Latency	~320ms
Model Size	7MB

The 75% precision means: if investigators review our top-ranked claims until they've found half of all fraud, they'll be right 3 out of 4 times. Compared to random selection (~6% fraud rate), that's a 12x improvement in investigator efficiency.

Key Takeaways

Combine model strengths: Use transformers for what they're good at (language), gradient boosting for what it's good at (tabular data)
Prefer parameter-efficient fine-tuning: LoRA gives you 97% of full fine-tuning at 1.5% of the cost
Encode domain knowledge: Monotonic constraints prevent embarrassing predictions and satisfy regulators
Always calibrate: Raw model scores aren't probabilities—calibrate before making threshold decisions
Optimize for production: A 10x faster model with 95% accuracy beats a slow model with 96% accuracy

Aegis is open source and available on GitHub. The training pipeline, including LoRA fine-tuning and XGBoost calibration, runs on Colab.

Self-Healing Voice Agents: Automating Prompt Remediation with AIOps

Fri, 31 Oct 2025 00:00:00 GMT

Note: This RSS feed strips out SVGs and embeds. You might want to read the post on the webpage here.

System that detects voice agent failures and fixes them without human intervention

October 31, 2025

Voice agents break in ways you don't expect. A prompt that handles "I need an oil change" perfectly in testing fails when a customer says "my car's making a weird noise, probably needs service." The intent is the same, but the phrasing is different enough to confuse the classifier.

This is manageable when you're running one agent for one location. You see the failure in logs, you tweak the prompt, you move on. But when the same agent runs across 50 dealerships, each generating hundreds of calls per day with their own regional phrasings and edge cases, manual monitoring becomes impossible. You can't have an engineer watching dashboards all day, and you certainly can't have one responding to 3 AM failures.

I built a system that closes this loop automatically. It detects when the agent starts failing, identifies what's going wrong, generates a better prompt, and deploys it—all without waking anyone up. The key insight is treating prompt degradation the same way we treat infrastructure failures: as something that can be detected via metrics and remediated via automation.

The Problem with Voice AI at Scale

The challenge isn't making a voice agent that works. The challenge is making one that keeps working as the world changes around it.

Consider what happens after you deploy a successful appointment-booking agent. Customers start calling with requests you didn't anticipate: "Can I get in tomorrow instead of Tuesday?" (a reschedule, not a new booking), "Do you guys do tire rotations?" (a service inquiry, not an appointment), "My husband's car, not mine" (a vehicle mismatch your slot-booking logic doesn't handle).

Each of these is a small failure. Individually, they're noise. But when they accumulate—when 15% of calls are failing instead of 5%—you have a systemic problem that requires prompt changes.

The traditional response is reactive. Someone notices the metrics are down. They pull call logs. They identify the failure pattern. They draft a new prompt. They test it locally. They deploy it. This cycle takes hours at best, days at worst. And during that time, customers are having bad experiences.

Automation can compress this cycle to minutes. The system watches its own metrics, detects when something breaks, diagnoses what's failing, generates a fix, and deploys it. The humans review the results in the morning, but the fires get put out overnight.

System Architecture

The solution requires three components working together:

The Voice Agent handles customer calls. It classifies intent, executes the appropriate conversation flow (booking, rescheduling, FAQ), and logs outcomes. Critically, it logs structured outcomes—not just "success" or "failure" but why it failed.

Keep is an open-source AIOps platform. It ingests metrics from the voice agent, applies anomaly detection (e.g., "conversion rate dropped 20% in the last hour"), and triggers remediation workflows. Keep handles the detection side: knowing something is wrong.

The GEPA Optimizer handles remediation. It receives failed call data from Keep, identifies patterns in the failures, generates an improved prompt, and deploys it as a new version. GEPA handles the fix side: doing something about the problem.

The three form a closed loop. Metrics flow from the agent to Keep. When thresholds are breached, Keep triggers the optimizer. The optimizer generates a new prompt and deploys it back to the agent. The loop runs continuously without human intervention.

┌─────────────────────────────────────────────────────────────────────────────┐│                          VOICE AGENT LAYER                                   ││                                                                              ││  ┌────────────────────┐      ┌────────────────────┐      ┌────────────────┐  ││  │  Intent Classifier │      │  Conversation      │      │ Outcome Logger │  ││  │  ─────────────     │ ───▶ │  Flows             │ ───▶ │ ───────────    │  ││  │  • Booking         │      │  • BookingFlow     │      │ • success      │  ││  │  • Reschedule      │      │  • RescheduleFlow  │      │ • failure_type │  ││  │  • FAQ             │      │  • FAQFlow         │      │ • transcript   │  ││  └────────────────────┘      └────────────────────┘      └───────┬────────┘  │└──────────────────────────────────────────────────────────────────┼───────────┘                                                                   │ metrics                                                                   ▼┌─────────────────────────────────────────────────────────────────────────────┐│                          AIOPS LAYER (Keep)                                  ││                                                                              ││  ┌────────────────────┐      ┌────────────────────┐      ┌────────────────┐  ││  │  Metric Ingestion  │      │  Anomaly Detection │      │ Workflow       │  ││  │  ─────────────     │ ───▶ │  ─────────────     │ ───▶ │ Trigger        │  ││  │  • conversion_rate │      │  "Rate dropped     │      │ • fetch fails  │  ││  │  • avg_call_time   │      │   20% in 1 hour"   │      │ • call webhook │  ││  └────────────────────┘      └────────────────────┘      └───────┬────────┘  │└──────────────────────────────────────────────────────────────────┼───────────┘                                                                   │ alert + data                                                                   ▼┌─────────────────────────────────────────────────────────────────────────────┐│                          OPTIMIZATION LAYER (GEPA)                           ││                                                                              ││  ┌────────────────────┐      ┌────────────────────┐      ┌────────────────┐  ││  │  Failure Analysis  │      │  Prompt Generation │      │ Version Store  │  ││  │  ─────────────     │ ───▶ │  ─────────────     │ ───▶ │ ───────────    │  ││  │  NO_SLOTS: 18      │      │  Current + fails   │      │ v1 → v2 → v3   │  ││  │  CONFIDENCE: 8     │      │  + objectives      │      │                │  ││  │  → derive goals    │      │  → LLM → v4        │      │ is_active: v4  │  ││  └────────────────────┘      └────────────────────┘      └───────┬────────┘  │└──────────────────────────────────────────────────────────────────┼───────────┘                                                                   │      ┌────────────────────── deploys v4 ◀─────────────────────────┘      │      ▼┌─────────────────────────────────────────────────────────────────────────────┐│  CLOSED LOOP: Agent uses new prompt → metrics recover → Keep auto-resolves  │└─────────────────────────────────────────────────────────────────────────────┘

The architecture separates concerns cleanly. The agent focuses on handling calls well. Keep focuses on detecting problems. The optimizer focuses on generating fixes. Each component is independently testable and replaceable.

Structured Failure Logging

The entire system depends on one thing: knowing why calls fail, not just that they fail.

A log entry that says "call failed" tells you nothing actionable. Did it fail because no appointment slots were available? Because the customer hung up? Because the agent couldn't understand the request? Each of these requires a different fix. Slot availability is a data problem, not a prompt problem. Customer hangups might indicate the greeting is too long. Intent confusion indicates the prompt needs better examples.

I defined a canonical set of failure reasons:

class FailureReason(str, Enum):    NO_SLOTS = "no_slots"    CUSTOMER_DISENGAGED = "customer_disengaged"    AGENT_CONFIDENCE_LOW = "agent_confidence_low"    UNKNOWN = "unknown"

Every call outcome includes the appropriate failure reason. When a customer hangs up mid-conversation, that's CUSTOMER_DISENGAGED. When the booking system has no available slots, that's NO_SLOTS. When the agent can't classify what the customer wants, that's AGENT_CONFIDENCE_LOW.

This structure propagates through the entire system. When the optimizer receives a batch of failed calls, it can aggregate by failure type: "Of these 30 failures, 18 were NO_SLOTS, 8 were AGENT_CONFIDENCE_LOW, and 4 were CUSTOMER_DISENGAGED." That breakdown tells you what to fix. If most failures are slot-related, the prompt probably needs to offer waitlist options more proactively. If most are confidence-related, the prompt needs better intent-handling instructions.

The key insight is that failure categorization is a design decision made at the agent level, but it pays dividends at the optimization level. You can't derive this structure from raw transcripts—the agent has to record it explicitly.

Deriving Optimization Objectives

Raw failure counts aren't useful to an LLM generating prompt improvements. "18 NO_SLOTS failures" doesn't tell the model what to do. The system needs to translate failure patterns into natural language objectives that guide generation.

This translation encodes domain knowledge about what each failure type implies:

def derive_objectives(failure_reasons: list[FailureReason]) -> list[str]:    counts = Counter(failure_reasons)    objectives = []    if counts.get(FailureReason.NO_SLOTS, 0) > 3:        objectives.append("Proactively offer waitlist when no immediate slots are available")        objectives.append("Suggest alternative dates before saying slots are unavailable")    if counts.get(FailureReason.CUSTOMER_DISENGAGED, 0) > 2:        objectives.append("Shorten the initial greeting to reach the customer's intent faster")        objectives.append("Use warmer, more engaging language to maintain customer attention")    if counts.get(FailureReason.AGENT_CONFIDENCE_LOW, 0) > 2:        objectives.append("Ask clarifying questions when the request is ambiguous")        objectives.append("Handle compound requests by addressing each part separately")    return objectives or ["Improve overall call resolution rate"]

This function is where human expertise enters the system. Someone who understands voice agent failure modes wrote these mappings. NO_SLOTS failures suggest the prompt should handle slot unavailability more gracefully—that's domain knowledge, not something the system can infer automatically.

The objectives become part of the prompt sent to the LLM for generation. Instead of asking "make this prompt better," the system asks "make this prompt better by offering waitlists when slots are unavailable and asking clarifying questions when requests are ambiguous." Specific objectives produce specific improvements.

The Optimization Workflow

When Keep detects a metric anomaly—say, conversion rate dropping below 80%—it triggers a workflow that calls the optimizer. The workflow is defined in YAML and executed by Keep's workflow engine:

workflow:  id: voice-ai-remediation  triggers:    - type: alert      filters:        - key: name          value: conversion_rate_drop  steps:    - name: fetch-failures      provider:        type: postgres      with:        query: |          SELECT transcript, summary, failure_reason          FROM call_outcomes          WHERE success = false AND created_at > NOW() - INTERVAL '1 hour'    - name: trigger-optimization      provider:        type: http      with:        url: "{{ env.GEPA_OPTIMIZER_URL }}/optimize"        method: POST        body:          alert_id: "{{ alert.id }}"          failed_calls: "{{ steps.fetch-failures.results }}"

The optimizer receives this payload and executes the generation workflow. It loads the current active prompt, derives objectives from the failure patterns, builds a generation prompt that includes the current prompt text plus failed call summaries plus objectives, calls the LLM to generate an improved version, and stores the result as a new prompt version.

The generation prompt looks like this:

Current prompt (version v3):[current prompt text]Recent failures (30 calls in the last hour):- Customer asked about tire rotation, agent failed to classify intent- Customer requested Tuesday instead of Monday, agent treated as new booking- Customer mentioned husband's car, agent didn't handle vehicle mismatch...Objectives:- Ask clarifying questions when the request is ambiguous- Handle rescheduling requests distinctly from new bookings- Confirm vehicle ownership before proceedingGenerate an updated prompt that addresses these failures while preserving what works.

The LLM generates improved prompt text based on this context. The optimizer parses the response, stores it as a new version (v4), and makes it active. The next call to the agent will use the new prompt.

Prompt Versioning

Prompts aren't mutable text files. They're versioned artifacts with full history.

Every optimization creates a new version. Version identifiers increase monotonically: v1, v2, v3, and so on. The database stores the complete history: what each version's text contained, when it was created, why it was created (the optimization notes), and whether it's currently active.

@dataclassclass PromptVersion:    id: int    version: str           # "v3"    content: str           # Full prompt text    notes: str             # "Addressed NO_SLOTS failures, added waitlist offer"    is_active: bool        # Only one version is active at a time    created_at: datetime

Versioning enables several capabilities:

Rollback: If v4 performs worse than v3, you can revert by changing which version is active. The old text still exists; you're just pointing at it again.

Audit trail: Every prompt change has a recorded reason. "Why does the prompt mention waitlists?" Look at v4's notes: it was generated to address NO_SLOTS failures. This is essential for debugging and for explaining decisions to stakeholders.

Comparison: You can diff v3 against v4 to see exactly what changed. Did the optimizer add a new section? Modify existing wording? Remove something that was working? The version history makes this visible.

The versioning overhead is minimal—a few database rows per optimization run—but the debugging and rollback capabilities are invaluable. You never want to be in a situation where you can't recover a previous prompt because someone overwrote it.

Scoring Before Deployment

Not every generated prompt is an improvement. LLMs hallucinate. They misunderstand objectives. They sometimes make prompts worse while trying to make them better.

The system scores new prompts before deploying them. The score evaluates multiple dimensions:

Coverage: Does the prompt address the failure types that triggered optimization? If we optimized because of AGENT_CONFIDENCE_LOW failures, the new prompt should include language about handling ambiguous requests.

Specificity: Are the instructions concrete or vague? "Be helpful" is vague. "Ask clarifying questions when the customer's request mentions multiple services" is specific.

Brevity: Longer isn't always better. Prompts that balloon in size may confuse the model or hit token limits. The score penalizes excessive length.

@dataclassclass ScoreBreakdown:    coverage: float    specificity: float    brevity: float    def total(self) -> float:        return self.coverage + self.specificity + self.brevity

The score determines whether the prompt gets deployed. If score.total() exceeds a minimum threshold, the new version goes active. If not, the optimization run is logged as attempted but the previous prompt remains active.

This prevents regressions. A bad LLM generation won't tank your production system because it won't pass the scoring threshold. The safeguard isn't perfect—scoring is based on the prompt text, not actual call outcomes—but it catches obvious failures like prompts that ignore the optimization objectives entirely.

Why Keep?

You could build the detection piece yourself. Poll a metrics endpoint, fire webhooks when thresholds are breached, done. So why use Keep?

Keep provides value beyond simple threshold alerting:

Deduplication: If conversion drops across 15 dealerships simultaneously, you don't want 15 separate alerts. Keep correlates related events into a single incident. The optimizer gets called once, not 15 times.

Enrichment: Keep workflows can query multiple data sources, call external APIs, and enrich alerts with context before calling the optimizer. "Conversion dropped" becomes "Conversion dropped, 80% of failures are AGENT_CONFIDENCE_LOW, mostly happening in Texas region."

Lifecycle management: When metrics recover after a prompt update, Keep can auto-resolve the incident. The full remediation cycle—detection, diagnosis, fix, recovery—is tracked as one incident with clear timestamps for each phase.

Workflow orchestration: The workflow definition is declarative YAML, not scattered code. Changing the remediation logic means editing a config file, not modifying application code.

Using established AIOps tooling for ML systems is an underappreciated pattern. Prompt degradation is as much an operations problem as a model problem. The same techniques that monitor database latency and API error rates apply to voice agent conversion rates.

Limitations and Future Work

The system has clear limitations that I'd address with more time:

No shadow testing: New prompts deploy immediately to all traffic. A safer approach would deploy to 10% of calls first, compare conversion rates against the old prompt, and only fully roll out if the new version performs better in production.

Single-shot optimization: The optimizer generates one prompt per run. Iterative refinement—generate a candidate, test it, refine based on results, repeat—would likely produce better outcomes than single-shot generation.

Heuristic scoring: The scoring evaluates prompt text characteristics, not actual call outcomes. True validation requires A/B testing against live traffic, which the current system doesn't do.

Manual rollback: If a new prompt performs worse, reverting requires manual intervention. Automated rollback triggered by metric degradation would complete the self-healing loop.

These are engineering improvements, not fundamental architectural changes. The closed-loop structure is sound; the implementation could be more sophisticated.

Results

The complete remediation cycle works like this:

Conversion rate drops below 80% → Keep fires alert
Keep workflow fetches recent failed calls from database
Workflow calls optimizer with failure data
Optimizer derives objectives from failure distribution
LLM generates updated prompt
Scoring passes threshold; new version deployed
Conversion recovers; Keep auto-resolves incident

Total cycle time: minutes, not hours.

Technology Stack

Python 3.11, Flask for the optimizer API
PostgreSQL for prompt version storage and call outcome logging
Keep (open-source) for alerting and workflow orchestration
Qwen via Together API for prompt generation
Docker Compose for local development and testing

Peeking Inside Diffusion Language Models: A LogitLens Analysis

Mon, 12 Jan 2026 00:00:00 GMT

Note: This RSS feed strips out SVGs and embeds. You might want to read the post on the webpage here.

October 25, 2025

Last Updated: January 12, 2026

mechInterp aiSafety

Bookmark for https://drive.google.com/file/d/12eHAXqjh-SzbwIxOuO8azonSMGlwnjgA/view?usp=sharing

This notebook explores the inner workings of diffusion-based language models (DLMs) like LLaDA using the LogitLens interpretability method. The analysis is inspired by two key sources:

Alessio de Voto's blog post, which provides an excellent guide to the inner workings of Diffusion Language Models.
The "Prophet" paper, which investigates verifying emergent properties in diffusion models and suggests that answers often converge early in the decoding timeline.

Unlike traditional left-to-right text generators, DLMs are bidirectional models that learn to fill in masked tokens through an iterative denoising process. LogitLens allows us to explore how a model's predictions evolve across its layers by applying the unembedding matrix (the language modeling head) to the hidden states of all intermediate layers.

Today, we're turning this lens on LLaDA to uncover how it builds its understanding layer by layer, providing insights into DLM functionality and highlighting differences with respect to autoregressive models.

Final Analysis: Two Axes of Convergence

By extending the LogitLens analysis across the full iterative decoding process, we can visually reconcile the findings of the "Prophet" paper (early convergence in time) with the layer-by-layer patterns seen in our initial analysis (late convergence in depth).

The results reveal that Diffusion Language Models exhibit two distinct types of convergence:

1. Convergence in Time (The "Prophet" Paper's Finding)

Examining the sequence of heatmaps from each timestep makes the "early convergence" phenomenon clear.

In the initial steps (t=0 to t≈7), the model's final prediction (the bottom row) is uncertain and incoherent.
Around the halfway mark (e.g., timestep t=8), the model's output stabilizes on a complete and correct sentence.
In all subsequent steps (t=9 to t=15), this final prediction remains largely unchanged, with the model only growing more certain (as shown by the brighter colors in the final layers).

This confirms the core finding that the final answer is often determined early in the iterative decoding process.

2. Convergence in Depth (The Logit Lens Finding)

Within any single heatmap—including those where the final answer has already converged—we still observe a "late convergence" pattern across the model's layers.

The initial layers (top rows) consistently show low certainty (dark colors) and often predict nonsensical or placeholder tokens.
The correct tokens and high certainty (bright yellow colors) only emerge in the final, deeper layers of the network.

Conclusion

This experiment reveals a more complete picture of how DLMs operate. They first establish a high-level, globally coherent "plan" for the full sentence early in the decoding timeline. Then, within each step of executing that plan, they refine the details from abstract representations (in early layers) to concrete tokens (in the final layers).

Geometry of Hidden State

Sun, 19 Oct 2025 00:00:00 GMT

Note: This RSS feed strips out SVGs and embeds. You might want to read the post on the webpage here.

October 19, 2025

aiSafety mechInterp

💡

TL;DR

Linear Subspaces (CoNLL 2021): Many linguistic features, from part-of-speech to dependency arcs, are encoded in low-dimensional linear subspaces within a model's hidden states. These subspaces are structured, distributed across many neurons, and-crucially-intervening on them causally changes model outputs.
Affine Relational Maps (ICLR 2024): The complex, multi-layer computation that maps a subject ("Paris") to an object ("France") for a given relation (e.g., capital of) can often be approximated by a simple affine transformation (Wx + b) at a specific layer and token position. This holds for roughly half of the common relations tested.
Causal Geometry (ICML 2024): The very notion of "linearity" or a "direction" is only meaningful if the geometry (the inner product) respects the causal structure of language. A proposed "causal inner product" provides a principled way to define these directions, unifying the goals of probing (reading concepts) and steering (writing concepts), and often revealing linear structure where naive metrics like cosine similarity fail.

The internal world of a large language model is a labyrinth of high-dimensional vectors. Each forward pass is a journey through this geometric space, where abstract concepts are manipulated by matrices and non-linear functions. For years, our best descriptions of this process were frustratingly opaque. But a recent thread of research suggests that within this complexity lies a surprising amount of simple, linear structure.

This isn't just a curiosity. If concepts and relations are encoded linearly, we can build more reliable tools to interpret what models know, predict how they will behave, and even perform surgical interventions to edit their knowledge or control their outputs. Three key papers-one on linguistic subspaces, one on relational maps, and one on the underlying geometry itself-form a powerful, unified narrative. This post unpacks their findings to explain what "linearity" really buys us.

What We Mean by “Linearity”

Before diving in, it’s critical to untangle a few distinct ideas often conflated under the banner of "linearity." The claim is not that transformers are secretly linear models. They are deeply non-linear. The claim is that we can often use linear tools to meaningfully describe and manipulate their internal representations.

Here are three increasingly specific senses of linearity:

Linear Separability: The weakest claim. This means that hidden states corresponding to different concepts (e.g., positive sentiment vs. negative sentiment) can be separated by a simple hyperplane. A linear classifier (a "probe") can be trained to read out the concept from the hidden state. This is a classic observation, but it only tells us that the information is present, not how it's used.
Linear Decoding: A stronger claim. This suggests that a concept can be read out directly from a hidden state r(l)r^{(l)}r(l) via a simple linear map, often just multiplication by a vector or matrix. For example, the log-probability of a token might be a linear function of the final hidden state via the unembedding matrix WUW^UWU. This implies the information is not just separable but also in a format that downstream components can easily access.
Affine Relational Mapping: The most specific claim, and the focus of recent work. This idea posits that for a given relation rrr (like country -> capital), there exists an affine map (Wr,brW_r, b_rWr,br) that transforms the hidden state of the subject $s$ into the hidden state of the object. That is, LREr(s)=Wrs+br\text{LRE}_r(s) = W_r s + b_rLREr(s)=Wrs+br approximates the model's internal computation.

Caveat: The choice of geometry-how we measure distance and angles-fundamentally changes whether these claims appear true. Most analyses default to the Euclidean inner product (and its cousin, cosine similarity). But as the work on the Linear Representation Hypothesis (LRH) shows, this can be misleading. A "causal" inner product, which respects the statistical independence of concepts in language, can reveal linear structure that is invisible to standard metrics, or show that apparent linearity was an artifact. Declaring the geometry is non-negotiable for making a rigorous claim.

Background & Notation

Let's establish a common language. We view a transformer through the lens of its residual stream. A hidden state r(l)∈Rdr^{(l)} \in \mathbb{R}^dr(l)∈Rd at layer lll and a specific token position represents the accumulated information up to that point. The model's final prediction depends on the final state, typically via the unembedding matrix WUW^UWU which maps r(L)r^{(L)}r(L) to logits over the vocabulary.

A conceptccc (e.g., gender, tense) can be represented by a linear subspaceSc⊂RdS_c \subset \mathbb{R}^dSc⊂Rd. We can isolate the component of a hidden state that lies in this subspace using a projection operator PcP_cPc.
A Linear Relational Embedding (LRE) for a relation rrr is an affine map LREr(s)=Wrs+br\text{LRE}_r(s) = W_r s + b_rLREr(s)=Wrs+br. It takes a subject representation $s$ (a hidden state at a certain layer/position) and approximates the corresponding object representation. This map is often estimated by taking a first-order Taylor approximation of the model's computation around a reference subject s0s_0s0.
The Causal Inner Product⟨⋅,⋅⟩causal\langle \cdot, \cdot \rangle_{\text{causal}}⟨⋅,⋅⟩causal is a non-Euclidean metric designed to make concepts that are semantically independent in the world (e.g., tense and sentiment) orthogonal in the representation space. It provides a formal basis for when talking about "directions" is meaningful and connects the geometry of representations to counterfactual changes in model inputs and outputs.

With these tools, we will now examine the evidence.

Paper A (CoNLL 2021): Linear Subspaces in Practice

The work by Hernandez and Andreas (2021) provided strong, early evidence that linear structures weren't just a convenient analytical tool, but a core aspect of how models like BERT organize linguistic information.

Their method involved identifying subspaces for various linguistic features, from simple ones like part-of-speech to complex ones like the grammatical role of a word in a dependency parse. They trained simple linear probes to find these subspaces and then tested their properties.

Key Findings:

Low-Dimensionality: Subspaces for many distinct linguistic features were surprisingly low-dimensional, often occupying just 1-10 dimensions out of the full 768 in BERT-base. This suggests information is compressed into specific, targeted geometric structures.
Hierarchical Structure: The subspaces reflected linguistic hierarchies. For example, the subspace for verb would contain the more specific subspaces for transitive verb and intransitive verb. This geometric containment mirrors the conceptual relationship.
Distributed Representation: These subspaces were not aligned with individual neurons. The basis vectors for a concept's subspace were distributed across many coordinates, confirming that features are stored collectively, not in single "grandmother cells."

Most importantly, the paper established a causal link between these subspaces and model behavior. They performed interventions by projecting a hidden state onto a chosen subspace or removing its component in that subspace. For example, by manipulating the subspace for a dependency arc, they could systematically change which words the model attended to and alter its masked language modeling predictions in a controlled way.

Caveat: This work demonstrates that linear signals are present and can be causally manipulated. It does not, however, prove that the model always uses that signal for every downstream task. A feature might be linearly encoded but ignored by later layers in certain contexts. The causal link is strong evidence for its functional role, but not a guarantee of its universal application.

Paper B (ICLR 2024): Linear Relation Decoding

While the CoNLL paper focused on properties of single words or tokens, Hernandez et al. (2024) investigated the geometry of relations between entities. Their central idea is the Linear Relational Embedding (LRE), which tests whether the complex computation mapping a subject to an object can be simplified to an affine transformation.

To find the relational map (Wr,br)(W_r, b_r)(Wr,br) for a relation like country -> capital, they compute the local Jacobian of the model's transformation from the subject's hidden state to the object's hidden state. By averaging this estimate over a few example subjects, they obtain a robust linear operator that approximates the relation.

Empirical Scope and Limits:

The team tested this method on 47 different relations, spanning factual knowledge, commonsense reasoning, linguistic patterns, and even social biases.

Successes: The linear approximation worked remarkably well for about half of the relations. For these, applying the learned LRE to a subject's hidden state could reliably produce the hidden state of the correct object, often for subjects unseen during the LRE's estimation.
Failures: The approximation consistently failed for relations where the set of possible objects is vast and arbitrary, such as company -> CEO or person -> birth city.

Intuition: This pattern suggests a compelling hypothesis. For relations with a constrained, structured output space (like capitals or languages), the transformer can "pre-compile" the computation into a nearly-linear map within a specific processing window (i.e., across a few layers). For open-ended relations, it must rely on more complex, non-linear search and retrieval mechanisms that cannot be captured by a single LRE.

This work also introduced the Attribute Lens, a practical application of LREs. By applying a pre-computed LRE (e.g., for country -> capital) to the hidden states at every token and every layer for an input like "The government of France announced," one can visualize exactly where and when the model computes the capital "Paris," even if it is never explicitly mentioned.

Paper C (ICML 2024): Definitions & the Causal Inner Product

The Linear Representation Hypothesis (LRH) paper by Park, Choe, and Veitch (2024) provides the theoretical bedrock for these empirical findings. It asks a fundamental question: what does it even mean for a concept to be "linearly represented"?

They provide two formal definitions grounded in counterfactuals:

Output-Space (Measurement): A concept has a linear representation if swapping the concept in the output (e.g., changing "he" to "she") corresponds to moving along a consistent direction in the hidden state space. This formalizes the idea behind linear probing.
Input-Space (Intervention): A concept has a linear representation if swapping the concept in the input context (e.g., changing "The doctor said he..." to "The doctor said she...") corresponds to moving along a consistent direction. This formalizes the idea behind model steering or vector addition.

The central insight is that these two definitions are unified by choosing the right geometry. The standard Euclidean inner product is a poor choice because it treats all dimensions equally. Language is not like that; some concepts are related, others are independent. The authors introduce the causal inner product, a metric derived from the model's unembedding matrix that makes semantically independent concepts orthogonal.

What the Causal Inner Product Buys Us:

Principled Directions: It provides a rigorous way to define "concept directions" that align with counterfactual meaning. A steering vector for gender should, ideally, be orthogonal to a steering vector for tense under this metric.
Improved Separability: The paper demonstrates empirically on LLaMA-2 that concepts that are muddled under cosine similarity become cleanly separable using the causal inner product. For instance, contexts about France and Spain might seem similar in Euclidean space, but the causal metric can better isolate the countryspecific components.
Unification: With the causal inner product, the "probe" direction (from the output-space definition) and the "steering" direction (from the input-space definition) become one and the same. This provides a single, coherent geometric object for both reading and writing a concept.

Experiments with LLaMA-2 confirm that linear representations meeting these causal definitions exist for a variety of concepts, and that the choice of inner product is not merely academic-it fundamentally determines whether these structures are found.

Synthesis - One Geometry, Many Tools

These three papers are not independent observations; they are three views of the same underlying phenomenon. Together, they paint a remarkably coherent picture of how linear geometry enables complex reasoning in non-linear models.

The subspace picture (CoNLL '21) establishes that concepts live in specific, low-dimensional linear manifolds and that these manifolds are causally potent.
The affine relational maps (ICLR '24) show that for many common knowledge-retrieval tasks, the multi-step computation from a subject's subspace to an object's subspace can be locally approximated by a simple affine hop.
The causal geometry (ICML '24) provides the rigorous mathematical language to justify it all. It explains why we can talk about directions and subspaces, tells us which metric to use, and unifies the acts of observing a concept (probing) and controlling it (steering).

The unified story is this: To find the simple linear structure hidden inside a transformer, you need to look at the right slice of the representation (the right layer and token position) and measure it with the right metric (the causal inner product). In that sweet spot, complex operations often resolve into simple geometric translations and projections. Outside of that slice, or with the wrong metric, the model's behavior remains stubbornly non-linear and opaque.

Practical Guidance

For researchers and practitioners working on interpretability, this body of work suggests a more responsible way to make and evaluate claims about linearity.

Disclose Your Geometry: Always state the inner product used for any analysis (e.g., standard Euclidean/cosine or a causal inner product). If you perform any normalization like centering or whitening, report it.
Sweep Layers and Positions: Linearity is often localized. Avoid cherry-picking a single layer. A full sweep across layers and token positions reveals where a linear approximation holds and where it breaks down.
Use Counterfactuals for Validation: Define directions using minimal counterfactual pairs as suggested by LRH. Validate that these directions are meaningful by performing causal interventions, such as subspace projections (CoNLL-style) or applying LREs to new subjects.
Analyze Failures: Documenting where linearity fails is as important as showing where it succeeds. Reporting out-of-scope relations or layers where the geometry is non-linear provides crucial constraints on the hypothesis.

My View

Opinion: Linear tools are at their best when treated as local approximations of a model's behavior. They are powerful instruments for generating hypotheses, performing quick diagnostic tests, and executing surgical edits on model knowledge. They are not, however, a complete theory of how transformers compute. Their power lies in simplifying a small piece of a complex puzzle, but we must always remember and declare the boundaries of that simplification.
Intuition: I suspect that transformers operate by dynamically linearizing sub-problems. When the model needs to map a country to its capital, it routes the subject representation through a sequence of blocks that, in aggregate, perform an approximately affine map. Once that attribute is computed and added to the residual stream, subsequent blocks can perform other, unrelated non-linear operations. The linearity is a temporary state in a much longer, piecewise computation.
Caveat: The most common mistake in this area is to confuse the presence of a linearly decodable signal with its use in the model's final decision. A probe can often detect a feature at an early layer, but that doesn't mean the model uses it. Only a causal experiment-an intervention that shows a change in the representation leads to a predictable change in output-can close that loop.

Limits & Open Questions

This is a fast-moving area, and our understanding is far from complete. Key open questions include:

Faithfulness Bounds: When, precisely, does the causal inner product lead to substantively different conclusions than standard cosine similarity? Are there classes of problems where the simpler metric is "good enough"?
Portability: Do linear subspaces and LREs transfer across model families (e.g., LLaMA to GPT) or model sizes? Is the geometric structure of knowledge universal, or is it an idiosyncratic feature of each training run?
Compositionality: Can we chain multiple linear operations? For example, can we compose the LRE for author -> book with the LRE for book -> main character to reason about an author's characters? Early evidence suggests this is difficult, likely because the non-linearities between steps are crucial.
Effect of Instruction Tuning: How does fine-tuning on human preferences (RLHF) and other alignment techniques warp the pristine geometric structures found in base models? Does it consolidate concepts into clearer subspaces or shatter them into context-dependent, non-linear fragments?

The geometry of hidden states is no longer a complete black box. By choosing the right lens, we are beginning to resolve its structure, finding straight lines and simple planes in a space we once thought was impenetrably complex.

References

Hernandez, E., & Andreas, J. (2021). Uncovering the Hierarchical Structure of Learned Language Representations. In Proceedings of the 25th Conference on Computational Natural Language Learning (CoNLL). Association for Computational Linguistics.
Hernandez, E., Sharma, A. S., Haklay, T., Meng, K., Wattenberg, M., Andreas, J., Belinkov, Y., & Bau, D. (2024). Linearity of Relation Decoding in Large Language Models. In The Twelfth International Conference on Learning Representations (ICLR).
Park, K., Choe, Y. J., & Veitch, V. (2024). The Linear Representation Hypothesis and the Causal Inner Product. In Proceedings of the 41st International Conference on Machine Learning (ICML). arXiv:2311.03658.

Causal Attribution & Patching (Part 2): Pathscopes

Sun, 19 Oct 2025 00:00:00 GMT

Note: This RSS feed strips out SVGs and embeds. You might want to read the post on the webpage here.

October 19, 2025

aiSafety mechInterp

💡

TL;DR

Fast Approximations Have Limits: Attribution Patching (AtP) and its successor, AtP*, offer scalable causal attribution through linearization. They are incredibly fast for ranking influential model components but are approximations and can fail, especially when model gradients are uninformative.
Natural Language Explanations: Patchscopes provides a unifying framework to inspect hidden states by using the model itself to generate natural-language descriptions of them. This approach helps probe traditionally opaque early layers and enables novel capabilities like cross-model explanations and multi-hop error correction.
Feature-Level Causal Graphs: Sparse Feature Circuits (SFCs) move beyond coarse components (heads, neurons) to discover causal subnetworks of more interpretable "features" learned by Sparse Autoencoders (SAEs). This enables a powerful editing technique, Shift, which modifies model generalization by ablating spurious features.
A Modern Workflow: The cutting edge of causal analysis combines these ideas: use Patchscopes for high-level hypotheses, AtP* for rapid ranking of influential components, and SFCs to dissect the underlying mechanism at the feature level and perform targeted edits.

Recap: From Part 1 to Part 2

In Part 1, we established causal patching as a core technique for mechanistic interpretability. By intervening on a model’s internal activations—running it on a "clean" input and patching those activations into a forward pass on a "corrupted" input—we can identify which components are sufficient to restore a desired behavior. This moves us from mere correlation to causal evidence, allowing us to ask not just what information a component encodes, but what mechanism it implements.

However, this approach faces two major limitations. First, our units of analysis—attention heads, MLP layers, or individual neurons—are often polysemantic, meaning they represent multiple, unrelated concepts simultaneously. This makes their role in a circuit hard to pin down. Second, the output of a patching experiment is typically a numeric score (e.g., a change in loss or logit difference), which tells us that a component matters but not what it represents.

This brings us to Part 2. We will explore two major advances that address these challenges. First, we will move from coarse, polysemantic components to fine-grained, interpretable "features" identified by Sparse Autoencoders (SAEs). Second, we will move beyond simple numeric probes to generating rich, natural-language (NL) descriptions of what a hidden state represents. The goal is to build a more precise, interpretable, and editable understanding of the causal mechanisms inside language models.

Background: Patching & Its Fast Approximations

At the heart of our analysis is Activation Patching (also known as causal tracing or interchange interventions). It provides a direct, causal measure of a component's importance. We take a clean prompt (e.g., "The Eiffel Tower is in Paris") and a corrupted prompt (e.g., "The Eiffel Tower is in Rome"), then run the model on the corrupted prompt while replacing a specific internal state, like an attention head's output, with the corresponding state from the clean run. If this single patch restores the model's correct output ("Paris"), we have strong evidence that the component is causally sufficient for that behavior.

While powerful, activation patching is computationally expensive. It requires a separate forward pass for every component we want to test. For a large model with thousands of attention heads and MLP layers across dozens of layers, this quickly becomes intractable.

This scalability problem motivated the development of Attribution Patching (AtP). AtP reframes the patching intervention through the lens of a first-order Taylor-series approximation. Instead of performing a full forward pass for each patch, it estimates the effect using gradients.

Let fθf_\thetafθ be our model, and let $m$ be a scalar metric we care about, such as the logit of the target token t⋆t^\start⋆. Let scleans_{\text{clean}}sclean be the activation of a component on a clean input and scorrs_{\text{corr}}scorr be the activation on a corrupted input. Activation patching measures the exact change in the metric, Δm\Delta mΔm. Attribution Patching approximates this change:

Δm≈∇sm∣scorr⊤ (sclean−scorr) \Delta m \approx \nabla_{s} m \big|{s{\text{corr}}}^\top \, (s_{\text{clean}} - s_{\text{corr}}) Δm≈∇smscorr⊤(sclean−scorr)

This calculation requires only one forward pass on the clean input (to get scleans_{\text{clean}}sclean), one forward pass on the corrupted input (to get scorrs_{\text{corr}}scorr), and one backward pass from the metric on the corrupted run (to get the gradient ∇sm\nabla_{s} m∇sm). With these three passes, we can estimate the effect of patching every component in the model simultaneously, offering a massive speedup.

However, this is an approximation. The paper introducing AtP* identified two key failure modes where this linearization breaks down:

QK-Dependence: The outputs of attention heads depend non-linearly on the query (Q) and key (K) vectors via the softmax function. AtP's linear approximation struggles here. AtP* addresses this by partially recomputing the attention pattern change instead of linearizing it.
Gradient Nullspace: Sometimes the gradient of the metric with respect to a component's activation is zero (or near-zero) on the corrupted run, even if the component is causally important. This leads to false negatives. AtP* mitigates this with a technique called "gradient dropout," which computes gradients on a modified run where some components are disabled.

Caveat
AtP and AtP* are powerful for rapidly identifying candidate components for a circuit. However, they are not foolproof. Their reliability depends on how well the model's local behavior can be approximated by a linear function. When gradients are small or zero, these methods can miss crucial components. It is best practice to use AtP/AtP* for an initial, wide search and then verify the top-scoring candidates with exact activation patching.

Patchscopes: A Unifying Framework for Inspecting Representations

While AtP helps us find where causality flows, it doesn't tell us what information is flowing. This is where Patchscopes comes in. Introduced by Ghandeharioun et al. (ICML 2024), Patchscopes is a framework that uses the language model itself to explain its internal representations in natural language.

The core idea is simple but profound: "patch" a hidden state from a source computation into a target computation that is designed to elicit a natural language description. For example, to understand what information is encoded at the final token of the prompt "The Colosseum is in the city of," we can take the residual stream vector at that position and patch it into a templated prompt like "The city is [MASK]." The model's subsequent prediction for "[MASK]" ("Rome") reveals the information stored in that vector.

Patchscopes unifies and extends several prior interpretability techniques:

Unifying Framework: The paper shows that methods like Logit Lens (projecting hidden states to the vocabulary) and Tuned Lens are special cases of the Patchscopes framework. This provides a common ground for understanding and comparing these methods.
Mitigating Early-Layer Gaps: Prior methods often failed to produce coherent explanations for representations in early layers of a network, as these states are not yet "aligned" with the final vocabulary space. Patchscopes can overcome this by patching a state into the same layer of a different computation, allowing the model's own later layers to process and "translate" the representation effectively.
Cross-Model Explanations: A fascinating capability introduced by Patchscopes is using a more powerful model (e.g., GPT-4) to explain the internal representations of a smaller, weaker model (e.g., GPT-2). This leverages the advanced reasoning and linguistic capabilities of the larger model to interpret the more primitive representations of the smaller one.
Multi-Hop Error Correction: Patchscopes can be used not just for inspection but for intervention. In multi-hop reasoning tasks where a model makes an intermediate error, one can patch in a corrected hidden state to steer the model back toward the correct final answer, demonstrating a direct application of the insights gained.

Intuition
Think of Patchscopes as providing human-readable labels for the intermediate variables in a model's computation. Instead of seeing a raw vector and a numeric attribution score, you get a natural language phrase like "the city of Rome" or "a plural subject." This provides a powerful conceptual scaffold that can guide further investigation. The patch queries themselves become testable hypotheses about what a given representation means and does.

Sparse Feature Circuits: Causally Implicated Feature Graphs

Patchscopes gives us better descriptions, but we are often still working with coarse components like the full residual stream or attention head outputs. The problem of polysemanticity remains. Sparse Feature Circuits (SFCs), introduced by Marks et al. (ICLR 2025), tackle this head-on by changing the fundamental unit of analysis from neurons or heads to features.

The approach relies on Sparse Autoencoders (SAEs). An SAE is a small neural network trained to reconstruct a model's internal activations (e.g., an MLP layer's output) using a highly sparse, learned dictionary of features. The goal is to find a basis where each dimension corresponds to a more monosemantic, or conceptually singular, concept. For example, instead of a neuron that activates for both Python code and English lists, an SAE might learn separate features for "Python [ token" and "start of an English list."

With SAEs in hand, SFCs build a causal graph of the model's behavior in three main steps:

Define Nodes and Edges: The nodes in the graph are the learned SAE features at each layer. The edges represent the flow of information between features in one layer and features in the next.
Estimate Indirect Effects (IE): To determine which nodes and edges are causally important, the method estimates their indirect effect on the final output. The indirect effect of a feature uuu on a metric mmm is the change in mmm that occurs when we intervene on uuu while keeping its upstream causes fixed. While exact calculation is complex, it can be approximated efficiently using methods like path-gradient products (a linear approximation) or Integrated Gradients (IG), which often performs better in practice.
Assemble the Circuit: By computing the IEs for all nodes and edges, we can filter for the most causally significant ones, revealing a sparse subnetwork—the feature circuit—responsible for the behavior.

This feature-level view enables a powerful new editing capability called Shift. The SFC paper demonstrates this on a classifier trained to determine a person's profession from a biography. When the training data has a spurious correlation (e.g., doctors are always male, nurses are always female), the model learns to rely on gender-associated features. With an SFC, these gender features can be identified and manually ablated from the circuit. The result is a model that no longer relies on the spurious feature and generalizes better to a balanced test set. This is a significant step beyond simple fact-editing, as it alters a model's underlying reasoning strategy.

Caveat
The quality and interpretability of a Sparse Feature Circuit are fundamentally dependent on the quality of the underlying SAEs. Training SAEs is an active area of research, and challenges remain in ensuring features are consistently monosemantic and that the SAE reconstruction is faithful to the original model's activations. The choice of approximation for indirect effects (linear vs. IG) also involves trade-offs between speed and accuracy, particularly in early layers where non-linearities can be more pronounced.

Synthesis: From Components to Features, from Numbers to Language

Patchscopes and Sparse Feature Circuits are not competing methods; they are complementary tools for dissecting a model's internal algorithm. Patchscopes offers a top-down, exploratory lens, while SFCs provide a bottom-up, mechanistic one.

Patchscopes asks: "What does this representation mean in natural language?" It frames our understanding.
SFCs ask: "Which specific features in this representation causally drive the output?" It quantifies the mechanism.

When used together, they create a powerful loop: propose a hypothesis about a mechanism, use Patchscopes to find representations that seem to correspond to concepts in that hypothesis, then use SFCs to validate the causal role of the features within those representations and assemble the full circuit.

When to choose which tool:

Exploration and Hypothesis Generation: Start with Patchscopes. Its natural language outputs are ideal for initial exploration and for communicating findings to a broader audience.
Scalable Triage: Use AtP/AtP* to quickly scan thousands of components (or features) to find the most influential ones for a specific behavior. This narrows the search space dramatically.
Mechanism Discovery and Editing: Once you have candidate components, use SFCs to decompose them into interpretable features and map their causal pathways. For editing, the Shift technique is the tool of choice for removing reliance on spurious features. For verifying a key feature's role, a targeted activation patch on that feature is the gold standard.

A Practical Protocol for Modern Causal Patching

Here is a conceptual workflow for conducting a causal analysis project that integrates these modern tools:

Define Behavior and Metric: Isolate a specific, reproducible behavior you want to understand (e.g., subject-verb agreement, factual recall, a reasoning failure). Define a clean input, a corrupted input, and a clear metric mmm (e.g., the logit difference between the correct and incorrect tokens).
Hypothesis Generation with Patchscopes: Use Patchscopes to probe key hidden states. For example, query the representation after the subject noun with a template like "The subject is [MASK]." Note which layers and token positions yield clear, plausible natural language explanations that align with a hypothesized mechanism.
Fast Triage with AtP\*: Run AtP* to get a global ranking of the importance of all components (attention heads, MLP layers) in the model for your defined task. This will highlight "hotspots" that deserve closer inspection.
Feature-Level Discovery with SFCs:
- Train or load pre-trained SAEs for the component types and layers identified as important by AtP*.
- Compute node and edge Indirect Effects (IEs) for the SAE features. This will rank the causal importance of each feature.
- Aggregate the top-scoring nodes and edges to form a sparse feature circuit.

The most effective and robust workflow today combines these techniques in a layered approach: Patchscopes for high-level conceptual framing, AtP* for rapid, scalable triage, and Sparse Feature Circuits for deep mechanistic analysis and editing. Each tool shores up the weaknesses of the others. Patchscopes makes the numeric outputs of AtP interpretable, AtP narrows the immense search space for SFCs, and SFCs provide the granular, editable substrate that fully realizes the promise of causal interventions. The gold standard for confirming a critical finding remains the targeted, exact activation patch on the component or feature in question.

Edit and Evaluate:
- Shift Edit: Manually inspect the features in your circuit. If you identify spurious features (e.g., a feature tracking gender when the task is profession classification), ablate them.
- Minimal Patch: Alternatively, perform a precise activation patch on only the highest-IE feature(s) to confirm their sufficiency.
- Test Generalization: Evaluate the edited model's performance, specificity (does the edit affect only the target behavior?), and generalization (does the edit hold across paraphrases and different contexts?).
Report Findings: Document your prompts, seeds, and thresholds. Include the natural language descriptions from Patchscopes, diagrams of the final feature circuit, and a careful analysis of failure modes (e.g., instances where linear approximations broke down).

Intuition
Many local behaviors in transformers linearize surprisingly well, which is why AtP/AtP* are so effective as a ranking tool. The complexity of the overall function emerges from the composition of many simpler, near-linear steps. The feature lens provided by SAEs is transformative because it refactors the model’s internal ontology into one that is more aligned with human concepts. This turns the question from "What did patching this opaque vector do?" into "What was the causal effect of the 'plural noun' feature?"

Limits & Open Questions

This frontier of research is moving fast, but significant challenges remain:

Faithfulness of Approximations: How can we establish tighter theoretical bounds on when linear approximations like AtP* and IE estimates are reliable?
SAE Quality and Scale: Training high-quality, perfectly monosemantic SAEs for all components in a frontier model is a massive engineering and research challenge. Standardization and best practices are still emerging.
Early-Layer Fidelity: While Patchscopes and better IE approximations have improved our view of early layers, these representations remain harder to interpret and analyze than those in later layers.
Porting Circuits: How well do feature circuits generalize across model sizes and families? Discovering universal circuits or principles of their construction is a key open goal.
Dataset-Level Corroboration: Current methods focus on specific input pairs. Bridging these instance-level causal findings to dataset-level statistics and phenomena (e.g., connecting a feature circuit to a specific pre-training data distribution) is a crucial next step for building a complete picture of model behavior.

References

Ghandeharioun, A., Caciularu, A., Pearce, A., Dixon, L., Geva, M., et al. (2024). Patchscopes: A Unifying Framework for Inspecting Hidden Representations of Language Models. In Proceedings of the 41st International Conference on Machine Learning (ICML). PMLR.
Marks, S., Rager, C., Michaud, E. J., Belinkov, Y., Bau, D., & Mueller, A. (2024). Sparse Feature Circuits: Discovering and Editing Interpretable Causal Graphs in Language Models. arXiv preprint arXiv:2405.00172. To appear at ICLR 2025.
Syed, F., et al. (2023). Attribution Patching. In Proceedings of the 7th BlackboxNLP Workshop on Analyzing and Interpreting Neural Networks for NLP.
Kramár, J., Lieberum, T., Nanda, N., & Shah, R. (2024). AtP*: Efficient and scalable methods for localizing LLM behaviour to components. arXiv preprint arXiv:2403.00745. to components. arXiv preprint arXiv:2403.00745.

Chain Of Thought in Diffusion Language Models

Sun, 19 Oct 2025 00:00:00 GMT

Note: This RSS feed strips out SVGs and embeds. You might want to read the post on the webpage here.

October 17, 2025

Last Updated: October 19, 2025

research mechInterp aiSafety

More updates soon!

Solving a maze using Diffusion Model

Sun, 19 Oct 2025 00:00:00 GMT

Note: This RSS feed strips out SVGs and embeds. You might want to read the post on the webpage here.

October 17, 2025

Last Updated: October 19, 2025

projects GenAI

Your browser does not support the video tag.

This post documents a small, didactic project I built to understand diffusion models end-to-end. It’s a from-scratch, minimal implementation of schedules, samplers, and a compact U-Net that learns to draw a valid path on a 2D maze by diffusing only the path channel while conditioning on fixed walls/start/end.

System at a glance

Data contract. On disk: NumPy arrays [4, H, W] in [0,1] ordered [walls, start, end, solution]. The loader emits:
- condition=[3, H, W] in [0,1] (walls/start/end)
- target=[1, H, W] scaled to [-1,1] (path to diffuse)
Model. Compact U-Net (GroupNorm + SiLU, residual conv blocks) with sinusoidal time embeddings. in_channels=4 (noisy path + 3 condition), out_channels=1 (predicted noise ε).
Diffusion core. Linear and cosine schedules implemented in-house with precomputed α_t, \bar{α}_t, posterior variance, forward factors, and DDIM coefficients for a uniform timestep ladder.
Samplers. DDPM (ancestral) and DDIM (fast, η=0 deterministic or η>0 stochastic), with classifier-free guidance via two passes (cond/uncond) and linear mixing.
Training. Sample t, add noise only to the path channel, concatenate condition (optionally dropped for CFG training), predict ε with MSE loss. AdamW + grad clip + EMA.
Inference. Solver prefers EMA weights, uses cosine schedule and DDIM by default, expects images with black=walls, green=start, red=end, white=free. Returns raw [-1,1] and thresholded [0,1] path maps, plus validation.
Evaluation. Validity (no wall overlap; start/end on path; BFS 4-connectivity), plus IoU/F1 vs. ground truth and optimality gap by path length.

Note on channels: Inference must use the same in_channels as training. If a checkpoint was trained with, say, 7 channels (e.g., extra aux features), construct the inference U-Net with in_channels=7 or weights won’t load.How it works (briefly)How it works (brief)How it works

As a compact learning project, this codebase made diffusion mechanics tangible: ε-prediction training, schedule design, sampler behavior, and guidance trade-offs. If you run the quick experiments above, you’ll see how each choice-cosine vs. linear, steps, CFG scale, η-maps to validity, quality, and speed in a way that’s easy to reason about and extend.

Causal Attribution & Patching (Part 1): Internal Mechanisms of Language Models

Sun, 19 Oct 2025 00:00:00 GMT

Note: This RSS feed strips out SVGs and embeds. You might want to read the post on the webpage here.

October 16, 2025

Last Updated: October 19, 2025

mechInterp #CMSC848R

💡

TL;DR

Causal vs. Correlational: Mechanistic interpretability aims to understand not just what information is present in a model's activations (correlation), but what information the model causally uses to produce its output. Patching and mediation are tools for establishing this causal link.
Mediation Analysis Identifies Pathways: Inspired by the social sciences, causal mediation analysis allows us to identify and quantify the specific components (attention heads, neurons) that transmit an effect from input to output, pinpointing where mechanisms like bias are implemented.
Mechanisms Can Be Surprisingly Linear: For a range of relational tasks (e.g., country capitals, verb tenses), large language models sometimes implement a simple vector-arithmetic update within their MLP/FFN layers, especially when retrieving facts from pre-training memory.
Causal Tracing & ROME Localize and Edit Facts: By systematically "patching" clean states into a corrupted model run, we can trace the flow of information for factual recall. This localizes facts to specific MLP modules, which can then be surgically edited with a rank-one update to change the model's knowledge, with testable effects on specificity and generalization.

What They Mean by Causal Attribution (vs. Correlational Probing)

For years, a dominant paradigm in interpretability has been correlational probing. A typical study might train a simple linear classifier (a "probe") on the internal activations of a large language model (LLM) to predict a certain property, like the part-of-speech of a token or the sentiment of a sentence. If the probe achieves high accuracy, the conclusion is that the property is "encoded" in the model's representations. (Probing)

This is an essential first step, but it leaves a critical question unanswered: just because information is present, does that mean the model actually uses it to make its next prediction?

Imagine a car's dashboard. A probe could easily "decode" the car's speed by looking at the radio's volume knob-perhaps the driver turns up the music on the highway. There is a correlation. But the volume knob does not cause the car to accelerate. The speedometer and the engine do. Causal attribution is the effort to find the speedometer, not just another correlated signal.

This is the distinction between information being merely encoded versus being causally implicated in the model's behavior. An internal state can contain linearly decodable information about a concept, but that information might be an artifact of the training data or a precursor to a different computation, rather than a direct input to the final decision.

As Vig et al. (2020) note in their work on causal mediation, many interpretability methods "show only that information is present in the model’s representations, not that it is used." Causal methods, by contrast, rely on interventions. The core idea is simple: if you think component $M$ is causally responsible for behavior $Y$, then changing or disabling $M$ should change $Y$. If it doesn't, your hypothesis is likely wrong, no matter how much information a probe can extract from $M$.

Background & Formalization

The Intervention Model: Corrupt vs. Clean Runs

The foundational tool for causal analysis in LLMs is the patching experiment. This involves running the model on at least two different inputs and swapping internal states between them.

The Clean Run: We run the model on a clean, factual input. For example, the prompt X=’The Eiffel Tower is located in’X = \text{'The Eiffel Tower is located in'}X=’The Eiffel Tower is located in’. We expect the model to complete this with the target token Y=" Paris"Y = \text{" Paris"}Y=" Paris". We save all the internal activations (residual stream states, attention outputs, MLP outputs) from this run. Let's call a generic clean state sss.
The Corrupted Run: We run the model on a different, "corrupted" input designed to produce an incorrect answer. This could be an input with a different subject (e.g., "The Colosseum is located in") or a nonsensical prompt where we inject noise into the embeddings. The goal is to create a computational path that does not lead to the target "Paris". Let's call a generic corrupted state s~\tilde{s}s~.

The experiment is to run the model on the corrupted input but, at a specific point in the forward pass, intervene and "patch" in a clean state from the original run. For instance, we might replace the output of the MLP module in layer l=15l=15l=15 from the corrupted run with the MLP output from the clean run, and then let the computation proceed.

We measure the effect of this intervention on the model's output logits. Specifically, we look at the change in the logit of the target token (e.g., "Paris"). We can formalize this as a restoration score:

Δlogit(t⋆)=logitrestore(M)(t⋆)−logitcorrupt(t⋆)\Delta \text{logit}(t^\star) = \text{logit}{\text{restore}(M)}(t^\star) - \text{logit}{\text{corrupt}}(t^\star)Δlogit(t⋆)=logitrestore(M)(t⋆)−logitcorrupt(t⋆)

Here, t⋆t^\start⋆ is our target token ("Paris"). logitcorrupt(t⋆)\text{logit}{\text{corrupt}}(t^\star)logitcorrupt(t⋆) is the (presumably low) logit for "Paris" on the corrupted run. logitrestore(M)(t⋆)\text{logit}{\text{restore}(M)}(t^\star)logitrestore(M)(t⋆) is the logit for "Paris" when we patch in the clean state of a specific component, or mediator, MMM.

A large, positive Δlogit\Delta \text{logit}Δlogit is strong evidence that the component MMM is causally implicated in producing the token "Paris" for the clean prompt. It means that restoring just this one piece of the computation was sufficient to recover much of the correct answer.

The Idea of Causal Mediation

Causal mediation analysis provides a formal language for these interventions. It originates from fields like epidemiology and social sciences, where researchers want to understand how a treatment causes an outcome. For example, how does an educational program (treatment XXX) lead to higher salaries (outcome YYY)? Is it by increasing a student's confidence, or by teaching them a specific skill (mediator MMM)?

The framework decomposes the Total Effect (TE) of XXX on YYY into two pathways:

Natural Direct Effect (NDE): The effect of $X$ on $Y$ that does not go through the mediator MMM.
Natural Indirect Effect (NIE): The effect of XXX on YYY that flows through the mediator MMM.

TE=NDE+NIETE = NDE + NIETE=NDE+NIE

In LLMs, the "treatment" XXX is the input prompt, and the "outcome" YYY is the output distribution. The "mediators" MMM are the internal components we hypothesize are on the causal path: individual neurons, attention heads, or entire MLP modules. By estimating the indirect effect (NIE) for each component, we can quantify how much of the model's final behavior is attributable to that specific part of the network.

Paper 1 - Causal Mediation in Language Models (NeurIPS 2020)

Vig et al. (2020) were among the first to formally apply causal mediation analysis to LLMs to locate the sources of a specific behavior: gender bias.

Goal: The researchers wanted to move beyond simply showing that gender is encoded in representations and instead identify which specific model components (attention heads and neuron activations) are causally responsible for producing biased predictions. For example, in a prompt like "The doctor said that," which components cause the model to assign a higher probability to "he" than "she"?
Method: They framed this as a mediation problem. The "treatment" was the gender of a subject in a sentence. The "outcome" was the probability distribution over subsequent gendered pronouns. They then systematically estimated the Natural Indirect Effect (NIE) for each attention head and neuron in the model, effectively measuring how much of the gender-bias effect was flowing through that component.
Result: The analysis revealed that gender bias was not a diffuse, entangled property of the entire network. Instead, the causal effects were highly concentrated in a small number of specific attention heads and neurons. For instance, in GPT-2, a handful of heads were consistently identified as "gender-mediating heads." This was a powerful demonstration that complex behaviors can be localized to specific, modular parts of the model.
Limits & Assumptions: Causal mediation relies on strong statistical assumptions, chiefly the assumption of no unobserved confounding between the mediator and the outcome. In a densely connected neural network, this assumption is difficult to formally satisfy. Furthermore, the identified components are specific to the model, its training data, and the particular task being analyzed. The results may not generalize perfectly to different prompts or model families.

Paper 2 - Linear Mechanisms in Practice (NAACL 2024)

While mediation analysis tells us where an effect is happening, it doesn't always tell us how. The work by Todd et al. (2024) on vector-arithmetic mechanisms provides a compelling hypothesis for the "how," suggesting that the underlying computation can be surprisingly simple and linear.

Claim: The paper demonstrates that for several common relational tasks, LLMs of various sizes (from 124M to 176B parameters) often implement the solution via a simple, additive update in their Feed-Forward Network (FFN), also known as MLP, layers. This is particularly true when retrieving knowledge stored during pre-training, as opposed to information present in the immediate context.
For example, to solve the task "Capital of France?", the model might process the representation for "France", r(France)r(\text{France})r(France), and an MLP layer will compute something akin to:
MLP(l)(r(France))≈r(France)+d⃗capitalMLP^{(l)}(r(\text{France})) \approx r(\text{France}) + \vec{d}_{\text{capital}}MLP(l)(r(France))≈r(France)+dcapital is a consistent direction in activation space that transforms a country's representation into its capital's. They find similar additive mechanisms for tasks like uppercasing words and changing verb tense.
Connection to Patching: This finding provides a powerful explanatory framework for why patching interventions work so well. If a task-relevant computation is just the addition of a specific vector d⃗\vec{d}d by an MLP module, then a causal patching experiment that restores the output of that MLP should have an enormous effect. It's not just restoring some abstract "information"; it's restoring the result of a precise, near-linear arithmetic operation that is critical to solving the task.

Intuition
Think of the MLP layers as a massive key-value store. The attention layers route the correct query (e.g., the representation of the last token of "Eiffel Tower") to the right MLP layer. That MLP layer then retrieves the associated "value" vector (e.g., a vector that points towards "Paris") and adds it to the residual stream. This additive update pushes the model's overall representation towards the correct answer.

Caveat: The authors are careful to note that this is not a universal mechanism. LMs are complex, non-linear systems. However, for a significant class of knowledge-retrieval and relational tasks, this simple linear update model appears to be a common and effective strategy implemented by the network.

Paper 3 - Causal Tracing & ROME (NeurIPS 2022 project)

The ROME project (Meng et al., 2022) provides a practical and powerful synthesis of these ideas, using causal interventions not just to locate information but to surgically edit it.

Causal Tracing: Finding the Path

To edit a fact, you first need to know where it lives. ROME introduces a method called causal tracing to do exactly this. It is a direct application of the corrupt-and-restore patching protocol described earlier, executed at scale.

Method: The model is run on a clean prompt (e.g., "The Eiffel Tower is in") and a corrupted prompt with random token embeddings. Then, the algorithm iterates through every single state in the clean forward pass (every layer, every token position, for both attention and MLP outputs). At each step, it patches that single clean state into the corrupted run and measures the restoration of the probability of the correct answer ("Paris").
Localization Results: By aggregating the results, a clear picture emerges. The causal effect for retrieving this fact is not distributed evenly. It is highly localized at:
1. MLP Modules: The effect is overwhelmingly concentrated in the MLP layers, not the attention layers.
2. Mid-Layers: The critical MLP modules are typically found in the middle layers of the network (e.g., layers 15-20 in a model like GPT-J).
3. Last Subject Token: The effect is localized to the computation happening at the position of the last token of the subject (the token "Tower" in our example).

This process creates a causal map of factual recall, pointing to a very specific set of modules responsible for storing and retrieving the fact.

ROME: Rewriting Facts with Rank-One Edits

Once a fact is localized to a specific MLP module, ROME (Rank-One Model Editing) performs a direct edit.

Method: An MLP layer computes y=Woutσ(Winx)y = W_{out} \sigma(W_{in} x)y=Woutσ(Winx). The ROME method treats this layer as a linear key-value memory. It solves for a minimal change to one of the weight matrices (e.g., WinW_{in}Win) that will map the subject's representation (the "key") to a new target representation (the "value"). This change can be formulated as a rank-one update, meaning the modification matrix ΔW\Delta WΔW has a rank of one. This is a very low-complexity and targeted change.
Evaluation: A successful edit isn't just about changing one fact. The ROME project introduces a crucial evaluation framework based on three criteria, tested on the CounterFact dataset:
1. Efficacy: Does the model now output the new, edited fact?
2. Specificity: Does the model still correctly answer questions about other, unrelated subjects? (i.e., did the edit cause collateral damage?).
3. Generalization: Does the edit generalize to paraphrases of the original prompt? (e.g., "Where is the Eiffel Tower located?").

The results show that these rank-one edits are remarkably effective, achieving high success across all three criteria.

Synthesis - A Unifying View of Patching

These three anchors are not disparate findings; they are three parts of the same story, painting a progressively clearer picture of how models implement knowledge.

Causal Mediation (NeurIPS 2020) provides the high-level framework. It gives us a principled, theory-grounded way to move from correlation to causation, allowing us to ask, "Which components are on the causal path for this behavior?" It's the wide-angle lens for localizing mechanisms.
Linear Mechanisms (NAACL 2024) provides a concrete hypothesis for the type of computation happening at the localized sites. It suggests that for many factual recall tasks, the complex non-linear MLP module is, in practice, implementing a simple, additive update. This explains why a linear, directional intervention can be so effective.
Causal Tracing and ROME (NeurIPS 2022) provides the applied toolkit. Causal tracing is a brute-force implementation of the patching philosophy that validates the localization hypotheses from mediation. ROME then leverages the insight about linear, localized mechanisms to perform a surgical, rank-one edit, demonstrating that our understanding is precise enough to be actionable.

Together, they support a "linear pathway + targeted intervention" picture. The model uses attention to route information to the right place, and a specific MLP module at a specific token position performs a vector-arithmetic operation to retrieve a fact stored in its weights.

This picture breaks down when tasks require more complex, non-linear reasoning, multi-step compositional logic, or when the primary mechanism is not a memory lookup but a dynamic computation heavily reliant on attention patterns over the local context.

Caveat
Be meticulous in your reporting. Document the exact prompts, model versions, layers, and token positions tested. Causal effects can be highly specific, and what holds for one model or prompt template may not hold for another. Beware of over-claiming based on a small set of examples.

My View

For practitioners, the most effective workflow seems to be hierarchical. Start with a broad, mediation-style analysis or causal trace to map the computational path and identify the critical few components. Then, use the linear/additive mechanism hypothesis from Todd et al. (2024) to design simple, directional patching experiments to quickly test hypotheses about what the component is doing. Reserve surgical, rank-one edits like ROME for cases where you need to make a durable change to the model's behavior and are prepared to rigorously evaluate the collateral damage.

Intuition:
A key insight unifying these papers is that many "memory" tasks in LLMs seem to be implemented locally and linearly. The FFNs, despite their non-linear activation function, can be driven into a regime where they behave like simple additive key-value stores. Causal tracing and ROME succeed because they correctly identify this implementation strategy and intervene at precisely the right layer and token position where this "memory lookup" occurs.

Caveat:
The success of these methods on factual recall should not be misinterpreted as evidence that the entire LLM is a simple, linear system. These techniques excel on atomized, single-hop tasks. They are less informative about multi-step reasoning, creative generation, or behaviors that emerge from complex interactions between many components. Furthermore, the causal pathways in a base model may be significantly altered by instruction tuning and RLHF, which layer on new, complex behaviors that may not be as easily localized.

Limits & Open Questions

This line of research, while incredibly promising, is still in its early stages. Key open questions remain:

Portability: How well do localized mechanisms transfer across model sizes and families? Is the fact "The Eiffel Tower is in Paris" stored in a similar location and way in GPT-J, Llama-3, and Claude 3?
Effect of Finetuning: How does instruction tuning or RLHF reshape these causal pathways? Does it create new, dedicated circuits for following instructions that override or modulate the base model's factual recall mechanisms? Are these new circuits as editable?
Compositionality: How can we move from editing single, atomic facts to editing compositional knowledge or complex rules? Editing "All birds can fly" is much harder than editing "The ostrich is a bird," as it requires a change that propagates through the model's conceptual space.
Beyond Factual Recall: What are the causal mechanisms for more abstract tasks, like identifying logical fallacies, summarizing a long document, or adopting a specific persona? Are they as localizable and editable?

Answering these questions will be critical as we move from understanding isolated mechanisms to building a more holistic, causal theory of how these powerful and complex models work.

References

Meng, K., Bau, D., Andonian, A., & Belinkov, Y. (2022). Locating and Editing Factual Associations in GPT. Project page based on work presented at NeurIPS 2022.

Todd, M. E., Mamou, J., & Lieber, O. (2024). Language Models Implement Simple Word2Vec-style Vector Arithmetic. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (NAACL).

Vig, J., Geiger, A., & Belinkov, Y. (2020). Investigating Gender Bias in Language Models Using Causal Mediation Analysis. In Advances in Neural Information Processing Systems (NeurIPS).

Training Data Attribution: Corroborative Methods

Fri, 17 Oct 2025 00:00:00 GMT

Note: This RSS feed strips out SVGs and embeds. You might want to read the post on the webpage here.

October 16, 2025

Last Updated: October 17, 2025

#CMSC848R mechInterp

💡

TL;DR

What is Corroborative Attribution? It is the process of finding evidence in a corpus that directly supports or entails a specific model behavior. Unlike contributive attribution, it does not estimate causal impact but instead asks: "Is there data that could plausibly explain this output?"
Why Does It Matter? Corroboration is a practical, scalable tool for auditing models. It helps us detect benchmark contamination, identify sources of toxicity or private information, and understand data provenance, especially when internal model access or full training logs are unavailable.
How to Do It At Scale: The What's In My Big Data? (WIMBD) project provides a concrete blueprint. By leveraging massive-scale "count and search" operations across over 35 terabytes of text, the authors uncovered high rates of duplication, PII, toxicity, and benchmark contamination in major pre-training corpora like RedPajama and C4.
The Deeper Motivation: As argued in You Are What You Eat, a model's internal structure and its resulting behaviors are fundamentally shaped by its training data. Corroborative audits are a necessary first step toward understanding this data-to-behavior pipeline, forming a foundation for safer and more aligned AI systems

What is “Corroborative Attribution”?

Training Data Attribution (TDA) is a broad term for methods that connect a model's behavior back to its training data. Within this field, a critical distinction exists between two primary goals: corroboration and contribution.

Corroborative attribution aims to find evidence. Given a specific model behavior-such as generating a piece of text, making a factual claim, or exhibiting a toxic stereotype-the goal is to search a candidate corpus and identify examples that directly support, contain, or entail that behavior. The central question is one of provenance and support: “Can I find a document in this dataset that substantiates what the model is saying?”

This stands in contrast to contributive attribution (Training Data Attribution: Contributive Methods). Contributive methods, like influence functions or TracIn, seek to estimate the causal impact of a training example on a model's parameters or a specific prediction. Their question is counterfactual: “How much would this prediction have changed if the model had not been trained on this specific example?”

While contributive methods offer deep insights into model learning, they are often computationally expensive and require significant access to the model's internals (e.g., gradients, checkpoints), which is frequently infeasible for practitioners working with proprietary, third-party models.

Corroborative attribution, on the other hand, is a powerful and practical first-line-of-defense auditing tool. It enables:

Data Provenance: Tracing model outputs to their likely sources for copyright, licensing, or fact-checking purposes.
Leakage & Contamination Audits: Systematically checking if evaluation benchmarks or sensitive test sets are present in the training data, which would invalidate reported performance metrics.
Safety Reviews: Identifying the sources of undesirable behaviors like toxicity, bias, or the memorization of personally identifiable information (PII).
Policy & Deployment Decisions: Providing concrete evidence to inform decisions about model deployment, even when the model itself is a black box.

Background & Formalization

We can formalize the goal of corroborative attribution as follows. Let BBB denote an observed model behavior, such as a generated sentence, a classification decision, or a factual claim. Let CCC be a candidate corpus, which could be a known pre-training dataset (like The Pile) or a suspected source of information (like a crawl of a specific website).

A corroborative method, or a corroborator, is a process that returns a support set Starg⊆CS_{\text{targ}} \subseteq CStarg⊆C. For each example z∈Stargz \in S_{\text{targ}}z∈Starg, we say that zzz supports the behavior $B$, which we denote as

z⇒Bz \Rightarrow Bz⇒B. This support relationship is established by an evidence relation RRR and a threshold τ\tauτ. Formally:

z∈Starg∧R(z,B)≥τ ⟹ z⇒Bz \in S_{\text{targ}} \wedge R(z, B) \geq \tau \implies z \Rightarrow Bz∈Starg∧R(z,B)≥τ⟹z⇒B

The evidence relation R(⋅,⋅)R(\cdot, \cdot)R(⋅,⋅) is the heart of the method and can take several forms, each with its own strengths and limitations:

Exact Match (REXACTR_{\text{EXACT}}REXACT): This is the simplest relation, checking if the string representing behavior BBB appears verbatim in document zzz. It is precise but brittle; it fails to capture even minor variations like changes in punctuation or capitalization.
Near-Duplicate Match (RNEARR_{\text{NEAR}}RNEAR): This relation uses fuzzy matching techniques (e.g., n-gram overlap, MinHash) to identify documents zzz that are nearly identical to BBB. It is more robust than exact matching but can still miss paraphrases.
Similarity Search (RSIMR_{\text{SIM}}RSIM): Here, we map both the behavior BBB and documents zzz into a high-dimensional embedding space (e.g., using a sentence-transformer). The relation is typically cosine similarity, RSIM(z,B)=cosine_sim(embed(z),embed(B))R_{\text{SIM}}(z, B) = \text{cosine\_sim}(\text{embed}(z), \text{embed}(B))RSIM(z,B)=cosine_sim(embed(z),embed(B)). This method is powerful for finding semantically related content but is sensitive to the choice of embedding model and the similarity threshold τ\tauτ. A low threshold can lead to many false positives (spurious connections).
Entailment Check (RENTR_{\text{ENT}}RENT): This is a more sophisticated logical check. Using a Natural Language Inference (NLI) model, we determine if a candidate document zzz textually entails the behavior BBB. For instance, if BBB is "Paris is the capital of France," a document zzz stating "The French government is based in Paris" would entail it. This helps filter out spurious semantic matches but depends heavily on the quality of the NLI model.

Caveat: Corroboration is not proof of training. Finding a document zzz in corpus CCC that supports behavior BBB demonstrates that the corpus contains information that could plausibly have led to that behavior. It does not definitively prove that the model used that specific document to produce the output. The same information could be present in many other documents or even in a different, unknown training corpus.

Paper 1 - What's In My Big Data? as a Corroborative Audit

The ICLR 2024 paper What's In My Big Data? (WIMBD) by Elazar et al. serves as a landmark example of a large-scale, systematic corroborative audit. Instead of focusing on a single model's behavior, the authors developed a platform to analyze the properties of the data itself, creating a foundation for countless future corroborative investigations.

Method: "Count + Search at Scale"

The core innovation of WIMBD is its ability to perform two fundamental operations-counting and searching-efficiently across massive text corpora. The authors analyzed ten major datasets, including C4, The Pile, RedPajama, and LAION-2B-en, totaling over 35 terabytes of text. By building specialized indices and using optimized search tools, they could execute complex queries on a standard compute node, making large-scale data audits accessible without requiring a supercomputer.

Key Findings as Corroborative Evidence

WIMBD's analyses directly enable corroborative workflows by characterizing the evidence available in these corpora. Their findings were striking:

High Duplication: The study revealed massive levels of duplication. For example, they found that approximately 50% of the documents in both RedPajama and the English portion of LAION-2B are duplicates. This is a critical finding for attribution; a behavior corroborated by a document present thousands of times is more likely to be memorized or heavily weighted by a model.
Benchmark Contamination: The authors systematically searched for test and evaluation data within the pre-training corpora. They found significant contamination for key benchmarks, including the Winograd Schema Challenge (WSC), and parts of the GLUE and SuperGLUE suites. This discovery corroborates the hypothesis that high performance on these benchmarks may be due to memorization rather than true generalization.
PII and Toxicity Prevalence: The platform was used to search for patterns matching personally identifiable information (email addresses, phone numbers) and to score documents for toxicity. The widespread presence of this content provides the necessary substrate for corroborating model behaviors that leak PII or generate toxic language.

WIMBD is a quintessential example of corroborative auditing because it links potential model behaviors (high benchmark scores, PII leakage) directly to evidence at the corpus level. It does not measure the influence of any single document but rather characterizes the sea of data from which these behaviors could plausibly emerge.

Intuition: The analyses in WIMBD effectively transform the opaque concept of "pre-training data" into a searchable, queryable substrate. This move from an unknown blob of text to an indexed and analyzed corpus is the foundational step that makes corroborative attribution possible. It allows a researcher to move from "the model said something weird" to "the model said something weird, and here are 10,000 documents in its likely training data that say the same thing."

Paper 2 - Why Corroboration Matters (You Are What You Eat)

The position paper You Are What You Eat argues that a deep understanding of the relationship between training data, a model's internal structure, and its downstream generalization is necessary for AI alignment. It posits that merely evaluating a model's behavior on test sets is insufficient to guarantee safety, as different internal structures can produce identical outputs on a given dataset but generalize in dangerously different ways.

The central thesis is that data shapes structure, and structure determines behavior. This provides a powerful motivation for investing in robust corroborative audits. If we want to build safe, aligned AI, we cannot treat the training data as an afterthought. We must understand its contents, its flaws, and its latent structures, as these will inevitably be reflected in the models we train.

Corroborative attribution is the most direct method we have for drawing this line from data to behavior. When we corroborate a model's toxic output by finding toxic source documents, or corroborate a factual error by finding source documents containing the same misinformation, we are empirically verifying the "You Are What You Eat" principle.

Opinion: For practitioners working on the front lines of AI safety and deployment, especially with limited access to model internals, corroboration is often the only tractable audit lever available. While mechanistic interpretability and contributive attribution are vital research directions, the ability to quickly search a corpus for evidence supporting a harmful behavior is a practical, scalable, and indispensable tool for risk assessment today.

Method Families & Practical Protocols (Corroborative)

A practical corroborative audit follows a structured protocol, moving from a behavior of interest to a final report. This protocol typically combines several families of methods.

Method Families

Exact & Near-Duplicate Matching: This is the go-to method for investigating suspected memorization, copyright infringement, or benchmark contamination. Using lexical tools (string matching) and fuzzy hashing (MinHash), one can quickly find documents that are verbatim or near-verbatim copies of the target behavior BBB. WIMBD's contamination checks for benchmarks like GLUE relied heavily on this approach.
Embedding & k-NN Retrieval: For finding semantic support, not just lexical copies, embedding-based retrieval is key. The process involves encoding the behavior $B$ into a vector and using an approximate nearest neighbor (k-NN) search index (like Faiss or ScaNN) to retrieve the top-k most similar documents from the corpus embeddings. The relevance of the results is then determined by applying a similarity threshold, RSIM≥τR_{\text{SIM}} \ge \tauRSIM≥τ.
Textual Entailment Checks: Semantic similarity can be misleading. A document might be about the same topic as BBB but contradict it. An entailment model can act as a crucial filter. After retrieving a set of candidates via embedding search, you can run an NLI model to check if a candidate document actually entails, contradicts, or is neutral with respect to the claim in BBB.
Corpus Health Metrics: Before investigating any specific behavior, it is useful to have global statistics about the corpus. The work of WIMBD exemplifies this: knowing the overall duplication rate, toxicity levels, and prevalence of PII helps prioritize which behaviors to investigate and provides context for the findings.

A Practical Protocol

A typical corroborative audit workflow can be structured as follows:

Define Behavior $B$: Start with a clear, specific behavior to investigate. This could be a verbatim quote from a model, a factual claim it made, or a pattern of biased completions.
Query Construction: Transform $B$ into one or more queries. For lexical search, this might be the string itself. For semantic search, it's the embedding of the string. For broader audits, it could be a regular expression (for PII) or a list of keywords.
Tiered Search: Begin with the cheapest, most precise methods first. Run exact and near-duplicate searches. If no strong evidence is found, proceed to broader semantic search (embedding/k-NN retrieval).
Filtering & Refinement: The output of a semantic search can be noisy. Filter the initial candidate set StargS_{\text{targ}}Starg using a relevance threshold τ\tauτ. For high-stakes audits, add a textual entailment (RENTR_{\text{ENT}}RENT) check to ensure the retrieved documents logically support the behavior.
Human-in-the-Loop Review: Automated metrics are not enough. A human expert should review the final, filtered set of supporting documents to assess their quality and relevance, and to identify subtle issues that automated tools might miss.
Reporting: The final report should be transparent. Disclose the corpora searched, the search methods used (REXACT,RSIMR_{\text{EXACT}}, R_{\text{SIM}}REXACT,RSIM, etc.), the thresholds (τ\tauτ) applied, and the final support set StargS_{\text{targ}}Starg. Crucially, include caveats about the limitations of the audit.

Caveat: The efficacy of any corroborative audit is limited by the accessibility of the training corpus. When auditing a model trained on a closed-source, proprietary dataset, one must rely on open-source proxies (e.g., using RedPajama as a proxy for the Llama training set). In such cases, the report must clearly state this limitation and frame the findings as plausible evidence rather than definitive proof.

Evaluation & Reporting

A rigorous corroborative audit is not just about finding examples; it's about evaluating the process itself. If a ground-truth dataset exists (e.g., a set of known memorized sequences and their sources), the performance of a corroborator can be measured with standard metrics:

Precision: What fraction of the retrieved support examples (StargS_{\text{targ}}Starg) are truly relevant?
Recall: What fraction of all possible relevant examples in the corpus CCC did the method successfully retrieve?

In practice, practitioners should perform ablations over the similarity threshold τ\tauτ to understand the trade-off between precision and recall. Reporting should be standardized and reproducible, including:

The exact corpora searched, with versioning information.
The query templates or construction methods used.
The software and models used for embedding, search, and entailment.
For audits involving human review, inter-annotator agreement scores should be reported to quantify the subjectivity of the task.

Tools like dataset cards and datasheets should be used not just for the training data itself, but for the outputs of corroborative audits, providing a transparent record of what was searched for and what was found.

My View

Opinion: Every high-stakes audit of a powerful language model should begin with a corroborative analysis. It is the most direct, scalable, and empirically grounded way to assess data-related risks. Only when this foundation is established, and when access permits, should teams move to more complex and computationally intensive contributive methods. Corroboration provides the "what" and "where"; contribution can then explore the "how much."

Intuition: The core argument of You Are What You Eat-that data shapes structure which in turn shapes behavior-is deeply compelling. If this causal chain holds, then examining the source data is not just one of many options; it is an indispensable first step. Ignoring the properties of the training data is akin to a doctor diagnosing a patient without asking about their diet.

Caveat: Corroborative methods have blind spots. They are biased towards finding evidence that is lexically or semantically close to the target behavior. They can over-credit popular sources that are highly duplicated and may fail to identify support from a collection of diverse, long-tail documents that, in aggregate, teach the model a concept through paraphrase and analogy. An audit report must always be transparent about the types of evidence it is designed to find and the types it is likely to miss.

Limits & Open Questions

Despite its practical utility, corroborative attribution faces several significant challenges and open research questions:

Coverage vs. Privacy: How can we conduct thorough audits when training corpora are increasingly proprietary and closed-source? Can we develop methods for privacy-preserving corroboration that allow model developers to answer queries about their data without revealing it?
Bridging to Contribution: Finding a supporting document is the first step. The next is understanding its impact. Future research should focus on integrating corroborative search (to find candidates) with lightweight contributive estimation (to rank candidates by influence).
Watermarking and Fingerprinting: Could we proactively embed invisible signals (watermarks) or unique phrases (fingerprints) into training corpora? This would transform corroboration from a search problem into a simple detection problem, making attribution far more reliable.
The Cost of Scale: While WIMBD demonstrated feasibility on a single node, analyzing future trillion-token datasets will require continued innovation in efficient indexing and search algorithms.

Ultimately, corroborative methods are a vital part of the toolkit for ensuring AI systems are safe, transparent, and accountable. They ground our understanding of model behavior in the concrete reality of the data they consume, reminding us that, in the world of AI, you are what you eat.

References

Pepin Lehalleur, S., Hoogland, J., Farrugia-Roberts, M., Wei, S., Gietelink Oldenziel, A., Wang, G., Carroll, L., & Murfet, D. (2025). You Are What You Eat -- AI Alignment Requires Understanding How Data Shapes Structure and Generalisation. arXiv preprint arXiv:2502.05475.
Elazar, Y., Bhagia, A., Magnusson, I., Ravichander, A., Schwenk, D., Suhr, A., Walsh, P., Groeneveld, D., Soldaini, L., Singh, S., Hajishirzi, H., Smith, N., & Dodge, J. (2024). What's In My Big Data? Proceedings of the International Conference on Learning Representations (ICLR).
Geva, M., Schuster, T., Berant, J., & Globerson, A. (2023). Connecting Pre-training and Fine-tuning with Attribution for T5. OpenReview.
Cohen-Wang, C., Chuang, I. L., & Zhang, C. (2023). Unifying Corroborative and Contributive Attributions in Large Language Models. OpenReview.

Training Data Attribution: Contributive Methods

Fri, 17 Oct 2025 00:00:00 GMT

Note: This RSS feed strips out SVGs and embeds. You might want to read the post on the webpage here.

October 16, 2025

Last Updated: October 17, 2025

#CMSC848R mechInterp

💡

TL;DR

What is Training Data Attribution (TDA)? TDA is a set of techniques for identifying which training examples are most responsible for a specific model behavior, such as a particular prediction or an overall performance metric. It moves beyond asking what a model predicts to asking why, in terms of its training data.
Why does it matter? As we move past chasing benchmark scores, we need tools for debugging, enhancing fairness, ensuring robustness, and understanding the origins of model capabilities. TDA provides a crucial link between the data we curate and the behaviors the models exhibit.
What are contributive methods? This post focuses on contributive TDA, which assigns a score to each training point reflecting its helpful or harmful influence on a model's prediction. We will cover four major families: Influence Functions, TracIn, Representer Points, and Data-Shapley.
What are the key takeaways? TDA is not a single, solved problem. Each method comes with its own set of mathematical assumptions, computational trade-offs, and failure modes. The most effective approach for a given problem depends on the model architecture, the available resources (e.g., gradients, checkpoints), and the specific question being asked.
What this post covers: A deep dive into the math, intuition, and practical caveats of the core contributive methods, with a final synthesis to guide practitioners on which method to choose and when.

The field's reliance on chasing leaderboard scores is slowly giving way to a more mature focus on model understanding, and for good reason: benchmark score-chasing isn’t enough; we need methods that link predictions to specific training examples. Without this connection, we are flying blind, unable to explain why a model fails on a critical edge case, reproduces a harmful bias, or succeeds at a surprising new task. Training Data Attribution (TDA) provides the analytical lens for forging this link.

What Counts as “Training Data Attribution”?

At its core, Training Data Attribution is the process of assigning a score to one or more training examples, z_i, that quantifies their responsibility for a model's behavior on a specific test point, z_test, or for its overall performance. It seeks to answer the question: "Which training data points made the model make this specific prediction?"

This is distinct from several related concepts. Data provenance simply tracks the origin of data, without assessing its impact on the model. Memorization detection identifies instances where the model has stored and reproduced training data verbatim, which is a specific type of influence but not the whole story. Dataset auditing typically involves high-level statistical analysis of a dataset (e.g., class balance, subgroup representation) rather than example-level attribution for a specific prediction.

The motivation for TDA has become increasingly urgent. Research such as "Impact of Pretraining Term Frequencies on Few-Shot Numerical Reasoning" (Razeghi et al., 2022) demonstrates that a model's emergent abilities are deeply connected to the statistical properties of its training corpus. For example, they show that the frequency of numbers in the pretraining data correlates with a model's ability to reason about those numbers. While this provides a corpus-level explanation, TDA offers the tools to drill down and ask: given this statistical backdrop, which specific examples were most instrumental in teaching the model a given concept or causing a particular failure?

Background & Scope

Before diving into methods, let's establish our notation and scope. We consider a training set D = {z_i = (x_i, y_i)}_{i=1}^n of n examples. Our model is parameterized by θ ∈ R^p. The loss on a single example z is ℓ(z, θ). The empirical risk, which the model aims to minimize, is R(θ) = (1/n) Σ_i ℓ(z_i, θ) + λΩ(θ), where Ω(θ) is a regularization term with weight λ. The Hessian of the risk is the matrix of second derivatives, H = ∇_θ^2 R(θ). Our goal is to understand the influence of a training point z_i on a prediction for a test point z_test.

TDA methods are most cleanly applied in a supervised learning context, such as fine-tuning a language model on a specific task. The dataset is well-defined, accessible, and of a manageable size. Applying TDA to the pretraining stage of a foundation model is significantly harder due to the immense scale and frequent inaccessibility of the pretraining corpus.

Different TDA methods also require different levels of access to the model and training process:

Some require full access to gradients and the ability to compute Hessian-vector products.
Some rely on having saved model checkpoints from various stages of training.
Others are only applicable to models with specific architectures, like a linear final layer.
The most computationally demanding methods may require retraining the model on many different subsets of the data.

Contributive Methods (Core Families)

Contributive methods form the backbone of TDA. They estimate the positive (helpful) or negative (harmful) impact of each training point on a specific outcome.

(a) Influence Functions

Method
Influence Functions, introduced to machine learning by Koh & Liang (2017), are a classic technique from robust statistics. The core idea is to approximate the effect of removing a training point by asking: "How would the model's optimal parameters θ change if we infinitesimally up-weighted a single training point z_i?" This change in parameters is then used to measure the change in loss on a test point z_test.

Math
The effect on the parameters θ from up-weighting a training point z_i by a small amount ε is given by the influence function:

dθ_ε / dε |_(ε=0) = - (1/n) H⁻¹ ∇_θ ℓ(z_i, θ)

This equation tells us that the change in parameters is proportional to the gradient of the training point's loss, scaled by the inverse Hessian H⁻¹. The Hessian measures the curvature of the loss landscape; inverting it translates a step in gradient space into a change in parameter space.

To find the influence of z_i on the loss at z_test, we use the chain rule. The influence score, IF(z_i → z_test), is defined as the change in the test loss:

IF(z_i → z_test) = - (1/n) ∇_θ ℓ(z_test, θ)ᵀ H⁻¹ ∇_θ ℓ(z_i, θ)

A positive score means that up-weighting z_i increases the test loss, implying z_i is harmful to the prediction on z_test. A negative score means z_i is helpful.

Assumptions
Influence Functions rest on strong assumptions. They require the loss function to be twice-differentiable and strictly convex, ensuring a unique minimum and an invertible Hessian. The method is formally derived as a first-order approximation, meaning it is most accurate for small perturbations and assumes that the model parameters are already at a local minimum.

Computation
The primary computational bottleneck is calculating and inverting the p x p Hessian, H, which is infeasible for models with millions or billions of parameters. Instead, we never form H explicitly. We compute the term H⁻¹ v (where v is a vector, like ∇_θ ℓ(z_i, θ)) using iterative algorithms like the conjugate gradient method. These algorithms only require access to Hessian-vector products (HVPs), which can be computed efficiently with automatic differentiation. A damping term is often added to the Hessian diagonal to improve stability.

Strengths and Failure Modes

Strengths: Influence Functions provide a fine-grained, theoretically grounded attribution score for each training point. They have been successfully used to identify mislabeled examples, understand dataset shortcuts, and debug model predictions.
Failure Modes: The convexity assumption is strongly violated by deep neural networks. The local quadratic approximation that underpins Influence Functions may not hold in the complex, non-convex loss landscapes of Transformers. This can make the resulting scores fragile and unreliable, especially for out-of-distribution training points.

Caveat: Influence Function scores for deep, non-convex models should be treated with caution. They can be a powerful signal, but their fragility means they should ideally be validated with more direct checks, such as leave-one-out retraining on a small subset of the most influential points. Think of them as a sophisticated heuristic, not ground truth.

(b) TracIn

Method
TracIn (Pruthi et al., 2020) offers a more direct, empirical approach. Instead of approximating the effect of removal from a single final model, TracIn calculates influence by integrating the impact a training point has over the entire training trajectory. The intuition is simple: if, during training, the gradient for z_i frequently points in the same direction as the gradient for z_test, then the updates driven by z_i were likely helpful for z_test.

Math
The TracIn score is the sum of dot products between the gradients of the test point and the training point, evaluated at various checkpoints θ_k saved during training, and scaled by the learning rate η_k:

TracIn(z_i → z_test) ≈ Σ_(k∈K) η_k ∇_θ ℓ(z_test, θ_k) ⋅ ∇_θ ℓ(z_i, θ_k)

A large positive score indicates that z_i consistently pushed the model in a direction that would also reduce the loss on z_test, making it an influential proponent. A large negative score indicates it was an opponent.

Assumptions
TracIn makes fewer theoretical assumptions than Influence Functions. It does not require convexity. Its primary practical assumption is the availability of a representative set of model checkpoints K spanning the training process. The quality of the attribution depends directly on the quality and frequency of these checkpoints.

Computation
TracIn is computationally more straightforward than Influence Functions. It requires saving model checkpoints periodically during training. The attribution step involves a forward and backward pass for each training point and the test point at each checkpoint to compute the gradients. The cost scales linearly with the number of training points and the number of checkpoints.

Strengths and Failure Modes

Strengths: Because it directly uses the training trajectory, TracIn is often more robust and reliable for deep, non-convex models than Influence Functions. It is conceptually simple and relatively easy to implement.
Failure Modes: The scores are highly dependent on the choice of checkpoints. If only the final checkpoints are used, TracIn can miss important influences from early in training. The score is also sensitive to the learning rate schedule and optimizer dynamics (e.g., momentum).

Intuition: TracIn's power comes from its direct connection to what the optimizer actually did. It follows the path taken, while Influence Functions analyze the local geometry at the destination. For the winding roads of non-convex optimization, the path often tells a more accurate story than a local map of the final location.

(c) Representer Points

Method
The Representer Theorem provides a way to express the solution of certain kernel methods as a linear combination of the training examples. Representer Point Selection (Yeh et al., 2018) adapts this idea to deep learning. Under specific architectural assumptions, it allows us to decompose a network's pre-activation output for a test point into a weighted sum of similarities to the training points. The weights, or "representer values," serve as the attribution scores.

Math
For a model with a final linear layer that minimizes an L2-regularized objective, the prediction function f(x) can be approximated as a linear combination of kernel similarities between the test input's representation φ(x) and each training input's representation φ(x_i):

f(x) ≈ Σ_i α_i φ(x_i)ᵀ φ(x)

The α_i are the representer values. A positive α_i means that z_i is an "excitatory" example, pushing the prediction in a positive direction. A negative α_i means it is "inhibitory." These values can be computed efficiently from the model's final-layer parameters and the loss gradients.

Assumptions
This method's applicability is constrained by its strong architectural assumptions: the model must be trained with L2 regularization on its final layer parameters, and the final output must be a linear function of the learned representations φ(x). While this holds for many standard classification networks, it may not apply directly to more complex architectures like Transformers without modification.

Computation
Once the model is trained, computing the representer values is typically fast. It involves a single backward pass to get gradients and then solving a system of linear equations that depends on the training set size. The main cost is in pre-computing and storing the representations φ(x_i) for all training data.

Strengths and Failure Modes

Strengths: The method provides directly interpretable signed attributions (excitatory vs. inhibitory), which is a powerful debugging tool. When its assumptions are met, it is computationally efficient.
Failure Modes: Its primary limitation is the strict set of architectural assumptions. It is not a general-purpose method for arbitrary deep learning models. The quality of the attribution also depends on the quality of the learned representations φ(x).

(d) Data-Shapley

Method
Data-Shapley (Ghorbani & Zou, 2019; Jia et al., 2019) brings a game-theoretic approach to TDA. It treats the training data points as players in a cooperative game where the goal is to achieve high model performance. The Shapley value is a concept from game theory that provides a principled way to fairly distribute the total "payout" (model utility) among the players. The Shapley value of a data point is its average marginal contribution to the model's performance across all possible subsets (coalitions) of the training data.

Math
Let U(S) be a utility function that measures the performance of a model trained on a subset of data S ⊆ D. The Shapley value ϕ_i for a data point z_i is:

ϕ_i = E_π [U(S_π ∪ {i}) - U(S_π)]

Here, the expectation is taken over all random permutations π of the training data, and S_π is the set of all data points preceding i in that permutation. In words, it's the average improvement in performance gained by adding z_i to a random subset of data.

Assumptions
Data-Shapley's main assumption is the choice of a meaningful utility function U. This could be validation accuracy, log-likelihood, or some other performance metric. The "fairness" of the attribution is only as good as the utility function is representative of the desired model behavior.

Computation
Computing exact Shapley values is computationally prohibitive, as it requires training the model on all 2^n subsets of the data. In practice, approximations are used. The most common is the Truncated Monte Carlo (TMC) estimation, which samples random permutations of the data and stops adding points once the marginal contributions stabilize. For specific models like k-Nearest Neighbors, efficient closed-form or near-closed-form variations like KNN-Shapley exist.

Strengths and Failure Modes

Strengths: Data-Shapley is founded on a solid, axiomatic basis (efficiency, symmetry, linearity, null player), making it a "fair" measure of contribution. It is particularly well-suited for dataset-level tasks like data valuation for monetization or identifying low-value points for dataset pruning.
Failure Modes: The computational cost is its greatest weakness, making it impractical for large models and datasets without significant approximations. The results can have high variance due to the stochasticity of both model training and the sampling process. Furthermore, the choice of utility function can drastically change the resulting attributions.

Opinion: Data-Shapley is the gold standard in theory but often a computational nightmare in practice. It is best reserved for situations where you need to make high-stakes decisions about dataset composition (e.g., pruning, valuation) and have the resources to invest in robust approximation. For routine prediction-level debugging, simpler methods are often more practical.

Bridge to LLMs

Applying these methods directly to large language model (LLM) pretraining is fraught with challenges. The sheer scale of the data, its often-proprietary nature, and the non-stationary training dynamics make full attribution analysis infeasible.

However, TDA is highly actionable and relevant for the fine-tuning stages of LLMs (e.g., supervised fine-tuning (SFT) or instruction tuning). In this regime, the datasets are curated, accessible, and orders of magnitude smaller. Here, TDA can help us understand:

Which instruction-tuning examples are most responsible for a specific alignment behavior?
Which documents in a domain-adaptation dataset are causing catastrophic forgetting of general abilities?
Why does the model fail on a specific query after being fine-tuned on our data?

The findings of Razeghi et al. (2022) provide the perfect motivation. We know that corpus-level statistics shape a model's core capabilities. Fine-tuning is our chance to precisely perturb those capabilities. TDA is the microscope that lets us see which examples in our fine-tuning data are driving those changes, for better or for worse.

Synthesis & Guidance

With four families of methods, how should a practitioner choose?

For quick, scalable debugging of specific predictions on deep models: Start with TracIn. It requires saving checkpoints but is empirically robust and sidesteps the theoretical pitfalls of Influence Functions in non-convex settings.
For finding mislabeled data or when a local approximation is sufficient: Use Influence Functions, but be prepared to validate the results. The HVP-based implementations scale well.
If your model has a linear readout and is trained with L2 regularization:Representer Points offer uniquely interpretable, signed attributions that can distinguish helpful from harmful examples.
For dataset-level valuation, pruning, or making policy decisions about data subsets: Use Data-Shapley approximations. This is the most principled but also the most expensive option.

A robust workflow often involves triangulation. If TracIn and Influence Functions both flag the same training example as highly influential for a given error, your confidence in that attribution should be much higher. Before taking a drastic action like removing data, seek consensus across multiple methods. When reporting results, be transparent about your assumptions: which checkpoints were used for TracIn, what damping was used for Influence Functions, or which utility function was chosen for Shapley.

Limits & Open Questions

Training Data Attribution is an active and evolving field of research. Major open questions remain:

Scalability: How can we develop TDA methods that are truly feasible for foundation model pretraining corpora?
Memorization vs. Contribution: What is the precise relationship between an example being influential and being memorized? An example can teach a generalizable concept without being memorized, and a memorized example might have little influence on other predictions.
Connection to Unlearning: Can we leverage TDA to perform more efficient and effective "machine unlearning," where the goal is to surgically remove a model's knowledge of specific data points for privacy or safety reasons?
Complex Objectives: How do we perform attribution in more complex training pipelines like Reinforcement Learning from Human Feedback (RLHF), where the "loss function" is itself a learned reward model?

TDA is more than an academic exercise; it is a fundamental tool for building more reliable, transparent, and debuggable machine learning systems. By connecting model behavior back to the data it learned from, we take a critical step from opaque prediction engines to understandable and trustworthy tools.

References

Ghorbani, A., & Zou, J. (2019). "Data-Shapley: Equitable Valuation of Data for Machine Learning". Proceedings of the 36th International Conference on Machine Learning (ICML).
Jia, R., et al. (2019). "Towards Efficient Data Valuation Based on the Shapley Value". Proceedings of the 22nd International Conference on Artificial Intelligence and Statistics (AISTATS).
Koh, P. W., & Liang, P. (2017). "Understanding Black-box Predictions via Influence Functions". Proceedings of the 34th International Conference on Machine Learning (ICML).
Pruthi, G., et al. (2020). "TracIn: An Influence-based Approach for Tracing Training Data". Advances in Neural Information Processing Systems (NeurIPS).
Razeghi, Y., et al. (2022). "Impact of Pretraining Term Frequencies on Few-Shot Numerical Reasoning". Findings of the Association for Computational Linguistics: EMNLP 2022.
Tenney, I. (2023). "A Hitchhiker's Guide to Training Data Attribution".
Yeh, C.-K., et al. (2018). "Representer Point Selection for Explaining Deep Neural Networks". Advances in Neural Information Processing Systems (NeurIPS).

Beyond the Leaderboard: New Frameworks for Understanding LLM Behavior

Fri, 17 Oct 2025 00:00:00 GMT

Note: This RSS feed strips out SVGs and embeds. You might want to read the post on the webpage here.

October 16, 2025

Last Updated: October 17, 2025

mechInterp #CMSC848R

💡

Notes from the McCoy, R. T., et al. (2024). "Embers of autoregression " [Link] and Holtzman, A., et al. (2025) and "Generative Models as a Complex Systems Science”. [Link]

TL;DR

Problem: Standard benchmarks tell us what an LLM can do, but not why it succeeds or fails. This limits our ability to reliably improve models and anticipate failures.
Framework 1 (Teleological Diagnostics): McCoy et al. (PNAS 2024) propose a diagnostic framework. Since LLMs are trained to maximize the log-likelihood of text, their accuracy on a task empirically co-varies with: (i) the probability of the task itself, (ii) the probability of the correct output string, and (iii) the probability of the input string.
Framework 2 (Complex Systems Science): Holtzman et al. (JSC 2025) argue for treating LLMs as complex systems with emergent behaviors (abilities not explicitly programmed). This reframes benchmarks as experimental probes for discovering and mapping these hidden capabilities.
Synthesis: These frameworks are complementary. The Complex Systems approach is for discovery (mapping what a model can do in general), while the Teleological approach is for diagnostics (explaining why a specific task failed).

1. Introduction: The Limits of Single-Score Evaluations

Current LLM evaluation is dominated by performance on benchmarks like MMLU or HumanEval. While these leaderboards provide a useful snapshot of capabilities, they collapse a model's complex behavior into a single score. This post explores two recent papers that offer frameworks for a deeper analysis, moving from "what score did it get?" to "why does it behave this way?".

2. Framework 1: A Diagnostic Toolkit for LLM Failures (McCoy et al., 2024)

2.1 Core Concept: Aligning Tasks with the Training Objective

LLMs are fundamentally sequence models trained to optimize the next-token log-likelihood, log pθ(output|input). McCoy et al. (PNAS 2024) propose that a model's performance on a downstream task is strongly influenced by how well that task aligns with this core objective. They use the analogy of human back pain: our spines are optimized for quadrupedalism, and the mismatch with bipedalism causes issues. Similarly, an LLM fails when a task creates a mismatch with its core training of predicting common text sequences.

2.2 Three Diagnostic Factors

The paper identifies three key empirical factors that correlate with model accuracy:

Task Probability P(task): The frequency of the task in the training data.
Output Probability P(output): The prior probability of the correct output string.
Input Probability P(input): The probability of the input string.

These are not used by the model internally but are powerful external diagnostics for predicting success or failure.

2.3 A Toy Example: The Shift Cipher

This directly relates to the "shift cipher" analyses in McCoy et al. They find models are far better at ROT13 than other shifts (e.g., ROT12). This is not because the model "understands" modular arithmetic better for the number 13, but because ROT13 is a common topic in web text (high P(task)), while ROT12 is virtually nonexistent (low P(task)).

2.4 Schematic: Accuracy Correlates with Output Probability

The original paper demonstrates a strong relationship between P(output) and model accuracy. We can illustrate this conceptually.

This schematic shows that model accuracy on a task tends to increase as the log-probability of the correct output string increases. For detailed methodology, see the logistic regression analyses in the original paper.

2.5 Code Demo: Measuring `P(output)` with GPT-2

This snippet demonstrates how a sentence with a common proper noun ("Instagram") is more "probable" to a model than one with a rare, fictional name ("Annathon").

import torchfrom transformers import GPT2LMHeadModel, GPT2Tokenizer# --- Setup & Reproducibility ---torch.manual_seed(42)tokenizer = GPT2Tokenizer.from_pretrained('gpt2')model = GPT2LMHeadModel.from_pretrained('gpt2')model.eval() # Set model to evaluation mode# --- Prepare Batched Inputs ---sentences = [    "First, she just posted to her Instagram story.", # High P(output)    "Sorry, Annathon writes to our Copyright Users." # Low P(output)]inputs = tokenizer(sentences, padding=True, return_tensors='pt')input_ids = inputs.input_idsattention_mask = inputs.attention_mask# --- Calculate Loss (Negative Log Likelihood) ---with torch.no_grad():    outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)    # The model's loss is the cross-entropy loss, which is the negative log-likelihood (NLL)    # averaged over the tokens.    nll = outputs.loss# We need to compute per-sequence loss as the batch-level loss is an average.per_sequence_nll = []for i in range(input_ids.shape[0]):    seq_ids = input_ids[i, :]    seq_labels = input_ids[i, :]    with torch.no_grad():        loss = model(input_ids=seq_ids.unsqueeze(0), labels=seq_labels.unsqueeze(0)).loss        per_sequence_nll.append(loss.item())# --- Print Results with Precise Metrics ---seq_lengths = [torch.sum(mask).item() for mask in attention_mask]print(f"Sentence 1: \\"{sentences[0]}\\"")print(f"  Sequence Length: {seq_lengths[0]}")print(f"  Avg Token NLL: {per_sequence_nll[0]:.4f}\\n")print(f"Sentence 2: \\"{sentences[1]}\\"")print(f"  Sequence Length: {seq_lengths[1]}")print(f"  Avg Token NLL: {per_sequence_nll[1]:.4f}\\n")print("Note on Tokenization: The lower probability for 'Annathon' is also affected by its")print("tokenization into subwords ('Ann', 'ath', 'on'), which can be less probable")print("as a sequence than the single token for 'Instagram'.")

3. Framework 2: LLMs as a Complex Systems Science (Holtzman et al., 2025)

3.1 Key Concept: Emergence in LLMs

Holtzman et al. (JSC 2025) argue we should treat LLMs as complex systems. In such systems, simple, local interactions between many components (neurons) give rise to complex, global behaviors (emergence) that are not explicitly designed. For LLMs, abilities like chain-of-thought reasoning are emergent.

3.2 A New Role for Benchmarks: Experimental Probes

This framework reframes benchmarks not as final exams, but as building blocks for experiments. The goal shifts from measuring a score to probing for the presence of a specific emergent behavior across diverse contexts.

3.3 An Example Probe: Testing for Deception

While Holtzman et al. provide the high-level framework, a concrete example of this style of analysis can be seen in other research. For instance, work on steganography (hiding messages) asks if a model can perform a task while also hiding information within its answer.

Attribution Note: The specific "steganography + GLUE" test is a hypothetical experiment inspired by the complex systems framework and real experiments from labs like Redwood Research (e.g., "Preventing Language Models from Hiding Their Reasoning," 2024).

3.4 Code Demo: Structuring a Behavioral Probe

This code doesn't run a real probe but outlines the experimental structure for testing an emergent skill like "topic adherence" across different benchmark tasks.

def run_benchmark_task(model, task_prompt):    """Placeholder for running a single instance of a benchmark task."""    # In reality, this would call an API or run a model forward pass    print(f"  Running task: '{task_prompt[:30]}...'")    return "This is the model's answer."def check_emergent_behavior(model_output, behavior_check_prompt):    """Placeholder for checking a behavior, e.g., using another LLM as a judge."""    # Example: Check if the output avoided a forbidden topic    print(f"  Checking behavior: '{behavior_check_prompt}'")    return "forbidden_topic" not in model_output.lower()# --- Experimental Design ---def run_probe_on_benchmark(model, benchmark_tasks, probe_behavior):    """    Tests for an emergent behavior across a suite of benchmark tasks.    Args:        model: The LLM to test.        benchmark_tasks: A list of prompts from a benchmark (e.g., MMLU, GLUE).        probe_behavior: A tuple containing the behavior instruction and the check.    """    print(f"Starting probe for behavior: '{probe_behavior['instruction']}'")    successes = 0    for task in benchmark_tasks:        # Combine the task with the behavioral instruction        full_prompt = f"{probe_behavior['instruction']}\n\nNow, please complete the following task:\n{task}"        output = run_benchmark_task(model, full_prompt)        if check_emergent_behavior(output, probe_behavior['check']):            successes += 1    print(f"\nResult: Behavior was present in {successes}/{len(benchmark_tasks)} tasks.")    return successes / len(benchmark_tasks)# --- Define the Experiment ---mock_mmlu_tasks = ["What is the capital of France?", "Solve for x: 2x + 3 = 7"]topic_adherence_probe = {    "instruction": "IMPORTANT: In your response, do not mention the word 'politics'.",    "check": "Does the response contain the word 'politics'?"}# Run itrun_probe_on_benchmark("mock_model", mock_mmlu_tasks, topic_adherence_probe)

4. Synthesis: The Analyst's Toolkit

These two frameworks are not in conflict; they are complementary tools for a deeper analysis of LLMs.

Feature	Teleological Approach (McCoy et al.)	Complex Systems Approach (Holtzman et al.)
Primary Goal	Diagnose specific failures.	Discover general, emergent capabilities.
Core Question	"Why did the model fail this time?"	"What hidden skills does this model possess?"
Use of Benchmarks	As a source of failure cases to analyze.	As building blocks for new experiments.
Direction	Top-down (training objective -> behavior)	Bottom-up (simple interactions -> emergence)

An effective LLM analyst uses both. They use the complex systems approach to map the landscape of what a model can do, and when they find a cliff where the model fails, they use the teleological approach to understand the geology of why that cliff exists.

5. Conclusion & Final Takeaways

Moving beyond leaderboards doesn't mean abandoning measurement, but enriching it. By combining diagnostic and exploratory frameworks, we can build a more robust, reliable, and ultimately more useful science of large language models.

Takeaway 1: LLM failures are often not random but are systematic consequences of their training objective (the "back pain" analogy).
Takeaway 2: The probabilities P(output) and P(task) are powerful predictors of model performance on a given task.
Takeaway 3: LLMs should be viewed as complex systems that exhibit emergent behaviors not explicitly coded into them.
Takeaway 4: Benchmarks can be repurposed from "final exams" into "experimental probes" to test for these emergent behaviors.
Takeaway 5: A mature approach to LLM analysis requires both a diagnostic toolkit (for explaining failures) and an exploratory one (for discovering capabilities).

References

McCoy, R. T., et al. (2024). "Embers of autoregression show how large language models are shaped by the problem they are trained to solve." Proceedings of the National Academy of Sciences. [Link]
Holtzman, A., et al. (2025). "Generative Models as a Complex Systems Science: How can we make sense of large language model behavior?" Journal of Social Computing. [Link]
Millière, R., et al. (2024). "Preventing Language Models from Hiding Their Reasoning." Redwood Research. [Link]

Inner Workings of a Transformer: Interpretability Techniques

Fri, 17 Oct 2025 00:00:00 GMT

Note: This RSS feed strips out SVGs and embeds. You might want to read the post on the webpage here.

Notes on the “Primer on teh Inner workings of the Transformer based LMs”

October 16, 2025

Last Updated: October 17, 2025

mechInterp #CMSC848R

💡

This is my deep dive into the sections 3 and 4 of Ferrando et al.'s "A Primer on the Inner Workings of Transformer-based Language Models," which connects the modern interpretability techniques to the model's fundamental additive architecture.

This post is a detailed walkthrough of two key families of interpretability techniques-Behavior Localization and Information Decoding-as presented in Sections 3 and 4 of the paper. We explore how to pinpoint which parts of a Transformer are responsible for a given behavior and how to understand the information those parts represent.

The Key is the Architecture: The Transformer's core design-summing component outputs into a residual stream and then linearly projecting that sum to get logits-is not just an implementation detail. It's the mathematical foundation that makes principled attribution and intervention possible.
Localization (Sec. 3): Techniques like Direct Logit Attribution (DLA) and activation patching directly exploit the model's additive structure. DLA accounts for how much each component pushed the final answer, while patching causally tests a component's importance by swapping its output.
Decoding (Sec. 4): Techniques like probing, linear feature analysis, and Sparse Autoencoders (SAEs) investigate what information is encoded in the model's activations. These methods help us build hypotheses about what concepts the model has learned, which can then be causally verified with localization tools.
A Unified Workflow: We will frame these two sections as a coherent workflow: start with the model's linear structure, use localization to find responsible components, use decoding to hypothesize what they represent, and iterate to build a grounded understanding of the model's internal algorithm.

Bridge from Section 2: Why This All Works

Before diving into specific techniques, we must ground our analysis in the mathematical structure of the Transformer, as laid out in Section 2 of the paper. At its heart, a Transformer is a sequence of layers, each containing attention and MLP components. Crucially, each of these components, which we can index with c, computes an update vector h_c and adds it to the residual stream.

Let r^(l) be the state of the residual stream at a specific token position after layer l. The final residual stream state, r_L, is the sum of the initial token/positional embedding and all subsequent component updates:
r_L = r^(0) + h_{attn,0} + h_{mlp,0} + ... + h_{attn,L-1} + h_{mlp,L-1}

For simplicity, we can express this as a sum over all components c in the computational graph:
r_L = ∑_c h_c

The model then produces logits, z, by applying a final linear transformation, the unembedding matrix W_U. The logit for a specific token y, denoted z_y, is the dot product of the final residual stream with the corresponding token's unembedding vector, u_y:
z_y = u_y^T r_L = u_y^T (∑_c h_c)

By the linearity of the dot product, this expands to:
z_y = ∑_c u_y^T h_c

This simple equation is the bedrock of modern mechanistic interpretability. It tells us that the final logit for any token is just a linear sum of contributions from every component in the model. This structure licenses two powerful modes of analysis :

Additive Accounting: We can directly measure the contribution of each component c to the logit z_y by calculating the term u_y^T h_c. This is the core idea behind Direct Logit Attribution.
Targeted Causal Intervention: We can test the causal importance of a component c by manipulating its output vector h_c. If we replace h_c with a different vector h_c~ (from a counterfactual input), the change in the final logit z_y will be precisely u_y^T (h_c~ - h_c), assuming no downstream components adapt to the change. This is the logic behind activation patching.

With this architectural foundation in mind, we can now explore the techniques in Sections 3 and 4 as principled methods for exploiting this structure.

Section 3: Behavior Localization - Finding Which Parts Matter

Behavior localization aims to identify which parts of the model-from individual input tokens to specific attention heads or MLP neurons-are responsible for a given output. These methods are our primary tools for moving from "the model predicted token Y" to "the model predicted token Y because head 7.3 attended to token X and FFN 12 modified the representation in this way."

3.1. Input Attribution: Tracing Behavior Back to Tokens

The most basic localization question is: which input tokens were most influential for a given prediction?

Gradient-based Attribution:
A standard approach in machine learning is to use gradients. If we have a function f(x) that maps an input x to a scalar output of interest (like the logit z_y), its first-order Taylor expansion f(x + Δ) ≈ f(x) + ∇_x f(x)^T Δ tells us that the gradient ∇_x f(x) represents the sensitivity of the output to infinitesimal perturbations of the input. We can aggregate these sensitivities per input token embedding to get a saliency score.

A common variant is gradient·input, which computes ∇_x f(x) ⊙ x, where ⊙ is the element-wise product. This approximates the contribution of each input feature to the final output score. However, raw gradients can be noisy and suffer from saturation. Techniques like SmoothGrad (averaging gradients over noisy inputs) and Integrated Gradients (integrating gradients along a path from a baseline input to the actual input) provide more robust estimates by mitigating local noise and accounting for non-linearities along the integration path .

Intuition: Gradient attribution asks, "If I could wiggle each input token embedding a tiny bit, which wiggles would have the biggest impact on the final logit?" It's a measure of local sensitivity.

Attention Weights vs. Context Mixing:
It is tempting to use attention weights as a direct measure of input importance. However, this is often misleading. An attention head could place a high weight on a token, but if the value vector produced at that token position is near-zero or orthogonal to what the rest of the model cares about, its actual influence is negligible. The weights only tell you where the information is being routed from, not what that information is or how it's used .

More faithful methods account for the full information flow through the attention head's OV-circuit (the part that computes Attention(Q,K)V). This can involve weighting attention scores by the norm of the value vectors or using more complex rollout algorithms that recursively trace influence through multiple layers. However, simple rollout techniques can misattribute importance by assuming linearity where it doesn't hold .

Contrastive Attribution:
Models don't just predict a token y in a vacuum; they predict it instead of other tokens. The softmax function amplifies differences between logits. Therefore, a more meaningful attribution target is often the logit difference, z_y - z_y~, where y~ is a plausible alternative token. By attributing this difference, we ask, "Which inputs caused the model to favor y over y~?" This aligns our analysis with the competitive nature of the model's output layer.

Caveat: Input attribution methods, especially gradient-based ones, provide a correlational signal, not a strictly causal one. They can be insensitive to causal model components, disagree with each other, and fail when perturbations push the model into out-of-distribution territory where its behavior is unpredictable. They are best used as a first-pass tool for hypothesis generation .

3.2. Model Component Importance: Pinpointing Heads and Neurons

While input attribution is useful, the core of mechanistic interpretability lies in understanding the internal components. Here, we leverage the Section 2 decomposition to its fullest.

Direct Logit Attribution (DLA):
As derived earlier, the total logit z_y is a sum of contributions from all components: z_y = ∑_c u_y^T h_c. The term u_y^T h_c is the Direct Logit Attribution (DLA) of component c to token y. It is the projection of the component's output vector h_c onto the unembedding direction u_y for that token. A large positive DLA means the component directly "pushed up" the logit for y.

Just as with input attribution, a contrastive analysis is often more powerful. The Contrastive Direct Logit Attribution for y vs. y~ is:
DLDA_c(y, y~) = u_y^T h_c - u_y~^T h_c = (u_y - u_y~)^T h_c

This measures how much component c pushed the logits in favor of y over y~. Because this method is a direct calculation based on the model's forward pass, it is computationally cheap and provides a complete, additive decomposition of the logits .

Intuition: DLA is the ultimate accounting tool. It opens up the final sum ∑_c h_c and shows you exactly how much each h_c term contributed to the final answer along the u_y direction.

Causal Interventions via Activation Patching:
DLA tells us about correlations in a single forward pass. To establish causality, we need to perform interventions. Activation patching is the canonical technique for this. The methodology is as follows:

Clean Run: Run the model on a "clean" input where it exhibits the behavior of interest (e.g., correctly translating a sentence). Store the internal activations at various points (e.g., the output of a specific attention head).
Corrupted Run: Run the model on a "corrupted" or counterfactual input where it does not exhibit the behavior (e.g., an input that would lead to a mistranslation).
Patching Run: Rerun the model on the corrupted input, but at a specific component c, intervene and replace its activation a_c with the activation stored from the clean run.
Measure the Effect: Observe the change Δ in the final output metric (e.g., the logit difference between the correct and incorrect translation). If Δ is large, meaning the behavior is restored, then component c is causally implicated in the behavior.

The core intervention is a <- a~, where a is the activation on the corrupted run and a~ is from the clean run. A large Δ on a metric S means the component is causally important for that metric.

Caveat: The choice of corrupted/counterfactual input is critical and can significantly affect the results. An overly simplistic ablation (e.g., zero-ablating a component) might push downstream components into an out-of-distribution regime, causing them to behave erratically. This concept, termed ecological validity, emphasizes the importance of using counterfactuals that keep the model in-distribution . Moreover, models can exhibit self-repair, where downstream components compensate for a patched or ablated component, masking its true importance.

Subspace Activation Patching and Distributed Interchange Intervention (DII):
Often, a feature or concept is not represented by a single neuron but by a direction or subspace within the high-dimensional activation space. Instead of patching an entire activation vector a, we can perform a more surgical intervention.

Given a subspace S (e.g., spanned by a single direction vector v), we can define a projector P_S onto that subspace. A subspace patch replaces only the part of the activation that lies in S:
a_patched = P_S a~ + (I - P_S) a

Here, we take the feature from the clean run (a~) and inject it into the corrupted run (a) while preserving everything in the orthogonal complement of S. This approach, also known as Distributed Interchange Intervention (DII), allows us to test the causal role of specific, distributed features rather than entire activation vectors . This is a powerful refinement of patching that aligns with the hypothesis that features are linear.

Finding Circuits: Edge and Path Patching:
Components don't just act independently; they form circuits. An early attention head might extract information that a later MLP processes. Edge patching extends activation patching to test the importance of the connection between two components. For example, to test the A -> B connection, one would patch the output of A only as it is read by B, leaving other components that read from A unaffected.

Path patching generalizes this to entire chains of components. However, the number of possible paths is combinatorially explosive. This has led to automated search algorithms like ACDC (Automatic Circuit Dissection and Correspondence) that greedily prune edges from the computational graph. Other methods like Edge Attribution Patching (EAP) use gradient-based approximations to estimate the effect of patching all possible edges more efficiently, though these linearized estimates can fail in the presence of strong non-linearities .

Section 4: Information Decoding - Discovering What is Represented

While Section 3 focuses on where a computation happens, Section 4 focuses on what information is being computed and represented. These techniques aim to put meaningful labels on the vectors and subspaces within the model.

4.1. Probing

A probe is a simple, supervised model (often linear) trained to predict a property of interest (e.g., part-of-speech, syntactic depth, presence of a specific concept) from a Transformer's internal activations. If a probe can be trained to high accuracy on the activations r^(l) at layer l, it suggests that information about that property is linearly decodable from the residual stream at that layer.

Caveat: This is a crucial distinction: a successful probe demonstrates that information is encoded, not necessarily that it is used by the model to solve its primary task. A powerful probe might simply be learning a complex function that the main model has no need for. To mitigate this, researchers use controls (e.g., training a probe on randomized labels to see what accuracy it can achieve by chance) and apply principles like Minimum Description Length (MDL) to penalize probe complexity, ensuring the discovered correlation is not spurious . A probe gives you a hypothesis; Section 3 tools are needed to test it.

4.2. The Linear Representation Hypothesis and Sparse Autoencoders (SAEs)

A central idea in modern interpretability is the linear representation hypothesis: many concepts and features the model learns are represented as directions or subspaces in activation space. This hypothesis is powerful because it connects directly to the linear readout mechanism from Section 2. If a concept C is represented by a direction v_C, then adding α * v_C to the residual stream should predictably shift the model's behavior towards expressing concept C.

This enables two key interventions:

Erasure: To test if a feature is necessary for a behavior (e.g., stereotypical gender bias), we can identify its direction v and project it out of the residual stream: r_new = r - P_v r. If the behavior disappears, the feature was causally necessary.
Steering: We can guide model generation at inference time by adding a feature vector to the residual stream, r_new = r + α * v, effectively "steering" the model towards a desired attribute or style.

The surprising effectiveness of these simple linear edits provides strong evidence for the linear representation hypothesis .

Unmixing Features with Sparse Autoencoders (SAEs):
A major challenge is that neurons are often polysemantic, meaning a single neuron activates for multiple, unrelated concepts. Furthermore, models may use superposition to represent more features than they have dimensions by storing them in a non-orthogonal basis.

Sparse Autoencoders (SAEs) are an unsupervised technique designed to tackle this. An SAE is a simple neural network trained to reconstruct a model's internal activations (e.g., r^(l)) from a much wider, "overcomplete" set of latent features. The key is a strong sparsity penalty (L1) on the latent activations. The objective is:
min_W,θ E[ ||r - W^T s||_2^2 + λ||s||_1 ]
where s = σ_θ(r) is the sparse code produced by the encoder, W is the decoder dictionary, and λ is the sparsity coefficient.

The columns of W represent learned "features." Because the code s is forced to be sparse, each activation r is reconstructed using only a few active features from this dictionary. These SAE features are often much more monosemantic (interpretable) than the original neuron basis. Once an SAE is trained, we can analyze its features using the full suite of Section 3 tools: calculate their DLA, patch them, and search for the inputs that maximally activate them .

Intuition: An SAE acts like a prism for activation vectors. It takes a dense, tangled vector where many features are superimposed and splits it into a sparse set of clean, interpretable feature directions.

Caveat: SAEs are not perfect. There is a trade-off between reconstruction fidelity and sparsity. Furthermore, the reconstruction error r - W^T s itself can contain important information, and causal analyses must account for it. Newer architectures like Gated SAEs improve this trade-off by separating the detection of a feature from its magnitude, leading to better fidelity for a given level of sparsity .

4.3. Decoding via the Vocabulary Interface

The unembedding matrix W_U is the bridge between the model's internal representation space and the human-readable vocabulary space. We can leverage this bridge to decode the content of internal states.

Logit Lens: The simplest application is the logit lens, where we apply the unembedding matrix directly to an intermediate residual stream r^(l) to get a premature prediction: z^(l) = W_U^T r^(l). This can reveal how the model's prediction evolves layer by layer. However, raw logit lenses often produce poor results because the distribution of r^(l) can be very different from that of the final residual stream r_L. Tuned lenses learn a simple affine transformation to correct for this distributional shift .
Patchscopes: This technique combines patching (Section 3) with a readout mechanism. The core idea is to patch an activation from a source context into a target context and then immediately decode it using a readout function (like the logit lens) to see what information it contains in that new context. This provides a more controlled way to understand what an activation vector "means" .
Logit Spectroscopy and SVD Views: We can analyze the linear operators within the model (like FFN weight matrices) by understanding their effect on the vocabulary. By taking the Singular Value Decomposition (SVD) of a weight matrix, we can find the input directions it amplifies the most. Projecting these directions into vocabulary space via W_U reveals which tokens or concepts the matrix is most sensitive to. This "logit spectroscopy" helps characterize the function of entire MLP layers or attention heads .
Maximally Activating Inputs: A straightforward way to understand a component (like a neuron or an SAE feature) is to search through a large dataset for the input examples that cause it to activate most strongly. This provides concrete examples of what the feature responds to. A riskier alternative is to use optimization to synthesize an input that maximizes the activation, but this can produce unnatural, "frankenstein" inputs that exploit adversarial patterns, leading to interpretability illusions .

Opinion: The ability of modern language models to generate their own natural language explanations for their features is tantalizing. However, as the primer correctly warns, these explanations should be treated as hypotheses, not ground truth. A model might generate a plausible-sounding but completely unfaithful explanation. Such claims must be rigorously tested using the causal intervention methods from Section 3.

Synthesis: A Coherent Workflow for Understanding Transformers

Sections 3 and 4 of the primer are not just a catalog of methods; they describe a unified, iterative process for reverse-engineering the algorithms learned by Transformers. The workflow, grounded in the Section 2 decomposition, looks like this:

Start with the Additive Structure (Sec. 2): The z_y = ∑_c u_y^T h_c equation is your entry point. It guarantees that the problem of understanding a logit can be broken down into understanding the contributions of individual components.
Localize with DLA and Patching (Sec. 3): For a specific behavior, use DLA to get a cheap, comprehensive ranking of which components contributed most to the outcome. Formulate a hypothesis (e.g., "Head 8.4 is responsible for copying the name"). Then, use activation patching to causally test this hypothesis. Does patching head 8.4's output from a clean run restore the correct name? If so, your hypothesis is supported.
Decode with Probes and SAEs (Sec. 4): Now that you've identified a causally important component, ask what information it represents. Is there a specific concept it seems to encode? Train a probe to see if that concept is linearly decodable. If the component's activations seem polysemantic, train an SAE on them to disentangle the underlying features.
Refine and Verify (Iterate): The features discovered by your SAE in step 3 are now new hypotheses. Re-enter the Section 3 workflow: calculate the DLA of these SAE features. Perform subspace patching on a single SAE feature direction to confirm its causal role. Use maximally activating examples to name the feature.

This loop-from localization to decoding and back to causal verification-allows us to move from coarse-grained observations about model behavior to fine-grained, validated claims about the specific algorithms the model has learned.

Limits and Open Questions

Despite this progress, significant challenges remain. The primer highlights several critical limitations that practitioners must keep in mind:

The Attention-as-Explanation Fallacy: Raw attention weights are not explanations. Always consider the full OV-circuit.
Distribution Shift in Patching: Causal interventions can be misleading if they push the model into out-of-distribution states. Careful design of counterfactuals is paramount.
Probe Faithfulness: Probes can discover spurious correlations. Their findings must be validated with causal methods.
SAE Quality: The utility of SAEs depends on finding the right balance of sparsity and reconstruction fidelity, and their learned features are not guaranteed to be perfectly monosemantic.
The Limits of Linearity: While the linear representation hypothesis has been incredibly fruitful, models are fundamentally non-linear. Linearized methods like EAP can fail, and complex reasoning may involve non-linear interactions between features that are difficult to analyze.

Part 4 Reverse Engineering Transformers: Induction Heads

Fri, 17 Oct 2025 00:00:00 GMT

Note: This RSS feed strips out SVGs and embeds. You might want to read the post on the webpage here.

October 16, 2025

Last Updated: October 17, 2025

mechInterp #CMSC848R

💡

Final part in the series:
Part 1: Reverse Engineering Transformers: Deconstructing Attention

Part 2: Reverse Engineering Transformers: Attention Heads & Circuits
Part 3 Reverse Engineering Transformers: Path Expansion

Transformer Circuits VI: Induction Heads

TL;DR

Induction heads are a two-head circuit that implements a simple but powerful form of in-context learning: repeating sequences.
The circuit's algorithm is: "find a previous occurrence of the current token, attend to the token after it, and copy that token's identity to the current position."
This is implemented via K-composition with a previous token head from an earlier layer.
The previous token head shifts information, making a token's Key vector at position t encode information about token[t-1].
The induction head's Query vector at position t+k can then match with the Key at position t, allowing it to attend to t and use its OV-circuit to copy token[t].
The discovery of induction heads shows that powerful, general capabilities like in-context learning can emerge from simple, mechanically understandable circuits.

Context & Motivation

One of the most remarkable abilities of large language models is in-context learning: the ability to perform a task from a few examples in a prompt, without any weight updates. If you show a model A -> B, C -> D, E ->, it will likely predict F. How does it do this?

The authors of the "Mathematical Framework" paper discovered that even tiny, two-layer models develop a surprisingly general version of this skill. When trained on random text, they learn to complete repeating sequences (e.g., ...[A][B]...[A] -> [B]). The circuit responsible for this is the induction head.

Understanding this circuit is a major milestone. It's not just a description of a pattern; it's a full, end-to-end reverse engineering of a key algorithmic capability. It demonstrates how the abstract building blocks we've discussed assemble into something greater than the sum of its parts.

Prereqs

A complete understanding of Parts I-V, especially head composition and the roles of W_Q, W_K, and W_V.

Core Idea in One Picture

The induction circuit consists of (at least) two heads. A "previous token head" in Layer 0 sets up the information, and the "induction head" in Layer 1 executes the core logic.

graph TD
    subgraph Input Sequence
        direction LR
        A1("... token[t-1]") --> B1("token[t]") --> C1("...") --> A2("token[t+k-1]") --> B2("token[t+k]") --> Next("???")
    end

    subgraph Layer 0: Prev Token Head (H0)
        B1 -- Attends to --> A1
        A1 -- Copies info --> B1_Stream[Stream at B1]
    end

    subgraph Layer 1: Induction Head (H1)
        B2 -- Query: "I am B" --> B1_Stream
        B1_Stream -- Key: "Prev token was A" --> B2
        B1_Stream -- Value: "I am B" --> B2_Stream[Stream at B2]

        B2 -- Attends to --> B1
        B1 --OV Circuit --> B2_Stream
    end

    B2_Stream -- Projects via W_U --> Logit_for_C["'C' logit is high"]

    style H0 fill:#ccf
    style H1 fill:#fcc
    style B1_Stream stroke:#00f,stroke-width:2px,stroke-dasharray: 2 2

Figure 1: The Induction Head circuit. At a token B (at t+k), the model wants to predict the next token. If another B appeared at t, the circuit finds it. It does this because a Layer 0 head (H0) copied information about A (token[t-1]) into the residual stream at B's position (t). H1 at t+k can now attend to B at t (by matching its Query for 'B' with 'B's content) and use its Value to copy the next token's identity, thereby predicting 'C'.

Walkthrough: The Algorithm Step-by-Step

Let's trace the full algorithm for the sequence ...[A] [B] [C]... [A] [B] -> ?, where the model should predict [C]. We are at position t+k, which contains the token [B].

Step 1: The Setup (Previous Token Head, `H0`)

Head Type: A "previous token head" in an early layer (let's say Layer 0).
QK-Circuit: This head has a very simple QK circuit. Its weights are configured such that the query at position t will always give the highest attention score to the key at position t-1.
OV-Circuit: This head's OV circuit is a "copying" circuit. It takes the embedding of the previous token (token[t-1]) and adds it to the residual stream at position t.
Result: After Layer 0, the residual stream at every position t, x_t, contains a representation of both token[t] (from the original embedding) and token[t-1] (added by H0).
- At position t, where the token is [B], the stream x_t now effectively says: "I am [B], and the token before me was [A]."

Step 2: The Core Logic (Induction Head, `H1`)

Now we move to a later layer (Layer 1). The induction head H1 executes the main algorithm.

Where does H1 attend?H1 needs to attend from the current position t+k (token [B]) back to the previous occurrence of [B] at position t.
The QK-Circuit (K-Composition in action):
- Query: The query vector q_{t+k} for H1 at the current position is formed from the stream, so it says "I am [B]."
- Key: This is the clever part. The key vector k_t at the earlier position t is formed from its stream x_t. And x_t contains information about token[t-1], which is [A]. So, k_t says something like "The token after me is [B]" or "My predecessor was [A]".
- The induction head's W_Q and W_K matrices are learned to match these specific signals. The query q_{t+k} looks for a key k_t that corresponds to a previous occurrence of the same token. The K-composition (the effect of H0 on H1's key) makes this possible. The Query "I am [B]" matches the Key "The token after me is [B]".
What does H1 move? The OV-Circuit:
- H1's attention is now focused on position t.
- Its W_V matrix is trained to read the information from the residual stream at t, x_t. Crucially, x_t still contains the original embedding for token[t], which is [B].
- However, the composed W_V^1 W_O^1 circuit of the induction head performs a specific transformation: it is trained to take the an embedding from the stream and effectively "shift it forward by one". It acts to copy the next token.
- Wait, how can it copy the next token, [C], by looking at [B]? The paper hypothesizes that the model leverages existing correlations. Since [B] is often followed by [C] in the training data, the model can learn a W_V W_O transformation that maps the embedding of [B] to a vector that boosts the logit for [C].

Self-correction/Refined View: A simpler and more common view of the induction head's OV-circuit is that it attends to the position t-1 (the position of A) and copies the value from there. The logic is:

H1 Query at t+k (B): Look for a previous B.
H1 Key at t+k-1 (A). Via K-Composition with H0, this Key says "The token after me is B".
H1 attends from t+k to t+k-1.
H1 Value reads from t+k-1's stream, which contains the embedding for A.
The OV Circuit of H1 moves this information to predict B.
(This seems to be a slight variation on the original paper's description, but is a common interpretation).

Let's stick to the paper's primary finding: the QK circuit attends to the previous instance of the same token, and the OV circuit copies the next token. The key is that W_Q and W_K are trained to "look for [token]" and H0's output lets them look for "[token] in the next position".

Implications & Limits

Implication: This is a concrete mechanism for in-context learning. Many of the impressive few-shot learning capabilities of LLMs may be powered by scaled-up, more abstract versions of this basic circuit.
Implication: It shows that general, algorithmic capabilities can emerge spontaneously from a simple predictive loss, and that these emergent algorithms can be reverse-engineered.
Limit: This is one specific circuit. It doesn't explain all forms of in-context learning (e.g., following complex instructions). However, it provides a foundational example.
Limit: The explanation relies on the linear approximation provided by path expansion; the real dynamics with LayerNorm are noisier and more complex, but the core signal is captured by this model.

Takeaways

Induction heads are a two-layer, two-head circuit that allows a model to continue repeating sequences.
They are a mechanistic explanation for a simple form of in-context learning.
The circuit combines a "previous token head" (to provide history) and an "induction head" (to use that history).
K-Composition is the critical mechanism that allows the induction head's query to find a match based on the previous token's identity.
The induction head's QK-circuit locates the previous occurrence of the current token.
Its OV-circuit acts to copy the token that followed it.
The existence of such a clean, effective, and discoverable circuit provides strong evidence for the validity of the circuits research agenda.

Transformer Circuits VII: A Conclusion and The Road Ahead

TL;DR

This series reverse-engineered Transformers using the "circuits" framework, moving from high-level architecture to specific, learned algorithms.
The core idea is to treat a trained model as a program to be decompiled, not an inscrutable black box.
We built a hierarchy of abstractions: the residual stream (the memory), independent heads (the operators), QK/OV circuits (the sub-routines), and composition (the program flow).
This framework allowed us to fully dissect the induction head circuit, a concrete mechanism for in-context learning.
Major frontiers remain: understanding FFN/MLP layers, tackling non-linearities like LayerNorm, dealing with superposition, and scaling these analyses to frontier models.

Context & Motivation

Over the last six posts, we've journeyed deep into the internals of Transformers. We began with a single, ambitious goal: to move beyond treating models as black boxes and instead reverse engineer the algorithms they learn during training. The "Mathematical Framework for Transformer Circuits" provided our roadmap.

We started with the broadest architectural concept-the residual stream-and progressively zoomed in, dissecting layers into heads, heads into sub-circuits, and finally, tracing the compositional paths that weave these components into complex programs.

This final post zooms back out. We will recap our journey, solidify the key insights of the circuits perspective, and honestly assess its limitations and the exciting frontiers that lie ahead for the field of mechanistic interpretability.

The Journey from Architecture to Algorithm: A Recap

Our entire analysis was built on a series of increasingly granular, yet complementary, mental models.

The Residual Stream as the Workspace: We first reframed the Transformer not as a sequential chain of transformations, but as a collaborative workspace. The residual stream is the central data bus where each component-embeddings, attention heads, FFNs-reads the current state and writes back an additive update. This "assembly line" view is the foundational mental model.
Attention Heads as Independent Operators: We then simplified the multi-head attention block, showing its output is mathematically equivalent to the sum of contributions from independent heads. This allowed us to isolate a single head as a valid unit of analysis.
QK and OV Circuits within a Head: Zooming in further, we split each head into two sub-modules: the QK-circuit that determines where to look (attention patterns) and the OV-circuit that determines what to move (information payloads). This let us separate the search and retrieval aspects of attention.
Composition as the Source of Power: With the building blocks defined, we explored how they connect. Composition-where heads in later layers operate on the outputs of earlier heads-is the source of depth's power. By analyzing Q-, K-, and V-composition, we developed a language to describe how multi-step algorithms are formed.
The Payoff: Induction Heads: Finally, we put it all together to explain a genuine, non-trivial algorithm learned by the model: the induction head circuit. We saw how a "previous token head" composes with a later "induction head" via K-composition to implement a general algorithm for sequence repetition-a form of in-context learning.

This final step is the proof-of-concept for the entire framework. It demonstrates that our abstractions are not just neat metaphors; they are precise enough to fully reverse-engineer a key capability of the model.

Implications & The Bigger Picture

The circuits perspective is more than just a set of analytical tricks; it's a paradigm for studying neural networks. Its core implications are:

Models Learn Interpretable Algorithms: The most profound takeaway is that models don't just learn a soup of statistical correlations. They appear to learn clean, specific, and often human-understandable algorithms made of composable parts.
A Path to Rigorous Safety and Auditing: If we can truly understand the circuits driving a model's capabilities, we can begin to audit them for flaws, biases, or unintended behaviors. This is a far more robust approach to safety than relying on behavioral testing alone.
A New Kind of Neuroscience: By analogy, MI is like a form of computational neuroscience for artificial minds. We are moving from lesion studies (ablations) to the detailed tracing of neural pathways to understand how cognition emerges from component interactions.

Limitations and the Road Ahead

For all its successes, the circuits framework is still in its infancy, and major challenges remain.

The Tyranny of LayerNorm: Our clean, linear story of path expansion is an approximation. Layer Normalization is a critical non-linear component that entangles the contributions of different paths. Future work must develop techniques that can handle these non-linearities more gracefully.
The Mystery of the FFN Layers: This series, like the foundational paper, is heavily attention-centric. The Feed-Forward (or MLP) layers, which consume half the model's parameters, are far less understood. While hypotheses exist (e.g., they act as key-value memories), a full mechanistic picture is still missing.
The Problem of Superposition: The "privileged basis" assumption-that individual neurons or directions correspond to clean features-often breaks. Models appear to represent more features than they have dimensions by storing them in overlapping linear combinations, a phenomenon called superposition. Decomposing these representations is a major open research problem.
Scaling to a Trillion Parameters: Manually finding circuits in a two-layer model is one thing. Doing so in a 100-layer, trillion-parameter frontier model is another. The central challenge for the field is developing automated or semi-automated techniques to scale circuit discovery.

Final Takeaways

The "circuits" perspective provides a powerful, hierarchical framework for reverse-engineering Transformers.
The core mental model is of components reading from and writing additive updates to a central residual stream.
This framework allows us to decompose the model into understandable parts: layers, heads, and QK/OV sub-circuits.
Depth enables composition, where components chain together to form multi-step algorithms like induction heads.
Induction heads provide a concrete, mechanistic explanation for an important form of in-context learning.
While powerful, this linear, attention-focused view is an approximation that must be expanded to handle non-linearities (LayerNorm), FFN layers, and superposition.
Mechanistic interpretability is a nascent but promising field that aims to make AI models truly understandable, not just superficially predictable.

This journey from architecture to algorithm is just the beginning. The tools and concepts laid out here are the foundation upon which the next generation of interpretability research is being built, aiming to one day make even the most complex AI systems transparent and trustworthy.

Part 3 Reverse Engineering Transformers: Path Expansion

Fri, 17 Oct 2025 00:00:00 GMT

Note: This RSS feed strips out SVGs and embeds. You might want to read the post on the webpage here.

October 16, 2025

Last Updated: October 17, 2025

mechInterp #CMSC848R

💡

Other parts in the series:
Part 1: Reverse Engineering Transformers: Deconstructing Attention
Part 2: Reverse Engineering Transformers: Attention Heads & Circuits
Part 4 Reverse Engineering Transformers: Induction Heads

Transformer Circuits IV: Path Expansion

TL;DR

Path expansion is a technique for decomposing a Transformer's full forward pass into a sum of individual end-to-end paths.
A "path" is a sequence of matrix multiplications representing information flow from an input token, through a series of attention heads and FFN layers, to a final output logit.
The total logit for a specific token is the sum of the contributions from all possible paths leading to it.
For this linear decomposition to work, we must ignore the non-linear LayerNorm operations, treating path expansion as a useful "linear approximation" of the true computation.
This technique allows us to move from analyzing what a single component does to analyzing how components compose into circuits that solve a task.

Context & Motivation

So far in our series, we have established a powerful mental model:

The Transformer is a sequence of read/write operations on a central residual stream.
The attention layer can be decomposed into a sum of independent head contributions.
Each head can be further split into a QK-circuit (where to look) and an OV-circuit (what to move).

We have successfully broken the model down into its smallest functional parts. The next logical question is: how do these parts combine across layers to implement complex algorithms? How does a "previous token head" in Layer 0 enable an "induction head" in Layer 5?

The answer lies in path expansion. By recursively expanding the definition of the residual stream, we can re-express the entire model's computation for a single output logit as a massive sum of individual end-to-end information paths. This allows us to attribute the model's final output to specific combinations of components, which we call "circuits."

Prereqs

Understanding of the residual stream, independent heads, and QK/OV circuits (Parts I-III).
Comfort with the idea of matrix composition (multiplying matrices in a sequence).

Core Idea in One Picture

A Transformer's output for a token is the sum of all the ways information can travel from the input embeddings to the final unembedding layer.

graph TD
    subgraph Input Embeddings
        T_A[Token A]
        T_B[Token B]
    end

    subgraph Layer 0
        H0_1[Head 0.1]
        F0[FFN 0]
    end

    subgraph Layer 1
        H1_1[Head 1.1]
        F1[FFN 1]
    end

    subgraph Output Logits
        U[Unembedding] --> Logit_C[Logit for 'C']
    end

    T_A -- Path 1 --> U
    T_A -- Path 2 --> H0_1 -- Path 2 --> U
    T_B -- Path 3 --> F0 -- Path 3 --> H1_1 -- Path 3 --> U

    style Path 1 stroke:#ff0000,stroke-width:2px,stroke-dasharray: 5 5
    style Path 2 stroke:#0000ff,stroke-width:2px,stroke-dasharray: 5 5
    style Path 3 stroke:#008000,stroke-width:2px,stroke-dasharray: 5 5```
*Figure 1: A conceptual illustration of path expansion. The final logit for token 'C' is a sum of contributions from many paths. Path 1 is a direct connection. Path 2 goes through one head. Path 3 goes from a different token through an FFN and another head. Each path corresponds to a product of weight matrices.*

Definitions & Setup

Let's trace the computation algebraically. The final logit for a specific token T is computed by taking the dot product of the final residual stream state x_final with the token's corresponding row in the unembedding matrix, W_U[T].

Logit(T) = x_final @ W_U[T]

We know that x_final is the sum of the initial embedding x_0 and all component updates:
x_final = x_0 + ∑_{l, h} Δ_{l,h}^{attn} + ∑_{l} Δ_{l}^{ffn}

If we ignore LayerNorm, we can substitute this in:
Logit(T) = (x_0 + ∑ Δ_{l,h}^{attn} + ∑ Δ_{l}^{ffn}) @ W_U[T]

Since the dot product distributes over addition, this becomes:
Logit(T) = (x_0 @ W_U[T]) + ∑ (Δ_{l,h}^{attn} @ W_U[T]) + ∑ (Δ_{l}^{ffn} @ W_U[T])

This is the first level of expansion. We've expressed the logit as a sum of contributions from each component's direct write to the final residual stream. But now we can expand each Δ term, because each component's output depends on the state of the stream before it, which is itself a sum of previous terms. This recursive expansion is path expansion.

A path is a term in this fully expanded sum. Each path corresponds to a unique sequence of components that information flows through.

Walkthrough: Expanding Two Simple Paths

Let's analyze the paths for predicting the token 'C' that follows the prompt "A B".

Path 1: The Direct "Bigram" Path

The simplest path is the information flowing directly from an input token's embedding to the output.

Path:Embedding('B') → Unembedding('C')
Description: This path represents the model's learned bigram statistics. How likely is 'C' to follow 'B' without any deeper context?
Algebraic Contribution: The contribution to the logit of 'C' from this path is:

\text{Contribution}1 = (\mathbf{W_E}[\text{'B'}] + \mathbf{W{pos}}) \cdot \mathbf{W_U}[\text{'C'}]

This is a single scalar value, computed from a product of two embedding vectors. This path exists for every input token.

Path 2: The "Skip-Trigram" Attention Path

Now consider a path that involves one attention head. Let's trace the information from token 'A' to the output at position 'B', which then predicts 'C'.

Path:Embedding('A') → L0 Head OV Circuit → Unembedding('C')
Description: This path models skip-trigrams (e.g., in "The dog ran", it helps predict "ran" based on "The" after seeing "dog"). The head at position 'B' attends to 'A' and moves information to 'B's residual stream.
Algebraic Contribution: Let's assume head h in layer 0 at position q=B is attending to source position k=A.
- The information moved by the OV circuit is (x_A @ W_V^h) @ W_O^h.
- This is projected by the unembedding: ((x_A @ W_V^h) @ W_O^h) @ W_U[C].
- The full contribution is this value multiplied by the attention probabilityA_{B ← A} paid from B to A.
  
  \text{Contribution}2 = A{B \leftarrow A} \cdot \left( (\mathbf{W_E}[\text{'A'}] + \mathbf{W_{pos}}) \mathbf{W_V^h} \mathbf{W_O^h} \mathbf{W_U}[\text{'C'}] \right)
  
  This shows how the matrices along a path compose. The term W_V^h W_O^h W_U forms a single composed operator that describes the end-to-end effect of this path.

The full logit for 'C' is Contribution_1 + Contribution_2 + contributions from all other paths (from 'B' through heads, from 'A' through FFNs, through multiple layers, etc.).

Implications & Limits

Implication: Finding Circuits. Path expansion is the theoretical tool that allows us to find circuits. The "Indirect Object Identification" circuit, for instance, is a set of paths involving specific heads that compose to perform that task. By summing the contributions of just the paths in a hypothesized circuit, we can see how much of the model's behavior on a specific task that circuit explains.
Implication: Explaining Logits, Not Probabilities. This is an analysis of the pre-softmax logits. Logits are additive, so this decomposition is mathematically sound (ignoring LayerNorm). You cannot do this with probabilities, which are the output of the non-linear softmax function. softmax(a+b) != softmax(a) + softmax(b).
Limit #1: LayerNorm Breaks Everything. As stated before, this entire framework relies on the linear additivity of the residual stream. LayerNorm is non-linear and breaks this assumption. LN(x_0 + Δ_1) is not LN(x_0) + LN(Δ_1). Therefore, path expansion is a linear approximation. For many circuits and models, this approximation holds up surprisingly well, but it is not the ground truth.
Limit #2: Combinatorial Explosion. The number of paths is exponential in the depth of the model. Fully expanding all paths for a large model is computationally intractable. In practice, researchers analyze the top-K most important paths or analyze contributions on a component-by-component basis, which is a coarser-grained version of path expansion.

Pitfalls

Forgetting the Linear Approximation. It is the most common mistake to treat path expansion results as the exact ground truth of the model's computation. Always remember the LayerNorm caveat.
Summing Probabilities. Never attempt to sum the "probability contributions" of different paths. The analysis is valid only for the additive logits.

Takeaways

Path expansion reframes a Transformer's computation as a sum over end-to-end information flow paths.
Each path has a value (its contribution to a final logit) and is composed of a product of weight matrices.
This technique is the key to moving from analyzing components to analyzing multi-component circuits.
It provides a granular way to attribute a model's output to specific model parameters and interactions.
The validity of simple path expansion rests on a linear approximation that ignores LayerNorm.
This analysis applies to logits, not probabilities, due to the additivity requirement.

Transformer Circuits V: The Power of Depth is Composition

TL;DR

The power of deep Transformers comes from composition, where components in later layers operate on the outputs of components in earlier layers.
An earlier head (H0) can influence a later head (H1) in three ways: Q-Composition, K-Composition, and V-Composition.
Q-Composition: The output of H0 affects H1's Query vector, changing what H1 is looking for.
K-Composition: The output of H0 affects H1's Key vector, changing how H1's source tokens appear to other queries.
V-Composition: The output of H0 affects H1's Value vector, changing what information H1 moves. This can create powerful "virtual attention heads."
These composition mechanisms are the building blocks of complex circuits, like the famous "induction heads."

Context & Motivation

In our last post on Path Expansion, we established that a Transformer's output can be viewed as a massive sum of information paths. However, this view can be misleading if it suggests paths are just independent, parallel computations. The most interesting behaviors in Transformers arise when these paths interact-when the output of one component directly becomes the input to another.

A single attention head in Layer 0 is limited. It can only implement simple algorithms like "attend to the previous token" or "attend to nouns." To build a program that does something complex, like "find the last time the current token appeared, and copy the token that came after it," you need multiple steps. You need depth.

Composition is the formal term for how these steps are chained together. Specifically, we'll examine how an attention head in an earlier layer (H0) can pass information through the residual stream to be read by a head in a later layer (H1), fundamentally altering H1's behavior.

Prereqs

Solid understanding of the residual stream, independent heads, QK/OV circuits, and the concept of path expansion (Parts I-IV).

Core Idea in One Picture

A head in Layer 0 (H0) writes an update to the residual stream. A head in Layer 1 (H1) then reads from this updated stream. Depending on which part of H1 reads the update (its Query, Key, or Value projection), we get a different type of composition.

graph TD
    subgraph Layer 0
        x0[Residual Stream @ L0] --> H0[Head H0]
        H0 -- Δ_0 --> x0_updated((+))
    end

    subgraph Layer 1
        x1[Residual Stream @ L1] --> H1[Head H1]
    end

    x0 --> x0_updated
    x0_updated --> x1

    subgraph H1
        direction LR
        x1 -- W_Q^1 --> Q1[Query]
        x1 -- W_K^1 --> K1[Key]
        x1 -- W_V^1 --> V1[Value]
    end

    linkStyle 0 stroke-width:0px
    linkStyle 4 stroke-width:0px

    subgraph Legend
        QComp[Q-Composition]
        KComp[K-Composition]
        VComp[V-Composition]
    end

    H0 -- "affects" --> Q1;
    H0 -- "affects" --> K1;
    H0 -- "affects" --> V1;

    style H0 fill:#ccf
    style H1 fill:#fcc
    style QComp fill:#f9d,stroke:#333
    style KComp fill:#d9f,stroke:#333
    style VComp fill:#9df,stroke:#333

Figure 1: The three types of head composition. An earlier head, H0, writes an update (Δ_0) to the stream. A later head, H1, reads this update. If the update influences H1's Query, it's Q-Composition; if its Key, K-Composition; if its Value, V-Composition.

Definitions & Setup

For simplicity, let's consider a two-layer, attention-only model. The residual stream entering Layer 1 is the sum of the initial embeddings (x_0) and the output of Layer 0 (Δ_{L0}), which is the sum of all Layer 0 head outputs. Let's focus on the interaction between one head H0 and one head H1.

The stream after H0 is x_1 = x_0 + Δ_0 (again, ignoring LayerNorm for our linear analysis). H1 now computes its Q, K, and V vectors from this x_1.

Q-Composition:H1's query is q_1 = x_1 @ W_Q^1 = (x_0 + Δ_0) @ W_Q^1. The term Δ_0 @ W_Q^1 represents the influence of H0 on H1's query.
K-Composition:H1's key is k_1 = x_1 @ W_K^1 = (x_0 + Δ_0) @ W_K^1. The term Δ_0 @ W_K^1 represents the influence of H0 on H1's key.
V-Composition:H1's value is v_1 = x_1 @ W_V^1 = (x_0 + Δ_0) @ W_V^1. The term Δ_0 @ W_V^1 represents the influence of H0 on H1's value.

Walkthrough: What Each Composition Type Does

Step 1: Q-Composition ("Changing What to Look For")

Imagine H0 is a "proper noun detector." When it sees a name like "Mary," it writes a +is_a_proper_noun feature into the residual stream at that position. H1 might be a head that wants to find the main verb of a sentence. Through Q-composition, H1 can learn to modify its query based on H0's output. Its query at the "Mary" token, influenced by Δ_0, might become something like "I am a proper noun, now I am looking for the verb I am the subject of." This allows H1's search to be context-dependent.

Step 2: K-Composition ("Changing How You Are Seen")

This is subtler but crucial. H0's output modifies the key of a token, changing how it "appears" to queries from other positions. The canonical example is in forming induction heads.

Let H0 be a previous token head. At every position t, it attends to position t-1 and copies its embedding. So, Δ_0 at position t contains information about token t-1.
Now consider H1, the induction head. It wants to find a previous occurrence of the current token. Let's say we're at token[t] = 'B' and we've seen "... A B ..." earlier in the sequence. H1 needs to attend to token[t-1] = 'A'.
How? At position t-1, H0 wrote information about token[t-2] = 'A'. H1's key at t-1 is thus modified by H0 to effectively broadcast "the token after me is 'A'". H1's query at t is looking for the token 'A'. The QK-dot-product will match, directing H1's attention to the correct previous token.

Step 3: V-Composition ("Changing What You Move")

This is perhaps the most powerful type, creating what the paper calls virtual attention heads. Here, H0 determines what information H1 will move. H1's attention pattern might direct it to attend to token A, but the information it moves can come from token B, as mediated by H0.

Let's expand the algebra for a V-composition path, where H1 at position q attends to position k_1, and H0 at k_1 attends to k_0.

H1's output: Output_1 = A_{q ← k_1} · (x_{k_1} @ W_V^1 @ W_O^1)
But x_{k_1} = x_0 + Δ_0 = x_0 + (A_{k_1 ← k_0} · (x_{k_0} @ W_V^0 @ W_O^0))
Plugging this in, the contribution to H1's output from this two-head path is:

Path Output=Aq←k1⋅Ak1←k0⋅(xk0@WV0WO0WV1WO1)\text{Path Output} = A_{q \leftarrow k_1} \cdot A_{k_1 \leftarrow k_0} \cdot ( \mathbf{x_{k_0}} @ \mathbf{W_V^0 W_O^0 W_V^1 W_O^1} )Path Output=Aq←k1⋅Ak1←k0⋅(xk0@WV0WO0WV1WO1)

This path moves information from x_{k_0} to position q via a two-step hop. It attends from q to k_1 and k_1 to k_0, but the matrices W_V^0 W_O^0 W_V^1 W_O^1 form a a new, composed OV circuit that operates directly on the information at x_{k_0}. We've created a "virtual head" that has the attention pattern of H1 but an OV circuit defined by both heads.

Implications & Limits

Implication: Head composition is the primary mechanism for building complex, multi-step algorithms in Transformers. It explains how models can perform tasks that are impossible for a single layer.
Implication: It allows us to reason about models as computational graphs where heads are nodes and composition links are edges, leading to a much richer understanding of model behavior.
Limit: This clean decomposition is, as always, complicated by LayerNorm, which mixes all parts of the residual stream together non-linearly. The real interactions are not quite as simple as adding up path contributions.
Limit: FFN/MLP layers also compose with heads and with each other. A full circuit analysis must account for their role, which is often less clear than that of attention heads.

Pitfalls

Viewing heads as independent agents.H1 and H0 are not collaborating consciously. They are a single system optimized end-to-end. H1 learns to use H0's output because the joint behavior is useful for minimizing the loss function. There is no guarantee H0's "purpose" is cleanly separable from H1's.
Assuming one type of composition per head pair. A single head pair (H0, H1) can and often does exhibit all three types of composition simultaneously. H0's output can influence H1's Q, K, and V paths all at once.

Takeaways

Composition is how Transformers build complex algorithms from simple components, and it is the key reason for the power of model depth.
The three types of composition-Q, K, and V-provide a vocabulary for describing how heads in different layers interact.
Q-Composition allows a head's search query to be context-dependent.
K-Composition allows a token's representation for search purposes to be context-dependent, enabling induction heads.
V-Composition allows a head to move information from sources it isn't directly attending to, creating powerful "virtual heads."
Analyzing these compositional circuits is a core activity in mechanistic interpretability.

Part 2: Reverse Engineering Transformers: Attention Heads & Circuits

Fri, 17 Oct 2025 00:00:00 GMT

Note: This RSS feed strips out SVGs and embeds. You might want to read the post on the webpage here.

October 16, 2025

Last Updated: October 17, 2025

mechInterp #CMSC848R

💡

Other Parts in the series:
Part 1: Reverse Engineering Transformers: Deconstructing Attention
Part 3 Reverse Engineering Transformers: Path Expansion
Part 4 Reverse Engineering Transformers: Induction Heads

Transformer Circuits II: Attention Heads as Independent Operators

The standard implementation of multi-head attention concatenates head outputs and uses a single large output projection (W_O).
We can mathematically reframe this as summing the outputs of independent heads, each with its own smaller output projection matrix (W_O^h).
This reframing is key for interpretability: it allows us to analyze the contribution of each attention head to the residual stream in isolation.
An attention head's "write" to the residual stream is its value-weighted attention pattern, projected back into d_model space by its personal W_O^h.
While heads in a layer compute in parallel, they are not fully independent; they read from the same source and their outputs may compete for capacity in the next layer's residual stream.

Context & Motivation

In our previous post, we established the residual stream as the central communication bus of a Transformer. We saw it as an assembly line where components "read" from the stream and "write" additive updates back to it. Our diagram represented "Attention" as a single block that performed one of these read-write operations.

Now, we need to look inside that block. A standard Transformer layer contains multi-head attention. The typical textbook explanation involves concatenating the outputs of all heads and projecting them back to the model dimension with a single output matrix, W_O. This implementation view is computationally efficient, but it's messy for interpretation. It mixes the outputs of all heads together, making it hard to ask: "What was the specific contribution of head 7?"

This post shows that we can adopt a more theoretically convenient, yet mathematically equivalent, perspective: the attention layer's total output is simply the sum of the outputs of each individual head. This allows us to treat each head as an independent operator writing to the residual stream, making our reverse engineering task dramatically simpler.

Prereqs

Familiarity with the concepts of Query, Key, and Value in attention.
Understanding of the residual stream as a communication bus (from Part I).

Core Idea in One Picture

The key is recognizing that the standard, monolithic attention output projection can be split into a sum of per-head contributions.

graph TD
    subgraph Standard View (Implementation)
        H1[Head 1 Out] --> C;
        H2[Head 2 Out] --> C;
        Hn[... Head N Out] --> C;
        C(Concat) --> WO1[W_O];
        WO1 --> LayerOut1[Layer Output];
    end

    subgraph Equivalent View (Interpretation)
        H1_i[Head 1 Out] --> WO_h1[W_O^1];
        H2_i[Head 2 Out] --> WO_h2[W_O^2];
        Hn_i[... Head N Out] --> WO_hn[W_O^n];
        WO_h1 --> S((+));
        WO_h2 --> S;
        WO_hn --> S;
        S --> LayerOut2[Layer Output];
    end

    LayerOut1 -- Mathematically Identical --> LayerOut2;

    style H1 fill:#f9f,stroke:#333
    style H2 fill:#f9f,stroke:#333
    style Hn fill:#f9f,stroke:#333
    style H1_i fill:#f9f,stroke:#333
    style H2_i fill:#f9f,stroke:#333
    style Hn_i fill:#f9f,stroke:#333

Figure 1: The standard implementation view (left) concatenates head outputs before a single projection. The interpretation view (right) shows this is equivalent to summing the outputs of heads projected by their own virtual output matrices.

Definitions & Setup

Let's formalize the two views. We assume a model with n_heads heads and a head dimension d_head. The model dimension is d_model, where typically d_model = n_heads * d_head. The input to the layer from the residual stream is x.

1. The Standard View (Implementation)
Each head h produces an output vector o_h of shape [seq_len, d_head]. These are concatenated along the last dimension to form a single large tensor O_cat of shape [seq_len, n_heads * d_head]. This tensor is then projected back to the residual stream's dimension by the output weight matrix W_O, which has shape [n_heads * d_head, d_model].

Output = Concat(o_1, o_2, ..., o_{n_heads}) @ W_O

2. The Theoretically Convenient View (Interpretation)
We can conceptually partition the large W_O matrix into n_heads smaller matrices, W_O^h, one for each head. Each W_O^h has shape [d_head, d_model].

The total output is then the sum of each head's output projected by its corresponding W_O^h:

Output=∑h=1nheadsoh@WOh\text{Output} = \sum_{h=1}^{n_{\text{heads}}} o_h @ W_O^hOutput=∑h=1nheadsoh@WOh

The output of a single head, o_h @ W_O^h, is that head's direct, additive contribution to the residual stream. This is exactly the "write update" we discussed in Part I.

Walkthrough

Step 1: Proof of Equivalence

Why are these two views identical? Let's look at the structure of the matrices. The Concat(o_1, ..., o_n) operation creates a wide matrix. The W_O matrix is correspondingly tall.

Let O_cat = [o_1 | o_2 | ...] be the row-wise concatenation. Let W_O be the block matrix composed of the W_O^h matrices stacked vertically:

WO=(WO1WO2⋮WOnheads)\mathbf{W_O} = \begin{pmatrix} \mathbf{W_O^1} \\ \mathbf{W_O^2} \\ \vdots \\ \mathbf{W_O^{n_{heads}}} \end{pmatrix}WO=WO1WO2⋮WOnheads

The matrix multiplication O_cat @ W_O is then:

[o1∣o2∣… ](WO1WO2⋮)=o1WO1+o2WO2+⋯=∑hohWOh[\mathbf{o_1} | \mathbf{o_2} | \dots] \begin{pmatrix} \mathbf{W_O^1} \\ \mathbf{W_O^2} \\ \vdots \end{pmatrix} = \mathbf{o_1 W_O^1} + \mathbf{o_2 W_O^2} + \dots = \sum_{h} \mathbf{o_h W_O^h}[o1∣o2∣…]WO1WO2⋮=o1WO1+o2WO2+⋯=∑hohWOh

(Note: This is a slight abuse of notation for clarity; the multiplication is properly defined by partitioning columns and rows.)

Step 2: A Minimal Experiment

We can verify this with a small PyTorch snippet. We'll simulate the outputs of two heads and show the two methods produce the same result.

import torch# Setup: 1 token, 2 heads, d_head=4, d_model=8seq_len, n_heads, d_head = 1, 2, 4d_model = n_heads * d_head# Random head outputso_1 = torch.randn(seq_len, d_head) # Output of head 1o_2 = torch.randn(seq_len, d_head) # Output of head 2# Big W_O matrix for the standard viewW_O = torch.randn(d_model, d_model)# 1. Standard View: Concatenate and multiplyo_cat = torch.cat([o_1, o_2], dim=-1) # Shape:output_standard = o_cat @ W_O# 2. Interpretation View: Split W_O and sumW_O1 = W_O[:d_head, :]  # Shape:W_O2 = W_O[d_head:, :]  # Shape:output_interp = (o_1 @ W_O1) + (o_2 @ W_O2)# Check if they are nearly identicalprint("Standard output:", output_standard)print("Interpretation output:", output_interp)print("Outputs are close:", torch.allclose(output_standard, output_interp))# Outputs are close: True

Step 3: Why This Reframing is Powerful

This equivalence is the foundation for analyzing individual heads. We can now cleanly state what a single head h does:

Read: It forms its Query, Key, and Value vectors by projecting from the LayerNorm'd residual stream: q = x_norm @ W_Q^h, k = x_norm @ W_K^h, v = x_norm @ W_V^h.
Compute: It calculates a weighted sum of Value vectors based on Query-Key similarity, producing its output o_h.
Write: It projects its output o_h via its personal output matrix W_O^h to create its total contribution Δ_h = o_h @ W_O^h, which is added directly to the residual stream.

This allows us to ask targeted questions like, "What information does head 8.3 read from the 'Paris' token to write to the 'capital' token?" and answer it by analyzing only the matrices W_Q^{8.3}, W_K^{8.3}, W_V^{8.3}, and W_O^{8.3}. We have successfully decomposed a layer-level operation into a sum of head-level operations.

Implications & Limits

Implication: The fundamental unit of analysis for attention inside a Transformer can be the individual head, not the entire layer. This makes reverse engineering tractable. We can study circuits composed of specific heads in different layers.
Limits: The heads within a layer are not truly independent.
1. Shared Input: They all read from the exact same normalized residual stream x_norm.
2. Shared Output: Their outputs are all summed into the same residual stream. A powerful update from one head could change the vector's direction or magnitude in a way that affects how it is processed by the LayerNorm in the next block.

Pitfalls

Confusing parallel computation with total independence. Just because heads are computed in parallel doesn't mean they don't influence each other's operating environment in subsequent layers. The additive nature of their outputs means they are in a "linear superposition" within the stream, but the subsequent LayerNorm is non-linear and will mix their signals.

Takeaways

A multi-head attention layer's output is mathematically equivalent to the sum of individual head outputs.
Each head can be thought of as having its own input (W_Q, W_K, W_V) and output (W_O) matrices.
This decomposition allows us to isolate and analyze the function of a single attention head.
A head's function is to read information from the residual stream, transform it via attention, and write an update back to the stream.
This "independent, additive" view is a cornerstone of the circuits framework for understanding Transformers.

Transformer Circuits III: The QK and OV Circuits Inside an Attention Head

An individual attention head's operation can be decomposed into two distinct sub-circuits: the QK-Circuit and the OV-Circuit.
The QK-Circuit determines the attention pattern-where the head looks. It uses Query and Key vectors to decide which source tokens to get information from.
The OV-Circuit determines the content being moved-what information the head moves. It uses Value vectors and the head's Output projection to specify the update written to the residual stream.
This separation is a powerful analytical tool, allowing us to categorize heads by their QK behavior (e.g., "previous token heads") and their OV behavior (e.g., "copying heads") separately.
The full effect of the OV-Circuit on the residual stream is captured by the composed matrix W_V @ W_O, which maps features from a source token's stream to the output update.

Context & Motivation

In our previous posts, we established that a Transformer layer's attention is best understood as a sum of contributions from independent heads, each writing an update to the residual stream. This was a crucial simplification, allowing us to isolate a single head for analysis.

Now, we zoom in further. What algorithm does a single head execute? An attention head performs a complex operation: for a given "query" token, it looks at all other "key" tokens, computes similarity scores, and then creates a weighted sum of their "value" vectors. This suggests two sub-tasks are happening at once:

Deciding which tokens are relevant (where to look).
Extracting and moving useful information from them (what to move).

The circuits framework makes this separation explicit, decomposing a head into a Query-Key (QK) circuit and an Output-Value (OV) circuit. By analyzing these two parts, we can achieve a much more precise understanding of a head's role in the model's overall algorithm.

Prereqs

Understanding the residual stream as a communication bus.
The view of attention layers as a sum of independent head outputs.
Basic knowledge of the Query, Key, and Value formulation of attention.

Core Idea in One Picture

An attention head can be split into two parallel information paths: one calculates attention scores, the other prepares the information to be moved. They only combine at the very end.

graph TD
    subgraph Attention Head h
        direction LR
        x_norm[Normed Residual Stream]

        subgraph QK-Circuit (Where to look)
            x_norm -- W_Q^h --> Q[Query];
            x_norm -- W_K^h --> K[Key];
            Q --> AS(Attention Scores);
            K --> AS;
        end

        subgraph OV-Circuit (What to move)
           x_norm -- W_V^h --> V[Value];
           V -- W_O^h --> Projected_V[Projected Value];
        end

        AS -- Softmax --> A[Attention Pattern: A];
        A --> Combine((@));
        Projected_V --> Combine;
        Combine --> Output[Output Δ_h];
    end

    Output --> RS_Update{+ Add to Residual Stream};
    ```

Figure 1: An attention head decomposed into its QK and OV circuits. The QK circuit produces the attention pattern `A`, which acts as a set of weights. The OV circuit produces the information to be moved. The final output is a weighted sum of the projected values

Definitions & Setup

Let's formalize the two circuits for a single head h. The input from the previous layer is the normalized residual stream, x_norm.

1. The QK-Circuit (Attention Pattern)

The QK-circuit's job is to compute a scalar attention score between a query token t_q and a key token t_k.

Inputs: Residual stream vectors x_q and x_k.
Matrices:W_Q^h (Query) and W_K^h (Key), both of shape [d_model, d_head].
Computation: The attention score s_{q,k} is computed via a dot product.

sq,k=(xqWQh)⋅(xkWKh)dheads_{q,k} = \frac{(\mathbf{x_q} \mathbf{W_Q^h}) \cdot (\mathbf{x_k} \mathbf{W_K^h})}{\sqrt{d_{head}}}sq,k=dhead(xqWQh)⋅(xkWKh)
Output: A sequence of scores, which are passed through a softmax function to create the final attention pattern A. This pattern is a set of probabilities summing to 1, indicating where the head "looks".

The QK-circuit is fundamentally a search and matching mechanism.

2. The OV-Circuit (Content Movement)

The OV-circuit's job is to determine what information is moved from a source token t_v and how it is written to the destination token's residual stream.

Input: The residual stream vector x_v from a source token.
Matrices:W_V^h (Value) of shape [d_model, d_head] and the head's effective output matrix W_O^h of shape [d_head, d_model].
Computation: For a single source token t_v, the "information payload" it offers is the projected value:

Payloadv=(xvWVh)WOh\text{Payload}_v = (\mathbf{x_v} \mathbf{W_V^h}) \mathbf{W_O^h} Payloadv=(xvWVh)WOh
Output: A d_modeldimensional vector. This is the update that will be added to the destination token's residual stream if this token receives 100% of the attention.

The OV-circuit is a feature extraction and transformation mechanism. The composed matrix W_V^h W_O^h can be seen as a single [d_model, d_model] operation that specifies "if you attend to this token, this is the update you will get."

Walkthrough

Step 1: Intuition in a Case Study

Consider the prompt "The Roman Empire fell in 476 AD. This date..." and a hypothetical head in a later layer processing the token "date".

QK-Circuit in action: We want this head to find the actual date. The QK circuit at the "date" token must learn to attend to tokens that are numbers and correspond to years.
- The Query vector from "date" (q = x_{date} W_Q) must be a vector that means something like "I'm looking for a year."
- The Key vector from "476" (k = x_{476} W_K) must be a vector that means "I am a year."
- The dot product q · k will be high, resulting in high attention to "476".
OV-Circuit in action: Once the head is attending to "476", what information should it move? It should probably move the semantic concept of "the year is 476".
- The OV-circuit takes the residual stream at "476" (x_{476}). This vector already contains information about the token being a number.
- It computes the information payload: (x_{476} W_V) W_O. This operation might read the "is a number: 476" feature from x_{476} and transform it into a more abstract feature vector like +references_year(476).
- This payload vector is then added to the residual stream at the "date" token, enriching its representation.

Step 2: How to Inspect These Circuits

This decomposition isn't just a metaphor; it gives us concrete matrices to inspect.

To analyze the QK-Circuit: We can study the matrix W_QK = W_Q^T W_K. This [d_model, d_model] matrix reveals the structure of attention patterns. A large positive value at (i, j) in this matrix means that if feature i is present at the query token and feature j is present at the key token, the attention score will be high. This allows us to find patterns like "attend to the previous token" or "attend to tokens that are nouns".
To analyze the OV-Circuit: We study the composed matrix W_OV = W_V W_O. This [d_model, d_model] matrix tells us what kind of information processing the head performs. We can ask: if the input vector x_v has a feature in direction d_1, what features does the output vector (x_v W_{OV}) contain? For a "copying" head, this matrix would look similar to the identity matrix, meaning it moves information without much transformation. For other heads, it might map "proper noun" features to "capital city" features.

Implications & Limits

Implication: A Taxonomy of Heads. This separation allows us to categorize heads. For example, many models learn "previous token heads" whose QK circuit reliably attends to the token at position - 1. However, these heads can have very different OV circuits: one might copy the previous token's embedding, while another might check if the previous token was a verb.
Limit: Not Fully Independent. The circuits are optimized jointly during training. The gradients flowing back to W_Q and W_K depend on how useful the information moved by the OV-circuit was for the final task. A QK-circuit has no reason to create a sharp, meaningful attention pattern if the corresponding OV-circuit provides useless information. They co-evolve.

Pitfalls

Analyzing W_V in isolation: A common error is to just look at W_V to understand what information is being moved. W_V projects into the d_headdimensional head-space, which is often uninterpretable. You must compose it with W_O to see the head's full effect on the d_modeldimensional residual stream.
Over-interpreting Q/K vectors: It is tempting to think of q and k vectors as having rich semantic meaning. It's often more accurate to think of them as specialized pointers or search keys, designed only for the purpose of dot-product matching. The actual semantic content is carried by the v vectors.

Takeaways

An attention head can be productively decomposed into a QK-circuit (where to look) and an OV-circuit (what to move).
The QK-circuit's behavior is summarized by the composed matrix W_Q^T W_K.
The OV-circuit's behavior is summarized by the composed matrix W_V W_O.
This decomposition provides a powerful framework for categorizing and analyzing the function of individual heads.
It simplifies the problem of understanding a head into two smaller, more focused problems.
While analytically useful, the two circuits are not truly independent as they are co-optimized during training.

Part 1: Reverse Engineering Transformers: Deconstructing Attention

Fri, 17 Oct 2025 00:00:00 GMT

Note: This RSS feed strips out SVGs and embeds. You might want to read the post on the webpage here.

Blog series on how understanding the residual stream and it’s components are important for MechInterp

October 16, 2025

Last Updated: October 17, 2025

mechInterp readings #CMSC848R

💡

Other parts in the series:
Part 2: Reverse Engineering Transformers: Attention Heads & Circuits
Part 3 Reverse Engineering Transformers: Path Expansion

Part 4 Reverse Engineering Transformers: Induction Heads

This is a 4 part blog series on my notes and understanding of the foundational paper Elhage, N. et al. (2021). A Mathematical Framework for Transformer Circuits

The blogs will be discussing the following ideas:

Reverse Engineering Transformers: The primary goal is to break down the complex computations within a transformer into understandable "circuits" or algorithmic patterns. This is analogous to reverse-engineering a compiled program to understand its source code.
The Residual Stream as a Communication Channel: The paper conceptualizes the residual stream not just as an embedding, but as a central communication channel. Different components of the transformer (like attention heads and MLP layers) read from and write information to this shared space.
Attention Heads as Independent, Additive Operations: The authors reframe attention layers as a collection of independent attention heads that operate in parallel and add their outputs to the residual stream. This is a more theoretically convenient way to think about them than the typical "concatenate and multiply" view.
Separating Attention Head Components: The paper proposes that attention heads can be broken down into two largely independent circuits:
- QK (Query-Key) Circuit: This determines the attention pattern, i.e., which tokens to move information from.
- OV (Output-Value) Circuit: This determines what information is moved and how it's written to the destination.
Path Expansion: A key analytical technique is to expand the transformer's computation into a sum of end-to-end paths from input tokens to output logits. This allows for a more granular analysis of how information flows and is transformed.
Composition of Attention Heads: In multi-layer transformers, the paper explores how attention heads can compose with each other. This composition is what gives deeper transformers their power. The three types of composition are:
- Q-Composition: A second-layer head's query is influenced by a first-layer head.
- K-Composition: A second-layer head's key is influenced by a first-layer head.
- V-Composition: A second-layer head's value is influenced by a first-layer head, creating "virtual attention heads".
Induction Heads: A specific and powerful circuit that emerges in two-layer (or deeper) models. These heads are a mechanism for in-context learning, allowing the model to recognize repeated sequences and continue them. They work by searching for previous occurrences of the current token and then attending to the next token in that sequence.
Circuits, Residual Streams, and Induction Heads

We start with the first two points in this blog

What is Reverse Engineering in Mechanistic Interpretability?

At its core, mechanistic interpretability (MI) is the scientific field of reverse engineering neural networks. The goal is to take a fully trained model, a complex and opaque mathematical function defined by billions of parameters, and understand the specific, human-understandable algorithms it has learned to perform its tasks.

Think of it this way:

A traditional computer program starts with human-written source code (like Python). It's then compiled into a low-level binary file that a machine can execute. A reverse engineer would take that binary and try to reconstruct the original logic and algorithms.
A neural network starts with a chosen architecture (the "virtual machine") and is trained on data. This process results in a set of optimized weights (the "program binary"). A mechanistic interpretability researcher takes those weights and tries to understand the algorithms they implement.

This analogy is quite deep and helps frame the entire endeavor.

Traditional Computing	Mechanistic Interpretability of Neural Networks
Program Binary	Network Parameters (Weights and Biases)
VM / Processor	Network Architecture (e.g., Transformer)
Program State / Memory	Layer Representations / Activations
Variable / Memory Location	Neuron / Feature Direction

Why is this Analogy So Powerful?

It Addresses the Curse of Dimensionality
Neural networks operate on incredibly high-dimensional inputs. Trying to understand a model by mapping out its behavior for every possible input is impossible. The reverse engineering approach sidesteps this by instead focusing on understanding the finite description of the program itself-the model's parameters. Just as a programmer can understand a large piece of software without running every possible input, an MI researcher aims to understand the finite set of weights that define the network's behavior.
It Sets Realistic Expectations
Reverse engineering a complex binary like a modern operating system is incredibly difficult and requires painstaking, detailed work. Similarly, we shouldn't expect to find a simple, one-sentence explanation for how a large language model works. Mechanistic interpretability is expected to be a slow, iterative, and challenging scientific process.
It Frames the Core Challenge-Understanding "Variables"
A computer program is made up of operations acting on variables. A statement like y = x + 5 is meaningless unless you know what x and y represent. A reverse engineer must figure out what each piece of the program's memory represents.
In neural networks, the "variables" are the activations within the model. The weights are the "instructions" that describe how previous activations affect later ones. To understand the weights, we must first understand the activations they operate on.

The Central Task: Decomposing Activations

The immediate serious challenge: activations are high-dimensional vectors. How can we possibly understand them?

This is where the concept of a privileged basis becomes critical. The hope is that a network's internal representations aren't just arbitrary directions in a vector space. Instead, MI researchers hypothesize and find evidence that networks often learn to align meaningful, independent "features" with the basis vectors of their activation space (i.e., with individual neurons or specific directions).

If a privileged basis exists:
We can decompose activations into independently understandable pieces. This is like figuring out that certain bytes in a program's memory correspond to a player's health, while other bytes correspond to their score. Each can be understood separately. This makes the task of reverse engineering tractable.
If a privileged basis doesn't exist:
And every "feature" is stored as some complex linear combination of all neurons (a phenomenon called superposition), the task becomes much harder. It would be like trying to understand a program where a player's health is determined by the 3rd bit of one byte, the 7th bit of another, and so on, all mixed up.

The papers you provided show that while networks don't necessarily converge to a single, unique basis across different training runs, the basis they do learn is not arbitrary. You can't just randomly rotate a layer's activations and expect the network to recover, which provides strong evidence that the learned basis directions are meaningful and privileged.

Summary

"Reverse engineering" is the guiding metaphor for mechanistic interpretability. It frames the goal as deciphering the learned algorithms from the model's weights and highlights that the central challenge is to find a way to break down the model's high-dimensional internal state (activations) into meaningful, understandable variables.

Transformer Circuits I: The Residual Stream as a Communication Channel

TL;DR

The residual stream is the central data pathway in a Transformer, acting like a communication bus or an assembly line.
It's a high-dimensional vector space (d_model) where information is stored and progressively refined.
Transformer components (attention, FFNs) don't replace the stream's content; they read from it and write additive updates back to it.
A "Read" is a linear projection from the stream (e.g., to create Q, K, V vectors).
A "Write" is a linear projection that is added back into the stream (x_new = x_old + update).
This additive structure is the key property that allows us to decompose the model's behavior, but non-linearities like LayerNorm complicate the pure linear view.

Context & Motivation

A Transformer is a deep stack of identical blocks. A natural question is: how does information from early layers persist and get refined by later layers? The answer lies in the residual connections that are present at every step.

To reverse-engineer a Transformer, we need a robust mental model of its architecture. The most powerful one, proposed by the "circuits" interpretability agenda, is to view the model not as a series of opaque transformations, but as a group of specialists collaborating on a shared workspace. That workspace is the residual stream.

Understanding the stream as a communication bus where components perform additive read/write operations is the first and most critical step. It provides the foundation for decomposing the model into understandable parts, allowing us to trace how specific features are computed from input to output.

Prereqs

Basic familiarity with the Transformer architecture (embeddings, attention, FFN/MLP layers).
Conceptual understanding of vectors and matrices (linear projections).
Knowledge of PyTorch/NumPy tensor shapes.

Core Idea in One Picture

Picture the residual stream as an assembly line. For each token, a "package" of information (a vector) moves down the line. At each station, a component reads the package, adds something new, and puts it back on the line.

graph TD
    A[Input + Pos Embedding] --> B[Residual Stream]

    subgraph Layer0
        B --> Attn0[Attention Head]
        Attn0 --> B
        B --> FFN0[FFN Layer]
        FFN0 --> B
    end

    B --> MoreLayers[...]

    subgraph FinalLayer
        MoreLayers --> B
        B --> Unembed[Unembedding]
        Unembed --> Logits[Output Logits]
    end

    style B fill:#e6f3ff,stroke:#333,stroke-width:2px

Figure 1: The residual stream as a central communication bus. Information for a token starts as an embedding. Each component (Attention, FFN) reads the current stream state and adds its own contribution back. The final, refined state is read by the unembedding matrix to produce logits.

Definitions & Setup

Let's formalize this. The residual stream is the state x that is passed from one layer to the next.

Shape: The stream is a tensor of shape [batch_size, seq_len, d_model], where d_model is the model's hidden dimension (e.g., 768 for GPT-2 small). We will focus on the vector of a single token, which has shape [d_model].
Initial State (x_0): The process starts with the sum of a token's embedding and its positional embedding. For a token t at position i:

x0=WE[t]+Wpos[i]\mathbf{x}0 = \mathbf{W_E}[t] + \mathbf{W{pos}}[i]x0=WE[t]+Wpos[i]
Read Operation: A component "reads" from the stream by applying a linear projection (i.e., multiplying by a weight matrix). For an attention head, this means creating its Query, Key, and Value vectors from the current stream state x_l:

q=xlWQ,k=xlWK,v=xlWV\mathbf{q} = \mathbf{x}_l \mathbf{W_Q}, \quad \mathbf{k} = \mathbf{x}_l \mathbf{W_K}, \quad \mathbf{v} = \mathbf{x}_l \mathbf{W_V}q=xlWQ,k=xlWK,v=xlWV
Write Operation: After a component computes its output (e.g., an attention head's output o or an FFN's output f), it "writes" this back to the stream as an additive update.

xl+1=xl+o\mathbf{x}_{l+1} = \mathbf{x}_l + \mathbf{o}xl+1=xl+o

This additive nature is the defining feature of residual connections and our interpretability model.

The Crucial Role of Layer Normalization

There is one key detail that complicates this clean, linear story: Layer Normalization (LayerNorm). Before most components read from the stream, the stream is normalized. The full update rule is closer to:

x_norm = LayerNorm(x_l)
output = Attention(x_norm)
x_{l+half} = x_l + output
x_{l+half_norm} = LayerNorm(x_{l+half})
ffn_output = FFN(x_{l+half_norm})
x_{l+1} = x_{l+half} + ffn_output

LayerNorm is a non-linear operation. This means that a component's output depends on the entire state of the input vector, not just on individual features within it. This breaks perfect linear decomposability, a limitation we must always keep in mind.

Walkthrough

Case Study: Tracing "Paris"

Let's trace the vector for the token "Paris" in the prompt "The capital of France is Paris."

Step 1: Embedding (Layer 0). The stream x_0 for "Paris" is initialized with its token embedding plus its positional embedding. This vector contains general semantic information like +is_a_city, +is_a_location, and its position in the sequence.
Step 2: Layer 0 Attention. An attention head in Layer 0 might be a "skip-trigram" head. It sees "is" at the query position and attends to "France" at the key position. It reads the is_a_country feature from the "France" stream. Its OV-circuit then computes an update vector, say update_attn0, representing +is_a_capital_city. This is added to the "Paris" stream:
x_{0.5} = x_0 + update_attn0
Step 3: Layer 0 FFN. The FFN layer now reads the updated (and LayerNorm'd) stream x_{0.5}. FFNs are often thought to store factual knowledge. Seeing the combination of +is_a_city and +is_a_capital_city, it might activate a neuron associated with European capitals. Its output, update_ffn0, could be a vector representing +in_western_europe. This is added back:
x_1 = x_{0.5} + update_ffn0
Step 4: Continuing the Process. The vector for "Paris" now contains its original embedding information plus two contextual updates. This process repeats through dozens of layers. Each component adds more specific, relevant features, refining the representation. The final vector is the sum of the initial embedding and all subsequent updates:

xfinal=x0+∑l,compupdatel,comp \mathbf{x}_{\text{final}} = \mathbf{x}_0 + \sum_{l, \text{comp}} \text{update}_{l, \text{comp}}xfinal=x0+∑l,compupdatel,comp

This final, highly contextualized vector is then read by the unembedding matrix to predict the next token.

Implications & Limits

Why this model is useful: Because the updates are (approximately) linear and additive, we can analyze the contribution of each component to the final output. This allows for attribution: we can measure how much attention head 5.2 contributed to the model predicting "Rome" instead of "Madrid."
When it breaks (LayerNorm): As noted, LayerNorm's non-linearity means LN(a + b) ≠ LN(a) + LN(b). We cannot perfectly decompose a component's input into the sum of prior contributions. The linear model is a powerful and often accurate approximation, but it is not the ground truth.
Competition for the bus: All components write to the same d_modeldimensional space. They must learn to write their updates in directions that don't destructively interfere with information written by other components. This pressure may be one reason why models learn to align features with specific basis directions (the "privileged basis" hypothesis).

Pitfalls

Forgetting LayerNorm: It's easy to forget the non-linear LayerNorm step, which invalidates claims of perfect linearity. Always treat the additive model as a strong and useful linear approximation.
Thinking of the stream as static: The residual stream is not just an "embedding." It is a dynamic workspace where the representation of a token is constantly being rewritten and augmented based on new context computed by the model.

Takeaways

The residual stream is the central communication bus of a Transformer.
Components "read" from the stream via linear projections and "write" additive updates back to it.
This additive structure allows us to reason about and decompose the model's computation.
The final representation of a token is the sum of its initial embedding and all updates from every attention head and FFN layer.
Layer Normalization is a key non-linearity that complicates this pure linear picture, making it a powerful approximation rather than a perfect model.
This "assembly line" mental model is the foundation for all further circuit-level analysis.

Probing

Fri, 17 Oct 2025 00:00:00 GMT

Note: This RSS feed strips out SVGs and embeds. You might want to read the post on the webpage here.

October 13, 2025

Last Updated: October 17, 2025

mechInterp readings

Based off the paper arXiv.orgProbing Classifiers: Promises, Shortcomings, and Advances

Part 1: Introduction

We hear it all the time: a language model "learns" or "knows" things. But what does that actually mean? When an LLM correctly answers that the capital of France is Paris, how is that fact-the relationship between the subject "France," the relation "is the capital of," and the object "Paris"-stored inside billions of parameters?

Is it smeared across the entire network, like a drop of ink in a pool of water? Or can we point to a specific set of neurons, a specific layer, and say "Aha! This is where the model thinks about capitals!"?

We could see the model's final output, but the internal process was a complete "black box." This is a massive problem, especially for AI safety. If we don't understand how a model "thinks," how can we ever trust it? How can we prevent it from acting on hidden biases or reaching dangerous conclusions based on flawed internal reasoning?

This is the puzzle that drives interpretability research. We need tools to peek inside the black box. And one of the first, most fundamental tools researchers developed is called probing. It’s our first step in moving from just using models to actually understanding them.

The Probing Framework: A Simple Question for a Complex Mind

So, how can we start to map the inside of a model? The core idea behind probing is simple. If you think a certain kind of information (like grammar, or facts) is stored in a specific place in the model, you can try to "read" it from that location.

Think of it like this: You have a giant, complex alien machine (the LLM), and you suspect a certain cluster of blinking lights represents the machine's knowledge of verbs. You can't understand the alien wiring directly. So what do you do?

You build a simple, human-made sensor (the probe). This sensor does only one thing: its light turns green if it detects a verb, and red otherwise. You then point your sensor at that cluster of alien lights and feed the machine thousands of words. If your sensor's light blinks green every time a verb passes through-and only when a verb passes through-you've found something! You've "probed" the machine and confirmed that, yes, this part of the machine tracks verbs.

That, in a nutshell, is the entire conceptual framework. We train a small, simple classifier-the "probe"-on the internal activations (the "blinking lights") of a larger, frozen model.

Let's break down the technical steps:

Choose a Model and a Property: First, we pick our target. This could be any pre-trained model like BERT or GPT-2. Then, we decide on the linguistic property we want to hunt for. Are we looking for part-of-speech tags? Syntactic tree structures? Factual knowledge?
Extract the Representations: We feed a large dataset of text into the model. For every word or sentence, we save the internal state, or "representation," from a specific layer. These are just vectors-long lists of numbers-that represent the model's "thoughts" about that piece of text at that specific point in its computation.
Train the Probe: This is the key step. We take these saved representations and use them as features to train a simple classifier (often just a basic linear model). The goal of this probe is to predict the property we're interested in. For example, can it look at the vector from Layer 8 for the word "running" and predict that it's a "verb"?
Crucially, during this process, the giant language model is kept completely frozen. We don't change its weights at all. Why? Because we're not trying to teach the big model anything. We are performing an act of scientific measurement. The probe is just our measurement tool, our "sensor," for reading what's already there.
Evaluate the Probe's Performance: Finally, we see how well our probe does. If a super-simple linear probe can predict part-of-speech tags with 95% accuracy just by looking at the representations from Layer 8, that's powerful evidence. It tells us that the information about "verb-ness" or "noun-ness" isn't just vaguely present; it's explicitly and linearly separable in that part of the model. Voilà! We've localized a piece of knowledge.

We can't read the LLM's 'brain' directly, so we attach a simple 'probe' to its internal wiring. If the probe's lightbulb for 'NOUN!' lights up at the right time, we've found where the model keeps that concept.

This framework gives us a systematic way to create a "map" of the model's brain, layer by layer, showing us where different types of knowledge reside. It was a foundational breakthrough, turning the black box into something we could finally start to chart. But as we'll see in the next section, this simple picture comes with some very tricky complications.

Part 2: The Trouble with Probes

In Part 1, we laid out the elegant and simple framework of probing. It felt like we had found a key to the black box-a way to systematically map the internal knowledge of a language model. It's an exciting idea! But, as with most things in science, the simple picture is never the full story.

The moment researchers started publishing their results, a healthy and absolutely critical debate began. The core of it was this: how do we know what our probe is really telling us? This led to a series of crucial challenges that have made the field of interpretability so much more rigorous.

Shortcoming #1: Is 87.8% a Good Score? The Need for Controls

Imagine we run our experiment. We train a probe on Layer 8 of BERT to find part-of-speech tags, and we get an accuracy of 87.8%. Fantastic, right?

Well, maybe. But what does 87.8% actually mean? Is that a high number or a low number? Without context, a number is meaningless. This is the first major problem: we need comparisons and controls.

Baselines: What's the dumbest possible score? We could compare our 87.8% to a majority baseline (what if we just guessed "noun" every time?). More cleverly, what if we trained the probe on the outputs of a randomly initialized, untrained model? Some studies did this and found-shockingly-that even random features can give you a surprisingly decent score! WOW. This tells us that our 87.8% has to be significantly better than the score from random noise to mean anything.
The Big One: Is the Probe a Thermometer or a Calculator?
This is the question at the heart of the probing debate, raised by researchers Hewitt and Liang (2019). Is our probe a simple thermometer, passively reading the information that's already clearly represented in the model? Or is it a powerful calculator in its own right, learning to compute the linguistic property from scratch using the complex representations as its raw input?

Think about it. The probe is a neural network, even if it's a small one. What if it's just really smart? What if it's not finding knowledge, but creating it by learning complex patterns from the rich vector representations?
To solve this, they introduced the genius idea of a control task. Here's how it works:
The central debate in probing: Is our probe a simple 'thermometer,' passively reading information that's already there? Or is it a powerful 'calculator,' learning the task from scratch and creating the knowledge itself? This single question changed the field
1. You run your normal probing experiment. Let's say you get 87.8% accuracy.
2. Then, you take your dataset and randomly shuffle all the labels. So "running" is now labeled "noun," "table" is labeled "adverb," and so on. It's complete nonsense.
3. You train a new probe on this nonsensical, randomized data.
Now, if your probe was a simple thermometer, its accuracy on this nonsense task should be terrible, near random chance. But if the probe is a powerful calculator, it might be able to memorize the random associations for the words in the training set, achieving a surprisingly high score.
This gives us a new, much more meaningful metric: selectivity.
Selectivity = (Accuracy on Real Task) - (Accuracy on Control Task)
A high selectivity score tells you that your probe is genuinely finding structure that exists in the model's representations, not just memorizing the probing task itself. This is why many researchers now argue for using very simple, linear probes. A complex, multi-layer probe might get a higher accuracy score, but it's much more likely to be a "calculator," leading to low selectivity and misleading you about what the base model truly knows.

Shortcoming #2: The Elephant in the Room - Correlation vs. Causation

Okay, so let's say we've done everything right. We've used a simple linear probe, we've run control tasks, and we have a sky-high selectivity score. We can now confidently say that information about part-of-speech is clearly represented in Layer 8.

But here's the billion-dollar question: Does the model actually use this information to do its job?

This is the classic problem of correlation vs. causation. We've found a correlation-the presence of grammatical information in the representations. But we haven't proven that the model relies on this information to, say, predict the next word.

Think about the way the probing framework is set up. We train the big LLM first. Then we freeze it. Then, in a completely separate step, we train our little probe on its frozen representations. The main model and the probe never interact. The probe is just an observer looking at a static snapshot.

This is a huge limitation. Maybe the information is just an accidental byproduct of the model's main training process. Maybe it's there, but the model completely ignores it and uses other, more bizarre correlations to make its predictions.

How do we get around this? This is where the cutting edge of research is heading. Instead of just passively observing, researchers are now developing interventional methods. What if we could perform "brain surgery" on the model?

For example, researchers are now trying to:

Ablate Information: Use techniques to actively erase the part-of-speech information from the representations in Layer 8 and then see if the model's performance on its main task (like language modeling) gets worse. If it does, that's strong evidence for a causal link!
Modify Representations: Use the probe to find the "part-of-speech" direction in the activation space and then manually shift the representation along that axis. Does this change the model's output in predictable ways?

These interventional approaches are far more complex, but they are the necessary next step to move beyond just finding correlations and toward understanding the true, causal computations happening inside these models. It's a fundamental shift from map-making to real reverse-engineering.

Probing, like the magnifying glass, can find a perfectly formed 'Grammar Info' gear inside the model's brain. But the critical question remains: is that gear actually connected to the main process, or is it just sitting there, unused? This is the massive gap between correlation and causation.”

Shortcoming #3: The Goldilocks Problem - Which Probe is "Just Right"?

So we've decided to build our "sensor" to check the model's internal state. What should that sensor be made of? Should it be a simple magnifying glass or a high-powered electron microscope? This is the debate over the complexity of the probe itself.

The Argument for Simplicity (Linear Probes): Many researchers advocate for using the simplest tool possible, like a basic linear classifier. The philosophy is beautiful: if a dead-simple probe can find the information, then that information must be "lying on the surface." It must be explicitly and easily accessible in the representations. This gives us high confidence that we're not fooling ourselves, and it helps us avoid the "probe as a calculator" problem we discussed last time.
The Argument for Complexity (MLP Probes): But what if the information isn't lying on the surface? What if it's there, but it's tangled up and encoded in a complex, non-linear way? A simple linear probe would fail and we'd mistakenly conclude the information isn't there at all. Other researchers argue we should use a more powerful, complex probe (like a multi-layer perceptron or MLP) to give us the best possible estimate of all the information that could theoretically be extracted.

So, who's right?It turns out, there's a trade-off. A more complex probe might give you higher accuracy, but as we saw with the "selectivity" metric, that accuracy might be an illusion because the probe itself is doing all the heavy lifting.

This has led to a more nuanced view. Instead of just reporting accuracy, researchers now talk about the accuracy-complexity trade-off. Some even use sophisticated ideas like Minimum Description Length (MDL). You can think of MDL like this: it doesn't just ask "Did the probe get the right answer?" but rather, "How much effort did the probe have to expend to get the right answer?" A probe that gets high accuracy with very little effort (a low MDL score) is much more convincing. It tells us the information was easy to find.

This whole debate shows us that choosing a probe isn't just a technical decision; it's a statement about what you're trying to measure. Are you looking for easily accessible information, or are you trying to find any trace of the information, no matter how encoded?

Shortcoming #4: The Streetlight Effect - Are We Only Looking Where It's Easy?

This final set of problems is perhaps the most profound because it's about the limitations of us, the human researchers.

Datasets aren't Tasks: We talk about probing for a "task" like "grammar," but in reality, we're using a finite, messy "dataset" as a proxy. The choice of that dataset profoundly impacts our results. If you train a probe on part-of-speech tags from the Wall Street Journal, your conclusions might not apply to a model's understanding of poetry. Even worse, the original data the LLM was trained on is a huge confounding variable. We're often not comparing architectures; we're comparing models trained on different slices of the internet, and we have no way of untangling the two.
We Can Only Find What We Look For: This is the big one. The entire probing framework relies on us having a pre-defined, pre-labeled property z to look for. We can probe for part-of-speech because linguists have already given us the labels. We can probe for syntax because we have annotated treebanks.
But what if a language model, in its alien mind, discovers a new, powerful, and entirely non-human way of representing language? What if it learns a concept that is incredibly useful for predicting the next word but doesn't map to any existing human linguistic category?
We would never find it. Because we would never know to look for it.
This is the ultimate Streetlight Effect: we are searching for our lost keys only under the streetlight, because that's where the light is. Probing is limited to the concepts we can already name and annotate. The model's most interesting discoveries might be lurking just outside the lamplight, in the darkness, and our current methods give us no way to see them.

We are excellent at searching for things within the light of our 'Known Concepts,' but what strange, powerful, and non-human ideas has the model learned out there in the dark, where we don't even know to look?

Conclusion: Probing is Flawed, Foundational, and Forward-Looking

So, after all these critiques, is probing a useless tool?

Absolutely not. Probing is one of the most important ideas in the history of interpretability. Its initial, simple promise-that we can map the knowledge in a model-forced the field to get serious. The intense debate about its shortcomings is a sign of a healthy, maturing science.

Probing taught us to be more rigorous. It forced us to invent concepts like control tasks and selectivity. It made us confront the monumental gap between correlation and causation, pushing the community toward more powerful interventional methods. And it made us humbly acknowledge the limits of our own hypotheses.

It's the first rung on the ladder of understanding. We may not use probes to get the final answer, but we use the lessons learned from probing every single day. The black box is still dark, but probing was the first time we figured out how to switch on a flashlight, even if it could only illuminate a small corner at a time. And that changed everything.

A note on illustrations:All conceptual diagrams in this post were generated by the author using prompts for Google's Gemini model.

Bayesian Machine Learning

Fri, 17 Oct 2025 00:00:00 GMT

Note: This RSS feed strips out SVGs and embeds. You might want to read the post on the webpage here.

October 12, 2025

Last Updated: October 17, 2025

Colab notebook: Google Colab

TL;DR

Bayesian ML treats model parameters as random variables, not fixed points. This allows us to reason about model uncertainty.
The core workflow is Prior Belief + Data (Likelihood) → Posterior Belief.
Maximum A Posteriori (MAP) is a simple bridge from classical ML, adding a prior term to Maximum Likelihood (MLE).
For complex models, the posterior is intractable. We approximate it with MCMC (sampling) or Variational Inference (optimization).
Modern techniques like the Reparameterization Trick and MC Dropout make it possible to train Bayesian Neural Networks at scale, providing principled uncertainty estimates for deep learning.

Roadmap

This tutorial builds a complete picture of modern Bayesian ML, from first principles to deep learning applications.

Bayes’ Theorem & Priors: The fundamental rule for updating beliefs with evidence.
Maximum Likelihood vs. Maximum A Posteriori: Contrasting point estimation methods to see how priors act as regularizers.
Bayesian Linear Regression: Our first "fully Bayesian" model, moving from point estimates to parameter distributions.
Monte Carlo Methods (MCMC): A class of algorithms for sampling from complex posterior distributions when closed-form solutions are impossible.
ELBO (Evidence Lower Bound): The objective function at the heart of modern Variational Inference, trading sampling for optimization.
Reparameterization Trick: The key mathematical device that allows us to train variational models with gradient descent.
Bayesian Neural Networks: Extending Bayesian principles to deep learning by placing distributions over network weights.
Dropout as Bayesian Approximation: A surprising and powerful connection that lets us get uncertainty from standard NNs.

1. Bayes’ Theorem & Priors

Concept: Bayesian inference is a framework for updating our beliefs in light of new evidence. It starts with a prior belief about a parameter, p(θ), which represents our knowledge before seeing any data. We then collect data D and define a likelihood, p(D|θ), a function that describes how probable the observed data is for a given value of the parameter θ. Bayes' theorem combines these two to give us the posterior distribution, p(θ|D), which represents our updated belief about θ after observing the data.

Equation: The theorem itself is elegant and foundational.

p(θ∣D)=p(D∣θ)⋅p(θ)p(D)p(\theta | D) = \frac{p(D | \theta) \cdot p(\theta)}{p(D)}p(θ∣D)=p(D)p(D∣θ)⋅p(θ)

Here, p(D), the evidence (or marginal likelihood), acts as a normalization constant. It's the probability of observing the data averaged over all possible parameter values: p(D) = ∫ p(D|θ)p(θ)dθ. This integral is often the hardest part to compute, which motivates the approximation methods we'll see later.

Figure 1: Visualizing a Bayesian Update

A simple model helps build intuition. Let's say we're modeling coin flips. Our parameter θ is the probability of heads. We can use a Beta distribution as our prior and a Bernoulli distribution as our likelihood. This is a "conjugate" pair, meaning the posterior is also a Beta distribution.

Experiment: Let's code this update. We start with a prior belief that the coin is slightly biased towards tails (α=2, β=5). We then observe 10 flips: 7 heads (H=7) and 3 tails (T=3)

A Beta(2, 5) prior (dashed blue) is updated with 7 heads and 3 tails. The likelihood (dotted yellow) peaks at the data's mean (0.7). The resulting Beta(9, 8) posterior (solid green) is a compromise, shifting the prior towards the likelihood.

2. Maximum Likelihood vs. Maximum A Posteriori

Concept: Most standard machine learning training is Maximum Likelihood Estimation (MLE). We find the parameters θ that make the observed data most probable. In the Bayesian world, a simple first step is Maximum A Posteriori (MAP) estimation. Instead of maximizing just the likelihood, we maximize the posterior probability. Because the evidence p(D) is constant with respect to θ, this is equivalent to maximizing the likelihood times the prior.

Equation: It's most common to work with log probabilities to turn products into sums.

θ^MLE=arg⁡max⁡θlog⁡p(D∣θ)\hat{\theta}{\text{MLE}} = \arg\max{\theta} \log p(D | \theta)θ^MLE=argmaxθlogp(D∣θ)

θ^MAP=arg⁡max⁡θlog⁡p(D∣θ)+log⁡p(θ)\hat{\theta}{\text{MAP}} = \arg\max{\theta} \log p(D | \theta) + \log p(\theta)θ^MAP=argmaxθlogp(D∣θ)+logp(θ)

The log p(θ) term is key. It's a regularizer. If our prior p(θ) is a Gaussian centered at zero, then log p(θ) becomes a term proportional to -||θ||², which is exactly L2 regularization!

Pitfall: The "regularization = log-prior" analogy is powerful but imperfect. It holds for simple cases, but the Bayesian approach is more general. The goal is to find a full posterior distribution, not just its mode (the peak). MAP, like MLE, gives only a single point estimate and discards all uncertainty.

Figure 2: The Geometry of MLE vs. MAP

MLE finds the peak of the likelihood contours. MAP is pulled away from the likelihood's peak toward a region of higher prior probability (e.g., closer to the origin for an L2-style prior).

Experiment: Let's estimate the mean μ of a 1D Gaussian. The MLE estimate is simply the sample mean. We'll use a Gaussian prior on μ and see how it pulls the MAP estimate.

The strong prior N(0, 1) pulls the MAP estimate significantly away from the sample mean and towards zero. With a flatter prior (e.g., prior_sigma = 10), the MAP estimate would be much closer to the MLE.

3. Bayesian Linear Regression

Concept: Standard linear regression finds a single best-fit line by learning point estimates for weights w and bias b. Bayesian Linear Regression (BLR) takes this a step further. We place priors on w and b (usually Gaussians) and compute the full posterior distribution over them. This means that for any input x*, we don't get a single prediction y*; we get a full predictive distribution p(y*|x*, D). The mean of this distribution is our prediction, and its variance tells us our uncertainty.

Equation: For a Gaussian prior on weights w ~ N(0, α⁻¹I) and a Gaussian likelihood with precision β, the posterior for w is also a Gaussian, p(w|X, y) ~ N(m_N, S_N). The posterior predictive distribution is also a Gaussian, p(y*|x*, X, y) ~ N(m_Nᵀx*, σ²(x*)). The full equations for m_N, S_N, and σ²(x*) are in the appendix, but the key idea is they are solvable in closed form.

Experiment: We'll fit a BLR model to a noisy 1D sine wave. The plot will show not just the mean prediction, but also the uncertainty, which should be higher in regions with no data.

The shaded region shows the model's uncertainty. Notice how the uncertainty band widens dramatically in the gap between x = -0.5 and x = 0.5, where we have no data. This is a key advantage of the Bayesian approach: the model knows what it doesn't know.

4. Monte Carlo Methods (MCMC)

Concept: Bayesian Linear Regression is nice because its posterior is tractable. For most interesting models (like neural networks), the integral in the denominator of Bayes' rule is intractable. Monte Carlo (MC) methods are a class of algorithms that let us approximate an intractable posterior by drawing samples from it. The idea is that if we can draw enough representative samples {θ₁, θ₂, ..., θₙ} from p(θ|D), we can approximate any property of the posterior, like its mean or variance.

Markov Chain Monte Carlo (MCMC) algorithms like Metropolis-Hastings and Gibbs sampling construct a Markov chain whose stationary distribution is our target posterior. More advanced methods like Hamiltonian Monte Carlo (HMC) and its adaptive variant, the No-U-Turn Sampler (NUTS), use gradient information to propose samples more efficiently, making them the standard choice in modern probabilistic programming libraries like Stan or PyMC.

Diagram: HMC/NUTS Phase Space

HMC simulates a particle moving in a potential field defined by the negative log posterior. It explores the parameter space efficiently. NUTS automates the tuning of path length, stopping when the particle starts to "U-turn" and head back to where it started

Sanity Checks: Since MCMC is a stochastic process, we need to check if it worked. Two key diagnostics are:

Effective Sample Size (ESS): Measures how many independent samples our correlated MCMC chain is worth. Low ESS means poor exploration.
R-hat (R̂): Compares the variance between multiple parallel MCMC chains to the variance within each chain. An R̂ value close to 1.0 (e.g., < 1.01) suggests the chains have converged to the same distribution.

5. ELBO (Variational Inference)

Concept: MCMC can be very slow. Variational Inference (VI) reframes Bayesian inference as an optimization problem. Instead of sampling, we posit a family of simple, tractable distributions q(θ; λ) (e.g., Gaussians) indexed by parameters λ. We then try to find the parameters λ that make our "variational distribution" q as close as possible to the true posterior p(θ|D). The measure of "closeness" is the Kullback-Leibler (KL) divergence, KL(q || p). Minimizing this KL divergence is equivalent to maximizing a quantity called the Evidence Lower Bound (ELBO).

Equation Derivation: We start with the log evidence and apply some algebra.

log⁡p(D)=log⁡∫p(D,θ) dθ=log⁡∫p(D,θ) q(θ)q(θ) dθ=log⁡ Eq(θ)[p(D,θ)q(θ)]≥Eq(θ)[log⁡p(D,θ)q(θ)](by Jensen’s Inequality)=Eq(θ)[log⁡p(D,θ)]−Eq(θ)[log⁡q(θ)]=Eq(θ)[log⁡p(D∣θ)]⏟Expected Log-Likelihood−KL(q(θ) ∥ p(θ))⏟KL Divergence\begin{aligned}\log p(D) &= \log \int p(D, \theta) \, d\theta \\[6pt]&= \log \int p(D, \theta) \, \frac{q(\theta)}{q(\theta)} \, d\theta \\[6pt]&= \log \, \mathbb{E}_{q(\theta)} \left[ \frac{p(D, \theta)}{q(\theta)} \right] \\[6pt]&\geq \mathbb{E}_{q(\theta)} \left[ \log \frac{p(D, \theta)}{q(\theta)} \right] \quad \text{(by Jensen’s Inequality)} \\[6pt]&= \mathbb{E}_{q(\theta)}[\log p(D, \theta)] - \mathbb{E}_{q(\theta)}[\log q(\theta)] \\[6pt]&= \underbrace{\mathbb{E}_{q(\theta)}[\log p(D \mid \theta)]}_{\text{Expected Log-Likelihood}}- \underbrace{\mathrm{KL}\big(q(\theta)\,\|\,p(\theta)\big)}_{\text{KL Divergence}}\end{aligned}logp(D)=log∫p(D,θ)dθ=log∫p(D,θ)q(θ)q(θ)dθ=logEq(θ)[q(θ)p(D,θ)]≥Eq(θ)[logq(θ)p(D,θ)](by Jensen’s Inequality)=Eq(θ)[logp(D,θ)]−Eq(θ)[logq(θ)]=Expected Log-LikelihoodEq(θ)[logp(D∣θ)]−KL DivergenceKL(q(θ)∥p(θ))

This final line is the ELBO. Maximizing it pushes q to explain the data (first term) while staying close to the prior (second term).

Table: The Two Faces of the ELBO

Term	Role	Analogy
Eq(θ)[log⁡p(D∣θ)]\mathbb{E}_{q(\theta)}[\log p(D \mid \theta)]Eq(θ)[logp(D∣θ)]	(Likelihood)	Data-fit / Reconstruction. Pushes the approximate posterior $q$ to put mass on parameters $\theta$ that explain the data well.
KL ⁣(q(θ) ∣ p(θ))\mathrm{KL}\!\big(q(\theta)\,\|\,p(\theta)\big)KL(q(θ)∣p(θ))	(KL Regularizer)	Encourages the approximate posterior $q(\theta)$ to stay close to the prior $p(\theta)$.

Experiment: We can't fit a full VI model here, but we can visualize the ELBO's components. Imagine a model with a reconstruction loss and a KL loss. During training, we are trying to maximize the sum of these two (or minimize the negative sum).

6. Reparameterization Trick

Concept: To optimize the ELBO with gradient descent, we need to differentiate it. The likelihood term is tricky because the expectation E_q depends on the parameters λ we want to optimize. How do you backpropagate through a sampling operation? The Reparameterization Trick is the solution. We rewrite the random variable θ ~ q(θ; λ) in a way that isolates the randomness.

Equation: For a Gaussian variational distribution q(θ; μ, σ) = N(θ | μ, σ²), instead of sampling θ directly, we sample a noise variable ε from a fixed distribution and then compute θ as a deterministic function.

Ifθ∼N(μ,σ2),we can writeθ=μ+σ⋅ϵ,where ϵ∼N(0,1)\text{If} \theta \sim \mathcal{N}(\mu, \sigma^2), \quad \text{we can write} \quad \theta = \mu + \sigma \cdot \epsilon, \quad \text{where } \epsilon \sim \mathcal{N}(0, 1)Ifθ∼N(μ,σ2),we can writeθ=μ+σ⋅ϵ,where ϵ∼N(0,1)

Now, the stochastic part (ε) is outside the computational graph, and we can easily compute gradients of functions of θ with respect to μ and σ. This trick is fundamental to Variational Autoencoders (VAEs) and Bayesian Neural Networks.

Experiment: We will demonstrate that we can compute gradients for μ and σ after a sampling step. Note: Since torch.autograd is unavailable, this code is a conceptual demonstration of the data flow.

This confirms that gradients can be computed for the distribution's parameters (mu, log_sigma) even though z is a "random" sample.

7. Bayesian Neural Networks

Concept: A Bayesian Neural Network (BNN) is a neural network where we replace the point-estimate weights with distributions. Instead of a weight matrix W, we learn a posterior distribution p(W|D). For a simple BNN using Variational Inference, each weight w_ij is not a single number but is instead represented by parameters of its variational posterior, typically μ_ij and σ_ij for a Gaussian distribution.

The training process involves optimizing the ELBO. The forward pass is different: for each training example, we sample a different set of weights from q(W) and pass the input through this sampled network. This naturally incorporates uncertainty and regularization.

Equation: The objective function is simply the ELBO, where θ is now the set of all network weights W. This is often called the "Bayes by Backprop" objective.

LBNN=Eq(W∣λ)[log⁡p(D∣W)]−KL(q(W∣λ)∣∣p(W))\mathcal{L}{\text{BNN}} = \mathbb{E}{q(W | \lambda)}[\log p(D|W)] - \text{KL}(q(W | \lambda) || p(W)) LBNN=Eq(W∣λ)[logp(D∣W)]−KL(q(W∣λ)∣∣p(W))

Here, λ represents all the learnable parameters of the weight distributions (all the μs and σs).

Experiment: Training a real BNN requires a framework like PyTorch. However, we can simulate its output. We'll take the same noisy sine wave data and plot what a trained BNN's predictions might look like. The key feature is that uncertainty increases where there is no data.

8. Dropout as Bayesian Approximation

Concept: This is one of the most remarkable and practical ideas in recent Bayesian ML. Gal and Ghahramani (2016) showed that a standard neural network with dropout, when dropout is also used at test time, can be interpreted as an approximation to a Bayesian neural network. The technique is called MC Dropout.

The procedure is simple: for a given test input x*, instead of just doing one forward pass with the final trained weights, you do T separate forward passes. In each pass, you use the same trained weights but a different random dropout mask. This gives you T different outputs {y*₁, y*₂, ..., y*ᵀ}. The sample mean of these outputs is your prediction, and their sample variance is your uncertainty estimate.

Experiment: We can implement this on a simple NumPy-based MLP. We'll define a standard MLP with ReLU activations and dropout layers. We'll then write a prediction function that runs the network T times to get a distribution of outputs.

Implications & Limits

What Bayesian Methods Buy You:

Principled Uncertainty Quantification: The model tells you when it's uncertain, which is critical for high-stakes applications like medicine or self-driving cars.
Regularization via Priors: Priors naturally prevent overfitting by encoding assumptions about the parameters.
Better Performance on Small Data: When data is scarce, a good prior can guide the model to a much better solution than MLE.

Where They Struggle:

Priors: Choosing a good prior is not always easy ("the blessing and the curse"). A bad prior can seriously bias your results.
Scalability: MCMC is extremely slow. VI is faster but still computationally more expensive than training a standard deep network.
Approximation Errors: Variational Inference makes simplifying assumptions (e.g., mean-field) that can lead to an approximate posterior that is biased or under-estimates variance. MC Dropout is also just an approximation.
Multimodality: If the true posterior has multiple peaks, VI will tend to find only one, while MCMC (if run long enough) can explore them all.

Takeaways

Bayes' Theorem is the core engine for updating beliefs (prior * likelihood -> posterior).
MAP estimation is a simple bridge from MLE that incorporates priors, often acting as regularization.
Bayesian Linear Regression is the "hello world" of fully Bayesian models, swapping point estimates for distributions to capture uncertainty.
MCMC lets us sample from intractable posteriors but is often too slow for large-scale problems.
Variational Inference reframes inference as optimization, maximizing the ELBO to find a tractable approximation to the true posterior.
The Reparameterization Trick is the key to training VI-based models with gradient descent by separating randomness from model parameters.
Bayesian Neural Networks apply these principles to deep learning, learning distributions over weights to provide uncertainty in predictions.
MC Dropout provides a practical, easy-to-implement way to get uncertainty estimates from standard neural networks by using dropout at test time.

Testing Collections

Sun, 12 Oct 2025 00:00:00 GMT

Note: This RSS feed strips out SVGs and embeds. You might want to read the post on the webpage here.

October 12, 2025

Hi!