autoLRP

Layer-wise relevance propagation on the PyTorch autograd graph. No model rewriting, no module names: wrap the input, run the model as it is, pick the output scalar, call .lrp().

import torch.nn as nn
import autoLRP
from autoLRP import LRPConfig, BASE

x = autoLRP.tensor(image)          # the input you want relevance for
out = model(x)
out[0, pred].lrp()                 # relevance of class `pred`
heatmap = x.relevance              # same shape as `image`

How it works

autoLRP.tensor returns a tensor subclass. While the model runs, a few ops are replaced by our own (add, sub, mean, sum, cumsum, softmax, fused attention) so that the backward graph keeps the values the rules need; their gradients are the native ones, so the graph is otherwise unchanged. .lrp() walks the autograd graph, attaches a label ("fact") to some nodes (for example, which operand of a product is the softmax weights), installs one hook per node that turns the arriving gradient into relevance, and runs backward. Whatever reaches the wrapped input is its relevance.

Configuration

Every rule-bearing node is addressed by its autograd name without the version digit, or by a fact an analyzer attached to it. BASE is the starting table:

>>> print(BASE)
{'AddmmBackward': 'epsilon', 'MmBackward': 'epsilon', 'ConvolutionBackward': 'epsilon',
 'BmmBackward': 'epsilon', 'MulBackward': 'proportional', 'DivBackward': 'proportional',
 'AddBackward': 'proportional', 'SubBackward': 'proportional',
 'statistic_operand': ('detach', {'by': 'statistic_operand'})}

Override entries on it, or use a preset:

LRPConfig(rule={**BASE, 'AddmmBackward': 'zplus'})
LRPConfig(rule={**BASE, 'ConvolutionBackward': ('gamma', {'gamma': 0.25})})
LRPConfig.composite()               # z+ on conv, epsilon elsewhere
LRPConfig(attn='attnlrp')           # epsilon products, Jacobian softmax
LRPConfig(attn='cplrp')             # attention weights treated as constants
LRPConfig(attn='uniform')

The config says exactly what runs. A key that is not a node name or a registered fact, a rule the key's family cannot run, and a node that no entry addresses are errors:

LRPConfig(rule={**BASE, 'linear': 'zplus'})
  ValueError: unknown rule key 'linear': not a node name [...]
LRPConfig(rule={**BASE, 'MulBackward': 'zbox'})
  ValueError: rule entry 'MulBackward'='zbox': 'zbox' is not a choice here. Choices: [...]

Rule tables, by family:

Names are positional: detach_lhs zeros the operand written on the left of that op, always. The one virtual name 'detach' takes by=<fact> and picks the side per node from the fact's value.

Which family a product node belongs to is decided by which of its operands come from the wrapped input, not by its name. One operand: the op is a linear layer whose weight is the other operand (a constant, a parameter, or any tensor you did not wrap). Two: the op is bilinear. Frozen models (requires_grad=False) work like trainable ones.

Softmax, layer norm and activation nodes are set by their own fields: softmax= (passthrough, jacobian, detach), layernorm= (identity, passthrough, yx, detach_std), activation= (passthrough, yx).

Reading back what ran

explain installs the hooks, records the config entry and the rule function at every node, removes the hooks again, and runs no backward. explain_summary prints one line per distinct combination:

from autoLRP import explain, explain_summary
rows = explain(model(autoLRP.tensor(x))[0, pred], LRPConfig(attn='cplrp'))
print(explain_summary([r for r in rows if r[2] != 'native gradient']))

count  node                      key                   what
    4  AddmmBackward0            AddmmBackward         epsilon
    2  AddBackward               AddBackward           residual_proportional
    2  NativeLayerNormBackward0  None                  layernorm=identity
    1  BmmBackward0              weights_operand       detach_lhs_bmm
    1  BmmBackward0              BmmBackward           epsilon_bmm
    1  DivBackward0              None                  passthrough (constant operand)
    1  SoftmaxBackward           None                  softmax=passthrough

key is the entry that addressed the node, what the function it picked; native gradient rows are shape ops, where the gradient is already the routing.

Facts and your own analyzers

Built-in facts: statistic_operand (the normalization statistic in a mul, div or sub, detached by BASE), weights_operand (the softmax weights in a bmm), input_conv (the first convolution). An analyzer is a function over all nodes that returns the nodes carrying its fact; its registered name is the config key. The value is the fact's value: True for a plain tag, a slot number (0 or 1) for a side that a ('detach', {'by': ...}) entry can read:

from autoLRP import register_analyzer

@register_analyzer('first_linear')
def first_linear(nodes):
    hits = [n for n in nodes if 'AddmmBackward' in n.name()]
    return {hits[-1]: True} if hits else {}

LRPConfig(rule={**BASE, 'first_linear': ('zbox', {'low': -3.0, 'high': 3.0})})

Conventions worth knowing

• A constant that multiplies, divides or negates passes relevance through unchanged. A constant that is added is a bias, and a bias absorbs its share (apply_bias_split), so a layer with a bias emits less than it receives.
• Fused scaled_dot_product_attention is decomposed into matmul, softmax, matmul by default; set_decompose_attention(False) keeps the fused node, which is handled by its own installer and gives the same relevance to 1e-13.
• An op with no installer runs its native gradient and warns once, if it lies on the path to the wrapped input.
• BASE with the default eps=1e-11 is LRP-0; on deep networks it can be numerically unstable, and the recipes (composite, gamma, z+) are what to use there.

Recipes and evaluation

autoLRP.bilrp(model, x_a, x_b) (second-order, similarity models), autoLRP.clrp(...) (contrastive), and autoLRP.eval with perturbation_curve, aopc, sanity_check_cascade, sensitivity_correlation.

Tests

python -m pytest -q

352 tests; two modules skip without zennit and the examples file. The notebooks under examples/showcase are the showcases; the four in extras/ run without downloads.