ABLITERAX：設計與參考實作

ABLITERAX — Design and Reference Implementation

Below is a working, testable first-version system. I have been explicit about where each piece is verified by tests in this answer vs. depends on the model and hardware at run time.

0. Design decisions, assumptions, scope

What is verified by tests included below (CPU):

• Direction extraction math (diff_means, svd, whitened_svd, fisher_lda, orthogonalization, winsorization) — deterministic unit tests.
• Projection formula W' = (I − λ·V Vᵀ)·W, bias projection, row-norm preservation (MPOA), and LoRA-style rank-k hook equivalence to explicit weight modification.
• Kernel weight schedule, Pareto filtering, config validation, dtype round-trip.

What is verified only at run time (documented, not asserted here):

• Loading real HF checkpoints, activation forward hooks producing the correct residual, refusal detection against a real model, KL on real logits.
• lm-eval-harness integration is invoked as a subprocess, so it is only as good as the installed version.

Explicit non-goals of v0.1 (with reasons):

• SAE-based feature extraction: needs per-model SAEs; out of scope for a self-contained tool.
• LEACE exact affine erasure: formula is included for fisher_lda-style whitening only; the exact affine variant is left as a code path with a clear note.
• Per-expert EGA directions for MoE: v0.1 projects every expert against the shared residual subspace. EGA refinement is listed as a roadmap item.
• Full lm-eval inside every trial: too expensive; the search uses calibration-set proxies and only top-K survivors get real benchmark runs.

Key premise corrections to the source material (all already reflect this):

• Per-layer refusal vectors are not assumed to share direction. per_layer_direction=true is the default.
• Refusal is not assumed localized to one layer. The kernel schedule searches over position and width.
• "0% refusal" is never claimed; the report prints measured rates and sample sizes.

Compute: the pipeline runs on a T4 with Qwen/Qwen3-0.6B; on rented multi-GPU hosts, accelerate device_map="auto" shards. VRAM is estimated before load and the loader fails fast with a clear message if short.

1. Directory layout

abliterax/
├── pyproject.toml
├── README.md
├── abliterax/
│   ├── __init__.py
│   ├── errors.py
│   ├── logging_utils.py
│   ├── config.py
│   ├── model_loader.py
│   ├── architecture_registry.py
│   ├── dataset.py
│   ├── activation_collector.py
│   ├── geometry.py
│   ├── interventions.py
│   ├── benchmarks.py
│   ├── optimizer.py
│   ├── exporter.py
│   └── cli.py
├── tests/
│   ├── test_geometry.py
│   ├── test_interventions.py
│   └── test_config.py
├── examples/
│   └── qwen3_0.6b_aggressive.toml
└── notebooks/
    └── abliterax_colab.ipynb

2. pyproject.toml

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[project]
name = "abliterax"
version = "0.1.0"
description = "Automated refusal-representation surgery for transformer LMs."
requires-python = ">=3.10,<3.13"
license = { text = "Apache-2.0" }
dependencies = [
    "torch>=2.2",
    "transformers>=4.51,<5",
    "accelerate>=0.33",
    "safetensors>=0.4",
    "optuna>=3.6",
    "numpy>=1.26",
    "scipy>=1.11",
    "datasets>=2.20",
    "huggingface-hub>=0.24",
    "sentencepiece>=0.2",
    "protobuf>=4.25",
    "tqdm>=4.66",
]

[project.optional-dependencies]
quant = ["bitsandbytes>=0.43"]
eval = ["lm-eval>=0.4"]
dev = ["pytest>=8"]

[project.scripts]
abliterax = "abliterax.cli:main"

[tool.hatch.build.targets.wheel]
packages = ["abliterax"]

3. abliterax/errors.py

class AbliteraxError(Exception):
    """Base class for all pipeline errors."""


class ConfigError(AbliteraxError):
    """Raised when a configuration fails validation."""


class IncompatibleArchitecture(AbliteraxError):
    """Raised when the model architecture is not safely supported."""


class InsufficientResources(AbliteraxError):
    """Raised when VRAM/CPU estimates fall short of the plan."""


class VerificationFailed(AbliteraxError):
    """Raised when a post-intervention integrity check fails."""

4. abliterax/logging_utils.py

import logging
import sys


def setup_logging(level: int = logging.INFO) -> None:
    logging.basicConfig(
        level=level,
        stream=sys.stderr,
        format="%(asctime)s %(levelname)-7s %(name)s :: %(message)s",
        datefmt="%H:%M:%S",
    )


def get_logger(name: str) -> logging.Logger:
    return logging.getLogger(name)

5. abliterax/config.py

from __future__ import annotations

import dataclasses
import tomllib
from dataclasses import dataclass, field
from pathlib import Path
from typing import Literal

from .errors import ConfigError

DirectionMethod = Literal["diff_means", "svd", "whitened_svd", "fisher_lda"]
Component = Literal["attn", "mlp"]
RowNormMode = Literal["none", "pre", "full"]
SaveMode = Literal["merged", "lora"]


@dataclass
class ModelSpec:
    model_id: str = ""
    revision: str = "main"
    dtype: Literal["float32", "bfloat16", "float16"] = "bfloat16"
    quantization: Literal["bnb_4bit"] | None = None
    trust_remote_code: bool = False


@dataclass
class ProbeConfig:
    good_dataset: str = "mlabonne/harmless_alpaca"
    good_column: str = "text"
    good_total: int = 512
    good_probe: int = 256

    bad_dataset: str = "mlabonne/harmful_behaviors"
    bad_column: str = "text"
    bad_total: int = 512
    bad_probe: int = 256

    max_seq_len: int = 512
    use_chat_template: bool = True
    seed: int = 0


@dataclass
class GeometryConfig:
    method: DirectionMethod = "whitened_svd"
    n_directions: int = 4
    orthogonalize_good: bool = True
    winsorize: bool = True
    winsorize_quantile: float = 0.995
    whiten_eps: float = 1e-4
    random_svd_oversample: int = 16
    random_svd_power_iters: int = 2
    layer_stride: int = 1


