Add LRP ZPlusRule (#88)

Author: adrhill

README.md

@@ -84,6 +84,7 @@ Currently, the following analyzers are implemented:
     ├── WSquareRule
     ├── FlatRule
     ├── ZBoxRule
+    ├── ZPlusRule
     ├── AlphaBetaRule
     └── PassRule
 ```

docs/literate/example.jl

@@ -117,6 +117,7 @@ mosaic(heatmap(batch, analyzer, 1); nrow=10)
 #     ├── WSquareRule
 #     ├── FlatRule
 #     ├── ZBoxRule
+#     ├── ZPlusRule
 #     ├── AlphaBetaRule
 #     └── PassRule
 # ```

src/ExplainableAI.jl

@@ -47,7 +47,7 @@ export LRP
 export AbstractLRPRule
 export LRP_CONFIG
 export ZeroRule, EpsilonRule, GammaRule, WSquareRule, FlatRule, PassRule
-export ZBoxRule, AlphaBetaRule
+export ZBoxRule, ZPlusRule, AlphaBetaRule
 export modify_input, modify_denominator
 export modify_param!, modify_layer!
 export check_model

src/lrp/composite_presets.jl

@@ -31,7 +31,7 @@ $(repr("text/plain", EpsilonPlus()))
 function EpsilonPlus(; epsilon=1.0f-6)
     return Composite(
         GlobalTypeRule(
-            ConvLayer => AlphaBetaRule(1.0f0, 0.0f0), # TODO: replace with ZPlusRule
+            ConvLayer => ZPlusRule(),
             Dense => EpsilonRule(epsilon),
             DropoutLayer => PassRule(),
             ReshapingLayer => PassRule(),
@@ -71,7 +71,7 @@ $(repr("text/plain", EpsilonPlusFlat()))
 function EpsilonPlusFlat(; epsilon=1.0f-6)
     return Composite(
         GlobalTypeRule(
-            ConvLayer => AlphaBetaRule(1.0f0, 0.0f0), # TODO: replace with ZPlusRule
+            ConvLayer => ZPlusRule(),
             Dense => EpsilonRule(epsilon),
             DropoutLayer => PassRule(),
             ReshapingLayer => PassRule(),

src/lrp/rules.jl

@@ -311,6 +311,38 @@ function lrp!(Rₖ, rule::AlphaBetaRule, layer::L, aₖ, Rₖ₊₁) where {L}
     return nothing
 end
 
+"""
+    ZPlusRule()
+
+LRP-``z^{+}`` rule. Commonly used on lower layers.
+
+Equivalent to `AlphaBetaRule(1.0f0, 0.0f0)`, but slightly faster.
+See also [`AlphaBetaRule`](@ref).
+
+# References
+- [1] $REF_BACH_LRP
+- [2] $REF_MONTAVON_DTD
+"""
+struct ZPlusRule <: AbstractLRPRule end
+function lrp!(Rₖ, rule::ZPlusRule, layer::L, aₖ, Rₖ₊₁) where {L}
+    require_weight_and_bias(rule, layer)
+    reset! = get_layer_resetter(rule, layer)
+
+    aₖ⁺ = keep_positive(aₖ)
+    aₖ⁻ = keep_negative(aₖ)
+
+    modify_layer!(Val(:keep_positive), layer)
+    out_1, pullback_1 = Zygote.pullback(layer, aₖ⁺)
+    reset!()
+    modify_layer!(Val(:keep_negative_zero_bias), layer)
+    out_2, pullback_2 = Zygote.pullback(layer, aₖ⁻)
+    reset!()
+
+    y_α = Rₖ₊₁ ./ modify_denominator(rule, out_1 + out_2)
+    Rₖ .= aₖ⁺ .* only(pullback_1(y_α)) + aₖ⁻ .* only(pullback_2(y_α))
+    return nothing
+end
+
 # Special cases for rules that don't modify params for extra performance:
 for R in (ZeroRule, EpsilonRule)
     for L in (DropoutLayer, ReshapingLayer)