ffca 0.1.0a1__py3-none-any.whl

ffca/__init__.py

@@ -0,0 +1,41 @@
+"""FFCA — Feature-Function Curvature Analysis for any PyTorch model.
+
+Quick start (Python):
+
+    from ffca import FFCAReport
+    from ffca.adapters import TabularAdapter
+
+    adapter = TabularAdapter(model, feature_names=cols)
+    report = FFCAReport(adapter, val_loader).run()
+    report.save("out/")
+
+Or via the CLI: ``ffca-report --help``.
+"""
+from .adapters import (
+    ChannelAdapter,
+    PixelAdapter,
+    TabularAdapter,
+    TransformerEmbeddingAdapter,
+    TransformerHeadAdapter,
+)
+from .checkpoint import CheckpointLoader
+from .core import FFCAModelAdapter, FFCASignature
+from .improvements_pkg import CauchyHVP, CoSensitivityGroups, TrustScore
+from .report import FFCAReport
+
+__version__ = "0.1.0a1"
+
+__all__ = [
+    "FFCAReport",
+    "FFCAModelAdapter",
+    "FFCASignature",
+    "TabularAdapter",
+    "PixelAdapter",
+    "ChannelAdapter",
+    "TransformerEmbeddingAdapter",
+    "TransformerHeadAdapter",
+    "CheckpointLoader",
+    "CauchyHVP",
+    "TrustScore",
+    "CoSensitivityGroups",
+]

ffca/adapters/__init__.py

@@ -0,0 +1,16 @@
+"""Built-in adapters covering the four v0.1.0 model families."""
+from .channel import ChannelAdapter
+from .pixel import PixelAdapter
+from .tabular import TabularAdapter
+from .transformer import (
+    TransformerEmbeddingAdapter,
+    TransformerHeadAdapter,
+)
+
+__all__ = [
+    "TabularAdapter",
+    "PixelAdapter",
+    "ChannelAdapter",
+    "TransformerEmbeddingAdapter",
+    "TransformerHeadAdapter",
+]

ffca/adapters/channel.py

