feat(explain/lrp): end-to-end model wiring + attention-pooling CP-LRP (Phase 2c)

Wire the LRP engines to whole PhysioEx models:
- pooling.py: CP-LRP for the softmax attention-pooling blocks
  (out = Σ_t softmax(score)_t·x_t, weights treated as constant → relevance
  through the value path). LRPAttentionPooling (SleepTransformer) and
  LRPAttentionLayer (SeqSleepNet, reused by lseq/proto), matched by class name.
- model.py: prepare_model_for_lrp() swaps LSTM/GRU→LRPLSTM/GRU,
  TransformerEncoder(Layer)→LRP versions, MultiheadAttention→LRPMHA, custom
  pooling→CP-LRP, and wraps remaining Linear/Conv/BatchNorm with LXT EpsilonRule
  and LayerNorm with IdentityRule. ModelLRP deep-copies the model, prepares it
  (disabling in-place ops that would corrupt saved tensors) and runs a
  target-seeded backward, returning relevance shaped like the input.
- LRPLSTM/LRPGRU now return (output, states) like nn.LSTM/GRU so the models'
  `out, _ = self.rnn(x)` calls keep working after the swap.

Verified on Sofia: 33/33 synthetic LRP tests, PLUS end-to-end on the real
models across all three families — sleeptransformer, seqsleepnet, tinysleepnet:
prepared-forward matches the original (Δ~1e-8) and relevance is finite with the
input's shape.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: Guido Gagliardi

physioex/explain/lrp/__init__.py

@@ -17,10 +17,34 @@
 from physioex.explain.lrp.attributor import LRP
 from physioex.explain.lrp.canonizers import default_canonizers
 from physioex.explain.lrp.composites import epsilon_composite, physioex_composite
+from physioex.explain.lrp.model import ModelLRP, prepare_model_for_lrp
+from physioex.explain.lrp.pooling import LRPAttentionLayer, LRPAttentionPooling
+from physioex.explain.lrp.recurrent import LRPGRU, LRPLSTM
+from physioex.explain.lrp.transformer import (
+    LRPMultiheadAttention,
+    LRPTransformerEncoder,
+    LRPTransformerEncoderLayer,
+    swap_transformer_layers,
+)
 
 __all__ = [
+    # CNN / Zennit composites (Phase 1)
     "LRP",
     "physioex_composite",
     "epsilon_composite",
     "default_canonizers",
+    # recurrent (Phase 2a)
+    "LRPLSTM",
+    "LRPGRU",
+    # transformer / attention (Phase 2b)
+    "LRPMultiheadAttention",
+    "LRPTransformerEncoderLayer",
+    "LRPTransformerEncoder",
+    "swap_transformer_layers",
+    # attention pooling (Phase 2c)
+    "LRPAttentionPooling",
+    "LRPAttentionLayer",
+    # whole-model wiring (Phase 2c)
+    "ModelLRP",
+    "prepare_model_for_lrp",
 ]

physioex/explain/lrp/model.py

