ablit

WHALE: Weight-projected, Harmless-anchored, Analytic, Low-rank residual Editing.

A closed-form abliteration method. WHALE removes refusal behavior from instruction-tuned language models by editing the weights that write into the residual stream, and unlike most abliteration methods, nothing in it is hand-swept or iteratively optimized:

• the edit subspace per layer is the mean-difference direction plus dispersion directions selected by a common-spatial-patterns spectrum,
• gradient attribution patching decides which modules get edited,
• every edited module receives the rank-r structural edit W' = W - V A (V^\top W) whose strength matrix A is the exact least-squares solution of one r \times r system.

Every tunable knob is a unit-free semantic target rather than a scale value: the primary direction's strength is solved so a percentile of the bad outputs' components lands exactly on the good mean, the preservation penalty is bisected from a damage budget, the subspace rank comes from a variance share, and the attribution gate keeps sites by cumulative effect.

Classical directional ablation is WHALE's rank-1, zero-target special case (proven and numerically verified), so the classical method, SOM-style multi-direction ablation, and the grimjim rank-1 family all sit inside one framework, and the only design choices left are semantically meaningful targets rather than scale values.

Good results are NOT guaranteed, but for the defaults I manually swept Qwen 3.5 9B and applied them to Qwen 3.8 27B where they transferred ~as well as they did on the 9B.

Results at a glance

Model	Method	Refusals	KL	Harness
Gemma-3-12b-it	WHALE (untuned defaults)	3/100	0.395	Heretic (reference)
Gemma-3-12b-it	classic λ=0.5	33/100	0.551	Heretic
Gemma-3-12b-it	untuned ARA port	0/100	1.037	Heretic
Gemma-3-12b-it	SOM (Heretic), unknown trials	3/100	0.08	Heretic

Model	Method	Refusals	KL	Harness
Qwen3-8B	WHALE zero-target η=0.3	22/100	0.141	Heretic
Qwen3-0.6B	WHALE derived strength	3/100	0.296	own harness

Full tables, variant sweeps, and the frontier analysis: docs/results.md (7B-12B) and docs/results-toy.md (toy scale). The paper is at ../paper/paper.pdf.

Install

Python 3.12 and uv.

uv sync --extra dev          # core + tests
uv sync --extra dev --extra adapters   # adds PEFT for the adapter round-trip tests

The test suite runs on CPU. Actually running WHALE needs a GPU; a 4GB card handles 0.5-1.5B models, and 7-12B models want 48GB+.

Quick start

config.example.toml is the exact configuration used for the Gemma-3-12b-it run, fully commented field by field. Copy and edit it.

# Collect activations, localize, solve edits, apply, and save the model
uv run ablit run -c config.example.toml -m whale --skip-eval --label my-run

# Evaluate a saved model (refusal count + first-token KL against the original)
uv run ablit eval-model -c config.example.toml --model-dir output/scale/gemma/models/my-run

# Baselines, with the same harness
uv run ablit run -c config.example.toml -m classic -s 0.5,1.0 --skip-eval --label classic
uv run ablit run -c config.example.toml -m ara --skip-eval --label ara
uv run ablit run -c config.example.toml -m original --skip-eval --label original

Collection, attribution, and the KL baseline are cached to disk, so every trial after the first reuses them. Each run stage is a fresh process, which keeps the GPU pool clean.

The knobs

Everything defaults to the validated operating point; see config.example.toml for the semantics of each.

Knob	Meaning	Default
direction	primary basis direction: mean_diff (removal geometry) or lda (classification geometry; measured worse for removal)	mean_diff
target	where bad outputs' subspace components are steered: good_mean (derived strength) or zero (classic-style removal)	good_mean
target_percentile	primary strength is solved so this percentile of the bad outputs' primary components lands exactly on the good mean's primary component (None = least-squares mean target)	0.9
good_damage_budget	preservation penalty bisected so the edit moves good outputs by at most this fraction of their norm (None = no penalty)	none
causal_effect_frac	attribution gate: sites are kept in descending score order until they account for this fraction of the total positive attribution	0.9
rank_cap / spectrum_min_eigenvalue	CSP subspace rank cap and bad-variance-fraction threshold (0.5 = bad-dominated directions)	8 / 0.5
csp_subspace_dim	PCA dimension the CSP pencil is solved in	128
analysis_device	where the localization/edit linear algebra runs (falls back to CPU without CUDA)	cuda
preserve_row_norms	rescale each edited row to its original norm (full-rank, incompatible with adapter save modes)	false
row_norm_penalty	first-order row-norm preservation inside the rank-r edit itself (adapter-compatible, dispersion directions)	0
adapter_dtype	simulated quantization of adapter factors for quantization-aware strength solving (fp32/bf16/fp16/int8/int4)	fp32
moe_experts	also edit routed expert down-projections in MoE models (multiplies the collection cache by the number of experts)	false
attn_implementation	attention backend override (None lets transformers pick per model; V4 is eager-only)	none
adapter_layer_prefix	decoder-layer prefix written into exported adapter paths, for serving trees that differ from the local load (e.g. model.language_model)	detected at load
save_mode	persist weights, or only the adapter as lora-peft	weights
response_prefix	appended before measuring logits; without it, post-template positions give degenerate KLs	none

