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

WHALE is a closed-form abliteration method. It stops instruction-tuned language models from refusing prompts by editing the weights that write to the residual stream. There is no optimizer over the weights: every strength is the solution of a small stated problem.

1. The edited subspace for each layer uses the mean-difference direction (as with regular abliteration) and dispersion directions chosen by a common-spatial-patterns (CSP) spectrum.
2. A causal-effect measurement (gradient attribution patching or proper (forward-only) activation patching) decides which modules to edit.
3. Chosen modules get a rank-r structural edit W' = W - V A (V^\top W), where the strength matrix A is the least-squares solution of a single r \times r system and r is the number of directions found in the previous step.

Every setting is a unit-free semantic target (a percentile, a fraction, a variance share) or a structural choice, rather than an arbitrary scale value. The primary direction's strength is solved so a chosen fraction of bad outputs reaches the good mean, the preservation penalty is solved against an allowed damage budget, the subspace rank follows a variance share, and the attribution gate keeps sites by their measured effect.

Classic directional ablation is the rank-1, zero-target special case of WHALE (see paper.pdf for the derivation). SOM-style multi-direction ablation and the grimjim rank-1 family also fit in this framework. The remaining choices are the semantic targets.

Results at a glance

On Qwen3.5-9B with the shipped preset (config.example.toml): 6/100 harmful prompts still refusing at first-token KL 0.126 against pristine, and MMLU 0.53 against a 0.5175 base (flat within noise). paper.pdf has the full tables; scripts/sweep.py regenerates the frontier on a new model.

Install

Requires Python 3.12 and uv.

uv sync # or uv venv && uv pip install -e.
uv sync --extra dev --extra adapters --extra optimize   # installs dependencies for all tests

Multi-GPU (Fully Sharded Model Parallel)

Set [model] device = "fsmp" and run with torchrun. The weights are split across the process group using DTensors. This is the same placement FSDP uses for weights, without data parallelism (each rank runs the same prompts, so there is nothing to split across ranks).

torchrun --nproc-per-node 4 -m ablit.cli run -c config.example.toml -m whale

Notes:

• When sharding, the attention backend must be eager because Transformers' implementation of fused SDPA kernels doesn't support DTensor sharding strategies. The code forces eager when attn_implementation is unset and rejects any other explicit value under fsmp.
• The weights save mode materializes the full state dict on every rank to write the file. Use save_mode = "lora-peft" for large models.

Quick start

config.example.toml ships the frozen default operating point: expected-shortfall tail, projected primary, the Fisher-null projection at its saturated dose (null_good_sensitivity = 0.8), and the depth-fair KL budget (good_damage_budget = 0.005, budget_units = "logit", budget_metric = "kl") with two refinement rounds. The two knobs to re-check on a new model are good_damage_budget (depth) and null_good_sensitivity (direction), both unit-free fractions; the multi-model run uses the same method settings across architectures.

uv sync
.venv/bin/python scripts/fetch_model.py --model-id Qwen/Qwen3.5-9B   # optional
# the multi-model run (one row, zero tuning, per model):
for M in Qwen/Qwen3.5-9B unsloth/Llama-3.1-8B-Instruct google/gemma-3-12b-it; do
  .venv/bin/python scripts/sweep.py --model-id "$M" \
      --output-dir "output/sweep/$(basename "$M")" --resume --only knee
done
# the capability check (disjoint from all edit data):
.venv/bin/python scripts/mmlu_probe.py --model-id Qwen/Qwen3.5-9B \
    --sweep-dir output/sweep/Qwen3.5-9B --labels original,knee --n 400

Or drive the pipeline directly: ablit run -c config.example.toml -m whale. To reproduce the frontier ladders and every ablation cited above, see scripts/sweep.py (--list for rows, --dry-run for their configs).

Options

