Add attention rollout interpretability method (Abnar & Zuidema 2020)

[fbonc]

Contributor: Felipe Amaral Bonchristiano (felipea5@illinois.edu)

Contribution Type: New interpretability method

Description:
Adds vanilla attention rollout (Abnar & Zuidema, "Quantifying Attention Flow
in Transformers," 2020, arXiv:2005.00928) as a new interpretability module,
AttentionRollout. Rollout is the canonical forward-only, gradient-free,
class-agnostic attention-flow baseline: it accounts for residual connections
( = 0.5·(A + I)), fuses heads by mean, and composes per-layer attention by
matrix product to produce per-token relevance. It complements the existing
CheferRelevance (gradient-weighted, class-specific) by providing the standard
baseline that gradient-based attention methods are measured against, which the
interpretability suite currently lacks.

The implementation reuses the existing attention-readout methods already on
PyHealth's attention models (set_attention_hooks, get_attention_layers,
get_relevance_tensor), so it requires no model-side changes. It is a
single new method file plus its export, tests, docs, and example registrations.

Files to Review:

• pyhealth/interpret/methods/attention_rollout.py: core implementation (AttentionRollout)
• pyhealth/interpret/methods/__init__.py: exports AttentionRollout
• tests/core/test_attention_rollout.py: synthetic-data unit tests (see Testing below)
• docs/api/interpret/pyhealth.interpret.methods.attention_rollout.rst: API documentation
• docs/api/interpret.rst: added to the Attribution Methods toctree
• examples/interpretability/{mp,los,dka}_{transformer,stageattn}_mimic4_interpret.py: AttentionRollout registered in the method comparison dicts alongside CheferRelevance

Quick note:
The actual bounty on the doc lists "Rollout Attention" and links arXiv:2012.09838,
which is Chefer et al., Transformer Interpretability Beyond Attention
Visualization (CVPR 2021), a gradient/LRP relevance method, not the rollout
paper. (The existing CheferRelevance implements the related Chefer et al. ICCV
2021 method, arXiv:2103.15679.) I read the bounty's intent from its name and from
the actual gap in the suite, as there was no gradient-free, class-agnostic baseline,
and implemented canonical rollout (Abnar & Zuidema 2020) rather than more
Chefer-style work. If the literal citation was intended, happy to redirect.

Key design decisions:

• Canonical rollout, not an enhanced variant. Default is mean head fusion +
  0.5·(A + I); alternative fusions and residual schemes are deferred to optional
  kwargs. Again, this module's value is fidelity to the baseline, not improving on it.
• Model compatibility via duck-typing, not isinstance(CheferInterpretable).
  The three readout methods are general attention readout, not Chefer-specific;
  __init__ checks hasattr and raises TypeError naming the missing methods.
  This keeps the PR to one new file with zero edits to the shared interface.
• target_class_idx accepted but ignored, documented as a no-op, so rollout is
  drop-in swappable with class-specific interpreters in existing pipelines.
• _map_to_input_shapes duplicated from CheferRelevance (rather than factored
  to a shared util) so attributions match the raw-input granularity the
  comprehensiveness/sufficiency metrics expect, while keeping this PR free of edits
  to chefer.py.

Proposed follow-up: extract a general AttentionInterpretable
interface and a shared shape-mapping helper that both AttentionRollout and
CheferRelevance depend on, removing the duck-typing and the duplicated
_map_to_input_shapes. Kept separate to avoid bundling a refactor of shared code
into a feature PR.

Testing: Unit tests use small synthetic data (create_sample_dataset, tiny
config, seeded) and run in well under a second with no network or credentials.
Beyond shape and dict-key checks, they assert the two correctness invariants:
(1) per-token relevance sums to 1 before input-shape expansion (the product of
row-stochastic matrices is row-stochastic), and (2) identity attention at every
layer yields an identity rollout. Construction-time errors (incompatible model,
unsupported head_fusion) are covered.

Note on verification: I am not yet MIMIC-credentialed, so end-to-end correctness
is established via the synthetic unit tests above; AttentionRollout is registered
in the MIMIC-IV comparison scripts for parity with the other methods but I have not
run those end-to-end myself.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[jhnwu3]

Thanks for this — the implementation, docs, and test suite are in good shape. The rationale in the description is clear and the correctness invariants in the tests (row-stochastic check, identity-attention test) are exactly what we want to see. A few minor things (header comments, trailing newline in attention_rollout.py, the unreachable ValueError branch in _fuse_heads) are non-blocking; feel free to clean them up or leave them.

