Require evaluation mode for attribution patching

tests/unit/tools/test_attribution_patching.py

@@ -78,6 +78,7 @@ def __init__(self, *, dtype: torch.dtype = torch.float32) -> None:
         )
         self.compatibility_mode = False
         self._weights_processed = False
+        self.forward_calls = 0
         self.embed = nn.Embedding(D_VOCAB, D_MODEL, dtype=dtype)
         nn.init.normal_(self.embed.weight, std=0.2)
         self.hook_embed = HookPoint()
@@ -88,6 +89,7 @@ def __init__(self, *, dtype: torch.dtype = torch.float32) -> None:
         self.ln_final = nn.Identity()
         self.unembed = nn.Linear(D_MODEL, D_VOCAB, bias=False, dtype=dtype)
         nn.init.normal_(self.unembed.weight, std=0.2)
+        self.eval()
 
     @property
     def hook_dict(self) -> dict[str, HookPoint]:
@@ -118,6 +120,7 @@ def to_tokens(self, prompt: str) -> torch.Tensor:
     def forward(
         self, tokens: torch.Tensor, return_type: str | None = "logits"
     ) -> torch.Tensor | None:
+        self.forward_calls += 1
         residual = self.hook_embed(self.embed(tokens))
         for block in self.blocks:
             residual = block(residual)