@@ -0,0 +1,145 @@
+"""ChannelAdapter — channel-level FFCA at any intermediate layer.
+
+Splice mechanism:
+  1. Resolve the named layer fresh on every forward (smoothing may have
+     swapped its identity).
+  2. Capture: register a one-shot forward_hook that records the activation
+     and unhooks itself.
+  3. Replace: register a hook that returns a pre-set leaf tensor as the
+     layer's output, so subsequent layers run from that leaf.
+
+Both modes resolve the layer by *name* every time, so even if `smooth()`
+replaces ReLU→Softplus mid-pipeline, the hook still attaches to the
+right (current) instance.
+"""
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from ..core.adapter import FFCAModelAdapter, find_layer
+from ..core.scalars import ScalarFn, predicted_class
+
+
+class ChannelAdapter(FFCAModelAdapter):
+    def __init__(
+        self,
+        model: nn.Module,
+        layer_name: str,
+        *,
+        scalar: ScalarFn | None = None,
+        reduction: str = "spatial_mean",
+    ):
+        super().__init__(model, scalar=scalar or predicted_class())
+        self.layer_name = layer_name
+        # Verify the layer exists; we re-resolve each call.
+        _ = find_layer(model, layer_name)
+        self.reduction = reduction
+        self._activation_shape: tuple[int, ...] | None = None
+        self.feature_names: list[str] | None = None
+        self.n_features = -1
+        self.feature_shape = ()
+        self._raw_input: torch.Tensor | None = None
+
+    # ---------------------------------------------------------- splice utilities
+    def _capture_activation(self, x: torch.Tensor) -> torch.Tensor:
+        """Run a forward pass, capture the named layer's activation, return it."""
+        captured: list[torch.Tensor] = []
+        layer = find_layer(self.model, self.layer_name)
+
+        def hook(module, inputs, output):
+            captured.append(output.detach())
+
+        handle = layer.register_forward_hook(hook)
+        try:
+            with torch.no_grad():
+                self.model(x)
+        finally:
+            handle.remove()
+        if not captured:
+            raise RuntimeError(
+                f"Layer {self.layer_name!r} was not invoked during forward; "
+                f"check the layer name is correct for this model's forward path"
+            )
+        return captured[0]
+
+    def _forward_with_replacement(self, x: torch.Tensor,
+                                   replacement: torch.Tensor) -> torch.Tensor:
+        """Run a forward pass with the named layer's output replaced."""
+        layer = find_layer(self.model, self.layer_name)
+
+        def hook(module, inputs, output):
+            return replacement
+
+        handle = layer.register_forward_hook(hook)
+        try:
+            return self.model(x)
+        finally:
+            handle.remove()
+
+    # ---------------------------------------------------------- shape probe
+    def _probe(self, batch):
+        x = batch[0] if isinstance(batch, (list, tuple)) else batch
+        x = x.to(device=self.device(), dtype=self.dtype())
+        act = self._capture_activation(x)
+        self._activation_shape = tuple(act.shape[1:])
+        if self.reduction == "spatial_mean":
+            C = act.shape[1]
+            self.n_features = C
+            self.feature_shape = (C,)
+            self.feature_names = [f"ch_{i}" for i in range(C)]
+        elif self.reduction == "none":
+            d = int(np.prod(act.shape[1:]))
+            self.n_features = d
+            self.feature_shape = tuple(act.shape[1:])
+            self.feature_names = [f"unit_{i}" for i in range(d)]
+        else:
+            raise ValueError(f"unknown reduction {self.reduction!r}")
+
+    # ---------------------------------------------------------- adapter API
+    def feature_input(self, batch) -> torch.Tensor:
+        x = batch[0] if isinstance(batch, (list, tuple)) else batch
+        x = x.to(device=self.device(), dtype=self.dtype())
+
+        if self._activation_shape is None:
+            self._probe(batch)
+
+        act = self._capture_activation(x)  # full activation tensor
+
+        if self.reduction == "spatial_mean":
+            if act.dim() == 4:           # (B, C, H, W)
+                feat = act.mean(dim=(2, 3))
+            elif act.dim() == 3:         # (B, C, L) or (B, T, C)
+                feat = act.mean(dim=-1)
+            else:
+                feat = act
+        else:
+            feat = act.reshape(act.size(0), -1)
+
+        leaf = feat.clone().detach().requires_grad_(True)
+        self._raw_input = x
+        return leaf
+
+    def scalar_output(self, leaf: torch.Tensor, batch) -> torch.Tensor:
+        # Re-inflate the leaf to the original activation shape and run the
+        # rest of the model with the replacement hook.
+        if self.reduction == "spatial_mean":
+            shape = self._activation_shape
+            if len(shape) == 3:
+                C, H, W = shape
+                injected = leaf.view(leaf.size(0), C, 1, 1).expand(-1, C, H, W).contiguous()
+            elif len(shape) == 2:
+                C, L = shape
+                injected = leaf.view(leaf.size(0), C, 1).expand(-1, C, L).contiguous()
+            else:
+                injected = leaf
+        else:
+            injected = leaf.reshape(leaf.size(0), *self._activation_shape)
+        out = self._forward_with_replacement(self._raw_input, injected)
+        return self._scalar(out, batch)
+
+    def channel_count(self) -> int:
+        if self._activation_shape is None:
+            raise RuntimeError("ChannelAdapter.feature_input must be called once first")
+        return self._activation_shape[0]

ffca/adapters/pixel.py