@@ -0,0 +1,122 @@
+"""End-to-end LRP for whole PhysioEx models (Phase 2c wiring).
+
+:func:`prepare_model_for_lrp` swaps every fused / attention block in a model for
+its relevance-carrying counterpart:
+
+* ``nn.LSTM`` / ``nn.GRU``               → :class:`LRPLSTM` / :class:`LRPGRU`
+* ``nn.TransformerEncoder(Layer)``       → the ``transformer`` LRP versions
+* ``nn.MultiheadAttention``              → :class:`LRPMultiheadAttention`
+* custom softmax pooling (matched by class name ``AttentionPooling`` /
+  ``AttentionLayer``) → the ``pooling`` CP-LRP versions
+* remaining leaf ``Linear`` / ``Conv`` / ``BatchNorm`` → LXT ``EpsilonRule``
+* standalone ``LayerNorm``               → LXT ``IdentityRule``
+
+:class:`ModelLRP` deep-copies the trained model, prepares it, then runs a forward
+plus a target-seeded backward, returning relevance shaped like the input.
+
+Requires the ``explain`` extra (``lxt``).
+"""
+
+from __future__ import annotations
+
+import copy
+
+import torch
+import torch.nn as nn
+
+from physioex.explain.lrp.pooling import LRPAttentionLayer, LRPAttentionPooling
+from physioex.explain.lrp.recurrent import LRPGRU, LRPLSTM
+from physioex.explain.lrp.transformer import (
+    LRPMultiheadAttention,
+    LRPTransformerEncoder,
+    LRPTransformerEncoderLayer,
+)
+
+# custom pooling classes are matched by name to avoid importing the model modules
+_POOLING_BY_NAME = {
+    "AttentionPooling": LRPAttentionPooling,
+    "AttentionLayer": LRPAttentionLayer,
+}
+
+
+def prepare_model_for_lrp(model: nn.Module, epsilon: float = 1e-6) -> nn.Module:
+    """In-place: replace fused/attention blocks with their LRP counterparts and
+    wrap remaining linear leaves with LXT rules.  Returns ``model``.
+
+    Call on a copy (see :class:`ModelLRP`) — it mutates the module tree.
+    """
+    from lxt.explicit.rules import EpsilonRule, IdentityRule
+
+    for name, child in list(model.named_children()):
+        cname = type(child).__name__
+        if isinstance(child, nn.LSTM):
+            setattr(model, name, LRPLSTM.from_torch(child))
+        elif isinstance(child, nn.GRU):
+            setattr(model, name, LRPGRU.from_torch(child))
+        elif isinstance(child, nn.TransformerEncoder):
+            setattr(model, name, LRPTransformerEncoder.from_torch(child))
+        elif isinstance(child, nn.TransformerEncoderLayer):
+            setattr(model, name, LRPTransformerEncoderLayer.from_torch(child))
+        elif isinstance(child, nn.MultiheadAttention):
+            setattr(model, name, LRPMultiheadAttention.from_torch(child))
+        elif cname in _POOLING_BY_NAME:
+            setattr(model, name, _POOLING_BY_NAME[cname](child, epsilon))
+        elif isinstance(child, (nn.Linear, nn.Conv1d, nn.Conv2d, nn.Conv3d)):
+            setattr(model, name, EpsilonRule(child, epsilon))
+        elif isinstance(child, nn.LayerNorm):
+            setattr(model, name, IdentityRule(child))
+        elif isinstance(child, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)):
+            # in eval mode BatchNorm is an affine-linear op → ε-LRP
+            setattr(model, name, EpsilonRule(child, epsilon))
+        else:
+            # not a leaf/known block: recurse into it
+            prepare_model_for_lrp(child, epsilon)
+    return model
+
+
+class ModelLRP(nn.Module):
+    """LRP attribution for a whole PhysioEx model (recurrent / transformer /
+    attention architectures).
+
+    Args:
+        model: trained model, ``(B, L, C, T[, F]) -> (B, L, n_classes)`` (or a
+            dict output — set ``output_key``).
+        in_index / out_index: sequence epoch and class to explain.
+        output_key: for models returning a dict (e.g. CoReSleep → ``"combined"``).
+        epsilon: ε for the ε-LRP rules.
+
+    ``forward(x)`` returns relevance shaped like ``x``.  The trained model is
+    deep-copied and left untouched.
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        in_index: int = 0,
+        out_index: int = 0,
+        output_key: str | None = None,
+        epsilon: float = 1e-6,
+    ):
+        super().__init__()
+        self.in_index = in_index
+        self.out_index = out_index
+        self.output_key = output_key
+        prepared = copy.deepcopy(model).eval()
+        self.model = prepare_model_for_lrp(prepared, epsilon)
+        # LRP saves tensors for the custom backward, so in-place ops (e.g.
+        # ReLU(inplace=True)) would corrupt them — disable them on the copy.
+        for m in self.model.modules():
+            if getattr(m, "inplace", False):
+                m.inplace = False
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x.detach().requires_grad_(True)
+        out = self.model(x)
+        if self.output_key is not None:
+            out = out[self.output_key]
+        seed = torch.zeros_like(out)
+        seed[:, self.in_index, self.out_index] = out[
+            :, self.in_index, self.out_index
+        ].detach()
+        out.backward(seed)
+        return x.grad

physioex/explain/lrp/pooling.py

@@ -0,0 +1,87 @@
+"""CP-LRP for softmax attention-pooling blocks.
+
+PhysioEx's custom attention poolings all collapse a sequence with softmax
+weights: ``out = Σ_t softmax(score(x))_t · x_t`` — ``AttentionPooling``
+(sleeptransformer), ``AttentionLayer`` (seqsleepnet, reused by lseqsleepnet and
+protosleepnet), and ``ChannelMixer.attn_pool`` (protosleepnet).  Their softmax
+leaks relevance just like attention, so the conservative-propagation (CP-LRP)
+rule applies: treat the pooling weights as **constant** and route relevance
+through the value path, which conserves exactly.
+
+Requires the ``explain`` extra.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+
+class _CPWeightedPool(torch.autograd.Function):
+    """``out = Σ_t w_t · x_t`` over the sequence dim, weights constant.
+
+    Backward (ε-LRP through the linear pooling with fixed ``w``):
+    ``R(x[t,d]) = x[t,d] · w[t] · R_out[d] / (out[d]+ε)``.
+    """
+
+    @staticmethod
+    def forward(ctx, x, weights, epsilon):
+        out = (x * weights).sum(dim=1)  # (B, D)
+        ctx.save_for_backward(x, weights, out)
+        ctx.epsilon = epsilon
+        return out
+
+    @staticmethod
+    def backward(ctx, relevance):
+        x, weights, out = ctx.saved_tensors
+        eps = ctx.epsilon
+        denom = out + torch.where(out >= 0, eps, -eps)
+        s = (relevance / denom).unsqueeze(1)  # (B, 1, D)
+        r_x = x * weights * s
+        return r_x, None, None
+
+
+def cp_weighted_pool(x, weights, epsilon=1e-6):
+    """Weighted sum over ``dim=1`` with the CP-LRP backward.  ``weights`` is
+    ``(B, T, 1)`` and must be detached (constant) by the caller."""
+    return _CPWeightedPool.apply(x, weights, epsilon)
+
+
+class LRPAttentionPooling(nn.Module):
+    """CP-LRP for ``sleeptransformer.AttentionPooling`` (and any pooling with a
+    ``self.attention`` MLP + softmax over ``dim=1``)."""
+
+    def __init__(self, orig: nn.Module, epsilon: float = 1e-6):
+        super().__init__()
+        self.orig = orig
+        self.epsilon = epsilon
+
+    def forward(self, x):
+        with torch.no_grad():
+            w = torch.softmax(self.orig.attention(x), dim=1)  # (B, T, 1)
+        return cp_weighted_pool(x, w, self.epsilon)
+
+
+class LRPAttentionLayer(nn.Module):
+    """CP-LRP for ``seqsleepnet.AttentionLayer`` (additive/Bahdanau pooling with
+    a manual softmax)."""
+
+    def __init__(self, orig: nn.Module, epsilon: float = 1e-6):
+        super().__init__()
+        self.orig = orig
+        self.epsilon = epsilon
+
+    def forward(self, x, r_alphas: bool = False):
+        with torch.no_grad():
+            B, S, H = x.size()
+            v = torch.tanh(
+                torch.matmul(x.reshape(B * S, H), self.orig.W_omega)
+                + self.orig.b_omega.reshape(1, -1)
+            )
+            vu = torch.matmul(v, self.orig.u_omega.reshape(-1, 1))
+            exps = torch.exp(vu).reshape(-1, S)
+            alphas = (exps / exps.sum(1, keepdim=True)).reshape(B, S, 1)
+        out = cp_weighted_pool(x, alphas, self.epsilon)
+        if r_alphas:
+            return out, alphas.reshape(B, S)
+        return out

