diff --git a/docs/api/interpret/pyhealth.interpret.methods.attention_rollout.rst b/docs/api/interpret/pyhealth.interpret.methods.attention_rollout.rst
index 35875f3d5..0770a0d03 100644
--- a/docs/api/interpret/pyhealth.interpret.methods.attention_rollout.rst
+++ b/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
 -----------
 
diff --git a/docs/api/interpret/pyhealth.interpret.methods.chefer.rst b/docs/api/interpret/pyhealth.interpret.methods.chefer.rst
index 12a77cb63..2e842a5f9 100644
--- a/docs/api/interpret/pyhealth.interpret.methods.chefer.rst
+++ b/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
 -----------
 
diff --git a/examples/interpretability/attention_capture.py b/examples/interpretability/attention_capture.py
new file mode 100644
index 000000000..c776619ad
--- /dev/null
+++ b/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()
diff --git a/pyhealth/interpret/api.py b/pyhealth/interpret/api.py
index f84b8d249..6181b47ef 100644
--- a/pyhealth/interpret/api.py
+++ b/pyhealth/interpret/api.py
@@ -257,185 +257,54 @@ def get_embedding_model(self) -> nn.Module | None:
         raise NotImplementedError
 
 
-class CheferInterpretable(Interpretable):
-    """Abstract interface for models supporting Chefer relevance attribution.
-
-    This is a subclass of :class:`Interpretable` and therefore
-    inherits the embedding-level interface (``forward_from_embedding``,
-    ``get_embedding_model``).  Models that implement this interface
-    automatically satisfy the general interpretability contract **and** the
-    Chefer-specific contract, so they work with both embedding-perturbation
-    methods (DeepLIFT, LIME, …) and gradient-weighted attention methods
-    (Chefer).
-
-    The Chefer algorithm works as follows:
-
-    1. **Forward + hook registration** — run the model while capturing
-       attention weight tensors and registering backward hooks so their
-       gradients are stored.
-    2. **Backward** — back-propagate from a one-hot target class through
-       the logits.
-    3. **Relevance propagation** — for every feature key, iterate over
-       attention layers, compute gradient-weighted attention
-       (``clamp(attn * grad, min=0)``), and accumulate into a relevance
-       matrix ``R`` via ``R += cam @ R``.
-    4. **Attribution extraction** — extract the final per-token
-       attribution from ``R`` (e.g. read the CLS row, or the
-       last-valid-timestep row, possibly with reshaping).
-
-    Steps 1, 3-b and 4 are model-specific; the rest is generic.  This
-    interface captures exactly those model-specific pieces.
-
-    Inherited from ``InterpretableModelInterface``
-    -----------------------------------------------
-    forward_from_embedding(**kwargs) -> dict[str, Tensor]
-        Forward pass starting from pre-computed embeddings.
-    get_embedding_model() -> nn.Module | None
-        Access the embedding / feature-extraction stage.
-
-    Additional (Chefer-specific) methods
-    -------------------------------------
-    set_attention_hooks(enabled) -> None
-        Toggle attention map capture and gradient hook registration.
-    get_attention_layers() -> dict[str, list[tuple[Tensor, Tensor]]]
-        Paired (attn_map, attn_grad) for each attention layer, keyed by
-        feature key.
-    get_relevance_vector(R, **data) -> dict[str, Tensor]
-        Reduce relevance matrices to per-token attribution vectors.
-
-    Attributes
-    ----------
-    feature_keys : list[str]
-        The feature keys from the task's ``input_schema`` (e.g.
-        ``["conditions", "procedures"]``).  Already provided by
-        :class:`~pyhealth.models.base_model.BaseModel`.
-
-    Notes
-    -----
-    *  ``set_attention_hooks(True)`` must be called **before** the forward
-       pass, and ``get_attention_layers`` must be called **after** the
-       forward + backward passes, because attention maps are populated
-       during forward and gradients during backward.
-    *  The interface intentionally does **not** prescribe how hooks are
-       registered internally — ``nn.MultiheadAttention`` with
-       ``register_hook``, manual ``save_attn_grad`` callbacks, or explicit
-       QKV computation all work as long as the getter methods return the
-       right tensors.
-
-    Examples
-    --------
-    Minimal skeleton for a new model:
-
-    >>> class MyAttentionModel(BaseModel, CheferInterpretableModelInterface):
-    ...     # feature_keys is inherited from BaseModel
-    ...
-    ...     def forward_from_embedding(self, **kwargs):
-    ...         # ... prediction head from pre-computed embeddings ...
-    ...
-    ...     def get_embedding_model(self):
-    ...         return self.embedding_layer
-    ...
-    ...     def set_attention_hooks(self, enabled):
-    ...         self._register_hooks = enabled
-    ...
-    ...     def get_attention_layers(self):
-    ...         result = {}
-    ...         for key in self.feature_keys:
-    ...             result[key] = [
-    ...                 (blk.attention.get_attn_map(),
-    ...                  blk.attention.get_attn_grad())
-    ...                 for blk in self.encoder[key].blocks
-    ...             ]
-    ...         return result
-    ...
-    ...     def get_relevance_vector(self, R, **data):
-    ...         return {key: r[:, 0] for key, r in R.items()}
+class AttentionInterpretable(Interpretable):
+    """Interface for models exposing attention maps without requiring autograd.
+
+    Implementations must capture detached maps under ``torch.no_grad()`` when
+    ``capture_gradients=False``. Layers are ordered from input to output and
+    grouped by feature key. ``get_relevance_tensor`` reduces the resulting
+    relevance matrices to model-specific token scores.
+
+    Capture is sequential: overlapping captured forwards/backwards on the same
+    model are unsupported. Finish a backward pass before starting another
+    forward, otherwise a previous graph's hook may overwrite newer results.
+
+    Examples:
+        >>> from pyhealth.interpret.api import AttentionInterpretable
+        >>> issubclass(AttentionInterpretable, Interpretable)
+        True
     """
 
     @abstractmethod
-    def set_attention_hooks(self, enabled: bool) -> None:
-        """Toggle attention hook registration for subsequent forward passes.
-
-        When ``enabled=True``, the next call to ``forward()`` (or
-        ``forward_from_embedding()``) must:
-
-        1. Store attention weight tensors so they are retrievable via
-           :meth:`get_attention_layers`.
-        2. Register backward hooks on those tensors so that after
-           ``.backward()`` the corresponding gradients are also stored.
-
-        When ``enabled=False``, subsequent forward passes should **not**
-        capture attention maps or register gradient hooks, restoring the
-        model to its normal (faster) execution mode.
-
-        Parameters
-        ----------
-        enabled : bool
-            ``True`` to start capturing attention maps and registering
-            gradient hooks; ``False`` to stop.
-
-        Typical implementations set an internal flag that the model's
-        forward method checks::
-
-            def set_attention_hooks(self, enabled):
-                self._attention_hooks_enabled = enabled
-
-        And inside the forward / encoder logic::
-
-            if self._attention_hooks_enabled:
-                attn.register_hook(self.save_attn_grad)
+    def set_attention_hooks(
+        self, enabled: bool, *, capture_gradients: bool = False
+    ) -> None:
+        """Configure capture for subsequent forwards.
+
+        ``enabled=False`` disables both map and gradient capture regardless of
+        ``capture_gradients``, without immediately erasing the last result or
+        removing hooks needed by an outstanding backward pass. The next
+        uncaptured attention forward clears both cached tensors.
+
+        With capture enabled, each forward replaces the map and clears the old
+        gradient. Gradient capture always implies map capture. Attention-only
+        implementations must support ``capture_gradients=False`` without
+        registering backward hooks or requiring autograd.
         """
         ...
 
     @abstractmethod
     def get_attention_layers(
         self,
-    ) -> dict[str, list[tuple[torch.Tensor, torch.Tensor]]]:
-        """Return (attention_map, attention_gradient) pairs for all feature keys.
-
-        Must be called **after** ``set_attention_hooks(True)``,
-        a ``forward()`` call, and a subsequent ``backward()`` call so
-        that both attention maps and their gradients are populated.
-
-        Returns
-        -------
-        dict[str, list[tuple[torch.Tensor, torch.Tensor]]]
-            A dictionary keyed by ``feature_keys``.  Each value is a list
-            with one ``(attn_map, attn_grad)`` tuple per attention layer,
-            ordered from the first (closest to input) to the last
-            (closest to output).
-
-            Each tensor may have shape:
-
-            * ``[batch, heads, seq, seq]`` — multi-head (will be
-              gradient-weighted-averaged across heads by Chefer).
-            * ``[batch, seq, seq]`` — already head-averaged.
-
-            ``attn_map`` and ``attn_grad`` in the same tuple must have
-            the same shape.
-
-        Examples
-        --------
-        A model with stacked ``TransformerBlock`` layers per feature key:
-
-        >>> def get_attention_layers(self):
-        ...     return {
-        ...         key: [
-        ...             (blk.attention.get_attn_map(),
-        ...              blk.attention.get_attn_grad())
-        ...             for blk in self.transformer[key].transformer
-        ...         ]
-        ...         for key in self.feature_keys
-        ...     }
-
-        A model with a single MHA layer per feature key:
-
-        >>> def get_attention_layers(self):
-        ...     return {
-        ...         key: [(self.stagenet[key].get_attn_map(),
-        ...                self.stagenet[key].get_attn_grad())]
-        ...         for key in self.feature_keys
-        ...     }
+    ) -> dict[str, list[tuple[torch.Tensor | None, torch.Tensor | None]]]:
+        """Return ordered (map, gradient) pairs, keyed by feature.
+
+        Maps are available after a captured forward; gradients only after a
+        gradient-captured forward and backward. Either element can be ``None``.
+        Available pairs have matching shapes: ``[batch, heads, seq, seq]``.
+        Chefer also accepts head-averaged ``[batch, seq, seq]`` pairs; Rollout
+        requires the head dimension. Disabling capture preserves these results
+        until the next forward.
         """
         ...
 
@@ -447,7 +316,7 @@ def get_relevance_tensor(
     ) -> dict[str, torch.Tensor]:
         """Reduce relevance matrices to per-token attribution vectors.
 
