Add PassRule (#76)

* Add `PassRule`
* Clean up rule doc strings and add references

Author: adrhill

README.md

@@ -57,9 +57,11 @@ Currently, the following analyzers are implemented:
 ├── SmoothGrad
 ├── IntegratedGradients
 └── LRP
-    ├── LRPZero
-    ├── LRPEpsilon
-    └── LRPGamma
+    ├── ZeroRule
+    ├── EpsilonRule
+    ├── GammaRule
+    ├── PassRule
+    └── ZBoxRule
 ```
 
 One of the design goals of ExplainableAI.jl is extensibility.

docs/src/api.md

@@ -24,8 +24,9 @@ InterpolationAugmentation
 ## Rules
 ```@docs
 ZeroRule
-GammaRule
 EpsilonRule
+GammaRule
+PassRule
 ZBoxRule
 ```
 

src/ExplainableAI.jl

@@ -42,7 +42,7 @@ export LRP
 # LRP rules
 export AbstractLRPRule
 export LRP_CONFIG
-export ZeroRule, EpsilonRule, GammaRule, ZBoxRule
+export ZeroRule, EpsilonRule, GammaRule, ZBoxRule, PassRule
 export modify_input, modify_denominator
 export modify_param!, modify_layer!
 export check_model

src/lrp_rules.jl

@@ -72,8 +72,8 @@ Inplace-modify parameters before computing the relevance.
 @inline modify_param!(rule, param) = nothing # general fallback
 
 # Useful presets:
-modify_param!(::Val{:mask_positive}, p) = (p .= max.(zero(eltype(p)), p), return nothing)
-modify_param!(::Val{:mask_negative}, p) = (p .= min.(zero(eltype(p)), p), return nothing)
+modify_param!(::Val{:mask_positive}, p) = p .= max.(zero(eltype(p)), p)
+modify_param!(::Val{:mask_negative}, p) = p .= min.(zero(eltype(p)), p)
 
 # Internal wrapper functions for bias-free layers.
 @inline modify_bias!(rule::R, b) where {R} = modify_param!(rule, b)
@@ -101,18 +101,29 @@ end
 """
     ZeroRule()
 
-Constructor for LRP-0 rule. Commonly used on upper layers.
+LRP-0 rule. Commonly used on upper layers.
+
+# References
+[1]: S. Bach et al., On Pixel-Wise Explanations for Non-Linear Classifier Decisions by
+    Layer-Wise Relevance Propagation
 """
 struct ZeroRule <: AbstractLRPRule end
 @inline check_compat(::ZeroRule, layer) = nothing
 
+# Optimization to save allocations since weights don't need to be reset:
+get_layer_resetter(::ZeroRule, layer) = Returns(nothing)
+
 """
     EpsilonRule([ϵ=1.0f-6])
 
-Constructor for LRP-``ϵ`` rule. Commonly used on middle layers.
+LRP-``ϵ`` rule. Commonly used on middle layers.
 
 Arguments:
 - `ϵ`: Optional stabilization parameter, defaults to `1f-6`.
+
+# References
+[1]: S. Bach et al., On Pixel-Wise Explanations for Non-Linear Classifier Decisions by
+    Layer-Wise Relevance Propagation
 """
 struct EpsilonRule{T} <: AbstractLRPRule
     ϵ::T
@@ -121,13 +132,19 @@ end
 modify_denominator(r::EpsilonRule, d) = stabilize_denom(d, r.ϵ)
 @inline check_compat(::EpsilonRule, layer) = nothing
 
+# Optimization to save allocations since weights don't need to be reset:
+get_layer_resetter(::EpsilonRule, layer) = Returns(nothing)
+
 """
     GammaRule([γ=0.25])
 
-Constructor for LRP-``γ`` rule. Commonly used on lower layers.
+LRP-``γ`` rule. Commonly used on lower layers.
 
 Arguments:
-- `γ`: Optional multiplier for added positive weights, defaults to 0.25.
+- `γ`: Optional multiplier for added positive weights, defaults to `0.25`.
+
+# References
+[1]: G. Montavon et al., Layer-Wise Relevance Propagation: An Overview
 """
 struct GammaRule{T} <: AbstractLRPRule
     γ::T
@@ -140,15 +157,34 @@ function modify_param!(r::GammaRule, param::AbstractArray{T}) where {T}
 end
 @inline check_compat(rule::GammaRule, layer) = require_weight_and_bias(rule, layer)
 
+"""
+    PassRule()
+
+Pass-through rule. Passes relevance through to the lower layer.
+Supports reshaping layers.
+"""
+struct PassRule <: AbstractLRPRule end
+function lrp!(Rₖ, ::PassRule, layer, aₖ, Rₖ₊₁)
+    if size(aₖ) == size(Rₖ₊₁)
+        Rₖ .= Rₖ₊₁
+    end
+    Rₖ .= reshape(Rₖ₊₁, size(aₖ))
+    return nothing
+end
+# No extra checks as reshaping operation will throw an error if layer isn't compatible:
+@inline check_compat(::PassRule, layer) = nothing
+
 """
     ZBoxRule(low, high)
 
-Constructor for LRP-``z^{\\mathcal{B}}``-rule.
-Commonly used on the first layer for pixel input.
+LRP-``z^{\\mathcal{B}}``-rule. Commonly used on the first layer for pixel input.
 
 The parameters `low` and `high` should be set to the lower and upper bounds of the input features,
 e.g. `0.0` and `1.0` for raw image data.
 It is also possible to provide two arrays of that match the input size.
+
+## References
+[1]: G. Montavon et al., Explaining nonlinear classification decisions with deep Taylor decomposition
 """
 struct ZBoxRule{T} <: AbstractLRPRule
     low::T
@@ -194,6 +230,13 @@ for R in (ZeroRule, EpsilonRule)
     @eval lrp!(Rₖ, ::$R, ::ReshapingLayer, aₖ, Rₖ₊₁) = (Rₖ .= reshape(Rₖ₊₁, size(aₖ)))
 end
 
+# Special cases for rules that don't modify params for extra performance:
+for R in (ZeroRule, EpsilonRule)
+    for L in (DropoutLayer, ReshapingLayer)
+        @eval lrp!(Rₖ, ::$R, l::$L, aₖ, Rₖ₊₁) = lrp!(Rₖ, PassRule(), l, aₖ, Rₖ₊₁)
+    end
+end
+
 # Fast implementation for Dense layer using Tullio.jl's einsum notation:
 for R in (ZeroRule, EpsilonRule, GammaRule)
     @eval function lrp!(Rₖ, rule::$R, layer::Dense, aₖ, Rₖ₊₁)

src/types.jl

@@ -6,7 +6,7 @@ const ConvLayer = Union{Conv} # TODO: DepthwiseConv, ConvTranspose, CrossCor
 const DropoutLayer = Union{Dropout,typeof(Flux.dropout),AlphaDropout}
 
 """Union type for reshaping layers such as `flatten`."""
-const ReshapingLayer = Union{typeof(Flux.flatten), typeof(Flux.MLUtils.flatten)}
+const ReshapingLayer = Union{typeof(Flux.flatten),typeof(Flux.MLUtils.flatten)}
 
 """Union type for max pooling layers."""
 const MaxPoolLayer = Union{MaxPool,AdaptiveMaxPool,GlobalMaxPool}