physioex.explain.lrp

Layer-wise Relevance Propagation (LRP) for PhysioEx.

Conservation-based attribution (Bach et al. 2015). Two entry points:

• LRP — Zennit composites for feed-forward / CNN models (ε, γ, w², z-box, flat; BatchNorm canonization).

• ModelLRP — whole-model LRP for architectures with LSTM/GRU (Arras signal-take), attention / transformers (CP-LRP, Ali et al. 2022) and PhysioEx’s softmax poolings; leaves get ε-LRP, norms/activations the identity rule, residual + the proportional rule.

Both seed the target logit so Σ R ≈ f_c(x); use return_report=True / check_conservation() to see how much biases absorb, and audit_lrp_coverage() to list blocks left on plain autograd.

Requires the optional explain extra:

pip install "physioex[explain]"

class physioex.explain.lrp.ConservationReport(target, relevance_sum, ratio, absorbed)

Bases: object

Per-sample conservation summary.

Parameters:

• target (Tensor)

• relevance_sum (Tensor)

• ratio (Tensor)

• absorbed (Tensor)

target

the explained logit f_c(x) per sample, (B,).

Type:

torch.Tensor

relevance_sum

Σ R over all input elements per sample, (B,).

Type:

torch.Tensor

ratio

Σ R / f (1.0 = exact conservation).

Type:

torch.Tensor

absorbed

f − Σ R — relevance absorbed by biases / unruled ops.

Type:

torch.Tensor

class physioex.explain.lrp.EpsilonRule(module, epsilon=1e-06)

Bases: _RuleWrapper

ε-LRP wrapper for a single-input module linear in its input.

Parameters:

• module (nn.Module)

• epsilon (float)

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class physioex.explain.lrp.IdentityRule(module)

Bases: _RuleWrapper

Identity-rule wrapper: forward is the module’s, backward passes relevance.

Parameters:

module (nn.Module)

forward(x, *args, **kwargs)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class physioex.explain.lrp.LRP(model, in_index=0, out_index=0, composite=None, canonizers=None)

Bases: Module

LRP attribution over a trained feed-forward PhysioEx model (Zennit).

Parameters:

• model – trained model, (B, L, C, T) -> (B, L, n_classes) or (B, n_classes). Evaluated in eval() mode.

• in_index – sequence epoch to explain (central-epoch models emit L=1).

• out_index – class-logit index to explain.

• composite – a Zennit composite; default physioex_composite() (ε dense / γ conv / w² first layer) with the BatchNorm canonizer.

• canonizers – used only when composite is None.

Example

>>> relevance = LRP(model, out_index=2)(x)
>>> relevance, report = LRP(model, out_index=2)(x, return_report=True)

forward(x, return_report=False)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

• x (Tensor)

• return_report (bool)

class physioex.explain.lrp.LRPAttentionLayer(orig, epsilon=1e-06)

Bases: _PoolAdapter

CP-LRP for seqsleepnet.AttentionLayer (additive pooling, manual softmax).

Parameters:

• orig (nn.Module)

• epsilon (float)

forward(x, r_alphas=False)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

r_alphas (bool)

class physioex.explain.lrp.LRPAttentionPooling(orig, epsilon=1e-06)

Bases: _PoolAdapter

CP-LRP for sleeptransformer.AttentionPooling (self.attention MLP, softmax over dim=1, (B,T,D) -> (B,D)).

Parameters:

• orig (nn.Module)

• epsilon (float)

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class physioex.explain.lrp.LRPChannelMixer(orig, epsilon=1e-06)

Bases: _PoolAdapter

CP-LRP for protosleepnet.ChannelMixer: constant modality embedding → per-channel mcy logits → dropout (eval) → residual x + mixer(x) (proportional split; the mixer itself is swapped by prepare) → softmax channel pooling (CP-LRP). forward(x, zero_emb) -> (h, mcy_logits).

Parameters:

• orig (nn.Module)

• epsilon (float)

forward(x, zero_emb)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class physioex.explain.lrp.LRPGRU(input_size: int, hidden_size: int, num_layers: int = 1, bias: bool = True, batch_first: bool = False, dropout: float = 0.0, bidirectional: bool = False, device=None, dtype=None)

class physioex.explain.lrp.LRPGRU(*args, **kwargs)

Bases: GRU