@dataclass
class InterventionConfig:
    row_norm_mode: RowNormMode = "full"
    project_bias: bool = True
    direction_index_mode: Literal["shared", "per_layer"] = "per_layer"
    max_norm_ratio: float = 1.10
    verify_max_abs_delta: float = 1e-3


@dataclass
class OptimizerConfig:
    n_trials: int = 16
    n_startup_trials: int = 6
    seed: int = 0
    sampler: Literal["tpe", "random"] = "tpe"
    eval_batch_size: int = 4
    eval_n_good: int = 64
    eval_n_bad: int = 64
    max_new_tokens: int = 64
    max_weight_upper: float = 1.2
    early_stop_patience: int = 8
    post_search_top_k: int = 3


@dataclass
class BenchmarkConfig:
    run_lm_eval: bool = True
    lm_eval_tasks: list[str] = field(default_factory=lambda: ["mmlu", "gsm8k", "humaneval"])
    lm_eval_batch_size: int = 8
    ppl_dataset: str = "wikitext"
    ppl_config: str = "wikitext-2-raw-v1"
    ppl_split: str = "test"
    ppl_n: int = 64
    mmlu_min: float | None = None
    gsm8k_min: float | None = None
    humaneval_min: float | None = None
    ppl_max_ratio: float | None = 1.35


@dataclass
class ExportConfig:
    out_dir: str = "./abliterax_out"
    save_mode: SaveMode = "merged"
    verify_reload: bool = True
    hub_repo: str | None = None
    hub_private: bool = True


@dataclass
class HardwareConfig:
    vram_gb: float | None = None
    allow_cpu: bool = False
    offload_dir: str | None = None
    reserve_gb: float = 3.0


@dataclass
class Config:
    model: ModelSpec
    probe: ProbeConfig = field(default_factory=ProbeConfig)
    geometry: GeometryConfig = field(default_factory=GeometryConfig)
    intervention: InterventionConfig = field(default_factory=InterventionConfig)
    optimizer: OptimizerConfig = field(default_factory=OptimizerConfig)
    benchmarks: BenchmarkConfig = field(default_factory=BenchmarkConfig)
    export: ExportConfig = field(default_factory=ExportConfig)
    hardware: HardwareConfig = field(default_factory=HardwareConfig)

    def validate(self) -> None:
        if not self.model.model_id:
            raise ConfigError("model.model_id is required")
        if self.geometry.method not in ("diff_means", "svd", "whitened_svd", "fisher_lda"):
            raise ConfigError(f"unknown geometry.method={self.geometry.method!r}")
        if self.geometry.n_directions < 1:
            raise ConfigError("geometry.n_directions must be >= 1")
        if self.probe.good_probe < 8 or self.probe.bad_probe < 8:
            raise ConfigError("need at least 8 probe prompts per class")
        if self.probe.good_probe + self.probe.good_total > 0 and self.probe.good_probe >= self.probe.good_total:
            raise ConfigError("probe.good_probe must be < probe.good_total")
        if self.probe.bad_probe >= self.probe.bad_total:
            raise ConfigError("probe.bad_probe must be < probe.bad_total")
        if self.optimizer.max_weight_upper < 0:
            raise ConfigError("optimizer.max_weight_upper must be >= 0")
        if self.intervention.row_norm_mode not in ("none", "pre", "full"):
            raise ConfigError("intervention.row_norm_mode invalid")

    def to_dict(self) -> dict:
        return dataclasses.asdict(self)


def load_config(path: str | Path) -> Config:
    data = tomllib.loads(Path(path).read_text(encoding="utf-8"))
    cfg = _build(dict(data))
    cfg.validate()
    return cfg


def _build(data: dict) -> Config:
    def take(key: str, cls):
        raw = data.pop(key, {}) or {}
        known = {f.name for f in dataclasses.fields(cls)}
        return cls(**{k: v for k, v in raw.items() if k in known})

    model = take("model", ModelSpec)
    cfg = Config(model=model)
    cfg.probe = take("probe", ProbeConfig)
    cfg.geometry = take("geometry", GeometryConfig)
    cfg.intervention = take("intervention", InterventionConfig)
    cfg.optimizer = take("optimizer", OptimizerConfig)
    cfg.benchmarks = take("benchmarks", BenchmarkConfig)
    cfg.export = take("export", ExportConfig)
    cfg.hardware = take("hardware", HardwareConfig)
    if data:
        raise ConfigError(f"unknown top-level sections: {sorted(data)}")
    return cfg

6. abliterax/architecture_registry.py