@@ -0,0 +1,66 @@
+"""PixelAdapter — input-level FFCA on image/tensor inputs.
+
+Use for: any model that consumes a (C, H, W) image and you want pixel-level
+explanations. The feature axis is flattened C·H·W pixels.
+
+Adds an `fbr()` helper to compute the Foreground/Background interaction
+Ratio used in shortcut-learning diagnostics (e.g. Waterbirds): proportion
+of total interaction concentrated in a center foreground box vs the
+surrounding background ring. FBR < 0.5 suggests background-shortcut.
+"""
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from ..core.adapter import FFCAModelAdapter
+from ..core.scalars import ScalarFn, predicted_class
+
+
+class PixelAdapter(FFCAModelAdapter):
+    def __init__(
+        self,
+        model: nn.Module,
+        *,
+        input_shape: tuple[int, int, int],  # (C, H, W)
+        scalar: ScalarFn | None = None,
+    ):
+        super().__init__(model, scalar=scalar or predicted_class())
+        self.feature_shape = tuple(input_shape)
+        self.n_features = int(np.prod(input_shape))
+        C, H, W = input_shape
+        # Don't name every pixel by default (too many); CLI uses indices
+        self.feature_names = [f"px_{c}_{h}_{w}" for c in range(C)
+                              for h in range(H) for w in range(W)]
+
+    def feature_input(self, batch) -> torch.Tensor:
+        x = batch[0] if isinstance(batch, (list, tuple)) else batch
+        x = x.to(device=self.device(), dtype=self.dtype())
+        return x.clone().detach().requires_grad_(True)
+
+    def scalar_output(self, x: torch.Tensor, batch) -> torch.Tensor:
+        out = self.model(x)
+        return self._scalar(out, batch)
+
+    # --- spatial helpers ---------------------------------------------------
+    def reshape_to_image(self, per_pixel_score: np.ndarray) -> np.ndarray:
+        """Map a (d,) per-pixel score back to its (C, H, W) layout."""
+        return per_pixel_score.reshape(self.feature_shape)
+
+    def fbr(self, per_pixel_interaction: np.ndarray, fg_frac: float = 0.5) -> float:
+        """Foreground/Background interaction ratio.
+
+        fg_frac: side-length fraction defining the center foreground box.
+        Returns mean(fg) / (mean(fg) + mean(bg)). Values < 0.5 hint at a
+        background shortcut.
+        """
+        C, H, W = self.feature_shape
+        img = per_pixel_interaction.reshape(self.feature_shape)
+        py = int(H * (1 - fg_frac) / 2)
+        px = int(W * (1 - fg_frac) / 2)
+        fg = img[:, py:H - py, px:W - px].mean()
+        bg_sum = img.sum() - img[:, py:H - py, px:W - px].sum()
+        bg_count = img.size - img[:, py:H - py, px:W - px].size
+        bg = bg_sum / max(bg_count, 1)
+        return float(fg / (fg + bg)) if (fg + bg) > 0 else float("nan")

ffca/adapters/tabular.py

