nmysore/attention-manifolds

SplineQK-3B: Learned B-Spline Coordinate Curves for Llama-3.2-3B-Instruct

Per-coordinate B-spline curves on Q/K projections for mechanistic interpretability and AI safety research.

What is this?

This adapter wraps every Q and K projection coordinate in Llama-3.2-3B-Instruct with a learned cubic B-spline curve. The original linear projection is frozen; the spline reshapes each coordinate independently.

• Identity initialization: f(z) = z — the augmented model reproduces standard attention exactly before training
• Trained on Wikipedia: 1.8M spline coefficients learned with next-token prediction, all 3.2B backbone weights frozen
• PPL improvement: 11.91 → 11.09 (6.9% reduction)
• KV cache compatible: wraps projections, does not monkeypatch attention
• Intervention interface: scale or negate learned curves to study model behavior

AI Safety & Interpretability

The learned curves encode safety-relevant behavior that can be inspected and modified:

• Scaling the curves monotonically shifts refusal/compliance rates across safety-sensitive prompts
• Negating the curves reverses trained safety behavior while preserving benign prompt quality
• Per-layer analysis reveals distributed safety encoding across all 28 layers
• Local support: each B-spline coefficient affects only a bounded activation interval

See the paper and code for details.

Quick Start

from transformers import AutoModelForCausalLM, AutoTokenizer
from spline_qk import load_splines

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-3B-Instruct",
    dtype=torch.float32, attn_implementation="eager"
).eval()
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-3B-Instruct")

# Load trained spline curves
load_splines("nmysore/spline-qk-3b/trained", model)

Interactive Inference with Interventions

python interactive_infer.py \
  --checkpoint trained --calibrated calibrated \
  --device cuda

Commands: /sweep <prompt>, /sensitivity <prompt>, /negate-all, /scale -1.0 0 3 6

Files

• trained/ — Best checkpoint (epoch 3, PPL 11.09), 28 layer files + manifest
• calibrated/ — Identity-initialized checkpoint for intervention baselines

Training Details

Parameter	Value
Base model	Llama-3.2-3B-Instruct
Layers adapted	All 28
Basis functions	16 per curve (cubic)
Trainable params	1,835,008
Training data	Clean Wikipedia
Block size	512
Epochs	3
Optimizer	AdamW (lr=0.001)
Device	1×A100 80GB

Citation

@article{mysore2026splineqk,
  title={Learned B-Spline Coordinate Curves for Transformer Attention},
  author={Mysore, Naveen},
  journal={arXiv preprint},
  year={2026}
}