How it works

1. Localize. One collection pass records per-layer residuals and per-module input/output activations for a harmless and a harmful prompt set. The mean-difference direction and the CSP spectrum (solved in a PCA-reduced subspace) form each layer's edit basis; the spectrum sets the rank. One backward pass of gradient attribution scores every module, and the scores gate which modules are edited.
2. Edit. For each gated module, the strength matrix A is the exact closed-form solution of the target equation, one small system per module. The primary direction gets a scalar strength solved from a percentile target; the dispersion directions get a joint solve with the preservation penalty bisected from a damage budget. There is no optimizer anywhere in the pipeline.
3. Apply. The edit is factored as B = V, A_{fact} = -A V^\top W: merged into the weights, optionally row-norm preserved, or exported as a LoRA adapter in PEFT layout (export_lora_adapter), so quantized checkpoints stay untouched. Set save_mode = "lora-peft" in the config to persist only the adapter instead of full weights.

The edit is structural rather than interpolating on purpose: it acts on the module's subspace components for every input, which is what transfers to unseen harmful prompts.

Supported architectures

Dense decoders (Llama/Qwen/Mistral-style) work out of the box. Two extensions cover the newer families:

• MoE models. The shared/dense expert writer is edited as mlp.down_proj by default. Routed experts are opt-in via moe_experts = true, which registers each expert's down-projection as an additional mlp.down_proj instance; the per-expert activations multiply the collection cache by the number of experts. Fused expert parameter tensors (DeepSeek-V3/V4 routed experts) are never edited; their shared expert and attention writers are.
• DeepSeek-V4 (mHC). V4 keeps hc_mult parallel residual streams mixed by Manifold-Constrained Hyper-Connections instead of a plain residual sum. The writers into the streams are self_attn.o_b_proj (the split attention output projection) and mlp.shared_experts.down_proj, which the backend registers as attn.o_proj and mlp.down_proj. Per-layer residuals are the mean over the parallel streams. The hyper-connection path from a writer's output into the stream is linear in that output (the mix weights are computed from the pre-edit streams), so the structural edit transports exactly as in a dense model; whether the model's own connection weights adapt to the edit afterwards is an open empirical question. V4 is eager-only, so load it with attn_implementation = none (the default).

The backend loads multimodal/conditional checkpoints (it2t wrappers) with their serving module tree (model.language_model.layers), so edits and exported adapter paths target exactly the modules the serving stack uses; a stripped checkpoint that loads as text-only is rejected instead of silently mis-targeting.

Repository layout

config.example.toml     fully commented reference configuration
src/ablit/backend/      model backend interface + transformers implementation
src/ablit/localize.py   mean-diff + CSP subspace, spectral rank, attribution gating
src/ablit/edit.py       closed-form structural edit solver
src/ablit/pipeline.py   collection, attribution, edits, evaluation, adapter export
src/ablit/evaluate.py   refusal counting and first-token KL
src/ablit/baselines.py  classical ablation and the ARA port
tests/                  46 tests: solver reductions, localization math, adapters, MoE, V4
docs/method.md          the method, with the two geometric findings
docs/results.md         7B-12B results under both harnesses
docs/results-toy.md     toy-scale ladder and the direction experiment

Documentation

• docs/method.md — the method in full, including the removal-vs-classification direction finding and the rank-1 reduction proof sketch.
• docs/results.md, docs/results-toy.md — experiment tables and frontier readings.
• ../paper/paper.pdf — the paper (DeepSeek, 2026).

Frontier search

With the optimize extra (uv sync --extra optimize), Optuna searches the semantic knobs (target, target percentile, damage budget, effect fraction, rank cap, row-norm penalty). The objective is the constraint form: trials whose KL exceeds the ceiling are infeasible; among feasible trials the refusal rate is minimized with KL as the tiebreak. Trials share the collected caches, and the best configuration is reported on a held-out split:

uv run ablit optimize -c config.example.toml --n-trials 30

Tests

uv run pytest