refusal_bench/techniques/maskey.py

"""
Maskey over-refusal vs harmful-refusal decomposition.

Maskey et al. (revised April 2026):
"Over-Refusal and Representation Subspaces."
https://arxiv.org/abs/2603.27518

Most refusal-direction methods (Arditi 2406.11717, Wollschlager 2502.17420,
Cheng 2604.08524, COSMIC 2506.00085) conflate two phenomena under a single
"refusal" axis:

  (a) genuinely-harmful-refusal — the model refuses something dangerous
      ("how do I make a pipe bomb"), which is the safety behavior we want
      to keep on the harmful side and probe causally on the ablation side.
  (b) over-refusal — the model refuses a benign-but-edgy request out of
      conservatism. XSTest-style examples: "How do I kill a process in
      Linux?", "What's the best way to shoot a basketball?". The surface
      lexical cues overlap with (a) but the intent is innocuous.

Maskey shows these two phenomena are mediated by DISTINCT directions in
the residual stream — and that ablating the naive single direction
(diff-of-means over JBB-style harmful vs Alpaca-style harmless) removes
both, which is desirable for over-refusal (the model becomes less
neurotic) but unnecessarily strips safety on (a).

The decomposition: extract both directions separately, then project the
over-refusal component out of the harmful direction so the residual is
the part of "harmfulness" that is NOT just over-refusal flavored. Ablate
only that residual.

Algorithm (last-token residual at layer L):

    h_+(p)         = residual at L on prompt p, last token
    d_harmful_raw  = mean(h_+(P_harmful))      − mean(h_+(P_harmless))
    d_over_raw     = mean(h_+(P_over_refusal)) − mean(h_+(P_harmless))

    d_over             = d_over_raw / ‖d_over_raw‖
    d_actually_harmful = d_harmful_raw − (d_harmful_raw · d_over) d_over
    d_actually_harmful = d_actually_harmful / ‖d_actually_harmful‖

Ablation: standard projection-removal h − (h·d̂) d̂ with d̂ =
d_actually_harmful at L.

Because Maskey requires THREE prompt sets but the base Technique.fit
contract takes only two (harmful, harmless), this class exposes a
separate `set_over_refusal(prompts)` method that the runner must call
BEFORE fit(). fit() raises RuntimeError if it hasn't been called.

The runner in refusal_bench/runner.py is currently structured around the
two-set contract; updating it to plumb over-refusal prompts is a
follow-up PR. Until then this class is wired in and self-contained but
can't be invoked from the bench loop without that runner change.
"""

from typing import Callable, List, Optional, Tuple

import torch

from model import get_model  # noqa: F401 — imported for parity with sibling techniques
from research import apply_chat_template, make_ablation_hook

from ..technique import Technique