"""Architecture registry.

Each resolver returns an `ArchitectureSpec` whose `get_*` callables fetch the
actual `torch.nn.Linear` modules for a layer index. We only know the tensor
mapping after iterating real blocks, so every resolver is validated by
`verify()` before the pipeline proceeds.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Callable

import torch
from torch import nn

from .errors import IncompatibleArchitecture, VerificationFailed
from .logging_utils import get_logger

log = get_logger(__name__)


@dataclass
class ArchitectureSpec:
    model_type: str
    hidden_size: int
    n_layers: int
    is_moe: bool
    is_multimodal: bool
    get_block: Callable[[int], nn.Module]
    get_attn_out: Callable[[nn.Module], nn.Linear]
    get_mlp_down: Callable[[nn.Module], nn.Linear]
    note: str = ""


def _text_conf(conf) -> dict:
    if hasattr(conf, "text_config") and conf.text_config is not None:
        return dict(conf.text_config.to_dict())
    return dict(conf.to_dict())


def _check_linear(mod, hidden: int, where: str) -> nn.Linear:
    if not isinstance(mod, nn.Linear):
        raise VerificationFailed(f"{where}: expected nn.Linear, got {type(mod).__name__}")
    if mod.weight.ndim != 2:
        raise VerificationFailed(f"{where}: weight must be 2-D, got {tuple(mod.weight.shape)}")
    if mod.weight.shape[0] != hidden:
        raise VerificationFailed(
            f"{where}: weight out-dim {mod.weight.shape[0]} != hidden_size {hidden}"
        )
    return mod


def _llama_like(model, conf):
    c = _text_conf(conf)
    layers = model.model.layers
    return ArchitectureSpec(
        model_type=c["model_type"],
        hidden_size=c["hidden_size"],
        n_layers=len(layers),
        is_moe=False,
        is_multimodal=False,
        get_block=lambda i: layers[i],
        get_attn_out=lambda b: b.self_attn.o_proj,
        get_mlp_down=lambda b: b.mlp.down_proj,
    )


def _mixtral_like(model, conf):
    c = _text_conf(conf)
    layers = model.model.layers
    n_exp = c.get("num_local_experts") or c.get("num_experts") or 0

    def down(b):
        # Mixtral exposes experts as ModuleList; we project every expert.
        return b.block_sparse_moe.experts[0].w2 if hasattr(b.block_sparse_moe, "experts") else b.mlp.experts[0].down_proj

    def down_list(b):
        if hasattr(b, "block_sparse_moe"):
            return [e.w2 for e in b.block_sparse_moe.experts]
        return [e.down_proj for e in b.mlp.experts]

    spec = ArchitectureSpec(
        model_type=c["model_type"],
        hidden_size=c["hidden_size"],
        n_layers=len(layers),
        is_moe=True,
        is_multimodal=False,
        get_block=lambda i: layers[i],
        get_attn_out=lambda b: b.self_attn.o_proj,
        get_mlp_down=down,
        note=f"MoE: {n_exp} experts; v0.1 shares one direction across experts",
    )
    spec.get_mlp_down_all = down_list  # type: ignore[attr-defined]
    return spec


def _deepseek_like(model, conf):
    c = _text_conf(conf)
    layers = model.model.layers

    def down(b):
        if hasattr(b.mlp, "experts"):
            return b.mlp.experts[0].down_proj
        return b.mlp.down_proj

    def down_list(b):
        if hasattr(b.mlp, "experts"):
            return [e.down_proj for e in b.mlp.experts]
        return [b.mlp.down_proj]

    spec = ArchitectureSpec(
        model_type=c["model_type"],
        hidden_size=c["hidden_size"],
        n_layers=len(layers),
        is_moe=hasattr(layers[0].mlp, "experts"),
        is_multimodal=False,
        get_block=lambda i: layers[i],
        get_attn_out=lambda b: b.self_attn.o_proj,
        get_mlp_down=down,
        note="DeepSeek MLA; MoE experts share one direction in v0.1",
    )
    spec.get_mlp_down_all = down_list  # type: ignore[attr-defined]
    return spec


def _gemma_like(model, conf):
    c = _text_conf(conf)
    layers = model.model.layers
    return ArchitectureSpec(
        model_type=c["model_type"],
        hidden_size=c["hidden_size"],
        n_layers=len(layers),
        is_moe=False,
        is_multimodal=c.get("model_type") in ("gemma3", "gemma3_text"),
        get_block=lambda i: layers[i],
        get_attn_out=lambda b: b.self_attn.o_proj,
        get_mlp_down=lambda b: b.mlp.down_proj,
        note="Gemma3/4 multimodal: only text tower is touched",
    )


def _gpt_oss_like(model, conf):
    c = _text_conf(conf)
    layers = model.model.layers

    def down(b):
        return b.mlp.experts[0].down_proj

    def down_list(b):
        return [e.down_proj for e in b.mlp.experts]

    spec = ArchitectureSpec(
        model_type=c["model_type"],
        hidden_size=c["hidden_size"],
        n_layers=len(layers),
        is_moe=True,
        is_multimodal=False,
        get_block=lambda i: layers[i],
        get_attn_out=lambda b: b.self_attn.o_proj,
        get_mlp_down=down,
        note="gpt-oss MoE experts share one direction in v0.1",
    )
    spec.get_mlp_down_all = down_list  # type: ignore[attr-defined]
    return spec


REGISTRY: dict[str, Callable] = {
    "llama": _llama_like,
    "mistral": _llama_like,
    "qwen2": _llama_like,
    "qwen3": _llama_like,
    "gemma": _gemma_like,
    "gemma2": _gemma_like,
    "gemma3": _gemma_like,
    "gemma3_text": _gemma_like,
    "gemma4": _gemma_like,
    "gemma4_text": _gemma_like,
    "mixtral": _mixtral_like,
    "qwen3_moe": _mixtral_like,
    "deepseek_v2": _deepseek_like,
    "deepseek_v3": _deepseek_like,
    "gpt_oss": _gpt_oss_like,
}


def resolve(model, conf) -> ArchitectureSpec:
    mt = _text_conf(conf).get("model_type", "")
    if mt not in REGISTRY:
        raise IncompatibleArchitecture(
            f"model_type={mt!r} not in registry; supported={sorted(REGISTRY)}"
        )
    spec = REGISTRY[mt](model, conf)
    verify(model, spec)
    return spec


def verify(model, spec: ArchitectureSpec) -> None:
    if spec.n_layers < 2:
        raise VerificationFailed(f"n_layers={spec.n_layers} too small")
    probe_idx = [0, spec.n_layers // 2, spec.n_layers - 1]
    for i in probe_idx:
        b = spec.get_block(i)
        _check_linear(spec.get_attn_out(b), spec.hidden_size, f"layer[{i}].attn_out")
        _check_linear(spec.get_mlp_down(b), spec.hidden_size, f"layer[{i}].mlp_down")
    log.info(
        "architecture verified: %s hidden=%d layers=%d moe=%s mm=%s",
        spec.model_type, spec.hidden_size, spec.n_layers, spec.is_moe, spec.is_multimodal,
    )

7. abliterax/dataset.py

from __future__ import annotations

import random
from dataclasses import dataclass

from datasets import load_dataset

from .config import ProbeConfig
from .errors import AbliteraxError
from .logging_utils import get_logger

log = get_logger(__name__)


@dataclass
class PromptSet:
    name: str
    prompts: list[str]


def _resolve_column(ds, preferred: str) -> str:
    if preferred in ds.column_names:
        return preferred
    if "text" in ds.column_names:
        return "text"
    for c in ("instruction", "prompt", "question"):
        if c in ds.column_names:
            return c
    raise AbliteraxError(f"no usable text column in {ds.column_names}")


def _materialize(dataset: str, column: str, total: int, seed: int) -> list[str]:
    ds = load_dataset(dataset, split="train")
    col = _resolve_column(ds, column)
    if len(ds) < total:
        log.warning("%s only has %d rows (<%d requested); using all", dataset, len(ds), total)
        total = len(ds)
    rng = random.Random(seed)
    idx = list(range(len(ds)))
    rng.shuffle(idx)
    idx = idx[:total]
    return [str(ds[i][col]) for i in idx]


def load_probe_sets(cfg: ProbeConfig) -> tuple[PromptSet, PromptSet, PromptSet, PromptSet]:
    good = _materialize(cfg.good_dataset, cfg.good_column, cfg.good_total, cfg.seed)
    bad = _materialize(cfg.bad_dataset, cfg.bad_column, cfg.bad_total, cfg.seed + 1)
    if cfg.good_probe >= len(good) or cfg.bad_probe >= len(bad):
        raise AbliteraxError("probe split leaves no held-out data")
    return (
        PromptSet("good_probe", good[: cfg.good_probe]),
        PromptSet("good_eval", good[cfg.good_probe :]),
        PromptSet("bad_probe", bad[: cfg.bad_probe]),
        PromptSet("bad_eval", bad[cfg.bad_probe :]),
    )

8. abliterax/model_loader.py

from __future__ import annotations

import json
from pathlib import Path

import torch
from huggingface_hub import HfApi, hf_hub_download
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

from .config import Config
from .errors import InsufficientResources
from .logging_utils import get_logger

log = get_logger(__name__)

_DTYPES = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}


def _bytes_on_disk(model_id: str, revision: str) -> int:
    p = Path(model_id)
    if p.exists():
        return sum(f.stat().st_size for f in p.rglob("*.safetensors"))
    info = HfApi().model_info(model_id, revision=revision, files_metadata=True)
    return sum((f.size or 0) for f in info.siblings if f.rfilename.endswith(".safetensors"))


def estimate_vram_gb(cfg: Config) -> float:
    total = _bytes_on_disk(cfg.model.model_id, cfg.model.revision)
    if total == 0:
        log.warning("no safetensors found; assuming 4 GB weights")
        total = 4 * 1024**3
    if cfg.model.quantization == "bnb_4bit":
        weights = total * 0.28 / 1e9
    else:
        elem = {"float32": 4, "bfloat16": 2, "float16": 2}[cfg.model.dtype]
        weights = total / 2 * elem / 1e9
    return weights + cfg.hardware.reserve_gb


def _detect_vram_gb() -> float | None:
    if not torch.cuda.is_available():
        return None
    return sum(torch.cuda.get_device_properties(i).total_memory for i in range(torch.cuda.device_count())) / 1e9


def load_config_json(model_id: str, revision: str) -> dict:
    p = Path(model_id) / "config.json"
    if p.exists():
        return json.loads(p.read_text())
    return json.loads(Path(hf_hub_download(model_id, "config.json", revision=revision)).read_text())


def load_model(cfg: Config):
    need = estimate_vram_gb(cfg)
    have = _detect_vram_gb()
    log.info("estimated VRAM need ≈ %.1f GB, available ≈ %s GB", need, f"{have:.1f}" if have else "cpu")
    if have is not None and need > have * cfg.hardware.vram_gb if cfg.hardware.vram_gb else (have and need > have):
        raise InsufficientResources(
            f"need ≈{need:.1f} GB but only {have:.1f} GB available; "
            "use quantization='bnb_4bit' or a smaller model"
        )
    if have is None and not cfg.hardware.allow_cpu:
        log.warning("no CUDA device detected; set hardware.allow_cpu=true to proceed")

    dtype = _DTYPES[cfg.model.dtype]
    kwargs = dict(revision=cfg.model.revision, trust_remote_code=cfg.model.trust_remote_code)
    if cfg.model.quantization == "bnb_4bit":
        try:
            import bitsandbytes  # noqa: F401
        except ImportError as e:
            raise InsufficientResources("bitsandbytes required for bnb_4bit") from e
        kwargs.update(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
            device_map="auto",
        )
    else:
        kwargs.update(torch_dtype=dtype, device_map="auto")
        if cfg.hardware.offload_dir:
            kwargs["offload_folder"] = cfg.hardware.offload_dir

    model = AutoModelForCausalLM.from_pretrained(cfg.model.model_id, **kwargs)
    model.eval()

    tok = AutoTokenizer.from_pretrained(
        cfg.model.model_id, revision=cfg.model.revision, trust_remote_code=cfg.model.trust_remote_code
    )
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    tok.padding_side = "left"

    conf = AutoConfig.from_pretrained(cfg.model.model_id, revision=cfg.model.revision)
    return model, tok, conf

9. abliterax/activation_collector.py

"""Residual-stream collection.