@@ -0,0 +1,40 @@
+"""TabularAdapter — input-level FFCA on dense feature vectors.
+
+Use for: MLPs, tabular transformers, scikit-style wrappers over an MLP.
+The feature axis is the d input columns.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from ..core.adapter import FFCAModelAdapter
+from ..core.scalars import ScalarFn, predicted_class
+
+
+class TabularAdapter(FFCAModelAdapter):
+    def __init__(
+        self,
+        model: nn.Module,
+        *,
+        feature_names: list[str] | None = None,
+        n_features: int | None = None,
+        scalar: ScalarFn | None = None,
+    ):
+        super().__init__(model, scalar=scalar or predicted_class())
+        if feature_names is None and n_features is None:
+            raise ValueError("provide feature_names or n_features")
+        self.feature_names = feature_names
+        self.n_features = n_features if n_features is not None else len(feature_names)
+        self.feature_shape = (self.n_features,)
+        if feature_names is None:
+            self.feature_names = [f"feature_{i}" for i in range(self.n_features)]
+
+    def feature_input(self, batch) -> torch.Tensor:
+        x = batch[0] if isinstance(batch, (list, tuple)) else batch
+        x = x.to(device=self.device(), dtype=self.dtype())
+        return x.clone().detach().requires_grad_(True)
+
+    def scalar_output(self, x: torch.Tensor, batch) -> torch.Tensor:
+        out = self.model(x)
+        return self._scalar(out, batch)

ffca/adapters/transformer.py

@@ -0,0 +1,236 @@
+"""Transformer adapters for FFCA — embedding-level and attention-head.
+
+Both adapters work on any HuggingFace `AutoModel` / `AutoModelForCausalLM`
+without modification. They auto-detect the input-embedding attribute and
+the per-layer attention block by walking common name conventions:
+
+  GPT-2 / DistilGPT2  : model.transformer.wte                          (embeddings)
+                         model.transformer.h[L].attn                   (attention block)
+  BERT / DistilBERT   : model.bert.embeddings.word_embeddings
+                         model.bert.encoder.layer[L].attention.self
+  LLaMA / Mistral     : model.model.embed_tokens
+                         model.model.layers[L].self_attn
+  Generic fallback    : the first nn.Embedding found in the model
+
+If auto-detection fails, pass `embedding_module=` / `attention_layer=` by
+hand. Both adapters use the splice trick (forward_hook re-resolved by name
+each call) so the package's automatic ReLU→Softplus / MaxPool→AvgPool /
+flash-SDP→math-SDPA smoothing does not break them.
+"""
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from ..core.adapter import FFCAModelAdapter
+from ..core.scalars import ScalarFn, predicted_class
+
+
+# ---------------------------------------------------------------- locator
+def _find_embedding(model: nn.Module) -> nn.Embedding:
+    """Return the input-token embedding module of an HF transformer."""
+    for path in (
+        "transformer.wte",            # GPT-2 family
+        "model.embed_tokens",         # LLaMA, Mistral
+        "model.decoder.embed_tokens", # OPT, BART
+        "bert.embeddings.word_embeddings",
+        "embeddings.word_embeddings",
+        "shared",                     # T5
+    ):
+        obj = model
+        ok = True
+        for part in path.split("."):
+            if not hasattr(obj, part):
+                ok = False; break
+            obj = getattr(obj, part)
+        if ok and isinstance(obj, nn.Embedding):
+            return obj
+    # Fallback: first nn.Embedding in the module tree
+    for m in model.modules():
+        if isinstance(m, nn.Embedding):
+            return m
+    raise AttributeError("No nn.Embedding found in the model")
+
+
+def _find_attention_layer(model: nn.Module, layer_idx: int = -1) -> nn.Module:
+    """Return the attention module at the requested encoder layer index."""
+    for path in (
+        ("transformer.h", "attn"),                         # GPT-2
+        ("model.layers", "self_attn"),                     # LLaMA
+        ("bert.encoder.layer", "attention.self"),
+        ("encoder.layer", "attention.self"),
+        ("model.encoder.layers", "self_attn"),
+        ("model.decoder.layers", "self_attn"),
+    ):
+        block_path, attr_path = path
+        obj = model
+        ok = True
+        for part in block_path.split("."):
+            if not hasattr(obj, part):
+                ok = False; break
+            obj = getattr(obj, part)
+        if not ok or not hasattr(obj, "__getitem__"):
+            continue
+        try:
+            block = obj[layer_idx]
+        except (IndexError, TypeError):
+            continue
+        attr_obj = block
+        for p in attr_path.split("."):
+            if not hasattr(attr_obj, p):
+                attr_obj = None; break
+            attr_obj = getattr(attr_obj, p)
+        if attr_obj is not None and isinstance(attr_obj, nn.Module):
+            return attr_obj
+    raise AttributeError(
+        f"Could not auto-locate attention block at layer index {layer_idx}; "
+        f"pass attention_layer= explicitly"
+    )
+
+
+def _last_token_logit_scalar(out: torch.Tensor, batch=None) -> torch.Tensor:
+    """Default LLM scalar: max logit at the last sequence position."""
+    if hasattr(out, "logits"):
+        last = out.logits[:, -1, :]
+    elif hasattr(out, "last_hidden_state"):
+        last = out.last_hidden_state[:, -1, :]
+    else:
+        last = out
+    return last.max(dim=-1).values.sum()
+
+
+# ---------------------------------------------------------------- embedding adapter
+class TransformerEmbeddingAdapter(FFCAModelAdapter):
+    """Treat token × hidden-dim input embeddings as the feature axis.
+
+    For a sequence of length T and hidden size H, this yields T·H features.
+    Suitable for understanding which token position / which hidden dim drives
+    the model's prediction.
+
+    Args:
+        model           : HF model (AutoModel / AutoModelForCausalLM)
+        seq_len, hidden : sequence length and hidden size of the input
+        scalar          : optional ScalarFn; defaults to last-token max-logit
+        embedding_module: optional override for the nn.Embedding to differentiate
+    """
+    def __init__(
+        self,
+        model: nn.Module,
+        seq_len: int,
+        hidden: int,
+        *,
+        scalar: ScalarFn | None = None,
+        embedding_module: nn.Embedding | None = None,
+    ):
+        super().__init__(model, scalar=scalar or _last_token_logit_scalar)
+        self.seq_len = seq_len
+        self.hidden = hidden
+        self.n_features = seq_len * hidden
+        self.feature_shape = (seq_len, hidden)
+        self.feature_names = [f"t{t}_h{h}" for t in range(seq_len)
+                              for h in range(hidden)]
+        self._emb = embedding_module or _find_embedding(model)
+
+    def feature_input(self, batch) -> torch.Tensor:
+        ids = batch["input_ids"] if isinstance(batch, dict) else batch[0]
+        ids = ids.to(self.device())
+        with torch.no_grad():
+            embs = self._emb(ids)
+        return embs.clone().detach().requires_grad_(True)
+
+    def scalar_output(self, embs: torch.Tensor, batch) -> torch.Tensor:
+        kw = {}
+        if isinstance(batch, dict):
+            if "attention_mask" in batch:
+                kw["attention_mask"] = batch["attention_mask"].to(self.device())
+        out = self.model(inputs_embeds=embs, **kw)
+        return self._scalar(out, batch)
+
+
+# ---------------------------------------------------------------- head adapter
+class TransformerHeadAdapter(FFCAModelAdapter):
+    """Per-attention-head pooled activations from a chosen encoder layer.
+
+    For an attention block with H heads of dim D, the feature axis is H·D.
+    The mean-pool over tokens turns the (B, T, H·D) attention output into a
+    (B, H, D) feature tensor that FFCA can differentiate.
+
+    Args:
+        model           : HF model
+        n_heads, head_dim: from the model's config
+        layer_idx        : which layer (negative indexes from the end; -1 = last)
+        attention_layer  : optional override for the attention module
+        scalar           : optional ScalarFn
+    """
+    def __init__(
+        self,
+        model: nn.Module,
+        n_heads: int,
+        head_dim: int,
+        *,
+        layer_idx: int = -1,
+        attention_layer: nn.Module | None = None,
+        scalar: ScalarFn | None = None,
+    ):
+        super().__init__(model, scalar=scalar or _last_token_logit_scalar)
+        self.n_heads = n_heads
+        self.head_dim = head_dim
+        self.layer_idx = layer_idx
+        self.n_features = n_heads * head_dim
+        self.feature_shape = (n_heads, head_dim)
+        self.feature_names = [f"h{h}_d{d}" for h in range(n_heads)
+                              for d in range(head_dim)]
+        self._attn = attention_layer or _find_attention_layer(model, layer_idx)
+        self._batch_ids: torch.Tensor | None = None
+        self._batch_attn: torch.Tensor | None = None
+
+    def feature_input(self, batch) -> torch.Tensor:
+        ids = batch["input_ids"] if isinstance(batch, dict) else batch[0]
+        attn_mask = (batch.get("attention_mask")
+                     if isinstance(batch, dict) else None)
+        ids = ids.to(self.device())
+        if attn_mask is not None:
+            attn_mask = attn_mask.to(self.device())
+
+        captured: list[torch.Tensor] = []
+        def hook(module, inputs, output):
+            t = output[0] if isinstance(output, tuple) else output
+            captured.append(t.detach())
+        h = self._attn.register_forward_hook(hook)
+        try:
+            with torch.no_grad():
+                kw = {"input_ids": ids}
+                if attn_mask is not None:
+                    kw["attention_mask"] = attn_mask
+                self.model(**kw)
+        finally:
+            h.remove()
+        if not captured:
+            raise RuntimeError("attention hook never fired; "
+                                "check `attention_layer` is on the forward path")
+        c = captured[0]
+        # Pool over tokens, shape → (B, n_heads, head_dim)
+        flat = c.mean(dim=1).reshape(c.size(0), self.n_heads, self.head_dim)
+        leaf = flat.clone().detach().requires_grad_(True)
+        self._batch_ids = ids
+        self._batch_attn = attn_mask
+        return leaf
+
+    def scalar_output(self, leaf: torch.Tensor, batch) -> torch.Tensor:
+        T = self._batch_ids.size(1)
+        injected = leaf.reshape(leaf.size(0), 1, self.n_heads * self.head_dim) \
+                       .expand(-1, T, -1).contiguous()
+        def hook(module, inputs, output):
+            if isinstance(output, tuple):
+                return (injected,) + output[1:]
+            return injected
+        handle = self._attn.register_forward_hook(hook)
+        try:
+            kw = {"input_ids": self._batch_ids}
+            if self._batch_attn is not None:
+                kw["attention_mask"] = self._batch_attn
+            out = self.model(**kw)
+        finally:
+            handle.remove()
+        return self._scalar(out, batch)

ffca/checkpoint.py

@@ -0,0 +1,58 @@
+"""CheckpointLoader — yield a fresh adapter for each saved state_dict.
+
+For v0.1.0 we support plain `torch.save(model.state_dict(), path)` files.
+Lightning / HF / SafeTensors formats can be added with detect-by-extension
+in v0.2.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Callable, Iterator, Sequence
+
+import torch
+import torch.nn as nn
+
+
+class CheckpointLoader:
+    """Iterate (epoch_label, model) over a list of saved state_dicts.
+
+    Args:
+        model_factory: () -> fresh nn.Module instance (same architecture as
+                       what was saved).
+        checkpoints:   list of paths (or (label, path) pairs).
+        device:        target device for the loaded model.
+    """
+
+    def __init__(
+        self,
+        model_factory: Callable[[], nn.Module],
+        checkpoints: Sequence[str | tuple[str, str]],
+        device: str | torch.device = "cpu",
+    ):
+        self.model_factory = model_factory
+        self.checkpoints = [
+            (Path(c).stem, Path(c)) if isinstance(c, str) else (str(c[0]), Path(c[1]))
+            for c in checkpoints
+        ]
+        self.device = torch.device(device)
+
+    def __len__(self) -> int:
+        return len(self.checkpoints)
+
+    def __iter__(self) -> Iterator[tuple[str, nn.Module]]:
+        for label, path in self.checkpoints:
+            model = self.model_factory()
+            state = torch.load(path, map_location=self.device, weights_only=False)
+            # Accept either {state_dict} or just a state_dict
+            if isinstance(state, dict) and "state_dict" in state:
+                state = state["state_dict"]
+            try:
+                model.load_state_dict(state, strict=False)
+            except Exception as e:
+                raise RuntimeError(
+                    f"Failed to load checkpoint {path}: {e}. "
+                    f"Ensure model_factory() returns the same architecture."
+                )
+            model.to(self.device)
+            model.eval()
+            yield label, model