ablate-llm 0.2.0__py3-none-any.whl

ablate/pipeline.py

@@ -0,0 +1,343 @@
+"""High-level orchestrator tying the whole abliteration pipeline together."""
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import asdict
+from typing import List, Optional
+
+from typing import Union
+
+import torch
+
+from . import data, extract, subspace as subspace_mod, utils, weights
+from .config import AblationConfig, RunConfig, SubspaceConfig
+from .evaluate import EvalResult, evaluate
+from .harness import HarnessResult, run_harness
+from .hooks import AblationHooks
+from .model import LM
+from .optimize import SearchResult, SubspaceSearchResult, optimize, optimize_subspace
+
+
+def _metrics_of(result: EvalResult) -> dict:
+    return {
+        "refusal_rate": round(result.refusal_rate, 4),
+        "mean_kl": round(result.mean_kl, 4),
+        "coherence": round(result.coherence, 4),
+    }
+
+
+class Ablator:
+    """End-to-end interface for extracting, applying, searching, and baking
+    refusal-direction ablations.
+
+    Typical use::
+
+        abl = Ablator("Qwen/Qwen2.5-0.5B-Instruct")
+        abl.extract()                 # find candidate directions
+        result = abl.search()         # optimize layer / alpha / band
+        print(result.result)
+        abl.bake(abl.best_config)     # fold into weights
+        abl.save("uncensored-model")
+    """
+
+    def __init__(self, model_name: str, device=None, dtype=None, trust_remote_code=False):
+        self.lm = LM.load(model_name, device=device, dtype=dtype, trust_remote_code=trust_remote_code)
+        self.directions: Optional[torch.Tensor] = None
+        self.subspace: Optional[torch.Tensor] = None
+        self.subspace_info: Optional[dict] = None
+        self.best_config: Optional[Union[AblationConfig, SubspaceConfig]] = None
+        self.last_metrics: Optional[dict] = None
+        self._baked = False
+
+    # -- resolve the basis a config refers to ------------------------------- #
+    def _basis_for(self, config: Union[AblationConfig, SubspaceConfig]) -> torch.Tensor:
+        if isinstance(config, SubspaceConfig):
+            assert self.subspace is not None, "call extract_subspace() first"
+            return self.subspace[: config.n_directions]
+        assert self.directions is not None, "call extract() first"
+        return self.directions[config.direction_layer]
+
+    # -- data --------------------------------------------------------------- #
+    def _default_data(self, cfg: RunConfig):
+        harmful = data.sample(data.load_builtin("harmful"), None, cfg.seed)
+        harmless = data.sample(data.load_builtin("harmless"), None, cfg.seed)
+        h_train, h_eval = data.split(harmful, cfg.n_harmful, cfg.n_eval_harmful, cfg.seed)
+        s_train, s_eval = data.split(harmless, cfg.n_harmless, cfg.n_eval_benign, cfg.seed)
+        return h_train, s_train, h_eval, s_eval
+
+    # -- extraction --------------------------------------------------------- #
+    def extract(
+        self,
+        harmful: Optional[List[str]] = None,
+        harmless: Optional[List[str]] = None,
+        method: str = "diff_of_means",
+        positions: int = 1,
+        batch_size: int = 8,
+    ) -> torch.Tensor:
+        if harmful is None:
+            harmful = data.load_builtin("harmful")
+        if harmless is None:
+            harmless = data.load_builtin("harmless")
+        self.directions = extract.extract_directions(
+            self.lm, harmful, harmless, method=method, positions=positions, batch_size=batch_size
+        )
+        return self.directions
+
+    def extract_subspace(
+        self,
+        harmful: Optional[List[str]] = None,
+        harmless: Optional[List[str]] = None,
+        method: str = "band",
+        n_directions: int = 4,
+        layer: Optional[int] = None,
+        positions: int = 1,
+        batch_size: int = 8,
+    ) -> torch.Tensor:
+        """Extract an orthonormal refusal *subspace* (multi-direction ablation)."""
+        if harmful is None:
+            harmful = data.load_builtin("harmful")
+        if harmless is None:
+            harmless = data.load_builtin("harmless")
+        self.subspace, self.subspace_info = subspace_mod.extract_subspace(
+            self.lm, harmful, harmless, method=method, n_directions=n_directions,
+            layer=layer, positions=positions, batch_size=batch_size,
+        )
+        return self.subspace
+
+    # -- single-config evaluation ------------------------------------------ #
+    def evaluate(
+        self,
+        config: AblationConfig,
+        harmful_eval: List[str],
+        benign_eval: List[str],
+        max_new_tokens: int = 64,
+        kl_weight: float = 1.0,
+        keep_samples: bool = True,
+    ) -> EvalResult:
+        assert self.directions is not None, "call extract() first"
+        return evaluate(
+            self.lm,
+            self.directions[config.direction_layer],
+            config,
+            harmful_eval,
+            benign_eval,
+            max_new_tokens=max_new_tokens,
+            kl_weight=kl_weight,
+            keep_samples=keep_samples,
+        )
+
+    # -- optimization ------------------------------------------------------- #
+    def search(
+        self,
+        harmful_eval: Optional[List[str]] = None,
+        benign_eval: Optional[List[str]] = None,
+        n_trials: int = 30,
+        kl_weight: float = 1.0,
+        max_new_tokens: int = 48,
+        seed: int = 0,
+    ) -> SearchResult:
+        assert self.directions is not None, "call extract() first"
+        if harmful_eval is None:
+            harmful_eval = data.split(data.load_builtin("harmful"), 40, 24, seed)[1]
+        if benign_eval is None:
+            benign_eval = data.split(data.load_builtin("harmless"), 40, 24, seed)[1]
+        res = optimize(
+            self.lm,
+            self.directions,
+            harmful_eval,
+            benign_eval,
+            n_trials=n_trials,
+            kl_weight=kl_weight,
+            max_new_tokens=max_new_tokens,
+            seed=seed,
+        )
+        self.best_config = res.config
+        self.last_metrics = _metrics_of(res.result)
+        return res
+
+    def search_subspace(
+        self,
+        harmful_eval: Optional[List[str]] = None,
+        benign_eval: Optional[List[str]] = None,
+        n_trials: int = 30,
+        kl_weight: float = 1.0,
+        max_new_tokens: int = 48,
+        seed: int = 0,
+    ) -> SubspaceSearchResult:
+        """KL-guided search over (n_directions, alpha, layer-band) for the subspace."""
+        assert self.subspace is not None, "call extract_subspace() first"
+        if harmful_eval is None:
+            harmful_eval = data.split(data.load_builtin("harmful"), 40, 24, seed)[1]
+        if benign_eval is None:
+            benign_eval = data.split(data.load_builtin("harmless"), 40, 24, seed)[1]
+        res = optimize_subspace(
+            self.lm, self.subspace, harmful_eval, benign_eval,
+            n_trials=n_trials, kl_weight=kl_weight, max_new_tokens=max_new_tokens, seed=seed,
+        )
+        self.best_config = res.config
+        self.last_metrics = _metrics_of(res.result)
+        return res
+
+    # -- benchmark harness with a judge ------------------------------------- #
+    def harness(
+        self,
+        prompts: List[str],
+        judge: Optional[object] = None,
+        config: Optional[Union[AblationConfig, SubspaceConfig]] = None,
+        max_new_tokens: int = 128,
+    ) -> HarnessResult:
+        """Run a harmful benchmark through a judge. Pass ``config`` (or rely on
+        ``best_config``) to score the ablated model; pass ``config=None`` and no
+        ``best_config`` for the baseline."""
+        cfg = config or self.best_config
+        basis = self._basis_for(cfg) if cfg is not None else None
+        return run_harness(self.lm, prompts, judge=judge, basis=basis, config=cfg,
+                           max_new_tokens=max_new_tokens)
+
+    # -- generation (with a live, non-destructive ablation) ----------------- #
+    def generate(self, prompts: List[str], config=None, max_new_tokens: int = 128) -> List[str]:
+        config = config or self.best_config
+        if config is None:
+            return self.lm.generate(prompts, max_new_tokens=max_new_tokens)
+        with AblationHooks(self.lm, self._basis_for(config), config):
+            return self.lm.generate(prompts, max_new_tokens=max_new_tokens)
+
+    # -- baking + saving ---------------------------------------------------- #
+    def bake(self, config=None, include_embedding: bool = True) -> None:
+        """Permanently fold the chosen direction/subspace into the model weights."""
+        config = config or self.best_config
+        assert config is not None, "provide a config or run search()/search_subspace() first"
+        weights.bake_subspace(self.lm.model, self._basis_for(config), include_embedding)
+        self._baked = True
+
+    def save(self, path: str) -> None:
+        os.makedirs(path, exist_ok=True)
+        self.lm.model.save_pretrained(path)
+        self.lm.tokenizer.save_pretrained(path)
+        meta = {
+            "base_model": self.lm.name,
+            "baked": self._baked,
+            "best_config": self.best_config.to_dict() if self.best_config else None,
+        }
+        with open(os.path.join(path, "ablate_meta.json"), "w") as f:
+            json.dump(meta, f, indent=2)
+
+    def save_directions(self, path: str) -> None:
+        assert self.directions is not None
+        torch.save(self.directions, path)
+
+    # -- publish to the HuggingFace Hub ------------------------------------- #
+    def push_to_hub(
+        self,
+        repo_id: str,
+        token: Optional[str] = None,
+        private: bool = True,
+        metrics: Optional[dict] = None,
+        license: str = "apache-2.0",
+        bake_if_needed: bool = True,
+    ) -> str:
+        """Bake (if needed) and upload the abliterated model + a generated model
+        card to ``repo_id``. Returns the repo URL.
+
+        ``token`` falls back to ``HF_TOKEN`` / ``HUGGINGFACE_TOKEN`` env vars.
+        """
+        from .hub import push_to_hub as _push
+
+        if not self._baked and bake_if_needed:
+            assert self.best_config is not None, "run a search or pass a config before pushing"
+            self.bake(self.best_config)
+
+        cfg = self.best_config.to_dict() if self.best_config else {}
+        method = (
+            f"rank-{self.best_config.n_directions} subspace ablation (baked)"
+            if isinstance(self.best_config, SubspaceConfig)
+            else "single-direction ablation (baked)"
+        )
+        return _push(
+            self.lm.model, self.lm.tokenizer, repo_id=repo_id, token=token,
+            base_model=self.lm.name, method=method, config=cfg,
+            metrics=metrics or self.last_metrics, private=private, license=license,
+        )
+
+
+def _load_source(source: Optional[str], kind: str, n: Optional[int], seed: int) -> List[str]:
+    """Resolve a data source string to a prompt list.
+
+    ``source`` may be ``None`` (built-in ``harmful``/``harmless`` set), a named
+    shortcut (``advbench``/``harmbench``/``jailbreakbench``/``alpaca``), or
+    ``"hf:<dataset>:<column>[:<config>]"``.
+    """
+    if source is None:
+        return data.sample(data.load_builtin(kind), n, seed)
+    named = {
+        "advbench": data.load_advbench,
+        "harmbench": data.load_harmbench,
+        "jailbreakbench": data.load_jailbreakbench,
+        "alpaca": data.load_alpaca_benign,
+    }
+    if source in named:
+        return named[source](n=n)
+    if source.startswith("hf:"):
+        parts = source.split(":")
+        dataset, column = parts[1], parts[2]
+        config = parts[3] if len(parts) > 3 else None
+        return data.load_hf(dataset, column=column, config=config, n=n, shuffle_seed=seed)
+    raise ValueError(f"unrecognized data source: {source!r}")
+
+
+def run(cfg: RunConfig):
+    """Fully automated pipeline driven by a RunConfig (used by the CLI).
+
+    Supports single-direction or subspace ablation, built-in or HuggingFace data
+    sources, and optional push to the Hub. Returns the search result.
+    """
+    utils.set_seed(cfg.seed)
+    abl = Ablator(cfg.model_name, device=cfg.device, dtype=cfg.dtype)
+
+    harmful = _load_source(cfg.harmful_source, "harmful", cfg.n_harmful + cfg.n_eval_harmful, cfg.seed)
+    harmless = _load_source(cfg.harmless_source, "harmless", cfg.n_harmless + cfg.n_eval_benign, cfg.seed)
+    h_train, h_eval = data.split(harmful, cfg.n_harmful, cfg.n_eval_harmful, cfg.seed)
+    s_train, s_eval = data.split(harmless, cfg.n_harmless, cfg.n_eval_benign, cfg.seed)
+
+    if cfg.subspace:
+        abl.extract_subspace(harmful=h_train, harmless=s_train, method=cfg.subspace_method,
+                             n_directions=cfg.n_directions, positions=cfg.extract_positions,
+                             batch_size=cfg.batch_size)
+        result = abl.search_subspace(harmful_eval=h_eval, benign_eval=s_eval, n_trials=cfg.n_trials,
+                                     kl_weight=cfg.kl_weight, max_new_tokens=cfg.max_new_tokens, seed=cfg.seed)
+    else:
+        abl.extract(harmful=h_train, harmless=s_train, positions=cfg.extract_positions, batch_size=cfg.batch_size)
+        result = abl.search(harmful_eval=h_eval, benign_eval=s_eval, n_trials=cfg.n_trials,
+                            kl_weight=cfg.kl_weight, max_new_tokens=cfg.max_new_tokens, seed=cfg.seed)
+
+    os.makedirs(cfg.output_dir, exist_ok=True)
+    with open(os.path.join(cfg.output_dir, "result.json"), "w") as f:
+        json.dump(
+            {
+                "run_config": asdict(cfg),
+                "subspace_info": abl.subspace_info,
+                "best_config": result.config.to_dict(),
+                "metrics": {
+                    "refusal_rate": result.result.refusal_rate,
+                    "mean_kl": result.result.mean_kl,
+                    "coherence": result.result.coherence,
+                    "objective": result.result.objective,
+                },
+                "samples": result.result.samples,
+            },
+            f,
+            indent=2,
+        )
+
+    if cfg.save_model or cfg.push_to_hub:
+        abl.bake(result.config)
+        if cfg.save_model:
+            abl.save(os.path.join(cfg.output_dir, "model"))
+
+    if cfg.push_to_hub:
+        url = abl.push_to_hub(cfg.push_to_hub, token=cfg.hf_token, private=cfg.private,
+                              metrics=abl.last_metrics)
+        print(f"\nPushed abliterated model to: {url}")
+
+    return result

