Add compatibility checks for LRP rule & layer combinations (#75)

Changes to functionality:
* Add `check_compat` mechanism
* Modify all layers that have weights and biases
* Allow `ZBoxRule` on all layers, but throw layer compatibility error
* Remove named LRP constructors

Changes to tests, docs and benchmarks:
* Update rule tests to check for compat errors 
* Remove `TestWrapper` tests and benchmarks
* Remove refs for rule & layer combinations deprecated by `check_compat`
* Update `LRPCustom` preset in docs and VGG tests

Author: adrhill

benchmark/benchmarks.jl

@@ -25,7 +25,7 @@ end
 algs = Dict(
     "Gradient" => Gradient,
     "InputTimesGradient" => InputTimesGradient,
-    "LRPZero" => LRPZero,
+    "LRPZero" => LRP,
     "LRPCustom" => LRPCustom, #modifies weights
     "SmoothGrad" => model -> SmoothGrad(model, 10),
     "IntegratedGradients" => model -> IntegratedGradients(model, 10),
@@ -46,17 +46,6 @@ for (name, alg) in algs
     SUITE["VGG"][name]["analyze"] = @benchmarkable analyze($(img), $(analyzer))
 end
 
-# Rules benchmarks – use wrapper to trigger AD fallback
-struct TestWrapper{T}
-    layer::T
-end
-(w::TestWrapper)(x) = w.layer(x)
-modify_layer!(rule::R, w::TestWrapper) where {R} = modify_layer!(rule, w.layer)
-get_layer_resetter(rule::R, w::TestWrapper) where {R} = get_layer_resetter(rule, w.layer)
-get_layer_resetter(::ZeroRule, w::TestWrapper) = Returns(nothing)
-get_layer_resetter(::EpsilonRule, w::TestWrapper) = Returns(nothing)
-lrp!(Rₖ, rule::ZBoxRule, w::TestWrapper, aₖ, Rₖ₊₁) = lrp!(Rₖ, rule, w.layer, aₖ, Rₖ₊₁)
-
 # generate input for conv layers
 insize = (64, 64, 3, 1)
 in_dense = 500
@@ -67,8 +56,6 @@ layers = Dict(
     "MaxPool" => (MaxPool((3, 3); pad=0), aₖ),
     "Conv" => (Conv((3, 3), 3 => 2), aₖ),
     "Dense" => (Dense(in_dense, out_dense, relu), randn(T, in_dense, 1)),
-    "WrappedDense" =>
-        (TestWrapper(Dense(in_dense, out_dense, relu)), randn(T, in_dense, 1)),
 )
 rules = Dict(
     "ZeroRule" => ZeroRule(),

docs/literate/advanced_lrp.jl

@@ -24,13 +24,13 @@ input = reshape(x, 28, 28, 1, :);
 # For this purpose, we create an array of rules that matches the length of the Flux chain:
 rules = [
     ZBoxRule(0.0f0, 1.0f0),
-    GammaRule(),
-    GammaRule(),
-    EpsilonRule(),
     EpsilonRule(),
+    GammaRule(),
     EpsilonRule(),
     ZeroRule(),
     ZeroRule(),
+    ZeroRule(),
+    ZeroRule(),
 ]
 
 analyzer = LRP(model, rules)
@@ -60,18 +60,27 @@ function modify_param!(::MyGammaRule, param)
 end
 
 # We can directly use this rule to make an analyzer!
-analyzer = LRP(model, MyGammaRule())
+rules = [
+    ZBoxRule(0.0f0, 1.0f0),
+    EpsilonRule(),
+    MyGammaRule(),
+    EpsilonRule(),
+    ZeroRule(),
+    ZeroRule(),
+    ZeroRule(),
+    ZeroRule(),
+]
+analyzer = LRP(model, rules)
 heatmap(input, analyzer)
 
-# We just implemented our own version of the ``γ``-rule in 4 lines of code!
-# The outputs match perfectly:
-analyzer = LRP(model, GammaRule())
-heatmap(input, analyzer)
+# We just implemented our own version of the ``γ``-rule in 4 lines of code.
+# The heatmap perfectly matches the previous one!
 
 # If the layer doesn't use weights `layer.weight` and biases `layer.bias`,
 # ExplainableAI provides a lower-level variant of [`modify_param!`](@ref)
 # called [`modify_layer!`](@ref). This function is expected to take a layer
 # and return a new, modified layer.
+# To add compatibility checks between rule and layer types, extend [`check_compat`](@ref).
 
 #md # !!! warning "Using modify_layer!"
 #md #
@@ -98,7 +107,7 @@ mylayer([1, 2, 3])
 # Let's append this layer to our model:
 model = Chain(model..., MyDoublingLayer())
 
-# Creating an LRP analyzer, e.g. `LRPZero(model)`, will throw an `ArgumentError`
+# Creating an LRP analyzer, e.g. `LRP(model)`, will throw an `ArgumentError`
 # and print a summary of the model check in the REPL:
 # ```julia-repl
 # ┌───┬───────────────────────┬─────────────────┬────────────┬────────────────┐
@@ -144,7 +153,7 @@ model = Chain(model..., MyDoublingLayer())
 LRP_CONFIG.supports_layer(::MyDoublingLayer) = true
 
 # Now we can create and run an analyzer without getting an error:
-analyzer = LRPZero(model)
+analyzer = LRP(model)
 heatmap(input, analyzer)
 
 #md # !!! note "Registering functions"
@@ -163,7 +172,7 @@ model = Chain(Flux.flatten, Dense(784, 100, myrelu), Dense(100, 10))
 # Once again, creating an LRP analyzer for this model will throw an `ArgumentError`
 # and display the following model check summary:
 # ```julia-repl
-# julia> analyzer = LRPZero(model3)
+# julia> analyzer = LRP(model3)
 # ┌───┬─────────────────────────┬─────────────────┬────────────┬────────────────┐
 # │   │ Layer                   │ Layer supported │ Activation │ Act. supported │
 # ├───┼─────────────────────────┼─────────────────┼────────────┼────────────────┤
@@ -187,7 +196,7 @@ model = Chain(Flux.flatten, Dense(784, 100, myrelu), Dense(100, 10))
 LRP_CONFIG.supports_activation(::typeof(myrelu)) = true
 
 # now the analyzer can be created without error:
-analyzer = LRPZero(model)
+analyzer = LRP(model)
 
 # ## How it works internally
 # Internally, ExplainableAI dispatches to low level functions
@@ -248,6 +257,7 @@ analyzer = LRPZero(model)
 # compute ``c`` from the previous equation as a VJP, pulling back ``s_{k}=R_{k}/z_{k}``:
 # ```julia
 # function lrp!(Rₖ, rule, layer, aₖ, Rₖ₊₁)
+#    check_compat(rule, layer)
 #    reset! = get_layer_resetter(layer)
 #    modify_layer!(rule, layer)
 #    ãₖ₊₁, pullback = Zygote.pullback(layer, modify_input(rule, aₖ))
@@ -256,8 +266,8 @@ analyzer = LRPZero(model)
 # end
 # ```
 #
-# You can see how `modify_layer!`, `modify_input` and `modify_denominator` dispatch on the
-# rule and layer type. This is how we implemented our own `MyGammaRule`.
+# You can see how `check_compat`, `modify_layer!`, `modify_input` and `modify_denominator`
+# dispatch on the rule and layer type. This is how we implemented our own `MyGammaRule`.
 # Unknown layers that are registered in the `LRP_CONFIG` use this exact function.
 
 # ### Specialized implementations

docs/src/api.md

@@ -36,6 +36,7 @@ modify_input
 modify_denominator
 modify_param!
 modify_layer!
+check_compat
 LRP_CONFIG.supports_layer
 LRP_CONFIG.supports_activation
 ```

src/ExplainableAI.jl

@@ -37,7 +37,7 @@ export AbstractXAIMethod
 export Gradient, InputTimesGradient
 export NoiseAugmentation, SmoothGrad
 export InterpolationAugmentation, IntegratedGradients
-export LRP, LRPZero, LRPEpsilon, LRPGamma
+export LRP
 
 # LRP rules
 export AbstractLRPRule

src/flux.jl

@@ -65,3 +65,11 @@ function strip_softmax(l::Conv)
 end
 
 has_weight_and_bias(layer) = hasproperty(layer, :weight) && hasproperty(layer, :bias)
+function require_weight_and_bias(rule, layer)
+    !has_weight_and_bias(layer) && throw(
+        ArgumentError(
+            "$rule requires linear layer with weight and bias parameters, got $layer."
+        ),
+    )
+    return nothing
+end

src/lrp.jl

@@ -41,11 +41,8 @@ function LRP(model::Chain, r::AbstractLRPRule; kwargs...)
     rules = repeat([r], length(model.layers))
     return LRP(model, rules; kwargs...)
 end
-# Additional constructors for convenience:
+# Additional constructors for convenience: use ZeroRule everywhere
 LRP(model::Chain; kwargs...) = LRP(model, ZeroRule(); kwargs...)
-LRPZero(model::Chain; kwargs...) = LRP(model, ZeroRule(); kwargs...)
-LRPEpsilon(model::Chain; kwargs...) = LRP(model, EpsilonRule(); kwargs...)
-LRPGamma(model::Chain; kwargs...) = LRP(model, GammaRule(); kwargs...)
 
 # The call to the LRP analyzer.
 function (analyzer::LRP)(

src/lrp_rules.jl

@@ -1,12 +1,10 @@
 # https://adrhill.github.io/ExplainableAI.jl/stable/generated/advanced_lrp/#How-it-works-internally
 abstract type AbstractLRPRule end
 
-# TODO: support all linear layers that use properties `weight` and `bias`
-const WeightBiasLayers = (Dense, Conv)
-
 # Generic LRP rule. Since it uses autodiff, it is used as a fallback for layer types
 # without custom implementations.
 function lrp!(Rₖ, rule::R, layer::L, aₖ, Rₖ₊₁) where {R<:AbstractLRPRule,L}
+    check_compat(rule, layer)
     reset! = get_layer_resetter(rule, layer)
     modify_layer!(rule, layer)
     ãₖ₊₁, pullback = Zygote.pullback(layer, modify_input(rule, aₖ))
@@ -18,6 +16,7 @@ end
 # To implement new rules, define the following custom functions:
 #   * `modify_input(rule, input)`
 #   * `modify_denominator(rule, d)`
+#   * `check_compat(rule, layer)`
 #   * `modify_param!(rule, param)` or `modify_layer!(rule, layer)`,
 #     the latter overriding the former
 #
@@ -36,6 +35,17 @@ Modify denominator ``z`` for numerical stability on the forward pass.
 """
 @inline modify_denominator(rule, d) = stabilize_denom(d, 1.0f-9) # general fallback
 
+"""
+    check_compat(rule, layer)
+
+Check compatibility of a LRP-Rule with layer type.
+
+## Note
+When implementing a custom `check_compat` function, return `nothing` if checks passed,
+otherwise throw an `ArgumentError`.
+"""
+@inline check_compat(rule, layer) = require_weight_and_bias(rule, layer)
+
 """
     modify_layer!(rule, layer)
 
@@ -45,15 +55,12 @@ propagation.
 ## Note
 When implementing a custom `modify_layer!` function, `modify_param!` will not be called.
 """
-modify_layer!(rule, layer) = nothing
-for L in WeightBiasLayers
-    @eval function modify_layer!(rule::R, layer::$L) where {R}
-        if has_weight_and_bias(layer)
-            modify_param!(rule, layer.weight)
-            modify_bias!(rule, layer.bias)
-        end
-        return nothing
+function modify_layer!(rule::R, layer::L) where {R,L}
+    if has_weight_and_bias(layer)
+        modify_param!(rule, layer.weight)
+        modify_bias!(rule, layer.bias)
     end
+    return nothing
 end
 
 """
@@ -97,6 +104,22 @@ end
 Constructor for LRP-0 rule. Commonly used on upper layers.
 """
 struct ZeroRule <: AbstractLRPRule end
+@inline check_compat(::ZeroRule, layer) = nothing
+
+"""
+    EpsilonRule([ϵ=1.0f-6])
+
+Constructor for LRP-``ϵ`` rule. Commonly used on middle layers.
+
+Arguments:
+- `ϵ`: Optional stabilization parameter, defaults to `1f-6`.
+"""
+struct EpsilonRule{T} <: AbstractLRPRule
+    ϵ::T
+    EpsilonRule(ϵ=1.0f-6) = new{Float32}(ϵ)
+end
+modify_denominator(r::EpsilonRule, d) = stabilize_denom(d, r.ϵ)
+@inline check_compat(::EpsilonRule, layer) = nothing
 
 """
     GammaRule([γ=0.25])
@@ -115,20 +138,7 @@ function modify_param!(r::GammaRule, param::AbstractArray{T}) where {T}
     param .+= γ * relu.(param)
     return nothing
 end
-
-"""
-    EpsilonRule([ϵ=1.0f-6])
-
-Constructor for LRP-``ϵ`` rule. Commonly used on middle layers.
-
-Arguments:
-- `ϵ`: Optional stabilization parameter, defaults to `1f-6`.
-"""
-struct EpsilonRule{T} <: AbstractLRPRule
-    ϵ::T
-    EpsilonRule(ϵ=1.0f-6) = new{Float32}(ϵ)
-end
-modify_denominator(r::EpsilonRule, d) = stabilize_denom(d, r.ϵ)
+@inline check_compat(rule::GammaRule, layer) = require_weight_and_bias(rule, layer)
 
 """
     ZBoxRule(low, high)
@@ -146,45 +156,44 @@ struct ZBoxRule{T} <: AbstractLRPRule
 end
 
 # The ZBoxRule requires its own implementation of relevance propagation.
-for L in WeightBiasLayers
-    function lrp!(Rₖ, rule::ZBoxRule, layer, aₖ, Rₖ₊₁)
-        T = eltype(aₖ)
-        l = zbox_input_augmentation(T, rule.low, size(aₖ))
-        h = zbox_input_augmentation(T, rule.high, size(aₖ))
-        reset! = get_layer_resetter(rule, layer)
+function lrp!(Rₖ, rule::ZBoxRule, layer::L, aₖ, Rₖ₊₁) where {L}
+    require_weight_and_bias(rule, layer)
+    reset! = get_layer_resetter(rule, layer)
 
-        # Compute pullback for W, b
-        aₖ₊₁, pullback = Zygote.pullback(layer, aₖ)
+    l = zbox_input(aₖ, rule.low)
+    h = zbox_input(aₖ, rule.high)
 
-        # Compute pullback for W⁺, b⁺
-        modify_layer!(Val{:mask_positive}, layer)
-        aₖ₊₁⁺, pullback⁺ = Zygote.pullback(layer, l)
-        reset!()
+    # Compute pullback for W, b
+    aₖ₊₁, pullback = Zygote.pullback(layer, aₖ)
 
-        # Compute pullback for W⁻, b⁻
-        modify_layer!(Val{:mask_negative}, layer)
-        aₖ₊₁⁻, pullback⁻ = Zygote.pullback(layer, h)
-        reset!()
+    # Compute pullback for W⁺, b⁺
+    modify_layer!(Val{:mask_positive}, layer)
+    aₖ₊₁⁺, pullback⁺ = Zygote.pullback(layer, l)
+    reset!()
 
-        y = Rₖ₊₁ ./ modify_denominator(rule, aₖ₊₁ - aₖ₊₁⁺ - aₖ₊₁⁻)
-        Rₖ .= aₖ .* only(pullback(y)) - l .* only(pullback⁺(y)) - h .* only(pullback⁻(y))
-        return nothing
-    end
+    # Compute pullback for W⁻, b⁻
+    modify_layer!(Val{:mask_negative}, layer)
+    aₖ₊₁⁻, pullback⁻ = Zygote.pullback(layer, h)
+    reset!()
+
+    y = Rₖ₊₁ ./ modify_denominator(rule, aₖ₊₁ - aₖ₊₁⁺ - aₖ₊₁⁻)
+    Rₖ .= aₖ .* only(pullback(y)) - l .* only(pullback⁺(y)) - h .* only(pullback⁻(y))
+    return nothing
 end
 
-const ZBOX_BOUNDS_MISMATCH = "ZBoxRule bounds should either be scalar or match input size."
-function zbox_input_augmentation(T, A::AbstractArray, in_size)
-    size(A) != in_size && throw(ArgumentError(ZBOX_BOUNDS_MISMATCH))
+zbox_input(in::AbstractArray{T}, c::Real) where {T} = fill(convert(T, c), size(in))
+function zbox_input(in::AbstractArray{T}, A::AbstractArray) where {T}
+    @assert size(A) == size(in)
     return convert.(T, A)
 end
-zbox_input_augmentation(T, c::Real, in_size) = fill(convert(T, c), in_size)
 
-# Other special cases that are dispatched on layer type:
-const LRPRules = (ZeroRule, EpsilonRule, GammaRule, ZBoxRule)
-for R in LRPRules
+# Special cases for rules that don't modify params for extra performance:
+for R in (ZeroRule, EpsilonRule)
+    @eval get_layer_resetter(::$R, l) = Returns(nothing)
     @eval lrp!(Rₖ, ::$R, ::DropoutLayer, aₖ, Rₖ₊₁) = (Rₖ .= Rₖ₊₁)
     @eval lrp!(Rₖ, ::$R, ::ReshapingLayer, aₖ, Rₖ₊₁) = (Rₖ .= reshape(Rₖ₊₁, size(aₖ)))
 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ₖ₊₁)
@@ -196,7 +205,3 @@ for R in (ZeroRule, EpsilonRule, GammaRule)
         return nothing
     end
 end
-
-# Rules that don't modify params can optionally be added here for extra performance
-get_layer_resetter(::ZeroRule, l) = Returns(nothing)
-get_layer_resetter(::EpsilonRule, l) = Returns(nothing)