physioex/explain/lrp/recurrent.py

@@ -146,16 +146,17 @@ def _layer_dir(self, x, w_ih, w_hh, b_ih, b_hh, reverse: bool):
             # h' = o⊙tanh(c') ; source is tanh(c') → identity bwd to c'
             h = mul_signal_take(o, _st_act(c, torch.tanh(c)))
             outs[t] = h
-        return torch.stack(outs, dim=0)
+        return torch.stack(outs, dim=0), h, c  # (T,B,H), final (h, c)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: torch.Tensor):
+        """Returns ``(output, (h_n, c_n))`` like ``nn.LSTM``."""
         if self.batch_first:
             x = x.transpose(0, 1)  # (B, T, in) -> (T, B, in)
-        num_dir = 2 if self.bidirectional else 1
         layer_in = x
+        h_states, c_states = [], []
         for layer in range(self.num_layers):
             suff = f"_l{layer}"
-            fwd = self._layer_dir(
+            fwd, hf, cf = self._layer_dir(
                 layer_in,
                 self._p("weight_ih" + suff),
                 self._p("weight_hh" + suff),
@@ -164,7 +165,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
                 reverse=False,
             )
             if self.bidirectional:
-                bwd = self._layer_dir(
+                bwd, hb, cb = self._layer_dir(
                     layer_in,
                     self._p("weight_ih" + suff + "_reverse"),
                     self._p("weight_hh" + suff + "_reverse"),
@@ -173,12 +174,16 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
                     reverse=True,
                 )
                 layer_in = torch.cat([fwd, bwd], dim=-1)  # (T, B, 2H)
+                h_states += [hf, hb]
+                c_states += [cf, cb]
             else:
                 layer_in = fwd
+                h_states.append(hf)
+                c_states.append(cf)
         out = layer_in
         if self.batch_first:
             out = out.transpose(0, 1)  # (T, B, dir*H) -> (B, T, dir*H)
-        return out
+        return out, (torch.stack(h_states, 0), torch.stack(c_states, 0))
 
 
 # ---------------------------------------------------------------------------
@@ -256,15 +261,17 @@ def _layer_dir(self, x, w_ih, w_hh, b_ih, b_hh, reverse: bool):
             one_minus_z = 1.0 - z
             h = add2(mul_signal_take(one_minus_z, n), mul_signal_take(z, h))
             outs[t] = h
-        return torch.stack(outs, dim=0)
+        return torch.stack(outs, dim=0), h  # (T,B,H), final h
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: torch.Tensor):
+        """Returns ``(output, h_n)`` like ``nn.GRU``."""
         if self.batch_first:
             x = x.transpose(0, 1)
         layer_in = x
+        h_states = []
         for layer in range(self.num_layers):
             suff = f"_l{layer}"
-            fwd = self._layer_dir(
+            fwd, hf = self._layer_dir(
                 layer_in,
                 self._p("weight_ih" + suff),
                 self._p("weight_hh" + suff),
@@ -273,7 +280,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
                 reverse=False,
             )
             if self.bidirectional:
-                bwd = self._layer_dir(
+                bwd, hb = self._layer_dir(
                     layer_in,
                     self._p("weight_ih" + suff + "_reverse"),
                     self._p("weight_hh" + suff + "_reverse"),
@@ -282,9 +289,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
                     reverse=True,
                 )
                 layer_in = torch.cat([fwd, bwd], dim=-1)
+                h_states += [hf, hb]
             else:
                 layer_in = fwd
+                h_states.append(hf)
         out = layer_in
         if self.batch_first:
             out = out.transpose(0, 1)
-        return out
+        return out, torch.stack(h_states, 0)