Separate attention-map capture from gradient hooks

docs/api/interpret/pyhealth.interpret.methods.attention_rollout.rst

@@ -35,16 +35,53 @@ Key Features
 - **Class-agnostic**: Independent of the predicted/target class (``target_class_idx`` is accepted but ignored)
 - **Layer-wise composition**: Composes per-layer attention as ``rollout = Â_L @ ... @ Â_1`` with the residual correction ``Â = 0.5 * (A + I)``
 - **Distribution over tokens**: Because each ``Â`` is row-stochastic, so is their product; per-token relevance sums to 1 (before the input-shape expansion)
-- **Model-agnostic by duck-typing**: Works with any model exposing the attention-readout methods ``set_attention_hooks``, ``get_attention_layers`` and ``get_relevance_tensor`` (currently :class:`~pyhealth.models.Transformer` and :class:`~pyhealth.models.StageAttentionNet`), not just one named model
+- **Model interface**: Works with models implementing :class:`~pyhealth.interpret.api.AttentionInterpretable`, including Transformer and StageAttentionNet.
 
 Usage Notes
 -----------
 
 1. **Batch size**: For interpretability, use ``batch_size=1`` to get per-sample explanations.
-2. **Do not wrap in** ``torch.no_grad()``: Although rollout is gradient-free in its math, the shared attention-readout plumbing registers a gradient hook on the attention tensors during the forward pass, so calling ``attribute(**batch)`` inside ``torch.no_grad()`` raises a ``RuntimeError``. Call it under the default (grad-enabled) context; no backward pass is performed.
-3. **Model compatibility**: Works with any model that exposes ``set_attention_hooks``, ``get_attention_layers`` and ``get_relevance_tensor`` — not restricted to the Transformer. Incompatible models raise ``TypeError`` at construction.
+2. **No gradients required**: Rollout runs under ``torch.no_grad()`` and can be called inside an existing no-grad context. It captures detached attention maps without registering backward hooks or changing parameter gradients.
+3. **Model compatibility**: Models must implement ``AttentionInterpretable``. Incompatible models raise ``TypeError`` at construction.
 4. **Class specification**: ``target_class_idx`` is accepted for API compatibility but ignored, since rollout is class-agnostic.
 
+Attention capture interfaces
+----------------------------
+
+``AttentionInterpretable`` provides forward-only capture; its implementations
+must capture detached maps under ``torch.no_grad()`` without requiring autograd.
+``GradientInterpretable`` extends it with backward capture for Chefer.
+Both interfaces live in ``pyhealth.interpret.api``.
+
+.. important::
+
+   Custom models must explicitly inherit the appropriate interface and accept
+   ``set_attention_hooks(enabled, *, capture_gradients=...)``. Rollout passes
+   ``capture_gradients=False`` and Chefer passes ``True``. The compatibility alias
+   ``CheferInterpretable = GradientInterpretable`` preserves imports and type
+   checks, but does not adapt obsolete one-argument implementations: these raise
+   ``TypeError`` identifying ``capture_gradients``. Method names alone no longer
+   establish Rollout compatibility.
+
+Transformer and StageAttentionNet retain ``capture_gradients=True`` as the
+one-argument default. ``enabled=False`` always disables both kinds of capture,
+regardless of the gradient argument. At the attention layer, map capture is
+``capture_attention or register_hook`` and gradient capture is ``register_hook``.
+The new low-level ``capture_attention`` argument is keyword-only; existing
+positional ``register_hook`` calls retain their meaning.
+
+``get_attention_layers()`` returns feature-keyed ordered pairs of optional
+``(attention_map, attention_gradient)`` tensors. Rollout requires maps shaped
+``[batch, heads, seq, seq]``; Chefer also accepts head-averaged maps. Gradients
+are available only after a gradient-enabled forward and backward.
+Disabling capture preserves the latest results and outstanding backward hooks.
+The next attention forward replaces the map and clears the gradient; an
+uncaptured forward clears both. Finish any backward pass before starting a new
+forward: overlapping captured forwards/backwards on one model are unsupported.
+
+See ``examples/interpretability/attention_capture.py`` for a synthetic example
+running Rollout under ``torch.no_grad()`` followed by Chefer on the same model.
+
 Quick Start
 -----------
 

docs/api/interpret/pyhealth.interpret.methods.chefer.rst

@@ -33,9 +33,23 @@ Usage Notes
 
 1. **Batch size**: For interpretability, use batch_size=1 to get per-sample explanations
 2. **Gradients required**: Do not use within ``torch.no_grad()`` context
-3. **Model compatibility**: Only works with PyHealth's Transformer model
+3. **Model compatibility**: Requires ``GradientInterpretable``, implemented by Transformer and StageAttentionNet. Incompatible models raise ``ValueError``.
 4. **Class specification**: You can specify a target class or use the predicted class
 
+Capture requirements
+--------------------
+
+Chefer explicitly requests ``capture_gradients=True``. Empty attention-layer
+lists, missing maps or gradients, and mismatched map/gradient shapes raise
+``RuntimeError`` with the affected feature key before relevance propagation.
+
+``CheferInterpretable`` remains an alias for ``GradientInterpretable``. Custom
+models must accept the new ``capture_gradients`` keyword; old one-argument
+implementations raise ``TypeError``. See
+:doc:`pyhealth.interpret.methods.attention_rollout` for the interface migration
+and sequential capture lifecycle. A complete synthetic demonstration is in
+``examples/interpretability/attention_capture.py``.
+
 Quick Start
 -----------
 

examples/interpretability/attention_capture.py

@@ -0,0 +1,43 @@
+"""Run gradient-free Rollout and gradient-based Chefer on one synthetic model.
+
+Run: pixi run -e test python examples/interpretability/attention_capture.py
+"""
+
+import torch
+
+from pyhealth.datasets import create_sample_dataset, get_dataloader
+from pyhealth.interpret.methods import AttentionRollout, CheferRelevance
+from pyhealth.models import Transformer
+
+
+def main():
+    torch.manual_seed(42)
+    dataset = create_sample_dataset(
+        samples=[
+            {"patient_id": "p0", "visit_id": "v0", "codes": ["A", "B"], "label": 1},
+            {"patient_id": "p1", "visit_id": "v1", "codes": ["B"], "label": 0},
+        ],
+        input_schema={"codes": "sequence"},
+        output_schema={"label": "binary"},
+    )
+    model = Transformer(dataset=dataset, embedding_dim=8, heads=2, num_layers=2)
+    batch = next(iter(get_dataloader(dataset, batch_size=2, shuffle=False)))
+
+    with torch.no_grad():
+        rollout = AttentionRollout(model).attribute(**batch)
+    for attention, gradient in model.get_attention_layers()["codes"]:
+        assert attention is not None and not attention.requires_grad
+        assert gradient is None
+
+    # Complete each interpretation before starting the next captured forward.
+    chefer = CheferRelevance(model).attribute(**batch)
+    for attention, gradient in model.get_attention_layers()["codes"]:
+        assert attention is not None and gradient is not None
+        assert attention.shape == gradient.shape
+
+    print("Rollout:", rollout["codes"].tolist())
+    print("Chefer:", chefer["codes"].tolist())
+
+
+if __name__ == "__main__":
+    main()