Cell-level GRU carrying LRP relevance; drop-in for a trained nn.GRU.

forward(x) returns (output, h_n) like nn.GRU. PyTorch gate order is [r, z, n] with n = tanh(W_in x + b_in + r ⊙ (W_hn h + b_hn)) and h' = (1−z)⊙n + z⊙h; sources are n, h and the recurrent pre-activation gated by r.

forward(x, hx=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class physioex.explain.lrp.LRPLSTM(input_size: int, hidden_size: int, num_layers: int = 1, bias: bool = True, batch_first: bool = False, dropout: float = 0.0, bidirectional: bool = False, proj_size: int = 0, device=None, dtype=None)

class physioex.explain.lrp.LRPLSTM(*args, **kwargs)

Bases: LSTM

Cell-level LSTM carrying LRP relevance; drop-in for a trained nn.LSTM.

Build with from_torch(). forward(x) returns (output, (h_n, c_n)) like nn.LSTM; an explicit initial state hx is not supported (raises), dropout is ignored (eval-time attribution).

forward(x, hx=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class physioex.explain.lrp.LRPMultiheadAttention(embed_dim, num_heads, epsilon=1e-06)

Bases: Module

LRP-instrumented multi-head attention (self or cross), value-path CP-LRP.

Reproduces nn.MultiheadAttention (batch_first=True, same q/k/v embed dims, need_weights=False). forward(query, key, value) returns the attention output tensor.

Parameters:

• embed_dim (int)

• num_heads (int)

• epsilon (float)

attention_weights(query, key)

Averaged-over-heads attention matrix (B, Tq, Tk) (no grad).

forward(query, key, value)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class physioex.explain.lrp.LRPMultiheadAttentionModule(attn)

Bases: Module

nn.MultiheadAttention-compatible adapter for standalone attention.

Host models call out, w = mha(query=, key=, value=, ...). Returns (out, weights_or_None); raises on attention masks (unsupported) instead of silently ignoring them.

Parameters:

attn (LRPMultiheadAttention)

forward(query, key, value, key_padding_mask=None, need_weights=True, attn_mask=None, average_attn_weights=True, is_causal=False)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class physioex.explain.lrp.LRPTransformerEncoder(layers, norm=None)

Bases: Module

LRP replacement for nn.TransformerEncoder (the container’s own forward introspects layers[0].self_attn.batch_first, so the whole stack is replaced). Applies the LRP layers in turn, then the optional final norm.

forward(x, mask=None, src_key_padding_mask=None, is_causal=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class physioex.explain.lrp.LRPTransformerEncoderLayer(epsilon=1e-06)

Bases: Module

LRP-instrumented nn.TransformerEncoderLayer (post- and pre-norm).

Build with from_torch(). Self-attention only; src_mask / src_key_padding_mask are not supported (raise).

Parameters:

epsilon (float)

forward(x, src_mask=None, src_key_padding_mask=None, is_causal=False)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class physioex.explain.lrp.ModelLRP(model, in_index=0, out_index=0, output_key=None, epsilon=1e-06, patch_residuals=True, strict=False, copy_model=True)

Bases: Module

LRP attribution for a whole PhysioEx model.

Parameters:

• model (nn.Module) – trained model (B, L, C, T[, F]) -> (B, L, n_classes) (or a dict — set output_key). Deep-copied; the original is untouched.

• out_index (int) – sequence epoch and class logit to explain.

• output_key (Optional[str]) – for dict outputs (CoReSleep → "combined").

• epsilon (float) – ε of every ε-rule in the prepared model.

• patch_residuals (bool) – redirect plain + between grad-carrying tensors in the model’s own forward to the proportional rule (recommended).

• strict (bool) – raise if a parametric leaf is left without an LRP rule.

• copy_model (bool) – set False to prepare model in place (saves memory).

• in_index (int)

• out_index

forward(x) returns relevance shaped like x (seeded with the target logit so Σ R ≈ f_c(x), minus what biases absorb); forward(x, return_report=True) also returns a ConservationReport.

forward(x, return_report=False)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

• x (Tensor)

• return_report (bool)

physioex.explain.lrp.audit_lrp_coverage(model, include_containers=False)

Names of parametric leaf modules that carry no LRP rule (their forward runs on plain autograd and propagates gradient, not relevance).

Container modules that own parameters and have children (e.g. a learned CLS token or positional embedding added to the stream) are additive constants w.r.t. the input and are only listed with include_containers=True for manual inspection.

Parameters:

• model (Module)

• include_containers (bool)

Return type:

List[str]

physioex.explain.lrp.check_conservation(explainer, x)

Run explainer (an ModelLRP or LRP) on x and report.

Parameters:

x (Tensor)

Return type:

ConservationReport

physioex.explain.lrp.cp_weighted_pool(x, weights, epsilon=1e-06)

Weighted sum over dim=1 with the CP-LRP backward; weights is (B, T, 1) and must be detached (constant).

physioex.explain.lrp.default_canonizers(model=None)

Return the default canonizer list for a PhysioEx model.

Currently merges sequential BatchNorm layers into their neighbouring linear / convolution modules — needed by e.g. tinysleepnet, sleepfm and tfc. Models without BatchNorm are unaffected (the canonizer simply finds nothing to merge).

Parameters:

model – unused for now; accepted so callers can pass the model and we can later dispatch model-specific canonizers (e.g. GroupNorm in neurolm or the protosleepnet prototype head).

Return type:

List

physioex.explain.lrp.epsilon_composite(epsilon=1e-06, canonizers=None)

A pure ε-LRP composite (ε on every linear/conv layer).

Unlike physioex_composite(), this uses no special first-layer rule, so it approximately conserves relevance (Σ R ≈ f(x)) — useful as the reference for the conservation sanity check in the test suite.

Parameters:

• epsilon (float)

• canonizers (List | None)

physioex.explain.lrp.physioex_composite(first_rule='w2', epsilon=1e-06, gamma=0.25, zbox_low=-3.0, zbox_high=3.0, canonizers=None)

Build the default PhysioEx LRP composite (ε dense / γ conv / w² first).

Parameters:

• first_rule (str) – input-layer rule — "w2" (default, unbounded signals), "zbox" (bounded inputs; needs zbox_low/zbox_high) or "flat" (uniform baseline).

• epsilon (float) – stabiliser for the ε-rule on dense layers.

• gamma (float) – positive-weight amplification for the γ-rule on conv layers.

• zbox_low (float) – bounds for the zbox first-layer rule.

• zbox_high (float) – bounds for the zbox first-layer rule.

• canonizers (List | None) – list of Zennit canonizers to apply (e.g. from physioex.explain.lrp.canonizers.default_canonizers()).

Returns:

A zennit.composites.SpecialFirstLayerMapComposite.

physioex.explain.lrp.prepare_model_for_lrp(model, epsilon=1e-06, strict=False, _memo=None)

In-place: swap fused / attention / registered blocks for their LRP counterparts and wrap the remaining leaves with LRP rules. Returns model. Call on a copy (see ModelLRP).

strict=True raises if any parametric leaf is left on plain autograd. Shared modules (the same object under several attributes) are replaced by a single shared LRP module.

Parameters:

• model (Module)

• epsilon (float)

• strict (bool)

• _memo (Dict[int, Module] | None)

Return type:

Module

physioex.explain.lrp.register_lrp_adapter(cls, factory)

Register factory(module, epsilon) -> nn.Module for instances of cls.

Use for custom blocks whose forward contains softmax pooling / attention or other non-module ops that plain autograd would attribute as gradient.

Parameters:

• cls (type)

• factory (Callable[[Module, float], Module])

Return type:

None

physioex.explain.lrp.st_identity(pre, act)

Straight-through nonlinearity (identity rule): value act, relevance → pre.

physioex.explain.lrp.swap_transformer_layers(module, epsilon=1e-06)

In-place: replace nn.TransformerEncoder / nn.TransformerEncoderLayer children with their LRP counterparts (weights copied; the source leaves are not modified). Returns module. Prefer prepare_model_for_lrp().

Parameters:

• module (Module)

• epsilon (float)

Return type:

Module

physioex.explain.lrp.target_seed(out, in_index, out_index)

Relevance seed selecting the target neuron with its logit value.

Seeding with f_c(x) (not a bare 1.0) makes total relevance conserve to the class evidence, Σ R ≈ f_c(x). Handles (B, L, n_classes) sequence outputs and (B, n_classes). Returns (seed, target_values).

Parameters:

• out (Tensor)

• in_index (int)

• out_index (int)