Refactor lrp! (#70)

* Make `Rₖ₊₁` first argument of `lrp!`, matching Bang convention

* Change rule keyword-arguments to default arguments

* Type stability fixes for `GammaRule`

* Add explanations for non-Julia programmers to docs

Author: adrhill

benchmark/benchmarks.jl

@@ -52,7 +52,7 @@ struct TestWrapper{T}
 end
 (w::TestWrapper)(x) = w.layer(x)
 modify_layer(r::AbstractLRPRule, w::TestWrapper) = modify_layer(r, w.layer)
-lrp!(rule::ZBoxRule, w::TestWrapper, Rₖ, aₖ, Rₖ₊₁) = lrp!(rule, w.layer, Rₖ, aₖ, Rₖ₊₁)
+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)
@@ -81,7 +81,7 @@ for (layername, (layer, aₖ)) in layers
     Rₖ₊₁ = layer(aₖ)
     for (rulename, rule) in rules
         SUITE["Layer"][layername][rulename] = @benchmarkable lrp!(
-            $(rule), $(layer), $(Rₖ), $(aₖ), $(Rₖ₊₁)
+            $(Rₖ), $(rule), $(layer), $(aₖ), $(Rₖ₊₁)
         )
     end
 end

docs/literate/advanced_lrp.jl

@@ -191,16 +191,20 @@ analyzer = LRPZero(model)
 # ## How it works internally
 # Internally, ExplainableAI dispatches to low level functions
 # ```julia
-# function lrp!(rule, layer, Rₖ, aₖ, Rₖ₊₁)
+# lrp!(Rₖ, rule, layer, aₖ, Rₖ₊₁)
 #     Rₖ .= ...
 # end
 # ```
-# These functions use the arguments `rule` and `layer` to dispatch
-# `modify_params` and `modify_denominator` on the rule and layer type.
-# They in-place modify a pre-allocated array of the input relevance `Rₖ`
-# based on the input activation `aₖ` and output relevance `Rₖ₊₁`.
+# These functions in-place modify a pre-allocated array of the input relevance `Rₖ`
+# (the `!` is a [naming convention](https://docs.julialang.org/en/v1/manual/style-guide/#bang-convention)
+# in Julia to denote functions that modify their arguments).
+
+# The correct rule is applied via [multiple dispatch](https://www.youtube.com/watch?v=kc9HwsxE1OY)
+# on the types of the arguments `rule` and `layer`.
+# The relevance `Rₖ` is then computed based on the input activation `aₖ` and the output relevance `Rₖ₊₁`.
+# Multiple dispatch is also used to dispatch `modify_params` and `modify_denominator` on the rule and layer type.
 #
-# Calling `analyze` then applies a forward-pass of the model, keeping track of
+# Calling `analyze` on a LRP-model applies a forward-pass of the model, keeping track of
 # the activations `aₖ` for each layer `k`.
 # The relevance `Rₖ₊₁` is then set to the output neuron activation and the rules are applied
 # in a backward-pass over the model layers and previous activations.
@@ -241,7 +245,7 @@ analyzer = LRPZero(model)
 # The default LRP fallback for unknown layers uses AD via [Zygote](https://github.com/FluxML/Zygote.jl).
 # For `lrp!`, we end up with something that looks very similar to the previous four step computation:
 # ```julia
-# function lrp!(rule, layer, Rₖ, aₖ, Rₖ₊₁)
+# function lrp!(Rₖ, rule, layer, aₖ, Rₖ₊₁)
 #     layerᵨ = modify_layer(rule, layer)
 #     c = gradient(aₖ) do a
 #             z = layerᵨ(a)
@@ -263,7 +267,7 @@ analyzer = LRPZero(model)
 # Reshaping layers don't affect attributions. We can therefore avoid the computational
 # overhead of AD by writing a specialized implementation that simply reshapes back:
 # ```julia
-# function lrp!(::AbstractLRPRule, ::ReshapingLayer, Rₖ, aₖ, Rₖ₊₁)
+# function lrp!(Rₖ, ::AbstractLRPRule, ::ReshapingLayer, aₖ, Rₖ₊₁)
 #     Rₖ .= reshape(Rₖ₊₁, size(aₖ))
 # end
 # ```
@@ -272,7 +276,7 @@ analyzer = LRPZero(model)
 #
 # We can even implement the generic rule as a specialized implementation for `Dense` layers:
 # ```julia