-        The Chefer algorithm builds a relevance matrix of shape
+        Attention attribution builds a relevance matrix of shape
         ``[batch, seq_len, seq_len]`` for each feature key.  This method
         reduces each matrix to a ``[batch, seq_len]`` vector by selecting
         the row corresponding to the classification position — giving the
@@ -475,12 +344,12 @@ def get_relevance_tensor(
         --------
         CLS-token model (e.g. Transformer) — row 0 for all keys:
 
-        >>> def get_relevance_vector(self, R, **data):
+        >>> def get_relevance_tensor(self, R, **data):
         ...     return {key: r[:, 0] for key, r in R.items()}
 
         Last-valid-timestep model (e.g. StageAttentionNet):
 
-        >>> def get_relevance_vector(self, R, **data):
+        >>> def get_relevance_tensor(self, R, **data):
         ...     result = {}
         ...     for key, r in R.items():
         ...         mask = self._get_mask(key, **data)
@@ -491,9 +360,39 @@ def get_relevance_tensor(
         """
         ...
 
-    # TODO: Add postprocess_attribution() when ViT support is ready.
-    # ViT models need to strip the CLS column, reshape the patch vector
-    # into a spatial [batch, 1, H, W] map, and optionally interpolate to
-    # the original image size.  For EHR models this is a no-op.  We can
-    # either fold this into extract_attribution() or add it as a separate
-    # optional method.
\ No newline at end of file
+
+class GradientInterpretable(AttentionInterpretable):
+    """Extend attention capture with gradients for Chefer relevance.
+
+    After a gradient-enabled forward and backward, each captured attention map
+    must have a corresponding gradient of the same shape. Interpreters select
+    the gradient mode explicitly. Existing models keep the one-argument default
+    enabled for compatibility.
+
+    ``CheferInterpretable`` is an alias of this class. Custom implementations
+    must accept the ``capture_gradients`` keyword; the alias does not adapt old
+    ``set_attention_hooks(enabled)`` implementations. Both interpreters pass
+    this keyword, so obsolete signatures raise ``TypeError``.
+
+    Examples:
+        >>> CheferInterpretable is GradientInterpretable
+        True
+        >>> issubclass(GradientInterpretable, AttentionInterpretable)
+        True
+    """
+
+    @abstractmethod
+    def set_attention_hooks(
+        self, enabled: bool, *, capture_gradients: bool = True
+    ) -> None:
+        """Configure maps and optional gradients using the parent lifecycle.
+
+        ``set_attention_hooks(True)`` preserves historical gradient capture;
+        ``set_attention_hooks(True, capture_gradients=False)`` captures only
+        maps, including under ``torch.no_grad()``. Disabling always stops both.
+        """
+        ...
+
+
+# Preserve historical imports, inheritance, and isinstance checks.
+CheferInterpretable = GradientInterpretable
diff --git a/pyhealth/interpret/methods/attention_rollout.py b/pyhealth/interpret/methods/attention_rollout.py
index fc4e709f5..f7c955cd6 100644
--- a/pyhealth/interpret/methods/attention_rollout.py
+++ b/pyhealth/interpret/methods/attention_rollout.py
@@ -17,6 +17,7 @@
 
 import torch
 
+from pyhealth.interpret.api import AttentionInterpretable
 from pyhealth.models.base_model import BaseModel
 from .base_interpreter import BaseInterpreter
 
@@ -35,27 +36,16 @@ class AttentionRollout(BaseInterpreter):
     It serves as the standard baseline that gradient-based attention methods
     are compared against.
 
-    .. note::
-        "Gradient-free" refers to the attribution **math**: no backward pass
-        is run and no gradients enter the rollout computation. It does **not**
-        mean the call is safe inside ``torch.no_grad()``. The shared
-        attention-readout plumbing registers a gradient hook on the attention
-        tensors during the forward pass, so running ``attribute(**batch)``
-        under ``torch.no_grad()`` raises a ``RuntimeError``. Call it under the
-        default (grad-enabled) context.
-
-    This interpreter works with any model that exposes the attention-readout
-    methods ``set_attention_hooks``, ``get_attention_layers``, and
-    ``get_relevance_tensor`` (currently :class:`~pyhealth.models.Transformer`
-    and :class:`~pyhealth.models.StageAttentionNet`). Compatibility is checked
-    by duck-typing in ``__init__`` rather than by requiring a named interface,
-    since these methods are general attention readout and not specific to any
-    one method.
+    This interpreter requires :class:`~pyhealth.interpret.api.AttentionInterpretable`
+    (currently Transformer and StageAttentionNet). It captures maps under
+    ``torch.no_grad()`` without registering backward hooks or changing parameter
+    gradients. Custom models must accept ``capture_gradients=False`` in
+    ``set_attention_hooks``; matching method names alone are insufficient.
 
     The algorithm, per feature key:
 
-    1. Enable attention hooks via ``model.set_attention_hooks(True)`` and run a
-       single forward pass (no backward pass).
+    1. Enable map capture with ``capture_gradients=False`` and run a single
+       forward pass under ``torch.no_grad()`` (no backward pass).
     2. Retrieve per-layer attention maps via ``model.get_attention_layers()``,
        discarding the gradient element of each ``(attn_map, attn_grad)`` pair.
     3. Fuse heads (mean) to get one ``[batch, seq, seq]`` matrix per layer.
@@ -123,25 +113,14 @@ def __init__(self, model: BaseModel, head_fusion: str = "mean"):
                 "Currently supported values: mean."
             )
 
-        required_methods = [
-            "set_attention_hooks",
-            "get_attention_layers",
-            "get_relevance_tensor",
-        ]
-        missing_methods = [m for m in required_methods if not hasattr(model, m)]
-
-        if missing_methods:
-            raise TypeError(
-                "AttentionRollout requires a model that exposes the attention "
-                "interpretability methods: "
-                f"{', '.join(required_methods)}. "
-                f"Missing: {', '.join(missing_methods)}."
-            )
+        if not isinstance(model, AttentionInterpretable):
+            raise TypeError("Model must implement AttentionInterpretable interface")
 
         super().__init__(model)
         self.head_fusion = head_fusion
 
 
+    @torch.no_grad()
     def attribute(
         self,
         target_class_idx: Optional[int] = None,
@@ -169,13 +148,11 @@ def attribute(
                 row-stochastic matrices).
 
         Note:
-            Do not call this method inside a ``torch.no_grad()`` context. Even
-            though rollout uses no gradients, enabling attention hooks registers
-            a gradient hook during the forward pass, which requires grad-enabled
-            tensors and otherwise raises a ``RuntimeError``.
+            Runs under ``torch.no_grad()`` and can also be called inside an
+            existing no-grad context. Only attention maps are captured.
         """
 
-        self.model.set_attention_hooks(True)
+        self.model.set_attention_hooks(True, capture_gradients=False)
         try:
             self.model(**data)
         finally:
diff --git a/pyhealth/interpret/methods/chefer.py b/pyhealth/interpret/methods/chefer.py
index 26ce6ffd2..bfe885010 100644
--- a/pyhealth/interpret/methods/chefer.py
+++ b/pyhealth/interpret/methods/chefer.py
@@ -2,7 +2,7 @@
 
 This module implements the Chefer et al. relevance propagation method for
 explaining transformer-family model predictions.  It relies on the
-:class:`~pyhealth.interpret.api.CheferInterpretable` interface — any model
+:class:`~pyhealth.interpret.api.GradientInterpretable` interface — any model
 that implements that interface is automatically supported.
 
 Paper:
@@ -18,9 +18,8 @@
 import torch
 import torch.nn.functional as F
 
-from pyhealth.interpret.api import CheferInterpretable
+from pyhealth.interpret.api import GradientInterpretable
 from pyhealth.models.base_model import BaseModel
-from pyhealth.interpret.api import CheferInterpretable
 from .base_interpreter import BaseInterpreter
 
 
@@ -66,7 +65,7 @@ class CheferRelevance(BaseInterpreter):
     """Chefer's gradient-weighted attention method for transformer interpretability.
 
     This interpreter works with **any** model that implements the
-    :class:`~pyhealth.interpret.api.CheferInterpretable` interface, which
+    :class:`~pyhealth.interpret.api.GradientInterpretable` interface, which
     currently includes:
 
     * :class:`~pyhealth.models.Transformer`
@@ -74,7 +73,7 @@ class CheferRelevance(BaseInterpreter):
 
     The algorithm:
 
-    1. Enable attention hooks via ``model.set_attention_hooks(True)``.
+    1. Enable attention hooks with ``capture_gradients=True``.
     2. Forward pass → capture attention maps and register gradient hooks.
     3. Backward pass from a one-hot target class.
     4. Retrieve ``(attn_map, attn_grad)`` pairs via ``model.get_attention_layers()``.
@@ -82,12 +81,12 @@ class CheferRelevance(BaseInterpreter):
     6. Reduce ``R`` to per-token vectors via ``model.get_relevance_tensor()``.
 
     Steps 1, 4 and 6 are delegated to the model through the
-    ``CheferInterpretable`` interface, making this class fully
+    ``GradientInterpretable`` interface, making this class fully
     model-agnostic.
 
     Args:
         model (BaseModel): A trained PyHealth model that implements
-            :class:`~pyhealth.interpret.api.CheferInterpretable`.
+            :class:`~pyhealth.interpret.api.GradientInterpretable`.
 
     Example:
         >>> from pyhealth.datasets import create_sample_dataset, get_dataloader
@@ -133,8 +132,8 @@ class CheferRelevance(BaseInterpreter):
 
     def __init__(self, model: BaseModel):
         super().__init__(model)
-        if not isinstance(model, CheferInterpretable):
-            raise ValueError("Model must implement CheferInterpretable interface")
+        if not isinstance(model, GradientInterpretable):
+            raise ValueError("Model must implement GradientInterpretable interface")
         self.model = model
 
     def attribute(
@@ -158,7 +157,7 @@ def attribute(
                 per-token attribution scores.
         """
         # --- 1. Forward with attention hooks enabled ---
-        self.model.set_attention_hooks(True)
+        self.model.set_attention_hooks(True, capture_gradients=True)
         try:
             logits = self.model(**data)["logit"]
         finally:
@@ -184,6 +183,17 @@ def attribute(
         # --- 4. Relevance propagation per feature key ---
         R_dict: dict[str, torch.Tensor] = {}
         for key, layers in attention_layers.items():
+            if not layers:
+                raise RuntimeError(f"No attention layers captured for feature '{key}'.")
+            for cam, grad in layers:
+                if cam is None or grad is None:
+                    raise RuntimeError(
+                        f"Missing attention map or gradient for feature '{key}'."
+                    )
+                if cam.shape != grad.shape:
+                    raise RuntimeError(
+                        f"Attention map and gradient shapes differ for feature '{key}'."
+                    )
             num_tokens = layers[0][0].shape[-1]
             R = (
                 torch.eye(num_tokens, device=device)
diff --git a/pyhealth/models/stagenet_mha.py b/pyhealth/models/stagenet_mha.py
index 637f2fdc5..0686bfef7 100644
--- a/pyhealth/models/stagenet_mha.py
+++ b/pyhealth/models/stagenet_mha.py
@@ -8,7 +8,7 @@
 from pyhealth.models import BaseModel
 from pyhealth.models.utils import get_last_visit
 from .transformer import MultiHeadedAttention
-from pyhealth.interpret.api import CheferInterpretable
+from pyhealth.interpret.api import GradientInterpretable
 
 from .embedding import EmbeddingModel
 
@@ -193,6 +193,8 @@ def forward(
         time: Optional[torch.Tensor] = None,
         mask: Optional[torch.Tensor] = None,
         register_hook: bool = False,
+        *,
+        capture_attention: bool = False,
     ) -> Tuple[torch.Tensor, ...]:
         """Forward propagation.
 
@@ -201,8 +203,8 @@ def forward(
             static: a tensor of shape [batch size, static_dim].
             mask: an optional tensor of shape [batch size, sequence len], where
                 1 indicates valid and 0 indicates invalid.
-            register_hook: whether to register a backward hook on attention
-                weights for gradient inspection.
+            register_hook: whether to capture maps and register a backward hook.
+            capture_attention: whether to capture maps without gradients.
 
         Returns:
             last_output: a tensor of shape [batch size, chunk_size*levels] representing the
@@ -245,7 +247,8 @@ def forward(
 
         seq_for_mha = hidden_seq.permute(1, 0, 2)  # [batch, time, hidden]
         attn_output = self.mha(
-            seq_for_mha, seq_for_mha, seq_for_mha, mask=attn_mask, register_hook=register_hook
+            seq_for_mha, seq_for_mha, seq_for_mha, mask=attn_mask, register_hook=register_hook,
+            capture_attention=capture_attention,
         )
         self.attn_map = self.get_attn_map()
         self.attn_gradients = None  # will be populated after backward if hooked
@@ -298,7 +301,7 @@ def forward(
         return last_output, output, distance
 
 
-class StageAttentionNet(BaseModel, CheferInterpretable):
+class StageAttentionNet(BaseModel, GradientInterpretable):
     """StageAttentionNet model.
 
     Paper: Junyi Gao et al. Stagenet: Stage-aware neural networks for health
@@ -407,6 +410,7 @@ def __init__(
         self.chunk_size = chunk_size
         self.levels = levels
         self._attention_hooks_enabled = False
+        self._attention_gradients_enabled = False
 
         # validate kwargs for StageNet layer
         if "input_dim" in kwargs:
@@ -477,7 +481,7 @@ def forward_from_embedding(
                 embed: (if embed=True in kwargs) the patient embedding.
         """
         # Support both the flag-based API and legacy kwarg-based API
-        register_attn_hook = self._attention_hooks_enabled
+        register_attn_hook = self._attention_gradients_enabled
         patient_emb = []
         distance = []
 
@@ -544,7 +548,8 @@ def forward_from_embedding(
 
             # Pass through StageNet layer with embedded features
             last_output, _, cur_dis = self.stagenet[feature_key](
-                value, time=time, mask=mask, register_hook=register_attn_hook
+                value, time=time, mask=mask, register_hook=register_attn_hook,
+                capture_attention=self._attention_hooks_enabled,
             )
 
             patient_emb.append(last_output)
@@ -633,16 +638,20 @@ def get_embedding_model(self) -> nn.Module | None:
         return self.embedding_model
 
     # ------------------------------------------------------------------
-    # CheferInterpretable interface
+    # GradientInterpretable interface
     # ------------------------------------------------------------------
 
-    def set_attention_hooks(self, enabled: bool) -> None:
+    def set_attention_hooks(
+        self, enabled: bool, *, capture_gradients: bool = True
+    ) -> None:
+        """Configure future capture; disabling preserves the latest results."""
         self._attention_hooks_enabled = enabled
+        self._attention_gradients_enabled = enabled and capture_gradients
 
     def get_attention_layers(
         self,
-    ) -> dict[str, list[tuple[torch.Tensor, torch.Tensor]]]:
-        return {  # type: ignore[return-value]
+    ) -> dict[str, list[tuple[torch.Tensor | None, torch.Tensor | None]]]:
+        return {
             key: [
                 (
                     cast(StageNetAttentionLayer, self.stagenet[key]).get_attn_map(),
diff --git a/pyhealth/models/transformer.py b/pyhealth/models/transformer.py
index cc0dfc5ca..9c8bbab0b 100644
--- a/pyhealth/models/transformer.py
+++ b/pyhealth/models/transformer.py
@@ -12,7 +12,7 @@
 from pyhealth.datasets import SampleDataset
 from pyhealth.models import BaseModel
 from pyhealth.models.embedding import EmbeddingModel
-from pyhealth.interpret.api import CheferInterpretable
+from pyhealth.interpret.api import GradientInterpretable
 
 # VALID_OPERATION_LEVEL = ["visit", "event"]
 
@@ -64,7 +64,16 @@ def forward(
 
 
 class MultiHeadedAttention(nn.Module):
-    """Multi-head attention wrapper used by the Transformer block."""
+    """Multi-head attention wrapper used by the Transformer block.
+
+    Examples:
+        >>> attention = MultiHeadedAttention(h=2, d_model=8)
+        >>> x = torch.randn(1, 3, 8)
+        >>> with torch.no_grad():
+        ...     output = attention(x, x, x, capture_attention=True)
+        >>> attention.get_attn_map().shape
+        torch.Size([1, 2, 3, 3])
+    """
 
     def __init__(self, h: int, d_model: int, dropout: float = 0.1):
         """Initialize the attention module.
@@ -126,6 +135,8 @@ def forward(
         value: torch.Tensor,
         mask: Optional[torch.Tensor] = None,
         register_hook: bool = False,
+        *,
+        capture_attention: bool = False,
     ) -> torch.Tensor:
         """Run multi-head attention with optional gradient capture.
 
@@ -134,12 +145,15 @@ def forward(
             key: Key tensor aligned with ``query``.
             value: Value tensor aligned with ``query``.
             mask: Optional boolean mask ``[batch, len_q, len_k]``.
-            register_hook: True to attach a backward hook saving gradients.
+            register_hook: True to capture maps and attach a gradient hook.
+            capture_attention: Capture maps without requiring gradients.
 
         Returns:
             torch.Tensor: Attention mixed representation ``[batch, len_q, hidden]``.
         """
 
+        self.attn_map = None
+        self.attn_gradients = None
         batch_size = query.size(0)
 
         # 1) Do all the linear projections in batch from d_model => h x d_k
@@ -153,15 +167,11 @@ def forward(
             mask = mask.unsqueeze(1)
         x, attn = self.attention(query, key, value, mask=mask, dropout=self.dropout)
 
-        if register_hook:
-            # Only store attn_map and hook during interpretability passes.
-            # Using .detach() gives an independent copy whose storage
-            # is NOT shared with the live graph, so the graph can be freed
-            # normally after .backward() without leaking GPU memory.
+        if capture_attention or register_hook:
+            # Detach captured maps so the cache does not retain the live graph.
             self.attn_map = attn.detach()
+        if register_hook:
             attn.register_hook(self.save_attn_grad)
-        else:
-            self.attn_map = None
         # 3) "Concat" using a view and apply a final linear.
         x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.d_k)
   
@@ -230,6 +240,13 @@ class TransformerBlock(nn.Module):
         hidden: hidden size of transformer.
         attn_heads: head sizes of multi-head attention.
         dropout: dropout rate.
+
+    Examples:
+        >>> block = TransformerBlock(hidden=8, attn_heads=2, dropout=0.0)
+        >>> with torch.no_grad():
+        ...     output = block(torch.randn(1, 3, 8), capture_attention=True)
+        >>> output.shape
+        torch.Size([1, 3, 8])
     """
 
     def __init__(self, hidden, attn_heads, dropout):
@@ -246,7 +263,7 @@ def set_activation_hooks(self, hooks) -> None:
         """Deprecated compatibility stub; no-op."""
         return None
 
-    def forward(self, x, mask=None, register_hook = False):
+    def forward(self, x, mask=None, register_hook=False, *, capture_attention=False):
         """Forward propagation.
 
         Args:
@@ -256,7 +273,13 @@ def forward(self, x, mask=None, register_hook = False):
         Returns:
             A tensor of shape [batch_size, seq_len, hidden]
         """
-        x = self.input_sublayer(x, lambda _x: self.attention(_x, _x, _x, mask=mask, register_hook=register_hook))
+        x = self.input_sublayer(
+            x,
+            lambda _x: self.attention(
+                _x, _x, _x, mask=mask, register_hook=register_hook,
+                capture_attention=capture_attention,
+            ),
+        )
         x = self.output_sublayer(x, lambda _x: self.feed_forward(_x, mask=mask))
         return self.dropout(x)
 
@@ -297,7 +320,8 @@ def set_activation_hooks(self, hooks) -> None:
         return None
 
     def forward(
-        self, x: torch.Tensor, mask: Optional[torch.Tensor] = None, register_hook: bool = False
+        self, x: torch.Tensor, mask: Optional[torch.Tensor] = None,
+        register_hook: bool = False, *, capture_attention: bool = False,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         """Forward propagation.
 
@@ -315,13 +339,13 @@ def forward(
         if mask is not None:
             mask = torch.einsum("ab,ac->abc", mask, mask)
         for transformer in self.transformer:
-            x = transformer(x, mask, register_hook)
+            x = transformer(x, mask, register_hook, capture_attention=capture_attention)
         emb = x
         cls_emb = x[:, 0, :]
         return emb, cls_emb
 
 
-class Transformer(BaseModel, CheferInterpretable):
+class Transformer(BaseModel, GradientInterpretable):
     """Transformer model for PyHealth 2.0 datasets.
 
     Each feature stream is embedded with :class:`EmbeddingModel` and encoded by
@@ -385,6 +409,7 @@ def __init__(
         self.num_layers = num_layers
         self.max_seq_len = max_seq_len
         self._attention_hooks_enabled = False
+        self._attention_gradients_enabled = False
 
         assert (
             len(self.label_keys) == 1
@@ -481,7 +506,7 @@ def forward_from_embedding(
                 embed: (if embed=True in kwargs) the patient embedding.
         """
         # Support both the flag-based API and legacy kwarg-based API
-        register_hook = self._attention_hooks_enabled
+        register_hook = self._attention_gradients_enabled
         patient_emb = []
 
         for feature_key in self.feature_keys:
@@ -516,7 +541,8 @@ def forward_from_embedding(
                 mask = self._mask_from_embeddings(value).to(self.device)
 
             _, cls_emb = self.transformer[feature_key](
-                value, mask, register_hook
+                value, mask, register_hook,
+                capture_attention=self._attention_hooks_enabled,
             )
             patient_emb.append(cls_emb)
 
@@ -606,16 +632,20 @@ def get_embedding_model(self) -> nn.Module | None:
         return self.embedding_model
 
     # ------------------------------------------------------------------
-    # CheferInterpretable interface
+    # GradientInterpretable interface
     # ------------------------------------------------------------------
 
-    def set_attention_hooks(self, enabled: bool) -> None:
+    def set_attention_hooks(
+        self, enabled: bool, *, capture_gradients: bool = True
+    ) -> None:
+        """Configure future capture; disabling preserves the latest results."""
         self._attention_hooks_enabled = enabled
+        self._attention_gradients_enabled = enabled and capture_gradients
 
     def get_attention_layers(
         self,
-    ) -> dict[str, list[tuple[torch.Tensor, torch.Tensor]]]:
-        return {  # type: ignore[return-value]
+    ) -> dict[str, list[tuple[torch.Tensor | None, torch.Tensor | None]]]:
+        return {
             key: [
                 (
                     cast(TransformerBlock, blk).attention.get_attn_map(),
diff --git a/tests/core/test_attention_capture.py b/tests/core/test_attention_capture.py
new file mode 100644
index 000000000..2210a3036
--- /dev/null
+++ b/tests/core/test_attention_capture.py
@@ -0,0 +1,219 @@
+"""Regression tests for separate attention-map and gradient capture."""
+
+import unittest
+from unittest.mock import patch
+
+import torch
+
+from pyhealth.datasets import get_dataloader
+from pyhealth.interpret.api import (
+    AttentionInterpretable,
+    CheferInterpretable,
+    GradientInterpretable,
+)
+from pyhealth.interpret.methods import AttentionRollout, CheferRelevance
+from pyhealth.models.stagenet_mha import StageNetAttentionLayer
+from pyhealth.models.transformer import (
+    MultiHeadedAttention,
+    TransformerBlock,
+    TransformerLayer,
+)
+from tests.core import test_attention_rollout, test_stagenet_mha
+
+
+class AttentionOnly(torch.nn.Module, AttentionInterpretable):
+    """Minimal custom model whose forward needs no autograd."""
+
+    feature_keys = ["codes"]
+
+    def __init__(self):
+        super().__init__()
+        self.enabled = False
+        self.layers = {}
+
+    def set_attention_hooks(self, enabled, *, capture_gradients=False):
+        self.enabled = enabled
+
+    def forward(self, **data):
+        self.layers = {"codes": [(torch.eye(2).reshape(1, 1, 2, 2), None)]}
+        return {"logit": torch.zeros(1, 1)}
+
+    def get_attention_layers(self):
+        return self.layers
+
+    def get_relevance_tensor(self, R, **data):
+        return {key: value[:, 0] for key, value in R.items()}
+
+
+class TestAttentionCapture(unittest.TestCase):
+    def models_and_batches(self):
+        for case_class in (
+            test_attention_rollout.TestAttentionRollout,
+            test_stagenet_mha.TestStageNetMHA,
+        ):
+            torch.manual_seed(42)
+            case = case_class()
+            case.setUp()
+            case.model.eval()
+            batch = next(iter(get_dataloader(case.dataset, batch_size=2)))
+            if case_class is test_stagenet_mha.TestStageNetMHA:
+                # Supply the supported explicit padding mask for raw token IDs.
+                for key in case.model.feature_keys:
+                    time, value = batch[key]
+                    batch[key] = (time, value, value.ne(0))
+            yield case.model, batch
+
+    def assert_captured(self, model, gradients):
+        for layers in model.get_attention_layers().values():
+            self.assertTrue(layers)
+            for attention, gradient in layers:
+                self.assertIsNotNone(attention)
+                self.assertFalse(attention.requires_grad)
+                self.assertIsNone(attention.grad_fn)
+                if gradients:
+                    self.assertIsNotNone(gradient)
+                    self.assertEqual(attention.shape, gradient.shape)
+                else:
+                    self.assertIsNone(gradient)
+
+    def test_alias_and_interfaces(self):
+        self.assertIs(CheferInterpretable, GradientInterpretable)
+        for model, _ in self.models_and_batches():
+            self.assertIsInstance(model, AttentionInterpretable)
+            self.assertIsInstance(model, GradientInterpretable)
+
+    def test_rollout_no_grad_and_existing_parameter_gradients(self):
+        for model, batch in self.models_and_batches():
+            with self.subTest(model=type(model).__name__):
+                for parameter in model.parameters():
+                    parameter.grad = torch.ones_like(parameter)
+                before = [p.grad.clone() for p in model.parameters()]
+                interpreter = AttentionRollout(model)
+                baseline = interpreter.attribute(**batch)
+                with torch.no_grad():
+                    result = interpreter.attribute(**batch)
+                for key in baseline:
+                    torch.testing.assert_close(result[key], baseline[key])
+                self.assert_captured(model, gradients=False)
+                for previous, parameter in zip(before, model.parameters()):
+                    self.assertTrue(torch.equal(previous, parameter.grad))
+
+    def test_disable_preserves_results_until_next_forward(self):
+        for model, batch in self.models_and_batches():
+            with self.subTest(model=type(model).__name__):
+                model.set_attention_hooks(True)
+                output = model(**batch)
+                model.set_attention_hooks(False, capture_gradients=True)
+                self.assert_captured(model, gradients=False)
+                output["logit"].sum().backward()
+                self.assert_captured(model, gradients=True)
+                with torch.no_grad():
+                    model(**batch)
+                for layers in model.get_attention_layers().values():
+                    for attention, gradient in layers:
+                        self.assertIsNone(attention)
+                        self.assertIsNone(gradient)
+
+    def test_forward_only_does_not_capture_gradients_after_backward(self):
+        for model, batch in self.models_and_batches():
+            model.set_attention_hooks(True, capture_gradients=False)
+            model(**batch)["logit"].sum().backward()
+            self.assert_captured(model, gradients=False)
+
+    def test_chefer_rollout_chefer_sequence(self):
+        for model, batch in self.models_and_batches():
+            with self.subTest(model=type(model).__name__):
+                chefer = CheferRelevance(model)
+                first = chefer.attribute(**batch)
+                self.assert_captured(model, gradients=True)
+                AttentionRollout(model).attribute(**batch)
+                self.assert_captured(model, gradients=False)
+                second = chefer.attribute(**batch)
+                self.assert_captured(model, gradients=True)
+                for key in first:
+                    torch.testing.assert_close(first[key], second[key])
+
+    def test_chefer_explicitly_requests_gradients(self):
+        for model, batch in self.models_and_batches():
+            original = model.set_attention_hooks
+
+            def false_default(enabled, *, capture_gradients=False):
+                original(enabled, capture_gradients=capture_gradients)
+
+            with patch.object(model, "set_attention_hooks", false_default):
+                CheferRelevance(model).attribute(**batch)
+                self.assert_captured(model, gradients=True)
+
+    def test_cleanup_after_forward_exception(self):
+        for model, batch in self.models_and_batches():
+            for interpreter_class in (AttentionRollout, CheferRelevance):
+                with patch.object(model, "forward", side_effect=RuntimeError("failed")):
+                    with self.assertRaisesRegex(RuntimeError, "failed"):
+                        interpreter_class(model).attribute(**batch)
+                with torch.no_grad():
+                    model(**batch)
+                for layers in model.get_attention_layers().values():
+                    self.assertTrue(all(pair == (None, None) for pair in layers))
+
+    def test_attention_only_and_method_only_models(self):
+        model = AttentionOnly()
+        with torch.no_grad():
+            result = AttentionRollout(model).attribute(codes=torch.ones(1, 2))
+        torch.testing.assert_close(result["codes"], torch.tensor([[1.0, 0.0]]))
+        with self.assertRaises(ValueError):
+            CheferRelevance(model)
+
+        class MethodOnly(torch.nn.Module):
+            set_attention_hooks = AttentionOnly.set_attention_hooks
+            get_attention_layers = AttentionOnly.get_attention_layers
+            get_relevance_tensor = AttentionOnly.get_relevance_tensor
+
+        with self.assertRaises(TypeError):
+            AttentionRollout(MethodOnly())
+
+    def test_obsolete_signature_requires_migration(self):
+        class Obsolete(AttentionOnly, GradientInterpretable):
+            def set_attention_hooks(self, enabled):
+                self.enabled = enabled
+
+        for interpreter_class in (AttentionRollout, CheferRelevance):
+            with self.assertRaisesRegex(TypeError, "capture_gradients"):
+                interpreter_class(Obsolete()).attribute(codes=torch.ones(1, 2))
+
+    def test_chefer_missing_or_mismatched_capture(self):
+        model, batch = next(self.models_and_batches())
+        attention = torch.ones(2, 2, 3, 3)
+        cases = (
+            [],
+            [(None, attention)],
+            [(attention, None)],
+            [(attention, attention[:, :, :1])],
+        )
+        for layers in cases:
+            with self.subTest(layers=layers):
+                with patch.object(
+                    model, "get_attention_layers", return_value={"codes": layers}
+                ):
+                    with self.assertRaisesRegex(RuntimeError, "codes"):
+                        CheferRelevance(model).attribute(**batch)
+
+    def test_low_level_positional_gradient_argument(self):
+        x = torch.randn(2, 3, 6, requires_grad=True)
+        mha = MultiHeadedAttention(2, 6)
+        block = TransformerBlock(6, 2, 0.0)
+        layer = TransformerLayer(6, heads=2)
+        stage = StageNetAttentionLayer(6, chunk_size=2, levels=3, num_heads=2)
+        calls = (
+            (mha, lambda: mha(x, x, x, None, True)),
+            (block.attention, lambda: block(x, None, True)),
+            (layer.transformer[0].attention, lambda: layer(x, None, True)[0]),
+            (stage.mha, lambda: stage(x, None, None, True)[0]),
+        )
+        for attention, call in calls:
+            call().sum().backward()
+            self.assertIsNotNone(attention.get_attn_map())
+            self.assertIsNotNone(attention.get_attn_grad())
+
+
+if __name__ == "__main__":
+    unittest.main()