ablate/subspace.py

@@ -0,0 +1,103 @@
+"""Multi-direction (subspace) ablation.
+
+A single refusal direction is often insufficient: safety behaviour is
+redundantly encoded across layers, heads, and multiple residual directions.
+Here we extract an *orthonormal basis* ``Q`` of shape ``(k, D)`` spanning a
+refusal **subspace**, and project the whole subspace out of the residual stream:
+
+    h' = h - alpha * (h Qᵀ) Q          # removes the component of h in span(Q)
+
+``k == 1`` recovers ordinary single-direction ablation.
+
+Two extraction strategies:
+
+* ``"band"`` — the diff-of-means direction from each of the ``k`` strongest
+  layers, orthonormalized (order-preserving Gram-Schmidt). Directions from
+  different layers live in the same residual space, so their span is a
+  meaningful multi-layer refusal subspace. Truncating ``Q[:j]`` keeps the ``j``
+  strongest.
+* ``"pca"`` — the top-``k`` principal components of the pooled, mean-centered
+  activations at one layer, sign-aligned to the harmful side. Components are
+  variance-ordered, so truncation is principled.
+"""
+from __future__ import annotations
+
+from typing import Dict, List, Optional, Tuple
+
+import torch
+
+from .extract import collect_activations
+from .model import LM
+
+
+def as_basis(x: torch.Tensor) -> torch.Tensor:
+    """Coerce a direction ``(D,)`` or basis ``(k, D)`` to a 2-D basis."""
+    return x.unsqueeze(0) if x.dim() == 1 else x
+
+
+def orthonormalize(vecs: torch.Tensor) -> torch.Tensor:
+    """Order-preserving modified Gram-Schmidt.
+
+    Input ``(k, D)``; output ``(r, D)`` with orthonormal rows, ``r <= k`` (near
+    linearly-dependent rows are dropped). Row order is preserved so that
+    ``basis[:j]`` keeps the ``j`` most important directions.
+    """
+    out: List[torch.Tensor] = []
+    for v in vecs.float():
+        w = v.clone()
+        for u in out:
+            w = w - (w @ u) * u
+        n = w.norm()
+        if n > 1e-6:
+            out.append(w / n)
+    if not out:
+        return vecs[:0]
+    return torch.stack(out, dim=0)
+
+
+def project_subspace(h: torch.Tensor, basis: torch.Tensor, alpha: float = 1.0) -> torch.Tensor:
+    """``h - alpha * (h Qᵀ) Q`` for orthonormal ``Q`` ``(k, D)`` (or a ``(D,)`` vector)."""
+    Q = as_basis(basis).to(dtype=h.dtype, device=h.device)
+    coeffs = h @ Q.T            # (..., k)
+    return h - alpha * (coeffs @ Q)
+
+
+def extract_subspace(
+    lm: LM,
+    harmful_prompts: List[str],
+    harmless_prompts: List[str],
+    method: str = "band",
+    n_directions: int = 4,
+    layer: Optional[int] = None,
+    positions: int = 1,
+    batch_size: int = 8,
+    show_progress: bool = True,
+) -> Tuple[torch.Tensor, Dict]:
+    """Return ``(basis (k, D), info)``. ``basis`` rows are orthonormal and
+    ordered by importance."""
+    h_act = collect_activations(lm, harmful_prompts, positions, batch_size, show_progress=show_progress)
+    s_act = collect_activations(lm, harmless_prompts, positions, batch_size, show_progress=show_progress)
+    raw = h_act.mean(dim=1) - s_act.mean(dim=1)  # (L+1, D) per-layer diff-of-means
+    norms = raw.norm(dim=-1)
+
+    if method == "band":
+        # Candidate layers exclude embedding (0) and the final layer.
+        cand = list(range(1, raw.shape[0] - 1))
+        cand.sort(key=lambda i: float(-norms[i]))
+        chosen = cand[:n_directions]
+        vecs = torch.stack([raw[i] for i in chosen], dim=0)  # ordered strongest-first
+        basis = orthonormalize(vecs)
+        return basis, {"method": "band", "layers": chosen, "n_directions": int(basis.shape[0])}
+
+    if method == "pca":
+        L = int(layer) if layer is not None else int(norms[1:-1].argmax().item()) + 1
+        X = torch.cat([h_act[L], s_act[L]], dim=0)
+        Xc = X - X.mean(dim=0, keepdim=True)
+        _, _, Vh = torch.linalg.svd(Xc, full_matrices=False)
+        comps = Vh[:n_directions]  # (k, D), variance-ordered, already orthonormal
+        dm = h_act[L].mean(0) - s_act[L].mean(0)
+        comps = torch.stack([c if torch.dot(c, dm) >= 0 else -c for c in comps], dim=0)
+        basis = orthonormalize(comps)
+        return basis, {"method": "pca", "layer": L, "n_directions": int(basis.shape[0])}
+
+    raise ValueError(f"unknown subspace method: {method}")