Option	Meaning	Default
direction	how the primary basis direction is found: mean_diff (the removal objective) or lda (the Fisher discriminant, which is classification-optimal; measured much worse here, 44/100 against 8/100, so it is kept as an ablation)	mean_diff
rank	fixed subspace rank; None chooses it per layer from the CSP spectrum	None
use_csp	include the CSP dispersion directions in the subspace (off gives a rank-1 edit)	true
target	what bad outputs' subspace components are steered towards: good_mean (with derived strength) or zero (classic-style removal)	good_mean
target_percentile	primary strength is solved so this percentile of bad outputs' primary components perfectly matches the good mean's primary component (None = pure least-squares mean target)	0.9
tail_statistic	how the target tail is measured: quantile (the p-th order statistic, can enter the flip regime with strengths above 1) or shortfall (mean of the top 1-p fraction; a stable tail statistic always solved in the capped regime 0 ≤ c ≤ 1)	quantile
cap_primary_strength	clip the primary strength to c ≤ 1, disabling the flip regime for the quantile rule (ignored for shortfall)	false
good_damage_budget	a preservation penalty designed so the edit shifts good outputs by no more than this fraction of their mean squared norm; bounds the entire edit (primary + dispersion), shrinking the primary strength if it alone would exceed the budget (None = no penalty)	None
budget_metric	what the budget constrains: norm (isotropic movement of good outputs) or kl (the movement weighted by per-prompt diagonal-Fisher damage weights from the pristine first-token logits - the quantity the KL metric measures, so one fraction means one damage level; no extra backward pass)	norm
budget_units	how the budget is allocated across depths: local (each module bound against its own good-output energy), logit (each module's gain-scaled damage, with the gain the mean squared attribution-gradient norm, bound against one global reference, so sites closer to the logits are clipped harder; the kl_normalized gate keeps its local-damage ranking), or price (same accounting, but the budget is pooled across modules and allocated through one scalar price: each site keeps s_k = min(1, beta_k / 2 lambda q_k) of its unbounded edit and non-positive-benefit sites drop out. Measured worse than logit at scale because the mean-linearized benefit ignores which prompts cross the refusal threshold; kept as an ablation). Non-local modes require gradient attribution	local
project_primary / project_primary_dims	orthogonalize the mean-diff primary against the top-K harmless principal directions before removal (projected abliteration absorbed into the closed form): removes only the harmful-specific component of the mean shift	false / 4
null_good_sensitivity	Fisher-null projection: per site, the gradients of the good prompts' top-token log-probability w.r.t. the site output form a sensitivity Fisher form (the per-site lift of the diagonal-Fisher KL weights); the top eigenvectors spanning this fraction of each site's sensitivity energy are projected out of the whole edit basis before solving, so removal happens only along directions the good behavior does not press on. A unit-free fraction; 0 = off (skips the extra pass). Acts on the direction degree of freedom, so it composes with the budget (a magnitude degree of freedom) the way project_primary does. Costs one attribution-style backward over a fixed 128-prompt subsample of the good set, per refinement round	0.0
gate_score	how sites are scored for the cumulative gating: mean_patch (effect of patching bad activations to the good mean), edit_effect (linearized effect of the actual derived edit), or kl_normalized (mean-patch effect per unit KL-weighted damage of the site's own edit - cheap-tipping sites gate first)	mean_patch
iterative_refine	apply each round's edits, re-collect, and re-solve from the edited model, composing rounds into one factored artifact (fixes the one-shot composition error; costs one collection + gating pass per round; incompatible with preserve_row_norms)	false
refine_iterations	number of solve-and-apply rounds when iterative_refine is enabled	2
use_causal_weights	gate sites by their measured causal effect score (off edits every site)	true
causal_effect_frac	gating: sites are kept in descending score order until they hit this fraction of total positive score	0.9
gating_method	how site-effect scores are measured: attribution (gradient attribution patching, one backward pass) or activation_patching (proper activation patching, forward-only re-runs with each site's activation swapped for the mean good activation; works on backends without a backward pass such as vLLM/sglang)	attribution
attribution.use_grad_drop	AtP* GradDrop: gate on the mean over per-layer "dropped" gradients instead of plain AtP (only GradDrop from AtP* is applied; the QK fix targets key/query nodes, which are never edited)	true
attribution.batch_size	prompts per forward in the attribution pass; every prompt keeps its own backward, so larger batches trade matmul FLOPs for amortized overhead (2-8 helps launch-bound devices/small models); batched runs match single-prompt exactly via explicit position ids	1
attribution.checkpointing	gradient checkpointing during attribution: on keeps only one layer's activations; off makes each backward a plain backward (~1.5× plain AtP, ~2× GradDrop) at the cost of keeping all activations	true
rank_cap / spectrum_min_eigenvalue	CSP subspace rank limit (note that this also caps the rank of the edit, and is useful for fixing the maximum rank of the edit) and bad-variance-fraction threshold (0.5 = bad-dominated directions)	8 / 0.5
csp_subspace_dim	PCA dimension where the CSP's matrix pencil is solved	128
ridge	ridge added to the strength-matrix systems (None = a small numerical guard)	None
analysis_device	where the localization/edit math runs (cuda, mps on Apple Silicon, or cpu; falls back to CPU when the requested accelerator is unavailable)	cuda
preserve_row_norms	rescale each edited row back to its original norm (full-rank edit to the applied delta, so incompatible with adapter save modes)	false
moe_experts	also edit routed expert down-projections in MoE models	false
components	canonical component names to edit	["mlp.down_proj", "attn.o_proj"]
attn_implementation	attention backend override (None lets transformers choose)	None
adapter_layer_prefix	decoder-layer prefix used in exported adapter paths, useful for serving systems with different tree structures	detected at load
save_mode	save full weights, or just a PEFT-compatible adapter as lora-peft	weights
response_prefix	added before measuring logits; prevents degenerate KLs from post-template positions (i.e. on reasoning tokens)	None
attribution.prefix	text appended after the chat template before scoring the refusal-vs-compliance gradient; keys the attribution cache	None
attribution.refusal_token / attribution.compliance_token	the two first tokens whose logit difference is the scored scalar	"Sorry" / "Sure"
attribution.max_tokens	cap each prompt to its last N tokens for the attribution pass (None = no cap)	None

The Default column is the library default (Config()). config.example.toml overrides several of these for the shipped preset: project_primary = true, tail_statistic = "shortfall", good_damage_budget = 0.005, budget_metric = "kl", budget_units = "logit", null_good_sensitivity = 0.8, rank_cap = 64, iterative_refine = true, and attribution.use_grad_drop = false.

How it works

1. Localize. A collection pass records layer residuals and module activations for a set of harmless and harmful prompts. Each layer's edit basis uses the mean-difference direction and the CSP spectrum (solved via PCA). Site-effect gating scores each module, deciding which ones actually get edited. The library default measurement is gradient attribution patching (a backward pass, linearized effect of swapping bad activations for mean good activations) with the AtP* GradDrop estimator, which averages per-layer "dropped" gradients so direct/indirect-effect cancellation cannot under-score the residual writers WHALE edits. gating_method = "activation_patching" switches to proper activation patching, which performs the swap in forward-only re-runs and re-measures the refusal-vs-compliance metric per site; that is the option for backends without a backward pass (vLLM, sglang). The shipped preset sets attribution.use_grad_drop = false to save the L backward passes per prompt.
2. Edit. For every selected module, we solve a small system to get the exact strength matrix A. We calculate a scalar strength for the primary direction based on a target percentile. The dispersion directions are solved together using a penalty based on an allowed damage budget. There are no optimizers used.
3. Apply. The edit is factored into B = V, A_{fact} = -A V^\top W. You can merge this directly into the weights, apply row-norm preservation, or export it as a LoRA adapter in PEFT format, which keeps quantized checkpoints untouched. Set save_mode = "lora-peft" to save just the adapter instead of the massive full model.

The edit is deliberately structural rather than interpolating: it modifies the module's subspace components for all inputs, which allows it to generalize to new harmful prompts.

With iterative_refine = true, each round's edits are applied, the prompt sets are re-collected from the edited model, and the closed forms are re-solved. Factor stacking is exact because each round's right factors are solved against the already-edited weights, so the rounds compose into one artifact.

Supported architectures

Standard dense decoders (like Llama, Qwen, and Mistral) work immediately. We added two extensions for newer model types:

• MoE models. The shared/dense expert writer is edited as mlp.down_proj by default. You can enable routed experts with moe_experts = true. Per-expert Linear modules are edited individually. Fused expert containers (like Qwen3.5-MoE or DeepSeek-V3/V4) are edited as a single block: we solve the strength on the routed-sum outputs and apply it to every expert using PEFT's target_parameters layout.
• DeepSeek-V4 (mHC). V4 uses multiple parallel residual streams mixed by Hyper-Connections instead of a standard residual sum. We edit the writers that feed into these streams (self_attn.o_b_proj and mlp.shared_experts.down_proj). We average the residuals across the parallel streams. Since the hyper-connection path is linear, the structural edit behaves exactly as it would in a dense model (whether the model's own connection weights adapt afterward is an open question). V4 is eager-only, so set attn_implementation = "eager" (or leave it unset and let the loader choose).

The backend correctly identifies multimodal checkpoints (it2t wrappers) and uses their specific serving module paths (e.g., model.language_model.layers). This ensures edits target the right modules, rejecting stripped text-only checkpoints instead of silently applying edits in the wrong place.

Backends

The backend interface (src/ablit/backend/base.py) isolates model-specific work: collecting activations, applying edits, generating, and measuring site effects. TransformersBackend implements it for HuggingFace transformers, including the gradient-based attribution pass. Activation patching is part of the same interface as the forward-only alternative, so an inference-only backend (vLLM, sglang) can implement collect() and activation_patch(), set supports_backward = false, and let the pipeline route gating through activation patching when gating_method = "activation_patching". The pipeline raises a clear error if attribution is requested from a backend without a backward pass. No such backend ships here yet.

Repository layout

config.example.toml     fully commented reference configuration
paper.pdf               the method paper
docs/papers.md          related work and reference papers
src/ablit/backend/      model backend interface + transformers implementation
src/ablit/localize.py   attribution, mean-diff + CSP subspace, spectral rank
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
scripts/                sweep, calibration, MMLU probe, roundtrip and tiny-model checks
tests/                  134 tests: solver reductions, localization math, patching, adapters, MoE, V4

Documentation

• paper.pdf: the method, including the geometric findings and the rank-1 reduction proof (DeepSeek, 2026).
• docs/papers.md: related work and reference papers.
• config.example.toml: the shipped preset, commented.

Frontier search

Install the optimize extra (uv sync --extra optimize) to use Optuna to search five method settings: target, target percentile, good damage budget, causal effect fraction, and rank cap. Settings outside that list keep the values from the config file. The objective filters out trials where the KL exceeds the acceptable limit; among valid trials it minimizes the refusal rate and uses KL to break ties. Trials share caches and are solved one-shot (iterative_refine is not applied during the search). The best configuration is tested on a held-out split:

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

Sweep script

scripts/sweep.py runs the full configuration sweep: the reference row, the project_primary and kl_normalized gate rows, the KL budget ladder, the depth-fair logit-units ladder, the pooled price allocation, the Fisher-null family, the frozen preset row (knee, identical to config.example.toml), a b = 1.0 control that must reproduce the core row, and non-iterative twins of the headline rows. rank_cap = 64 and the CSP defaults are pinned to match scripts/kl_calibration.py, which measures the budget's binding range on a model before you sweep it:

uv run python scripts/sweep.py --list              # row labels only
uv run python scripts/sweep.py --dry-run           # rows + effective config
uv run python scripts/sweep.py                     # everything
uv run python scripts/sweep.py --only core,proj,b3e-2,b1e0
uv run python scripts/sweep.py --resume            # carry completed rows, run the rest
uv run python scripts/kl_calibration.py            # binding shares first

Because caches are keyed by what actually changes them, all rows share the expensive base passes (collection, attribution, KL baselines); iterative rows each pay a tagged round-2 collect + attribution (the edited model state differs per config). Tagged caches also include the repo HEAD, so a commit invalidates exactly them; untagged caches are protected only by the manual cache-schema version, so rerun everything after a change to the collection, attribution, or baseline math. Results are written incrementally to <output-dir>/sweep_results.json and .md after every row, so a killed batch resumes with --resume (completed rows only; rerun everything after a src/ change, since carried rows keep their old environment), or a --only subset; --skip-done keeps a multi-model bash loop resumable.

Tests

uv sync --extra dev --extra adapters --extra optimize
uv run pytest