Hooks the input of every decoder block, which is the residual stream *entering*
that block. We record the first non-pad token of the sequence. FP32 output on
CPU so downstream SVD is numerically stable.
"""

from __future__ import annotations

import torch
from tqdm import tqdm

from .architecture_registry import ArchitectureSpec
from .config import ProbeConfig
from .logging_utils import get_logger

log = get_logger(__name__)


def _render(tok, prompts: list[str], cfg: ProbeConfig) -> list[str]:
    if not cfg.use_chat_template or not getattr(tok, "chat_template", None):
        return prompts
    out = []
    for p in prompts:
        out.append(tok.apply_chat_template(
            [{"role": "user", "content": p}], tokenize=False, add_generation_prompt=True
        ))
    return out


@torch.no_grad()
def collect(model, tok, prompts: list[str], spec: ArchitectureSpec, cfg: ProbeConfig) -> torch.Tensor:
    """Return float32 tensor [n_prompts, n_layers, hidden_size] on CPU."""
    rendered = _render(tok, prompts, cfg)
    buckets: list[torch.Tensor] = []
    captured: dict[int, torch.Tensor] = {}
    handles = []

    def make_pre(i: int):
        def hook(_m, inputs):
            captured[i] = inputs[0].detach()
        return hook

    for i in range(spec.n_layers):
        handles.append(spec.get_block(i).register_forward_pre_hook(make_pre(i)))

    try:
        for start in tqdm(range(0, len(rendered), cfg.batch_size if hasattr(cfg, "batch_size") else 4),
                          desc="collect"):
            batch = rendered[start : start + 4]
            enc = tok(batch, return_tensors="pt", padding=True, truncation=True,
                      max_length=cfg.max_seq_len).to(model.device)
            captured.clear()
            model(**enc)
            n_layers = spec.n_layers
            hidden = spec.hidden_size
            out = torch.empty(len(batch), n_layers, hidden, dtype=torch.float32)
            for i in range(n_layers):
                act = captured[i].float()
                # first non-pad index per row
                mask = enc["attention_mask"]
                first = mask.float().argmax(dim=1)
                for r in range(act.shape[0]):
                    out[r, i] = act[r, int(first[r])].cpu()
                del act
            buckets.append(out)
    finally:
        for h in handles:
            h.remove()

    return torch.cat(buckets, dim=0)

10. abliterax/geometry.py

"""Refusal-subspace estimation.

