Improve LRP backward pass (#143)

* Move inner LRP loop into `lrp_backward_pass!`

* Reduce amount of Unicode

* Add Tullio to benchmark dependencies

Author: adrhill

benchmark/Project.toml

@@ -4,6 +4,7 @@ ExplainableAI = "4f1bc3e1-d60d-4ed0-9367-9bdff9846d3b"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
 PkgBenchmark = "32113eaa-f34f-5b0d-bd6c-c81e245fc73d"
+Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
 
 [compat]
 BenchmarkTools = "1"

benchmark/benchmarks.jl

@@ -1,5 +1,6 @@
 using BenchmarkTools
 using LoopVectorization
+using Tullio
 using Flux
 using ExplainableAI
 using ExplainableAI: lrp!, modify_layer

src/lrp/lrp.jl

@@ -46,13 +46,21 @@ LRP(model::Chain; kwargs...) = LRP(model, Composite(ZeroRule()); kwargs...)
 # Construct Chain-/ParallelTuple of rules by applying composite
 LRP(model::Chain, c::Composite; kwargs...) = LRP(model, lrp_rules(model, c); kwargs...)
 
-get_activations(model, input) = [input, Flux.activations(model, input)...]
+get_activations(model, input) = (input, Flux.activations(model, input)...)
 
-function mask_output_neuron!(Rᴺ, aᴺ, ns::AbstractNeuronSelector)
-    fill!(Rᴺ, 0)
-    idx = ns(aᴺ)
-    Rᴺ[idx] .= 1
-    return Rᴺ
+function mask_output_neuron!(R_out, a_out, ns::AbstractNeuronSelector)
+    fill!(R_out, 0)
+    idx = ns(a_out)
+    R_out[idx] .= 1
+    return R_out
+end
+
+function lrp_backward_pass!(Rs, as, rules, layers, modified_layers)
+    # Apply LRP rules in backward-pass, inplace-updating relevances `Rs[k]` = Rᵏ
+    for k in length(layers):-1:1
+        lrp!(Rs[k], rules[k], layers[k], modified_layers[k], as[k], Rs[k + 1])
+    end
+    return Rs
 end
 
 # Call to the LRP analyzer
@@ -63,11 +71,7 @@ function (lrp::LRP)(
     Rs = similar.(as)                         # allocate relevances Rᵏ for all layers k
     mask_output_neuron!(Rs[end], as[end], ns) # compute relevance Rᴺ of output layer N
 
-    # Apply LRP rules in backward-pass, inplace-updating relevances `Rs[k]` = Rᵏ
-    for k in length(lrp.model):-1:1
-        lrp!(Rs[k], lrp.rules[k], lrp.model[k], lrp.modified_layers[k], as[k], Rs[k + 1])
-    end
-
+    lrp_backward_pass!(Rs, as, lrp.rules, lrp.model, lrp.modified_layers)
     extras = layerwise_relevances ? (layerwise_relevances=Rs,) : nothing
     return Explanation(first(Rs), last(as), ns(last(as)), :LRP, extras)
 end
@@ -77,10 +81,7 @@ function lrp!(Rᵏ, rules::ChainTuple, chain::Chain, modified_chain::ChainTuple,
     Rs = similar.(as)
     last(Rs) .= Rᵏ⁺¹
 
-    # Apply LRP rules in backward-pass, inplace-updating relevances `Rs[i]`
-    for i in length(chain):-1:1
-        lrp!(Rs[i], rules[i], chain[i], modified_chain[i], as[i], Rs[i + 1])
-    end
+    lrp_backward_pass!(Rs, as, rules, chain, modified_chain)
     return Rᵏ .= first(Rs)
 end
 
@@ -91,15 +92,15 @@ function lrp!(
     # according to the contribution aᵏ⁺¹ᵢ of branch i to the output activation aᵏ⁺¹:
     #   Rᵏ⁺¹ᵢ = Rᵏ⁺¹ .* aᵏ⁺¹ᵢ ./ aᵏ⁺¹ = c .* aᵏ⁺¹ᵢ
 
-    aᵏ⁺¹s = [l(aᵏ) for l in parallel.layers]     # aᵏ⁺¹ᵢ for each branch i
-    c = Rᵏ⁺¹ ./ stabilize_denom(sum(aᵏ⁺¹s))
-    Rᵏ⁺¹s = [c .* aᵏ⁺¹ᵢ for aᵏ⁺¹ᵢ in aᵏ⁺¹s]      # Rᵏ⁺¹ᵢ for each branch i
-    Rᵏs = [similar(aᵏ) for _ in parallel.layers] # pre-allocate output Rᵏᵢ for each branch i
+    aᵏ⁺¹_parallel = [layer(aᵏ) for layer in parallel.layers] # aᵏ⁺¹ᵢ for each branch i
+    c = Rᵏ⁺¹ ./ stabilize_denom(sum(aᵏ⁺¹_parallel))
+    Rᵏ⁺¹_parallel = [c .* a for a in aᵏ⁺¹_parallel]          # Rᵏ⁺¹ᵢ for each branch i
+    Rᵏ_parallel = [similar(aᵏ) for _ in parallel.layers]     # pre-allocate output Rᵏᵢ for each branch
 
-    for (Rᵏᵢ, ruleᵢ, layerᵢ, modified_layerᵢ, Rᵏ⁺¹ᵢ) in
-        zip(Rᵏs, rules, parallel.layers, modified_parallel, Rᵏ⁺¹s)
-        # In-place update Rᵏᵢ (and therefore Rᵏs)
-        lrp!(Rᵏᵢ, ruleᵢ, layerᵢ, modified_layerᵢ, aᵏ, Rᵏ⁺¹ᵢ)
+    for (Rᵏᵢ, rule, layer, modified_layer, Rᵏ⁺¹ᵢ) in
+        zip(Rᵏ_parallel, rules, parallel.layers, modified_parallel, Rᵏ⁺¹_parallel)
+        # In-place update Rᵏᵢ and therefore Rᵏ_parallel
+        lrp!(Rᵏᵢ, rule, layer, modified_layer, aᵏ, Rᵏ⁺¹ᵢ)
     end
-    return Rᵏ .= sum(Rᵏs)
+    return Rᵏ .= sum(Rᵏ_parallel)
 end