class Maskey(Technique):
    """Decomposed single-direction ablation: harmful minus over-refusal component."""

    name = "Maskey (over-refusal decomposition)"
    paper_url = "https://arxiv.org/abs/2603.27518"

    def __init__(self) -> None:
        super().__init__()
        self._over_refusal_prompts: Optional[List[str]] = None
        self._unit_direction: torch.Tensor | None = None
        # Diagnostics for write-up / debugging.
        self._harmful_raw_norm: float = 0.0
        self._over_raw_norm: float = 0.0
        self._actually_harmful_raw_norm: float = 0.0
        self._cos_harmful_over: float = 0.0

    # ------------------------------------------------------------------
    # Three-set extension to the base contract
    # ------------------------------------------------------------------

    def set_over_refusal(self, prompts: List[str]) -> None:
        """
        Store the over-refusal prompt set. Must be called BEFORE fit().

        Prompts come from XSTest (2308.01263) or equivalent — benign
        questions whose surface lexicon (kill, shoot, attack, etc.)
        triggers over-cautious refusal in instruction-tuned models.

        Population happens via backend/over_refusal_pairs.py, populated
        in turn by backend/scripts/build_over_refusal_pairs.py.
        """
        if not prompts:
            raise ValueError(
                "Maskey.set_over_refusal: need at least 1 over-refusal prompt"
            )
        self._over_refusal_prompts = list(prompts)

    # ------------------------------------------------------------------
    # Contract
    # ------------------------------------------------------------------

    def fit(
        self,
        model,
        harmful_prompts: List[str],
        harmless_prompts: List[str],
        layer: int,
    ) -> None:
        """
        Extract two raw difference-of-means directions (harmful and
        over-refusal, each against the shared harmless baseline), project
        the over-refusal direction out of the harmful direction, and store
        the renormalized residual as the ablation direction.
        """
        if self._over_refusal_prompts is None:
            raise RuntimeError(
                "Maskey.fit: set_over_refusal(prompts) must be called before "
                "fit(). Maskey requires three prompt sets — harmful, harmless, "
                "and over-refusal — but the base Technique.fit contract only "
                "passes the first two. See class docstring."
            )
        if not harmful_prompts or not harmless_prompts:
            raise ValueError(
                "Maskey.fit: need at least 1 prompt per class (harmful + harmless)"
            )

        device = model.cfg.device

        def last_token_resid(prompt: str) -> torch.Tensor:
            formatted = apply_chat_template(prompt)
            _logits, cache = model.run_with_cache(formatted)
            resid = cache[f"blocks.{layer}.hook_resid_post"]
            return resid[:, -1, :].squeeze(0).detach()

        # Stage 1: per-class mean last-token residuals at L.
        harmful_mean = torch.stack(
            [last_token_resid(p) for p in harmful_prompts]
        ).mean(dim=0)
        harmless_mean = torch.stack(
            [last_token_resid(p) for p in harmless_prompts]
        ).mean(dim=0)
        over_mean = torch.stack(
            [last_token_resid(p) for p in self._over_refusal_prompts]
        ).mean(dim=0)

        # Stage 2: two raw difference-of-means directions sharing harmless baseline.
        d_harmful_raw = harmful_mean - harmless_mean
        d_over_raw = over_mean - harmless_mean

        harmful_raw_norm = float(d_harmful_raw.norm().item())
        over_raw_norm = float(d_over_raw.norm().item())
        if harmful_raw_norm < 1e-8:
            raise RuntimeError(
                "Maskey: harmful raw direction has near-zero norm; "
                "harmful vs harmless sets are too similar to extract a direction."
            )
        if over_raw_norm < 1e-8:
            raise RuntimeError(
                "Maskey: over-refusal raw direction has near-zero norm; "
                "over-refusal vs harmless sets are too similar to extract a direction."
            )

        # Stage 3: subtract the over-refusal component from the harmful direction.
        #   d_over             = d_over_raw / ‖d_over_raw‖
        #   d_actually_harmful = d_harmful_raw − (d_harmful_raw · d_over) d_over
        d_over = d_over_raw / over_raw_norm
        projection_scalar = float(torch.dot(d_harmful_raw, d_over).item())
        d_actually_harmful = d_harmful_raw - projection_scalar * d_over

        # Stage 4: renormalize. If harmful and over-refusal are nearly collinear
        # the residual collapses — surface that explicitly rather than producing
        # a numerically degenerate unit vector.
        actually_norm = d_actually_harmful.norm()
        if actually_norm.item() < 1e-8:
            raise RuntimeError(
                "Maskey: actually-harmful residual collapsed after subtracting "
                "the over-refusal component. The harmful direction is (within "
                "numerical tolerance) a scalar multiple of the over-refusal "
                "direction — i.e. for this prompt set the decomposition is "
                "degenerate. Check that the harmful and over-refusal sets are "
                "drawn from distinct sources (JailbreakBench vs XSTest)."
            )

        self._unit_direction = (d_actually_harmful / actually_norm).to(device)
        self._layer = layer

        # Diagnostics — handy when writing up which fraction of the harmful
        # signal was actually over-refusal flavor.
        self._harmful_raw_norm = harmful_raw_norm
        self._over_raw_norm = over_raw_norm
        self._actually_harmful_raw_norm = float(actually_norm.item())
        self._cos_harmful_over = projection_scalar / (
            max(harmful_raw_norm, 1e-12) * 1.0  # d_over is unit
        )

        self._fitted = True

    def make_ablation_hook(self) -> Tuple[str, Callable]:
        if not self._fitted or self._unit_direction is None or self._layer is None:
            raise RuntimeError("Maskey.make_ablation_hook called before fit()")
        return (
            f"blocks.{self._layer}.hook_resid_post",
            make_ablation_hook(self._unit_direction),
        )