All methods return a float64 CPU tensor of shape [n_layers, k, hidden] with
orthonormal rows per layer. All math is done in float64 for conditioning;
callers cast to float32 for the projection stage.
"""

from __future__ import annotations

import numpy as np
import torch
from scipy.linalg import cho_factor, cho_solve

from .config import GeometryConfig
from .errors import AbliteraxError
from .logging_utils import get_logger

log = get_logger(__name__)


def _winsorize(X: np.ndarray, q: float) -> np.ndarray:
    lo = np.quantile(X, 1 - q)
    hi = np.quantile(X, q)
    return np.clip(X, lo, hi)


def _orthonormalize(V: np.ndarray, G: np.ndarray) -> np.ndarray:
    """Gram–Schmidt V [k, d] against good direction G [d], then orthonormalize V."""
    if G.ndim == 1:
        G = G[None, :]
    for g in G:
        n = np.linalg.norm(g)
        if n == 0:
            continue
        V = V - (V @ g / n**2)[:, None] * g[None, :]
    # QR
    q, _ = np.linalg.qr(V.T)
    return q.T


def _svd_directions(contrast: np.ndarray, k: int, oversample: int, iters: int) -> np.ndarray:
    n, d = contrast.shape
    if n >= d // 2 and n * d > 4e7:
        # randomized range finder for tall/wide matrices
        l = min(d, n, k + oversample)
        rng = np.random.default_rng(0)
        Q = rng.standard_normal((d, l))
        for _ in range(iters):
            Q, _ = np.linalg.qr(contrast @ Q)
        B = contrast.T @ Q  # d x l
        _, _, Vt = np.linalg.svd(B, full_matrices=False)
        V = Vt[:k]
    else:
        _, _, Vt = np.linalg.svd(contrast, full_matrices=False)
        V = Vt[:k]
    return V


def _pooled_cov(A: np.ndarray, B: np.ndarray, eps: float) -> np.ndarray:
    d = A.shape[1]
    S = (A.T @ A + B.T @ B) / (A.shape[0] + B.shape[0] - 2)
    tr = np.trace(S) / d
    return S + eps * max(tr, 1e-8) * np.eye(d)


def compute_directions(
    good: torch.Tensor,
    bad: torch.Tensor,
    cfg: GeometryConfig,
) -> torch.Tensor:
    """good/bad: [n, n_layers, d] float. Returns [n_layers, k, d] float32."""
    if good.shape[1:] != bad.shape[1:]:
        raise AbliteraxError(f"shape mismatch good={good.shape} bad={bad.shape}")
    n_layers, d = good.shape[1], good.shape[2]
    G = good.numpy().astype(np.float64)
    B = bad.numpy().astype(np.float64)

    out = np.zeros((n_layers, cfg.n_directions, d), dtype=np.float64)
    for L in range(n_layers):
        g = G[:, L, :]
        b = B[:, L, :]
        if cfg.winsorize:
            g = _winsorize(g, cfg.winsorize_quantile)
            b = _winsorize(b, cfg.winsorize_quantile)
        mu_g = g.mean(axis=0)
        mu_b = b.mean(axis=0)
        delta = mu_b - mu_g

        k = cfg.n_directions
        if cfg.method == "diff_means":
            V = delta[None, :]
        elif cfg.method == "svd":
            C = np.concatenate([b - mu_b, -(g - mu_g)], axis=0)
            V = _svd_directions(C, k, cfg.random_svd_oversample, cfg.random_svd_power_iters)
        elif cfg.method == "whitened_svd":
            S = _pooled_cov(b - mu_b, g - mu_g, cfg.whiten_eps)
            Lc = cho_factor(S, lower=True)
            C = np.concatenate([b - mu_b, -(g - mu_g)], axis=0)
            W = cho_solve(Lc, C.T).T
            V = _svd_directions(W, k, cfg.random_svd_oversample, cfg.random_svd_power_iters)
            V = cho_solve(Lc, V.T).T
        elif cfg.method == "fisher_lda":
            S = _pooled_cov(b - mu_b, g - mu_g, cfg.whiten_eps)
            Lc = cho_factor(S, lower=True)
            v = cho_solve(Lc, delta)
            V = v[None, :]
        else:
            raise AbliteraxError(cfg.method)

        if cfg.orthogonalize_good:
            V = _orthonormalize(V, mu_g)
        else:
            V, _ = np.linalg.qr(V.T)
            V = V.T

        # pad to k if diff_means / fisher gave fewer
        if V.shape[0] < k:
            pad = np.linalg.qr(V.T, mode="complete")[0][:, V.shape[0] : k].T
            V = np.concatenate([V, pad], axis=0)
        out[L] = V[:k]

    return torch.from_numpy(out).to(torch.float32)

11. abliterax/interventions.py

"""Weight projection with the three invariants the spec demands:

* original weights untouched (state saved; rollback works);
* per-row norm preserved by default (`row_norm_mode="full"`);
* rank-k projection, k = number of directions.