-# function lrp!(rule::AbstractLRPRule, layer::Dense, Rₖ, aₖ, Rₖ₊₁)
+# function lrp!(Rₖ, rule::AbstractLRPRule, layer::Dense, aₖ, Rₖ₊₁)
 #     ρW, ρb = modify_params(rule, get_params(layer)...)
 #     ãₖ₊₁ = modify_denominator(rule, ρW * aₖ + ρb)
 #     @tullio Rₖ[j] = aₖ[j] * ρW[k, j] / ãₖ₊₁[k] * Rₖ₊₁[k] # Tullio ≈ fast einsum
@@ -283,7 +287,7 @@ analyzer = LRPZero(model)
 # you can also implement your own `lrp!` function and dispatch
 # on individual rule types `MyRule` and layer types `MyLayer`:
 # ```julia
-# function lrp!(rule::MyRule, layer::MyLayer, Rₖ, aₖ, Rₖ₊₁)
+# function lrp!(Rₖ, rule::MyRule, layer::MyLayer, aₖ, Rₖ₊₁)
 #     Rₖ .= ...
 # end
 # ```

src/lrp.jl

@@ -65,7 +65,7 @@ function (analyzer::LRP)(
 
     # Backward pass through layers, applying LRP rules
     for (i, rule) in Iterators.reverse(enumerate(analyzer.rules))
-        lrp!(rule, layers[i], rels[i], acts[i], rels[i + 1]) # inplace update rels[i]
+        lrp!(rels[i], rule, layers[i], acts[i], rels[i + 1]) # inplace update rels[i]
     end
 
     return Explanation(

src/lrp_rules.jl

@@ -2,13 +2,13 @@
 abstract type AbstractLRPRule end
 
 # Generic LRP rule. Since it uses autodiff, it is used as a fallback for layer types without custom implementation.
-function lrp!(rule::R, layer::L, Rₖ, aₖ, Rₖ₊₁) where {R<:AbstractLRPRule,L}
-    lrp_autodiff!(rule, layer, Rₖ, aₖ, Rₖ₊₁)
+function lrp!(Rₖ, rule::R, layer::L, aₖ, Rₖ₊₁) where {R<:AbstractLRPRule,L}
+    lrp_autodiff!(Rₖ, rule, layer, aₖ, Rₖ₊₁)
     return nothing
 end
 
 function lrp_autodiff!(
-    rule::R, layer::L, Rₖ::T1, aₖ::T1, Rₖ₊₁::T2
+    Rₖ::T1, rule::R, layer::L, aₖ::T1, Rₖ₊₁::T2
 ) where {R<:AbstractLRPRule,L,T1,T2}
     layerᵨ = modify_layer(rule, layer)
     c::T1 = only(
@@ -23,21 +23,21 @@ function lrp_autodiff!(
 end
 
 # For linear layer types such as Dense layers, using autodiff is overkill.
-function lrp!(rule::R, layer::Dense, Rₖ, aₖ, Rₖ₊₁) where {R<:AbstractLRPRule}
-    lrp_dense!(rule, layer, Rₖ, aₖ, Rₖ₊₁)
+function lrp!(Rₖ, rule::R, layer::Dense, aₖ, Rₖ₊₁) where {R<:AbstractLRPRule}
+    lrp_dense!(Rₖ, rule, layer, aₖ, Rₖ₊₁)
     return nothing
 end
 
-function lrp_dense!(rule::R, l, Rₖ, aₖ, Rₖ₊₁) where {R<:AbstractLRPRule}
+function lrp_dense!(Rₖ, rule::R, l, aₖ, Rₖ₊₁) where {R<:AbstractLRPRule}
     ρW, ρb = modify_params(rule, get_params(l)...)
     ãₖ₊₁ = modify_denominator(rule, ρW * aₖ .+ ρb)
     @tullio Rₖ[j, b] = aₖ[j, b] * ρW[k, j] / ãₖ₊₁[k, b] * Rₖ₊₁[k, b]
     return nothing
 end
 
 # Other special cases that are dispatched on layer type:
-lrp!(::AbstractLRPRule, ::DropoutLayer, Rₖ, aₖ, Rₖ₊₁) = (Rₖ .= Rₖ₊₁)
-lrp!(::AbstractLRPRule, ::ReshapingLayer, Rₖ, aₖ, Rₖ₊₁) = (Rₖ .= reshape(Rₖ₊₁, size(aₖ)))
+lrp!(Rₖ, ::AbstractLRPRule, ::DropoutLayer, aₖ, Rₖ₊₁) = (Rₖ .= Rₖ₊₁)
+lrp!(Rₖ, ::AbstractLRPRule, ::ReshapingLayer, aₖ, Rₖ₊₁) = (Rₖ .= reshape(Rₖ₊₁, size(aₖ)))
 
 # To implement new rules, we can define two custom functions `modify_params` and `modify_denominator`.
 # If this isn't done, the following fallbacks are used by default:
@@ -75,7 +75,7 @@ Constructor for LRP-0 rule. Commonly used on upper layers.
 struct ZeroRule <: AbstractLRPRule end
 
 """
-    GammaRule(; γ=0.25)
+    GammaRule([γ=0.25])
 
 Constructor for LRP-``γ`` rule. Commonly used on lower layers.
 
@@ -84,16 +84,17 @@ Arguments:
 """
 struct GammaRule{T} <: AbstractLRPRule
     γ::T
-    GammaRule(; γ=0.25) = new{Float32}(γ)
+    GammaRule(γ=0.25f0) = new{Float32}(γ)
 end
 function modify_params(r::GammaRule, W, b)
-    ρW = W + r.γ * relu.(W)
-    ρb = b + r.γ * relu.(b)
+    T = eltype(W)
+    ρW = W + convert(T, r.γ) * relu.(W)
+    ρb = b + convert(T, r.γ) * relu.(b)
     return ρW, ρb
 end
 
 """
-    EpsilonRule(; ϵ=1f-6)
+    EpsilonRule([ϵ=1.0f-6])
 
 Constructor for LRP-``ϵ`` rule. Commonly used on middle layers.
 
@@ -102,7 +103,7 @@ Arguments:
 """
 struct EpsilonRule{T} <: AbstractLRPRule
     ϵ::T
-    EpsilonRule(; ϵ=1.0f-6) = new{Float32}(ϵ)
+    EpsilonRule(ϵ=1.0f-6) = new{Float32}(ϵ)
 end
 modify_denominator(r::EpsilonRule, d) = stabilize_denom(d, r.ϵ)
 
@@ -122,8 +123,8 @@ struct ZBoxRule{T} <: AbstractLRPRule
 end
 
 # The ZBoxRule requires its own implementation of relevance propagation.
-lrp!(r::ZBoxRule, layer::Dense, Rₖ, aₖ, Rₖ₊₁) = lrp_zbox!(r, layer, Rₖ, aₖ, Rₖ₊₁)
-lrp!(r::ZBoxRule, layer::Conv, Rₖ, aₖ, Rₖ₊₁) = lrp_zbox!(r, layer, Rₖ, aₖ, Rₖ₊₁)
+lrp!(Rₖ, r::ZBoxRule, layer::Dense, aₖ, Rₖ₊₁) = lrp_zbox!(Rₖ, r, layer, aₖ, Rₖ₊₁)
+lrp!(Rₖ, r::ZBoxRule, layer::Conv, aₖ, Rₖ₊₁) = lrp_zbox!(Rₖ, r, layer, aₖ, Rₖ₊₁)
 
 _zbox_bound(T, c::Real, in_size) = fill(convert(T, c), in_size)
 function _zbox_bound(T, A::AbstractArray, in_size)
@@ -135,7 +136,7 @@ function _zbox_bound(T, A::AbstractArray, in_size)
     return convert.(T, A)
 end
 
-function lrp_zbox!(r::ZBoxRule, layer::L, Rₖ::T1, aₖ::T1, Rₖ₊₁::T2) where {L,T1,T2}
+function lrp_zbox!(Rₖ::T1, r::ZBoxRule, layer::L, aₖ::T1, Rₖ₊₁::T2) where {L,T1,T2}
     T = eltype(aₖ)
     in_size = size(aₖ)
     l = _zbox_bound(T, r.low, in_size)

test/test_rules.jl

@@ -27,7 +27,7 @@ const RULES = Dict(
 
     layer = Dense(W, b, relu)
     R̂ₖ = similar(aₖ) # will be inplace updated
-    @inferred lrp!(rule, layer, R̂ₖ, aₖ, Rₖ₊₁)
+    @inferred lrp!(R̂ₖ, rule, layer, aₖ, Rₖ₊₁)
     @test R̂ₖ ≈ Rₖ
 
     ## Pooling layer
@@ -42,7 +42,7 @@ const RULES = Dict(
 
     layer = MaxPool((2, 2); stride=(1, 1))
     R̂ₖ = similar(aₖ) # will be inplace updated
-    @inferred lrp!(rule, layer, R̂ₖ, aₖ, Rₖ₊₁)
+    @inferred lrp!(R̂ₖ, rule, layer, aₖ, Rₖ₊₁)
     @test R̂ₖ ≈ Rₖ
 end
 
@@ -53,7 +53,7 @@ pseudorandn(dims...) = randn(MersenneTwister(123), T, dims...)
 ## Test individual rules
 @testset "modify_params" begin
     W, b = [1.0 -1.0; 2.0 0.0], [-1.0, 1.0]
-    ρW, ρb = @inferred modify_params(GammaRule(; γ=0.42), W, b)
+    ρW, ρb = @inferred modify_params(GammaRule(0.42), W, b)
     @test ρW ≈ [1.42 -1.0; 2.84 0.0]
     @test ρb ≈ [-1.0, 1.42]
 end
@@ -67,7 +67,7 @@ aₖ_dense = pseudorandn(ins_dense, batchsize)
 
 layers = Dict(
     "Dense_relu" => Dense(ins_dense, outs_dense, relu; init=pseudorandn),
-    "Dense_identity" => Dense(Matrix(I, outs_dense, ins_dense), false, identity),
+    "Dense_identity" => Dense(Matrix{Float32}(I, outs_dense, ins_dense), false, identity),
 )
 @testset "Dense" begin
     for (rulename, rule) in RULES
@@ -76,7 +76,7 @@ layers = Dict(
                 @testset "$layername" begin
                     Rₖ₊₁ = layer(aₖ_dense)
                     Rₖ = similar(aₖ_dense)
-                    @inferred lrp!(rule, layer, Rₖ, aₖ_dense, Rₖ₊₁)
+                    @inferred lrp!(Rₖ, rule, layer, aₖ_dense, Rₖ₊₁)
 
                     @test typeof(Rₖ) == typeof(aₖ_dense)
                     @test size(Rₖ) == size(aₖ_dense)
@@ -118,8 +118,8 @@ equalpairs = Dict( # these pairs of layers are all equal
                     @test Rₖ₊₁ == l2(aₖ)
                     Rₖ1 = similar(aₖ)
                     Rₖ2 = similar(aₖ)
-                    @inferred lrp!(rule, l1, Rₖ1, aₖ, Rₖ₊₁)
-                    @inferred lrp!(rule, l2, Rₖ2, aₖ, Rₖ₊₁)
+                    @inferred lrp!(Rₖ1, rule, l1, aₖ, Rₖ₊₁)
+                    @inferred lrp!(Rₖ2, rule, l2, aₖ, Rₖ₊₁)
                     @test Rₖ1 == Rₖ2
 
                     @test typeof(Rₖ1) == typeof(aₖ)
@@ -152,7 +152,7 @@ layers = Dict(
                 @testset "$layername" begin
                     Rₖ₊₁ = layer(aₖ)
                     Rₖ = similar(aₖ)
-                    @inferred lrp!(rule, layer, Rₖ, aₖ, Rₖ₊₁)
+                    @inferred lrp!(Rₖ, rule, layer, aₖ, Rₖ₊₁)
 
                     @test typeof(Rₖ) == typeof(aₖ)
                     @test size(Rₖ) == size(aₖ)
@@ -173,7 +173,7 @@ struct TestWrapper{T}
 end
 (w::TestWrapper)(x) = w.layer(x)
 modify_layer(r::AbstractLRPRule, w::TestWrapper) = modify_layer(r, w.layer)
-lrp!(rule::ZBoxRule, w::TestWrapper, Rₖ, aₖ, Rₖ₊₁) = lrp!(rule, w.layer, Rₖ, aₖ, Rₖ₊₁)
+lrp!(Rₖ, rule::ZBoxRule, w::TestWrapper, aₖ, Rₖ₊₁) = lrp!(Rₖ, rule, w.layer, aₖ, Rₖ₊₁)
 
 layers = Dict(
     "Conv" => (Conv((3, 3), 2 => 4; init=pseudorandn), aₖ),
@@ -188,7 +188,7 @@ layers = Dict(
                     wrapped_layer = TestWrapper(layer)
                     Rₖ₊₁ = wrapped_layer(aₖ)
                     Rₖ = similar(aₖ)
-                    @inferred lrp!(rule, wrapped_layer, Rₖ, aₖ, Rₖ₊₁)
+                    @inferred lrp!(Rₖ, rule, wrapped_layer, aₖ, Rₖ₊₁)
 
                     @test typeof(Rₖ) == typeof(aₖ)
                     @test size(Rₖ) == size(aₖ)