@@ -364,6 +367,7 @@ def __init__(self, d_model: int, n_heads: int, d_head: int, layer: int, dtype: t
         self.w_z = nn.Linear(d_model, n_heads * d_head, bias=False, dtype=dtype)
         self.w_o = nn.Linear(n_heads * d_head, d_model, bias=False, dtype=dtype)
         self.w_mlp = nn.Linear(d_model, d_model, bias=False, dtype=dtype)
+        self.dropout = nn.Dropout(p=0.5)
         for linear in (self.w_z, self.w_o, self.w_mlp):
             nn.init.normal_(linear.weight, std=0.2)
         self.hook_z = HookPoint()
@@ -374,7 +378,8 @@ def __init__(self, d_model: int, n_heads: int, d_head: int, layer: int, dtype: t
     def forward(self, residual: torch.Tensor) -> torch.Tensor:
         batch, seq, _ = residual.shape
         z = self.hook_z(self.w_z(residual).reshape(batch, seq, self.n_heads, self.d_head))
-        residual = residual + self.w_o(z.reshape(batch, seq, self.n_heads * self.d_head))
+        attn_out = self.w_o(z.reshape(batch, seq, self.n_heads * self.d_head))
+        residual = residual + self.dropout(attn_out)
         mlp_out = self.hook_mlp_out(self.w_mlp(residual))
         return residual + mlp_out
 
@@ -403,6 +408,7 @@ def __init__(self, *, dtype: torch.dtype = torch.float32) -> None:
         )
         self.compatibility_mode = False
         self._weights_processed = False
+        self.forward_calls = 0
         self.embed = nn.Embedding(D_VOCAB, D_MODEL, dtype=dtype)
         nn.init.normal_(self.embed.weight, std=0.2)
         self.hook_embed = HookPoint()
@@ -413,6 +419,7 @@ def __init__(self, *, dtype: torch.dtype = torch.float32) -> None:
         self.ln_final = nn.Identity()
         self.unembed = nn.Linear(D_MODEL, D_VOCAB, bias=False, dtype=dtype)
         nn.init.normal_(self.unembed.weight, std=0.2)
+        self.eval()
 
     @property
     def hook_dict(self) -> dict[str, HookPoint]:
@@ -425,6 +432,7 @@ def hook_dict(self) -> dict[str, HookPoint]:
     def forward(
         self, tokens: torch.Tensor, return_type: str | None = "logits"
     ) -> torch.Tensor | None:
+        self.forward_calls += 1
         residual = self.hook_embed(self.embed(tokens))
         for block in self.blocks:
             residual = block(residual)
@@ -513,6 +521,57 @@ def test_attribution_patch_raises_on_batch_size_mismatch() -> None:
         attribution_patch(model, clean, corrupt, _metric_fn(answer=1, wrong=2))
 
 
+def test_attribution_patch_rejects_model_in_training_mode() -> None:
+    model = _NodeGraphToyBridge()
+    model.train()
+    tokens = torch.tensor([[1, 2, 3]])
+
+    with pytest.raises(ValueError, match=r"model\.eval\(\)"):
+        attribution_patch(model, tokens, tokens, _metric_fn(answer=1, wrong=2))
+    assert model.training is True
+    assert model.forward_calls == 0
+
+
+def test_attribution_patch_rejects_nested_submodule_in_training_mode() -> None:
+    model = _NodeGraphToyBridge()
+    model.blocks[1].dropout.train()
+    tokens = torch.tensor([[1, 2, 3]])
+    assert model.training is False
+    assert model.blocks[1].training is False
+
+    with pytest.raises(ValueError, match=r"model\.eval\(\)"):
+        attribution_patch(model, tokens, tokens, _metric_fn(answer=1, wrong=2))
+    assert model.training is False
+    assert model.blocks[1].training is False
+    assert model.blocks[1].dropout.training is True
+    assert model.forward_calls == 0
+
+
+def test_attribution_patch_rejects_hidden_original_model_training_mode() -> None:
+    model = _NodeGraphToyBridge()
+    original_model = nn.Sequential(nn.Linear(D_MODEL, D_MODEL))
+    original_model.eval()
+    original_model[0].train()
+    model.__dict__["original_model"] = original_model
+    tokens = torch.tensor([[1, 2, 3]])
+
+    with pytest.raises(ValueError, match="original_model"):
+        attribution_patch(model, tokens, tokens, _metric_fn(answer=1, wrong=2))
+    assert original_model.training is False
+    assert original_model[0].training is True
+    assert model.forward_calls == 0
+
+
+def test_attribution_patch_identical_inputs_are_exactly_zero_in_eval_mode() -> None:
+    model = _NodeGraphToyBridge()
+    tokens = torch.tensor([[1, 2, 3]])
+
+    result = attribution_patch(model, tokens, tokens, _metric_fn(answer=1, wrong=2))
+
+    assert result.node_scores
+    assert all(score == 0.0 for score in result.node_scores.values())
+
+
 # ---------------------------------------------------------------------------
 # Commit 5 — linear-model reconstruction identity
 # ---------------------------------------------------------------------------
@@ -670,6 +729,7 @@ def __init__(self, *, dtype: torch.dtype = torch.float32) -> None:
                 for layer in range(N_LAYERS)
             ]
         )
+        self.eval()
 
 
 def test_nonlinear_node_scores_read_the_corrupt_run_gradient() -> None:

transformer_lens/tools/analysis/_model_state.py

@@ -0,0 +1,31 @@
+"""Shared model-state validation for analysis tools."""
+
+from typing import Any
+
+import torch
+
+
+def require_eval_mode(model: Any, *, operation: str) -> None:
+    """Reject training state anywhere in a wrapped model without mutating it."""
+    training_modules: dict[int, str] = {}
+    roots = (("", model), ("original_model", getattr(model, "original_model", None)))
+    for prefix, root in roots:
+        if not isinstance(root, torch.nn.Module):
+            continue
+        for name, module in root.named_modules():
+            if not module.training:
+                continue
+            qualified_name = ".".join(part for part in (prefix, name) if part)
+            training_modules.setdefault(id(module), qualified_name or "<root>")
+    if not training_modules:
+        return
+
+    names = list(training_modules.values())
+    preview = ", ".join(names[:3])
+    if len(names) > 3:
+        preview += f", and {len(names) - 3} more"
+    raise ValueError(
+        f"{operation} requires the model and all submodules to be in evaluation "
+        f"mode; found training mode at {preview}. Call model.eval() before running "
+        "the analysis."
+    )