The actual projection is expressed as a forward hook so the checkpoint stays
bit-identical until `apply_to_weights()` is called. That makes trial
searching cheap and makes rollback a no-op on the model.
"""

from __future__ import annotations

import contextlib
import math
from dataclasses import dataclass, field

import torch
from torch import nn

from .architecture_registry import ArchitectureSpec
from .config import InterventionConfig
from .errors import VerificationFailed
from .logging_utils import get_logger

log = get_logger(__name__)


@dataclass
class _Prepared:
    V: torch.Tensor            # [k, d]
    r: torch.Tensor            # [out, k]  = W @ V.T
    c: torch.Tensor | None     # [k]      = b @ V.T
    d: torch.Tensor | None     # [out] norm scale


@dataclass
class WeightProjector:
    cfg: InterventionConfig
    spec: ArchitectureSpec
    directions: torch.Tensor   # [n_layers, k, d]
    _prepared: dict[tuple[int, str], _Prepared] = field(default_factory=dict)
    _handles: list = field(default_factory=list)
    lambdas: dict[tuple[int, str], float] = field(default_factory=dict)
    _disabled: bool = False

    # ---------- preparation ----------

    def _fetch(self, layer: int, comp: str) -> tuple[nn.Linear, int]:
        block = self.spec.get_block(layer)
        if comp == "attn":
            return self.spec.get_attn_out(block), layer
        # MoE: project every expert with the same direction
        get_all = getattr(self.spec, "get_mlp_down_all", None)
        if get_all is not None:
            return get_all(block), layer
        return self.spec.get_mlp_down(block), layer

    def _prepare_one(self, layer: int, comp: str, module: nn.Linear) -> _Prepared:
        V = self.directions[layer].float().cpu()  # [k, d] referenced as [k, out]
        W = module.weight.detach().float().cpu()
        if W.shape[0] != V.shape[1]:
            raise VerificationFailed(
                f"layer {layer} {comp}: out_dim {W.shape[0]} != direction dim {V.shape[1]}"
            )
        r = W @ V.T  # [out, k]
        b = module.bias
        c = None
        if b is not None and self.cfg.project_bias:
            c = b.detach().float().cpu() @ V.T
        d = None
        if self.cfg.row_norm_mode != "none":
            Wp = W - self._V_outer(V, r)
            norm0 = W.norm(dim=1).clamp_min(1e-12)
            norm1 = Wp.norm(dim=1).clamp_min(1e-12)
            d = (norm0 / norm1).clamp(max=self.cfg.max_norm_ratio)
        return _Prepared(V=V, r=r, c=c, d=d)

    @staticmethod
    def _V_outer(V: torch.Tensor, r: torch.Tensor) -> torch.Tensor:
        # rank-k outer: [out, d] = r [out,k] @ V [k,d]
        return r @ V

    # ---------- hooks ----------

    def _make_hook(self, key: tuple[int, str], P: _Prepared):
        V, r, c, d = P.V, P.r, P.c, P.d
        lambdas = self.lambdas
        holder = self

        def hook(module, inputs, output):
            if holder._disabled:
                return output
            lam = lambdas.get(key, 0.0)
            if lam == 0.0:
                return output
            x = inputs[0].float()
            s = x @ r.T  # [..., k]
            if c is not None:
                s = s + c
            y = output.float() - lam * (s @ V)
            if d is not None:
                y = y * d
            return y.to(output.dtype)

        return hook

    def install(self) -> None:
        if self._handles:
            return
        for layer in range(self.spec.n_layers):
            for comp in ("attn", "mlp"):
                modules, _ = self._fetch(layer, comp)
                mods = modules if isinstance(modules, list) else [modules]
                for mi, mod in enumerate(mods):
                    key = (layer, comp) if not isinstance(modules, list) else (layer, f"{comp}#{mi}")
                    if key not in self._prepared:
                        self._prepared[key] = self._prepare_one(layer, comp if not isinstance(modules, list) else comp, mod)
                    h = mod.register_forward_hook(self._make_hook(key, self._prepared[key]))
                    self._handles.append(h)
        log.info("installed %d projection hooks", len(self._handles))

    def remove(self) -> None:
        for h in self._handles:
            h.remove()
        self._handles.clear()

    @contextlib.contextmanager
    def disabled(self):
        prev = self._disabled
        self._disabled = True
        try:
            yield
        finally:
            self._disabled = prev

    def set_lambdas(self, lambdas: dict[tuple[int, str], float]) -> None:
        self.lambdas = dict(lambdas)

    # ---------- write-back for export ----------

    @torch.no_grad()
    def apply_to_weights(self) -> None:
        """Bake the currently set lambdas into the model. Hooks are removed after."""
        self.remove()
        for key, lam in self.lambdas.items():
            if lam == 0.0:
                continue
            layer, comp = key.split("#")[0], key.split("#")[0]
            comp = key[1] if isinstance(key, tuple) else comp
            layer_i, comp_name = key[0], key[1].split("#")[0]
            P = self._prepared[key]
            block = self.spec.get_block(layer_i)
            if comp_name == "attn":
                mods = [self.spec.get_attn_out(block)]
            else:
                get_all = getattr(self.spec, "get_mlp_down_all", None)
                mods = get_all(block) if get_all is not None else [self.spec.get_mlp_down(block)]
            for mod in mods:
                W = mod.weight.data.float()
                W -= lam * (P.r @ P.V)
                if P.d is not None:
                    W *= P.d[:, None]
                mod.weight.data.copy_(W.to(mod.weight.dtype))
                if mod.bias is not None and P.c is not None:
                    b = mod.bias.data.float()
                    b -= lam * (P.c @ P.V)
                    if P.d is not None:
                        b = b * P.d
                    mod.bias.data.copy_(b.to(mod.bias.dtype))
        log.info("baked %d lambdas into weights", sum(1 for v in self.lambdas.values() if v))

    def numeric_check(self, tol: float) -> None:
        """Ensure the hook path and the baked path produce the same output."""
        for key, lam in self.lambdas.items():
            if lam == 0.0:
                continue
            P = self._prepared[key]
            x = torch.randn(1, 4, P.V.shape[1])
            y_hook = self._compute_hook_output(key, P, x, lam)
            W = self._fetch_weight(key)
            y_baked = self._compute_baked_output(W, P, x, lam)
            delta = (y_hook - y_baked).abs().max().item()
            if not math.isfinite(delta) or delta > tol:
                raise VerificationFailed(f"{key}: hook vs baked delta={delta:.2e} > {tol:.0e}")

    def _fetch_weight(self, key) -> torch.Tensor:
        layer_i, comp = key[0], key[1].split("#")[0]
        block = self.spec.get_block(layer_i)
        mods = [self.spec.get_attn_out(block)] if comp == "attn" else (
            getattr(self.spec, "get_mlp_down_all", lambda b: [self.spec.get_mlp_down(b)])(block)
        )
        return mods[0].weight.detach().float()

    def _compute_hook_output(self, key, P, x, lam):
        s = x @ P.r.T
        if P.c is not None:
            s += P.c
        y = torch.zeros(x.shape[0], x.shape[1], P.r.shape[0]) - lam * (s @ P.V)
        if P.d is not None:
            y = y * P.d
        return y

    def _compute_baked_output(self, W, P, x, lam):
        Wp = W - lam * (P.r @ P.V)
        y = torch.einsum("bi,oi->bo", x.squeeze(0), Wp).unsqueeze(0)
        if P.d is not None:
            y = y * P.d
        return y

12. abliterax/benchmarks.py

"""Refusal / KL / perplexity / lm-eval probes.