Two things before merge:

---

1. One doc spot still lists only Chefer

docs/why_pyhealth.rst line 178 currently reads:

Attention-based: Chefer relevance propagation for transformers

Since rollout is the canonical gradient-free attention baseline that Chefer-style methods are measured against, please add it here — e.g.:

Attention-based: Chefer relevance propagation and attention rollout for transformers

Everything else checks out: the __init__ export, the toctree entry, the API rst page, and the six MIMIC comparison scripts are all wired up correctly.

---

2. A runnable benchmark example

Every existing script that calls evaluate_attribution depends on credentialed full MIMIC-IV plus pre-trained checkpoints and hardcoded /shared/eng/... paths — which is exactly why you noted you couldn't verify end-to-end faithfulness. Could you add a self-contained example under examples/interpretability/ that uses the MIMIC-IV demo dataset (the ~100-patient open-access subset on PhysioNet, freely downloadable with no credentialing: https://physionet.org/content/mimic-iv-demo/)? It should train a small Transformer from scratch and run evaluate_attribution comparing AttentionRollout, CheferRelevance, and RandomBaseline on Comprehensiveness/Sufficiency.

Here's a complete starting point you can drop in as examples/interpretability/attention_rollout_benchmark_mimic4_demo.py:

"""Benchmark AttentionRollout on the MIMIC-IV demo dataset (no credentialing).

Download once from PhysioNet (open access, no credentialing required):
  https://physionet.org/content/mimic-iv-demo/
Extract and point --ehr_root at the resulting directory.

Run:
    python examples/interpretability/attention_rollout_benchmark_mimic4_demo.py \
        --ehr_root /path/to/mimic-iv-clinical-database-demo
"""
import argparse

import torch

from pyhealth.datasets import MIMIC4Dataset, get_dataloader, split_by_patient
from pyhealth.interpret.methods import AttentionRollout, CheferRelevance, RandomBaseline
from pyhealth.metrics.interpretability import evaluate_attribution
from pyhealth.models import Transformer
from pyhealth.tasks import MortalityPredictionMIMIC4
from pyhealth.trainer import Trainer


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--ehr_root",
        required=True,
        help="Path to the extracted mimic-iv-clinical-database-demo directory",
    )
    parser.add_argument("--device", default="cpu")
    args = parser.parse_args()

    torch.manual_seed(0)

    # 1. Load the open-access demo dataset (no credentialing required).
    base_dataset = MIMIC4Dataset(
        ehr_root=args.ehr_root,
        ehr_tables=[
            "patients",
            "admissions",
            "diagnoses_icd",
            "procedures_icd",
            "prescriptions",
        ],
    )

    # 2. Apply the plain-Transformer mortality task; train fresh (no checkpoints).
    sample_dataset = base_dataset.set_task(MortalityPredictionMIMIC4())
    print(f"Loaded {len(sample_dataset)} samples")

    # 3. Patient-level split and loaders.
    train_ds, val_ds, test_ds = split_by_patient(
        sample_dataset, [0.7, 0.1, 0.2], seed=42
    )
    train_loader = get_dataloader(train_ds, batch_size=16, shuffle=True)
    val_loader = get_dataloader(val_ds, batch_size=16, shuffle=False)
    test_loader = get_dataloader(test_ds, batch_size=16, shuffle=False)

    # 4. Small Transformer — already exposes the attention-readout methods
    #    both AttentionRollout and CheferRelevance rely on.
    model = Transformer(
        dataset=sample_dataset,
        embedding_dim=64,
        heads=2,
        num_layers=2,
        dropout=0.1,
    )

    # 5. Train briefly on the demo data.
    trainer = Trainer(model=model, device=args.device, metrics=["roc_auc"])
    trainer.train(
        train_dataloader=train_loader,
        val_dataloader=val_loader,
        epochs=5,
        monitor="roc_auc",
        monitor_criterion="max",
    )
    model.eval()

    # 6. Compare attention interpreters against the random floor.
    methods = {
        "random": RandomBaseline(model),
        "chefer": CheferRelevance(model),
        "rollout": AttentionRollout(model),
    }

    print(f"\n{'method':<10}{'comprehensiveness':>20}{'sufficiency':>16}")
    print("-" * 46)
    for name, method in methods.items():
        scores = evaluate_attribution(
            model,
            test_loader,
            method,
            metrics=["comprehensiveness", "sufficiency"],
            percentages=[25, 50, 99],
        )
        print(
            f"{name:<10}"
            f"{scores['comprehensiveness']:>20.4f}"
            f"{scores['sufficiency']:>16.4f}"
        )