transformer_lens/tools/analysis/attribution_patching.py

@@ -36,6 +36,8 @@
 
 import torch
 
+from transformer_lens.tools.analysis._model_state import require_eval_mode
+
 MetricFn = Callable[[torch.Tensor], torch.Tensor]
 NamesFilter = Union[str, Sequence[str], Callable[[str], bool], None]
 
@@ -460,6 +462,10 @@ def attribution_patch(
     reconstruction identity holds) and per-node scores are averaged across the batch
     before ranking.
 
+    The model and every submodule must be in evaluation mode. Separate clean and
+    corrupt forwards cannot produce meaningful activation differences if stochastic
+    training layers such as dropout remain active.
+
     Args:
         model: A ``TransformerBridge`` (or compatible) exposing ``cfg.n_layers``,
             ``hook_dict``, and ``hooks()``.
@@ -476,8 +482,9 @@ def attribution_patch(
 
     Raises:
         ValueError: if ``clean``/``corrupt`` are not 2D, hold a different number of
-            pairs, or a pair tokenizes to different lengths (activations must align
-            position-by-position).
+            pairs, a pair tokenizes to different lengths (activations must align
+            position-by-position), or the model or one of its submodules is in
+            training mode.
     """
     del config  # node granularity + ig_steps=1 only; enforced at construction.
 
@@ -498,6 +505,8 @@ def attribution_patch(
             "Attribution patching aligns activations position-by-position."
         )
 
+    require_eval_mode(model, operation="attribution_patch()")
+
     node_hook_names = _required_hook_names(int(model.cfg.n_layers))
     batch = int(clean.shape[0])
     totals: dict[Node, float] = {}

transformer_lens/tools/analysis/jacobian_lens.py

@@ -81,6 +81,7 @@
 from tqdm.auto import tqdm
 
 from transformer_lens.ActivationCache import ActivationCache
+from transformer_lens.tools.analysis._model_state import require_eval_mode
 from transformer_lens.tools.analysis.jacobian_lens_coordinate_patch import (
     CoordinatePatch,
     solve_coordinate_patch,
@@ -1701,7 +1702,7 @@ def fit(
                 or layer indices, or if no prompt was long enough to fit on.
         """
         _require_raw_bridge(model)
-        _require_eval_mode_for_fit(model)
+        require_eval_mode(model, operation="JacobianLens.fit()")
         if not isinstance(corpus, str) or not corpus.strip():
             raise ValueError("corpus must be a non-empty provenance identifier")
         n_layers = model.cfg.n_layers
@@ -1869,32 +1870,6 @@ def _require_raw_bridge(model: Any) -> None:
         )
 
 
-def _require_eval_mode_for_fit(model: Any) -> None:
-    """Reject stochastic training state without mutating the caller's model."""
-    training_modules: Dict[int, str] = {}
-    roots = (("", model), ("original_model", getattr(model, "original_model", None)))
-    for prefix, root in roots:
-        if not isinstance(root, torch.nn.Module):
-            continue
-        for name, module in root.named_modules():
-            if not module.training:
-                continue
-            qualified_name = ".".join(part for part in (prefix, name) if part)
-            training_modules.setdefault(id(module), qualified_name or "<root>")
-    if not training_modules:
-        return
-
-    names = list(training_modules.values())
-    preview = ", ".join(names[:3])
-    if len(names) > 3:
-        preview += f", and {len(names) - 3} more"
-    raise ValueError(
-        "JacobianLens.fit() requires the model and all submodules to be in "
-        f"evaluation mode; found training mode at {preview}. Call model.eval() "
-        "before fitting."
-    )
-
-
 def _validate_metadata(metadata: Dict[str, Any]) -> None:
     """Reject values that ``torch.load(weights_only=True)`` cannot reload."""