Layer-wise Relevance Propagation for physioex.explain (CNN, recurrent, attention models)

[guidogagl]

Layer-wise Relevance Propagation for physioex.explain

Adds a conservation-based attribution family (Bach et al. 2015) to PhysioEx, faithful across the whole model zoo — including the recurrent and transformer architectures that plain conv/dense LRP cannot handle.

What's in

Two entry points, one convention (target neuron seeded with its logit so Σ R ≈ f_c(x); return_report=True gives a per-sample ConservationReport):

Models	Entry point	Rules
Feed-forward / CNN (Tsinalis, Chambon2018)	LRP (Zennit composites)	ε dense · γ conv · w² first layer (unbounded EEG, not the pixel z-box) · BatchNorm canonized · αβ/flat/z-box configurable
Recurrent (TinySleepNet, SeqSleepNet, L-SeqSleepNet, ProtoSleepNet-seq)	ModelLRP	Arras signal-take LSTM/GRU at cell level (LRPLSTM/LRPGRU subclass nn.LSTM/GRU, forward bit-identical), gate_rule="signal_take"|"uniform"
Attention / transformer (SleepTransformer, CoReSleep, ProtoSleepNet-tf)	ModelLRP	CP-LRP attention (Ali et al. 2022; value path, exact conservation), identity rule on LayerNorm/activations, proportional rule on residuals, ε on projections/FFN

prepare_model_for_lrp rewrites a deep copy of the model (fused RNNs, TransformerEncoder/MultiheadAttention, PhysioEx's softmax poolings and the learnable filterbank via an adapter registry, ε on leaves, BatchNorm folding); plain + residuals in a model's own forward are redirected to the proportional rule at runtime (patch_residuals); audit_lrp_coverage / strict=True flag any parametric leaf left on plain autograd (gradient ≠ relevance). All primitives use a signed, dtype-safe stabiliser z + ε·sign z. No LXT dependency — the explain extra is zennit only.

Verification

Developed under a multi-agent review: four independent reviewers (recurrent math, attention math, wiring/API, methodological completeness) → every reported defect reproduced numerically on Sofia before fixing → adversarial verification of the fixes → three further blockers fixed. Highlights of what that caught: the seqsleepnet filterbank propagating gradient instead of relevance, plain residuals over-counting exactly 2×, LXT's sign-blind stabiliser (±1e8 on near-cancelling sums), fp16 promotion, CoReSleep's in-place residual write invalidating saved views.

• 91 tests (tests/explain/lrp/): forward-equivalence to the fused modules, exact conservation on bias-free blocks/models, rule/adapter wiring, robustness (no_grad, shared modules, idempotence, device/dtype), and all six real architectures end-to-end (incl. bimodal CoReSleep, L-SeqSleepNet, ProtoSleepNet with channel mixer) with full rule coverage. Green on Sofia and in CI (3.11/3.12).
• Strict -W docs build green; physioex.explain.lrp added to the API reference.
• Applied check on the pretrained seqsleepnet-phan (MASS, N3): examples/explain/lrp_seqsleepnet_n3.py compares Saliency / Input×Gradient / IG with LRP under different rule assignments and shows why ε must be scaled to the activations (1e-6 inflates 2–3×, 1e-2 conserves).

Docs

docs/pages/explain/explain.md: entry-point/rule matrix, how ModelLRP rewrites a model, how to read conservation and choose ε, unsupported cases (masks, batch_first=False, RNN initial states, proj_size, VQ path), and a literature note on LRP vs path-based attribution (Deep Taylor, IG, DeepLIFT/DeepSHAP, "integrated LRP").

Known limitations / follow-ups (not blocking)

• CP-LRP gives zero relevance to the query/key path by design; full AttnLRP (non-conserving) could be added as attn_rule="attnlrp".
• Bias relevance is absorbed (LRP-ε convention); an Arras δ=1 bias_mode and a relative-ε default for recurrent models are natural next steps.
• Phase 3 (foundation encoders: SleepFM easy, NeuroLM medium, BIOT/REVE blockers) and Phase 4 (DFT/STFT-LRP on the existing spectral layers) not started.
• DeepLIFT-style explicit-reference rule and head-level relevance would complete the toolkit.

🤖 Generated with Claude Code