if __name__ == "__main__":
    main()

Both rollout and chefer should land above random on comprehensiveness — that's the sanity check that confirms the metric path works end-to-end for AttentionRollout.

Happy to review once those two are in.

---

Generated by Claude Code

[jhnwu3]

Check comments. But you're almost done.

[fbonc]

Hey, just wanted to give an update and ask about some decisions that need to be made before finalizing the benchmark. Sorry about the long message or if any of this stuff is trivial - lots of stuff that I'm learning as I'm going.

First, while building the demo benchmark, I hit a preexisting bug in the metric code (not in rollout). Essentially, pyhealth/metrics/interpretability/base.py, _apply_ablation has two paths. The tuple (time, values) path (StageNet) already guards discrete dtype and ablates with x * (1 - mask).long(). But, the plain tensor path (which a vanilla transformer uses for its code features) had no such guard: it did x * (1 - mask), upcasting long to float, so the next forward pass crashes in nn.Embedding ("indices must be Long/Int, got Float"). This hits Chefer and RandomBaseline on plain transformers too, not just rollout. To fix I mirrored the tuple branch's handling into the plain branch and factored the shared part into a helper so the two paths can't drift apart again. Could definitely split this into its own PR if you'd rather keep this one to just the rollout file, but it's a small fix.

Next, a methodology question aobut the metrics and unsigned attributions (also encountered during benchmarking). So, the removal metric framework evaluates predicted-negative samples by negating the attributions (SampleClass.NEGATIVE). This assumes signed attributions (negative means supports class 0). But, both Chefer and rollout produce unsigned, nonnegative importance. This means that negating those doesn't surface class 0 features, it just kind of inverts the ranking, so predicted-negative scores worse than random. In practice, this means that only predicted-positive samples evaluate cleanly for those methods. New issue perhaps?

Evidence: negating unsigned attributions makes them lose to random

---

All 29 test patients evaluated, predicted-negatives sent through the SampleClass.NEGATIVE (attribution-negating) path:

method	comprehensiveness ↑	sufficiency ↓
random	0.2358	0.2265
chefer	0.1427	0.2775
rollout	0.1103	0.2705

Random beats both real methods on both metrics. The negation inverts the (unsigned) rankings for every predicted negative patient.

---

Finally, on the demo I trained the model lightly (per your "don't chase AUC" note on dc), and it is (relatively) well calibrated from what I found while trying to understand why I was getting weird comprehensiveness/sufficiency numbers: P(death) over the 29 patients spans between 0.004-0.75 with only one patient more than 0.5. Comprehensiveness comes out negative for all three methods, including random, and with an empty input ablation baseline of around 0.5+ (Inferred, not measured: comprehensiveness is negative even for the random baseline, which, from what I understand, can only happen if ablation drifts predictions upward toward a high empty record baseline (avg P(death) = 0.49, avg comprehensiveness = −0.07 => ablated = 0.56), which is high, most patients sit well below it and removing features moves P(death) the "wrong" way. This isn't really a rollout issue, it's more a model/demo/evaluation interaction. So, is there a config you use to get a clean positive-comprehensiveness population. Like, is oversampling/rebalancing the training set acceptable here? At natural class balance I only get n=1 confidently positive.

Evidence: conservative model + high ablation baseline => negative comprehensiveness

---

Predicted P(death) on the 29 test patients:
count: 29
min / mean / max: 0.0036 / 0.1842 / 0.7499

>= 0.5: 1 patient
top 10: [0.7499, 0.4938, 0.4787, 0.4712, 0.4217,
              0.3387, 0.2478, 0.2457, 0.2332, 0.2266]

Predicted-positives only, threshold 0.3, attributions as-is (no negation), n = 6 valid patients:

method	comprehensiveness ↑	sufficiency ↓
random	-0.0703	-0.0492
chefer	-0.0450	-0.0447
rollout	-0.0426	0.0008

Comprehensiveness is negative for all three (random included), and the two metrics disagree on ranking. This is consistent with the empty-input ablation baseline sitting above most patients' P(death), so ablation moves P(death) the wrong way. Marking all 29 patients positive gives the same all-negative pattern.

---

[jhnwu3]

lgtm. I think for follow-up research, we can definitely look into benchmarking this to get more reasonable metric results.

Thanks for the hard work! Tbh the interp module is definitely the most opaque thing we have currently, and can definitely be a bit confusing to work with at first.