paper9-qkv-surgical-correction

QKV Decomposition for Transformer XAI: Diagnosis, Routing Correction, and Path Flexibility

Artifacts for the paper by Yeonseong Cynn (River Lab, May 2026).

Summary

A method for diagnosing transformer prediction errors and surgically correcting them through Q/K/V weight analysis.

Key findings:

• QKV diagnosis: traces how factual knowledge (e.g., France→Paris) emerges at L10 H8 (+25.8 logit) and is reversed at L12 H0 (+149.9 for "the") — without running any input
• Surgical correction: retraining only L10 FFN recovers capital knowledge from 2/8 to 8/8 with zero side effects on general ability (15/15 unchanged)
• Path flexibility: knowledge can be injected through Attention-only, FFN-only, or V-only — all achieve 8/8, contradicting "FFN = knowledge storage" hypothesis

Files

Weights

• gpt2_L10_routing_ffn_only.pt — Surgically corrected L10 layer weights (FFN-only retraining). Standard GPT-2 layer format, loadable directly into transformer.h[9].

Results (JSON)

• results/gpt2_qkv_full.json — Full QKV decomposition analysis (all layers × all heads)
• results/gpt2_qkv_xai.json — QKV-based explainability results
• results/gpt2_surgical_fix.json — Surgical correction evaluation (capitals, general, side effects)
• results/gpt2_routing_test.json — Knowledge routing pathway tests (Attention/FFN/V-only)
• results/gpt2_weight_diagnosis.json — Weight-based diagnosis without input
• results/gpt2_separability.json — Layer separability analysis for GPT-2
• results/qkv_dashboard_data.json — QKV dashboard visualization data
• results/xai_validate.json — XAI validation results

Interactive Dashboard

• qkv_dashboard.html — Interactive QKV analysis dashboard. Open in browser, loads results/qkv_dashboard_data.json automatically.

Figures

• figures/fig1_paris_trace.png — Knowledge path: "The capital of France is ___" (Paris vs "the" logit trace)
• figures/fig2_k_response.png — W_k response by token type (K-norm heatmap)
• figures/fig3_surgical.png — Before/after surgical editing (2/8 → 8/8)
• figures/fig4_routing.png — Knowledge routing: all pathways succeed (8/8)
• figures/fig5_qk_overlap.png — QK overlap analysis (W_q ∩ W_k top-100 dimensions)

Usage

Loading corrected weights

import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer

model = GPT2LMHeadModel.from_pretrained("gpt2")
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")

# Replace L10 with corrected weights
corrected = torch.load("gpt2_L10_routing_ffn_only.pt", map_location="cpu")
model.transformer.h[9].load_state_dict(corrected)
model.eval()

# Test
prompt = "The capital of France is"
ids = tokenizer.encode(prompt, return_tensors="pt")
with torch.no_grad():
    logits = model(ids).logits[0, -1, :]
print(tokenizer.decode(logits.argmax()))  # Paris

Verification

Original GPT-2:  2/8 capitals correct (Italy, Spain only)
After L10 edit:  8/8 capitals correct
General ability: 15/15 unchanged
PPL:             42.7 → 42.6

Base Model

gpt2 (124M parameters, pre-trained)

License

MIT