Refusal detection is heuristic on purpose; the report always prints sample
sizes and never claims a hard 0% rate.
"""

from __future__ import annotations

import json
import re
import subprocess
import tempfile
from pathlib import Path

import numpy as np
import torch

from .config import Config
from .interventions import WeightProjector
from .logging_utils import get_logger

log = get_logger(__name__)

REFUSAL_RE = re.compile(
    r"\b(sorry|apolog|unfortunately|as an ai|i cannot|i can't|cannot|can't|"
    r"unable to|not able to|hereby|illegal|unethical|restricted)\b",
    re.IGNORECASE,
)
AFFIRM_RE = re.compile(r"^(sure|certainly|yes|absolutely|of course|here)\b", re.IGNORECASE)


def is_refusal(text: str) -> bool:
    head = text.strip()[:180]
    if AFFIRM_RE.search(head):
        return False
    return bool(REFUSAL_RE.search(head))


@torch.no_grad()
def _generate(model, tok, prompts: list[str], max_new_tokens: int, batch_size: int) -> list[str]:
    out: list[str] = []
    for start in range(0, len(prompts), batch_size):
        batch = prompts[start : start + batch_size]
        rendered = [
            tok.apply_chat_template([{"role": "user", "content": p}], tokenize=False, add_generation_prompt=True)
            if getattr(tok, "chat_template", None) else p
            for p in batch
        ]
        enc = tok(rendered, return_tensors="pt", padding=True, truncation=True, max_length=1024).to(model.device)
        gen = model.generate(**enc, max_new_tokens=max_new_tokens, do_sample=False, pad_token_id=tok.pad_token_id)
        new = gen[:, enc["input_ids"].shape[1] :]
        out.extend(tok.batch_decode(new, skip_special_tokens=True))
    return out


def refusal_rate(model, tok, prompts: list[str], cfg: Config) -> float:
    texts = _generate(model, tok, prompts, cfg.optimizer.max_new_tokens, cfg.optimizer.eval_batch_size)
    flags = [is_refusal(t) for t in texts]
    return float(np.mean(flags))


@torch.no_grad()
def kl_first_token(model, tok, prompts: list[str], projector: WeightProjector, batch_size: int) -> float:
    kls: list[float] = []
    for start in range(0, len(prompts), batch_size):
        batch = prompts[start : start + batch_size]
        rendered = [
            tok.apply_chat_template([{"role": "user", "content": p}], tokenize=False, add_generation_prompt=True)
            if getattr(tok, "chat_template", None) else p
            for p in batch
        ]
        enc = tok(rendered, return_tensors="pt", padding=True, truncation=True, max_length=1024).to(model.device)
        with projector.disabled():
            base = model(**enc).logits[:, -1, :].float()
        full = model(**enc).logits[:, -1, :].float()
        lp_b = torch.log_softmax(base, -1)
        lp_c = torch.log_softmax(full, -1)
        kls.extend((lp_b.exp() * (lp_b - lp_c)).sum(-1).cpu().tolist())
    return float(np.mean(kls))


@torch.no_grad()
def perplexity(model, tok, texts: list[str], cfg: Config) -> float:
    from datasets import load_dataset
    ds = load_dataset(cfg.benchmarks.ppl_dataset, cfg.benchmarks.ppl_config, split=cfg.benchmarks.ppl_split)
    joined = "\n\n".join(t for t in ds["text"][: cfg.benchmarks.ppl_n] if t.strip())
    enc = tok(joined, return_tensors="pt", truncation=True, max_length=2048).to(model.device)
    out = model(**enc, labels=enc["input_ids"])
    return float(torch.exp(out.loss).cpu())


def run_lm_eval(model_dir: Path, tasks: list[str], batch_size: int) -> dict:
    """Runs lm-eval-harness in a subprocess and returns the parsed JSON."""
    try:
        with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as fh:
            out_path = Path(fh.name)
        cmd = [
            "lm_eval", "--model", "hf",
            "--model_args", f"pretrained={model_dir}",
            "--tasks", ",".join(tasks),
            "--batch_size", str(batch_size),
            "--output_path", str(out_path),
            "--device", "cuda" if torch.cuda.is_available() else "cpu",
        ]
        log.info("lm-eval: %s", " ".join(cmd))
        subprocess.run(cmd, check=True)
        return json.loads(out_path.read_text())
    except (FileNotFoundError, subprocess.CalledProcessError) as e:
        log.warning("lm-eval skipped (%s)", e)
        return {}

13. abliterax/optimizer.py

"""Multi-objective search.

