Interpretability

Occlusion sensitivity and related probes.

polyweave.interpretability

Interpretability probes for PolyWeave layers and teachers.

Currently: occlusion sensitivity, used to distinguish additive from multiplicative (Sigma-Pi) features via their conjunctive AND-signature.

conjunction_index

conjunction_index(response_fn: ResponseFn, x: Tensor, group_a: Sequence[int] | Tensor, group_b: Sequence[int] | Tensor, *, baseline: float | Tensor = 0.0, eps: float = 1e-08) -> torch.Tensor

AND-signature index in [0, 1] per item.

index = (drop_a + drop_b - drop_ab) / (|drop_ab| + eps), clamped to [0, 1]:

• ~0 — additive feature (single drops add up to the joint drop, no interaction).
• ~1 — multiplicative / conjunctive feature (either factor alone already collapses the response, so the single drops far exceed the joint drop).

Returns [N]. Aggregate with .mean() for a per-feature scalar.

Source code in polyweave/interpretability/occlusion.py

@torch.no_grad()
def conjunction_index(
    response_fn: ResponseFn,
    x: torch.Tensor,
    group_a: Sequence[int] | torch.Tensor,
    group_b: Sequence[int] | torch.Tensor,
    *,
    baseline: float | torch.Tensor = 0.0,
    eps: float = 1e-8,
) -> torch.Tensor:
    """AND-signature index in ``[0, 1]`` per item.

    ``index = (drop_a + drop_b - drop_ab) / (|drop_ab| + eps)``, clamped to
    ``[0, 1]``:

    * **~0** — additive feature (single drops add up to the joint drop, no
      interaction).
    * **~1** — multiplicative / conjunctive feature (either factor alone already
      collapses the response, so the single drops far exceed the joint drop).

    Returns ``[N]``. Aggregate with ``.mean()`` for a per-feature scalar.
    """
    d = group_drops(response_fn, x, group_a, group_b, baseline=baseline)
    num = d["drop_a"] + d["drop_b"] - d["drop_ab"]
    idx = num / (d["drop_ab"].abs() + eps)
    return idx.clamp(0.0, 1.0)

group_drops

group_drops(response_fn: ResponseFn, x: Tensor, group_a: Sequence[int] | Tensor, group_b: Sequence[int] | Tensor, *, baseline: float | Tensor = 0.0) -> dict

Single- and joint-occlusion drops for two disjoint feature groups.

Returns a dict with per-item tensors drop_a, drop_b, drop_ab and interaction = drop_ab - drop_a - drop_b (each [N]). Indices select along the flattened feature axis.

Source code in polyweave/interpretability/occlusion.py

@torch.no_grad()
def group_drops(
    response_fn: ResponseFn,
    x: torch.Tensor,
    group_a: Sequence[int] | torch.Tensor,
    group_b: Sequence[int] | torch.Tensor,
    *,
    baseline: float | torch.Tensor = 0.0,
) -> dict:
    """Single- and joint-occlusion drops for two disjoint feature groups.

    Returns a dict with per-item tensors ``drop_a``, ``drop_b``, ``drop_ab`` and
    ``interaction = drop_ab - drop_a - drop_b`` (each ``[N]``). Indices select
    along the flattened feature axis.
    """
    a = torch.as_tensor(list(group_a) if not isinstance(group_a, torch.Tensor) else group_a, dtype=torch.long)
    b = torch.as_tensor(list(group_b) if not isinstance(group_b, torch.Tensor) else group_b, dtype=torch.long)
    base = _as_response(response_fn(x))
    drop_a = base - _as_response(response_fn(_occlude(x, a, baseline)))
    drop_b = base - _as_response(response_fn(_occlude(x, b, baseline)))
    ab = torch.cat([a, b])
    drop_ab = base - _as_response(response_fn(_occlude(x, ab, baseline)))
    return {
        "drop_a": drop_a,
        "drop_b": drop_b,
        "drop_ab": drop_ab,
        "interaction": drop_ab - drop_a - drop_b,
    }

occlusion_sensitivity_1d

occlusion_sensitivity_1d(response_fn: ResponseFn, x: Tensor, *, baseline: float | Tensor = 0.0, window: int = 1, stride: int = 1) -> torch.Tensor

Per-feature occlusion sensitivity map over the feature axis of x.

Parameters:

Name	Type	Description	Default
response_fn	ResponseFn	maps [N, ...] -> per-item response [N] (or a tensor reduced to [N] by mean over trailing dims).	required
x	Tensor	input batch [N, ...]; dims after the batch axis are flattened to a single feature axis of length F.	required
baseline	float | Tensor	value substituted into the occluded window (default 0.0).	0.0
window	int	width of the occluding window in features.	1
stride	int	step between successive window starts.	1

Returns:

Type	Description
Tensor	[N, P] tensor of drops (response(x) - response(x_occluded)), one
Tensor	column per window position P; large positive entries mark features the
Tensor	response most depends on.

Source code in polyweave/interpretability/occlusion.py

@torch.no_grad()
def occlusion_sensitivity_1d(
    response_fn: ResponseFn,
    x: torch.Tensor,
    *,
    baseline: float | torch.Tensor = 0.0,
    window: int = 1,
    stride: int = 1,
) -> torch.Tensor:
    """Per-feature occlusion sensitivity map over the feature axis of ``x``.

    Args:
        response_fn: maps ``[N, ...]`` -> per-item response ``[N]`` (or a tensor
            reduced to ``[N]`` by mean over trailing dims).
        x: input batch ``[N, ...]``; dims after the batch axis are flattened to a
            single feature axis of length ``F``.
        baseline: value substituted into the occluded window (default ``0.0``).
        window: width of the occluding window in features.
        stride: step between successive window starts.

    Returns:
        ``[N, P]`` tensor of *drops* (``response(x) - response(x_occluded)``), one
        column per window position ``P``; large positive entries mark features the
        response most depends on.
    """
    base = _as_response(response_fn(x))  # [N]
    F = x.reshape(x.shape[0], -1).shape[1]
    starts = list(range(0, F - window + 1, stride))
    if starts and starts[-1] != F - window:
        starts.append(F - window)
    cols = []
    for s in starts:
        idx = slice(s, s + window)
        occ = _as_response(response_fn(_occlude(x, idx, baseline)))
        cols.append(base - occ)
    return torch.stack(cols, dim=1) if cols else base.new_zeros(x.shape[0], 0)

occlusion_sensitivity_2d

occlusion_sensitivity_2d(response_fn: ResponseFn, x: Tensor, *, baseline: float | Tensor = 0.0, window: int = 3, stride: int = 1, relative: bool = False, eps: float = 1e-08) -> torch.Tensor

Spatial occlusion sensitivity for [N, C, H, W] inputs.

A window x window patch (all channels) is occluded at each spatial stride position. Returns an [N, Hout, Wout] map of response drops, the classic Zeiler-Fergus heatmap.

With relative=True each drop is divided by the un-occluded base response (drop / (base + eps)), giving the fraction of the response that depends on each region. This is what exposes the AND-signature visually: for a multiplicative detector every contributing region is ~100% critical (occluding it alone collapses the response), whereas for an additive detector each region is only partially critical.

Source code in polyweave/interpretability/occlusion.py

@torch.no_grad()
def occlusion_sensitivity_2d(
    response_fn: ResponseFn,
    x: torch.Tensor,
    *,
    baseline: float | torch.Tensor = 0.0,
    window: int = 3,
    stride: int = 1,
    relative: bool = False,
    eps: float = 1e-8,
) -> torch.Tensor:
    """Spatial occlusion sensitivity for ``[N, C, H, W]`` inputs.

    A ``window x window`` patch (all channels) is occluded at each spatial stride
    position. Returns an ``[N, Hout, Wout]`` map of response drops, the classic
    Zeiler-Fergus heatmap.

    With ``relative=True`` each drop is divided by the un-occluded base response
    (``drop / (base + eps)``), giving the *fraction* of the response that depends
    on each region. This is what exposes the AND-signature visually: for a
    multiplicative detector every contributing region is ~100% critical (occluding
    it alone collapses the response), whereas for an additive detector each region
    is only partially critical.
    """
    if x.ndim != 4:
        raise ValueError(f"expected [N, C, H, W], got shape {tuple(x.shape)}")
    N, C, H, W = x.shape
    base = _as_response(response_fn(x))  # [N]
    ys = list(range(0, H - window + 1, stride)) or [0]
    xs = list(range(0, W - window + 1, stride)) or [0]
    out = base.new_zeros(N, len(ys), len(xs))
    for i, yy in enumerate(ys):
        for j, xx in enumerate(xs):
            occ = x.clone()
            occ[:, :, yy:yy + window, xx:xx + window] = baseline
            out[:, i, j] = base - _as_response(response_fn(occ))
    if relative:
        out = out / (base.reshape(N, 1, 1).abs() + eps)
    return out