Faster LRP rules on Dense layers (#31)

* Add faster AD-less LRP implementation for Dense layer using Tullio.jl

Author: adrhill

Project.toml

@@ -11,13 +11,15 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 ColorSchemes = "3"
 Flux = "0.12"
 ImageCore = "0.8, 0.9"
 PrettyTables = "1"
+Tullio = "0.3"
 Zygote = "0.6"
 julia = "1.6"
 

src/ExplainabilityMethods.jl

@@ -6,6 +6,7 @@ using Zygote
 using ColorSchemes
 using ImageCore
 using Base.Iterators
+using Tullio
 
 using Markdown
 using PrettyTables

src/lrp_rules.jl

@@ -4,23 +4,27 @@
 # can be implemented by dispatching on the functions `modify_params` & `modify_denominator`,
 # which make use of the generalized LRP implementation shown in [1].
 #
-# If the relevance propagation falls outside of this scheme, a custom function
+# If the relevance propagation falls outside of this scheme, custom functions
 # ```julia
 # (::MyLRPRule)(layer, aₖ, Rₖ₊₁) = ...
+# (::MyLRPRule)(layer::MyLayer, aₖ, Rₖ₊₁) = ...
+# (::AbstractLRPRule)(layer::MyLayer, aₖ, Rₖ₊₁) = ...
 # ```
-# can be implemented. This is used for the ZBoxRule.
+# that return `Rₖ` can be implemented.
+# This is used for the ZBoxRule and for faster computations on common layers.
 #
 # References:
 # [1] G. Montavon et al., Layer-Wise Relevance Propagation: An Overview
-# [2] W. Samek et al., Explaining Deep Neural Networks and Beyond:
-#     A Review of Methods and Applications
+# [2] W. Samek et al., Explaining Deep Neural Networks and Beyond: A Review of Methods and Applications
 
 abstract type AbstractLRPRule end
 
 # This is the generic relevance propagation rule which is used for the 0, γ and ϵ rules.
 # It can be extended for new rules via `modify_denominator` and `modify_params`.
 # Since it uses autodiff, it is used as a fallback for layer types without custom implementation.
-function (rule::AbstractLRPRule)(layer, aₖ, Rₖ₊₁)
+(rule::AbstractLRPRule)(layer, aₖ, Rₖ₊₁) = lrp_autodiff(rule, layer, aₖ, Rₖ₊₁)
+
+function lrp_autodiff(rule, layer, aₖ, Rₖ₊₁)
     layerᵨ = _modify_layer(rule, layer)
     function fwpass(a)
         z = layerᵨ(a)
@@ -30,7 +34,16 @@ function (rule::AbstractLRPRule)(layer, aₖ, Rₖ₊₁)
     return aₖ .* gradient(fwpass, aₖ)[1] # Rₖ
 end
 
-# Special cases are dispatched on layer type:
+# For linear layer types such as Dense layers, using autodiff is overkill.
+(rule::AbstractLRPRule)(layer::Dense, aₖ, Rₖ₊₁) = lrp_dense(rule, layer, aₖ, Rₖ₊₁)
+
+function lrp_dense(rule, l, aₖ, Rₖ₊₁)
+    ρW, ρb = modify_params(rule, get_params(l)...)
+    ãₖ₊₁ = modify_denominator(rule, ρW * aₖ + ρb)
+    return @tullio Rₖ[j] := aₖ[j] * ρW[k, j] / ãₖ₊₁[k] * Rₖ₊₁[k]
+end
+
+# Other special cases that are dispatched on layer type:
 (::AbstractLRPRule)(::DropoutLayer, aₖ, Rₖ₊₁) = Rₖ₊₁
 (::AbstractLRPRule)(::ReshapingLayer, aₖ, Rₖ₊₁) = reshape(Rₖ₊₁, size(aₖ))
 
@@ -104,7 +117,10 @@ Commonly used on the first layer for pixel input.
 struct ZBoxRule <: AbstractLRPRule end
 
 # The ZBoxRule requires its own implementation of relevance propagation.
-function (rule::ZBoxRule)(layer::Union{Dense,Conv}, aₖ, Rₖ₊₁)
+(rule::ZBoxRule)(layer::Dense, aₖ, Rₖ₊₁) = lrp_zbox(layer, aₖ, Rₖ₊₁)
+(rule::ZBoxRule)(layer::Conv, aₖ, Rₖ₊₁) = lrp_zbox(layer, aₖ, Rₖ₊₁)
+
+function lrp_zbox(layer, aₖ, Rₖ₊₁)
     W, b = get_params(layer)
     l, h = fill.(extrema(aₖ), (size(aₖ),))
 

src/utils.jl

@@ -3,7 +3,8 @@
 
 Replace zero terms of a matrix `d` with `eps`.
 """
-function stabilize_denom(d; eps=1.0f-9)
+stabilize_denom(d::Real; eps=1.0f-9) = ifelse(d ≈ 0, d + sign(d) * eps, d)
+function stabilize_denom(d::AbstractArray; eps=1.0f-9)
     return d + ((d .≈ 0) + sign.(d)) * eps
 end