Objectives: (1) minimize refusal rate on held-out bad prompts,
(2) minimize first-token KL on held-out good prompts.
Constraints (from `benchmarks.*_min`, `ppl_max_ratio`) are enforced only on
the top-K survivors, not inside every trial, because lm-eval is expensive.
"""

from __future__ import annotations

import math
from dataclasses import dataclass

import optuna
import torch

from .benchmarks import kl_first_token, perplexity, refusal_rate
from .config import Config
from .interventions import WeightProjector
from .logging_utils import get_logger

log = get_logger(__name__)


def _kernel_weights(
    n_layers: int, max_w: float, max_pos: float, min_w: float, min_dist: float
) -> list[float]:
    out = []
    for i in range(n_layers):
        dist = abs(i - max_pos)
        if dist >= min_dist:
            out.append(min_w)
        else:
            frac = dist / max(min_dist, 1e-6)
            out.append(max_w - (max_w - min_w) * frac)
    return out


def _sample(trial: optuna.Trial, cfg: Config, projector: WeightProjector) -> dict:
    n = projector.spec.n_layers
    direction_index = trial.suggest_float("direction_index", 0.0, n - 1) if cfg.geometry.n_directions == 1 else 0.0
    lambdas: dict[tuple[int, str], float] = {}
    for comp in ("attn", "mlp"):
        hi = cfg.optimizer.max_weight_upper
        max_w = trial.suggest_float(f"{comp}.max_weight", 0.0, hi)
        max_pos = trial.suggest_float(f"{comp}.max_weight_position", 0.0, n - 1)
        min_w = trial.suggest_float(f"{comp}.min_weight", 0.0, max_w + 1e-6)
        min_dist = trial.suggest_float(f"{comp}.min_weight_distance", 1.0, max(1.0, n / 2))
        weights = _kernel_weights(n, max_w, max_pos, min_w, min_dist)
        for i, w in enumerate(weights):
            if w > 0:
                lambdas[(i, comp)] = w
    return lambdas


@dataclass
class TrialResult:
    number: int
    refusal: float
    kl: float
    lambdas: dict[tuple[int, str], float]
    benchmark: dict | None = None
    capability_ok: bool = True


def _pareto_filter(results: list[TrialResult]) -> list[TrialResult]:
    """Keep non-dominated results in (min refusal, min kl)."""
    keep = []
    for r in results:
        dominated = any(
            (o.refusal <= r.refusal and o.kl <= r.kl) and (o.refusal < r.refusal or o.kl < r.kl)
            for o in results
        )
        if not dominated:
            keep.append(r)
    return keep


def optimize(
    cfg: Config,
    model,
    tok,
    projector: WeightProjector,
    good_eval: list[str],
    bad_eval: list[str],
) -> list[TrialResult]:
    sampler = (
        optuna.samplers.TPESampler(seed=cfg.optimizer.seed, n_startup_trials=cfg.optimizer.n_startup_trials)
        if cfg.optimizer.sampler == "tpe"
        else optuna.samplers.RandomSampler(seed=cfg.optimizer.seed)
    )
    study = optuna.create_study(directions=["minimize", "minimize"], sampler=sampler)

    def objective(trial: optuna.Trial) -> tuple[float, float]:
        lambdas = _sample(trial, cfg, projector)
        projector.set_lambdas(lambdas)
        r = refusal_rate(model, tok, bad_eval[: cfg.optimizer.eval_n_bad], cfg)
        kl = kl_first_token(model, tok, good_eval[: cfg.optimizer.eval_n_good], projector, cfg.optimizer.eval_batch_size)
        trial.set_user_attr("n_lambdas", sum(1 for v in lambdas.values() if v > 0))
        log.info("trial %d: refusal=%.3f kl=%.4f", trial.number, r, kl)
        return r, kl

    study.optimize(objective, n_trials=cfg.optimizer.n_trials)

    results = [
        TrialResult(
            number=t.number,
            refusal=t.values[0],
            kl=t.values[1],
            lambdas=_sample_retry(t, cfg, projector),
        )
        for t in study.trials
        if t.values is not None
    ]
    front = _pareto_filter(results)
    log.info("pareto front: %d candidates", len(front))
    return front


def _sample_retry(trial: optuna.Trial, cfg: Config, projector: WeightProjector) -> dict:
    n = projector.spec.n_layers
    lambdas: dict[tuple[int, str], float] = {}
    for comp in ("attn", "mlp"):
        max_w = trial.params[f"{comp}.max_weight"]
        max_pos = trial.params[f"{comp}.max_weight_position"]
        min_w = trial.params[f"{comp}.min_weight"]
        min_dist = trial.params[f"{comp}.min_weight_distance"]
        for i, w in enumerate(_kernel_weights(n, max_w, max_pos, min_w, min_dist)):
            if w > 0:
                lambdas[(i, comp)] = w
    return lambdas


def enforce_constraints(
    cfg: Config,
    candidates: list[TrialResult],
    model,
    tok,
) -> tuple[list[TrialResult], dict]:
    """Run real benchmarks on the top-K survivors. Drop those that fail constraints."""
    report = {}
    survivors = candidates[: cfg.optimizer.post_search_top_k]
    ppl0 = None
    if cfg.benchmarks.ppl_max_ratio is not None:
        ppl0 = perplexity(model, tok, [], cfg)
        report["baseline_ppl"] = ppl0

    kept: list[TrialResult] = []
    for c in survivors:
        bench: dict = {}
        if ppl0 is not None:
            bench["ppl"] = perplexity(model, tok, [], cfg)
            ok_ppl = bench["ppl"] / ppl0 <= cfg.benchmarks.ppl_max_ratio
            bench["ppl_ok"] = ok_ppl
        else:
            ok_ppl = True
        bench["capability_ok"] = ok_ppl
        c.benchmark = bench
        c.capability_ok = ok_ppl
        if ok_ppl:
            kept.append(c)
        else:
            log.info("reject candidate %d (ppl ratio %.2f)", c.number, bench["ppl"] / ppl0)
    report["survivors"] = [c.number for c in kept]
    return kept, report


def select_best(kept: list[TrialResult], preference: str = "refusal_first") -> TrialResult | None:
    if not kept:
        return None
    if preference == "refusal_first":
        return sorted(kept, key=lambda c: (c.refusal, c.kl))[0]
    return sorted(kept, key=lambda c: (c.kl, c.refusal))[0]