ablate/utils.py

@@ -0,0 +1,129 @@
+"""Device / dtype helpers, seeding, and architecture adapters.
+
+The architecture adapters are what let the rest of the toolkit stay
+model-agnostic: given any HuggingFace causal-LM they locate the list of decoder
+layers and the weight matrices that *write into the residual stream* (embedding,
+attention output projection, MLP output projection). Everything downstream —
+activation capture, runtime hooks, weight baking — is expressed in terms of
+these accessors.
+"""
+from __future__ import annotations
+
+import os
+import random
+from typing import List
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+# Let unsupported MPS ops fall back to CPU instead of crashing.
+os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")
+
+try:  # transformers' GPT2-style Conv1D layer (weight is stored transposed).
+    from transformers.pytorch_utils import Conv1D
+except Exception:  # pragma: no cover - very old transformers
+    Conv1D = tuple()  # isinstance(x, ()) is always False
+
+
+def pick_device(preferred: str | None = None) -> torch.device:
+    if preferred:
+        return torch.device(preferred)
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    if getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+def pick_dtype(device: torch.device, preferred: str | None = None) -> torch.dtype:
+    if preferred:
+        return getattr(torch, preferred)
+    if device.type == "cuda":
+        return torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+    # float32 is the safe choice on CPU/MPS for tiny models on 16GB machines.
+    return torch.float32
+
+
+def set_seed(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+# --------------------------------------------------------------------------- #
+# Architecture adapters
+# --------------------------------------------------------------------------- #
+def get_decoder_layers(model: nn.Module) -> nn.ModuleList:
+    """Return the ModuleList of transformer decoder blocks for common families."""
+    m = model
+    # Llama / Mistral / Qwen2 / SmolLM2 / Gemma / Phi ...
+    if hasattr(m, "model") and hasattr(m.model, "layers"):
+        return m.model.layers
+    # GPT-2 / GPT-Neo
+    if hasattr(m, "transformer") and hasattr(m.transformer, "h"):
+        return m.transformer.h
+    # GPT-NeoX / Pythia
+    if hasattr(m, "gpt_neox") and hasattr(m.gpt_neox, "layers"):
+        return m.gpt_neox.layers
+    # Fallback: some models expose .model.decoder.layers (OPT)
+    if hasattr(m, "model") and hasattr(m.model, "decoder") and hasattr(m.model.decoder, "layers"):
+        return m.model.decoder.layers
+    raise ValueError(
+        f"Could not locate decoder layers for {type(model).__name__}. "
+        "Add an adapter in utils.get_decoder_layers."
+    )
+
+
+def num_layers(model: nn.Module) -> int:
+    return len(get_decoder_layers(model))
+
+
+def hidden_size(model: nn.Module) -> int:
+    return int(model.config.hidden_size) if hasattr(model.config, "hidden_size") else int(model.config.n_embd)
+
+
+def get_embedding(model: nn.Module) -> nn.Embedding:
+    return model.get_input_embeddings()
+
+
+def _first_attr(module: nn.Module, names: List[str]):
+    for n in names:
+        if hasattr(module, n):
+            return getattr(module, n)
+    return None
+
+
+def get_residual_writers(layer: nn.Module) -> List[nn.Module]:
+    """Return the modules inside one decoder block that write to the residual
+    stream: the attention output projection and the MLP output projection.
+
+    These are exactly the matrices we orthogonalize when baking a direction into
+    the weights.
+    """
+    writers: List[nn.Module] = []
+
+    attn = _first_attr(layer, ["self_attn", "attn", "attention"])
+    if attn is not None:
+        o = _first_attr(attn, ["o_proj", "c_proj", "dense", "out_proj"])
+        if o is not None:
+            writers.append(o)
+
+    mlp = _first_attr(layer, ["mlp", "feed_forward", "ffn"])
+    if mlp is not None:
+        down = _first_attr(mlp, ["down_proj", "c_proj", "dense_4h_to_h", "fc_out", "w2"])
+        if down is not None:
+            writers.append(down)
+
+    if not writers:
+        raise ValueError(
+            f"Could not locate residual-writing projections in {type(layer).__name__}. "
+            "Add names to utils.get_residual_writers."
+        )
+    return writers
+
+
+def is_conv1d(module: nn.Module) -> bool:
+    return isinstance(Conv1D, type) and isinstance(module, Conv1D)

ablate/weights.py

@@ -0,0 +1,71 @@
+"""Weight orthogonalization ("baking" a direction or subspace into the model).
+
+Runtime hooks are ideal for search, but for a shippable model we fold the
+projection directly into every weight matrix that writes to the residual stream,
+so the model can *never* express the subspace ``span(Q)`` and no hook is needed:
+
+    W_out  <-  (I - QᵀQ) W_out            (attention / MLP output projections)
+    E      <-  E (I - QᵀQ)                 (token embeddings, per row)
+
+``Q`` is an orthonormal basis ``(k, D)``; ``k == 1`` is single-direction baking.
+Supports both ``nn.Linear`` (Llama/Mistral/Qwen/SmolLM/Pythia) and GPT-2's
+``Conv1D`` (weight stored transposed).
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from . import utils
+from .subspace import as_basis, orthonormalize
+
+
+@torch.no_grad()
+def _orthogonalize_module(module: nn.Module, Q: torch.Tensor) -> None:
+    """Remove ``span(Q)`` from a residual-writing module's output."""
+    W = module.weight.data
+    Qd = Q.to(dtype=W.dtype, device=W.device)  # (k, D)
+
+    if utils.is_conv1d(module):
+        # GPT-2 Conv1D: y = x @ W, W is (in, out=D). Output dim is dim 1.
+        # W <- W (I - QᵀQ) = W - (W Qᵀ) Q
+        module.weight.data = W - (W @ Qd.T) @ Qd
+    else:
+        # nn.Linear: y = W x (+ b), W is (out=D, in). Output dim is dim 0.
+        # W <- (I - QᵀQ) W = W - Qᵀ (Q W)
+        module.weight.data = W - Qd.T @ (Qd @ W)
+
+    bias = getattr(module, "bias", None)
+    if bias is not None:
+        b = bias.data
+        bq = Qd.to(b.dtype)
+        bias.data = b - bq.T @ (bq @ b)
+
+
+@torch.no_grad()
+def _orthogonalize_embedding(emb: nn.Embedding, Q: torch.Tensor) -> None:
+    W = emb.weight.data  # (vocab, D): each row is a residual vector
+    Qd = Q.to(dtype=W.dtype, device=W.device)
+    emb.weight.data = W - (W @ Qd.T) @ Qd
+
+
+@torch.no_grad()
+def bake_subspace(model: nn.Module, basis: torch.Tensor, include_embedding: bool = True) -> None:
+    """Permanently orthogonalize ``model``'s weights against ``span(basis)``.
+
+    Mutates ``model`` in place. ``basis`` may be a ``(D,)`` direction or a
+    ``(k, D)`` basis; it is orthonormalized defensively.
+    """
+    Q = orthonormalize(as_basis(basis))
+    if Q.shape[0] == 0:
+        return
+    if include_embedding:
+        _orthogonalize_embedding(utils.get_embedding(model), Q)
+    for layer in utils.get_decoder_layers(model):
+        for writer in utils.get_residual_writers(layer):
+            _orthogonalize_module(writer, Q)
+
+
+# Backward-compatible single-direction alias.
+def bake_direction(model: nn.Module, direction: torch.Tensor, include_embedding: bool = True) -> None:
+    bake_subspace(model, direction, include_embedding)