Add AlphaBetaRule (#78)

Author: adrhill

README.md

@@ -63,11 +63,12 @@ Currently, the following analyzers are implemented:
     ├── WSquareRule
     ├── FlatRule
     ├── ZBoxRule
+    ├── AlphaBetaRule
     └── PassRule
 ```
 
 One of the design goals of ExplainableAI.jl is extensibility.
-Individual LRP rules like `ZeroRule`, `EpsilonRule`, `GammaRule` and `ZBoxRule` [can be composed][docs-composites] and are easily extended by [custom rules][docs-custom-rules].
+Individual LRP rules [can be composed][docs-composites] and are easily extended by [custom rules][docs-custom-rules].
 
 ## Roadmap
 In the future, we would like to include:

benchmark/benchmarks.jl

@@ -13,20 +13,11 @@ model = flatten_model(strip_softmax(vgg11.layers))
 T = Float32
 img = rand(MersenneTwister(123), T, (224, 224, 3, 1))
 
-# Benchmark custom LRP composite
-function LRPCustom(model::Chain)
-    return LRP(
-        model,
-        [ZBoxRule(zero(T), oneunit(T)), repeat([GammaRule()], length(model.layers) - 1)...],
-    )
-end
-
 # Use one representative algorithm of each type
 algs = Dict(
     "Gradient" => Gradient,
     "InputTimesGradient" => InputTimesGradient,
-    "LRPZero" => LRP,
-    "LRPCustom" => LRPCustom, #modifies weights
+    "LRP" => LRP,
     "SmoothGrad" => model -> SmoothGrad(model, 10),
     "IntegratedGradients" => model -> IntegratedGradients(model, 10),
 )
@@ -53,7 +44,6 @@ out_dense = 100
 aₖ = randn(T, insize)
 
 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)),
 )
@@ -62,6 +52,9 @@ rules = Dict(
     "EpsilonRule" => EpsilonRule(),
     "GammaRule" => GammaRule(),
     "ZBoxRule" => ZBoxRule(zero(T), oneunit(T)),
+    "FlatRule" => FlatRule(),
+    "WSquareRule" => WSquareRule(),
+    "AlphaBetaRule" => AlphaBetaRule(),
 )
 
 SUITE["Layer"] = BenchmarkGroup([k for k in keys(layers)])

docs/src/api.md

@@ -27,6 +27,7 @@ ZeroRule
 EpsilonRule
 GammaRule
 WSquareRule
+AlphaBetaRule
 FlatRule
 ZBoxRule
 PassRule

src/ExplainableAI.jl

@@ -42,7 +42,8 @@ export LRP
 # LRP rules
 export AbstractLRPRule
 export LRP_CONFIG
-export ZeroRule, EpsilonRule, GammaRule, WSquareRule, FlatRule, ZBoxRule, PassRule
+export ZeroRule, EpsilonRule, GammaRule, WSquareRule, FlatRule, PassRule
+export ZBoxRule, AlphaBetaRule
 export modify_input, modify_denominator
 export modify_param!, modify_layer!
 export check_model

src/lrp_rules.jl

@@ -251,7 +251,62 @@ function zbox_input(in::AbstractArray{T}, A::AbstractArray) where {T}
     return convert.(T, A)
 end
 
+"""
+    AlphaBetaRule(alpha, beta)
+    AlphaBetaRule([alpha=2.0], [beta=1.0])
+
+LRP-``\alpha\beta`` rule. Weights positive and negative contributions according to the
+parameters `alpha` and `beta` respectively. The difference `alpha - beta` must be equal one.
+Commonly used on lower layers.
+
+Arguments:
+- `alpha`: Multiplier for the positive output term, defaults to `2.0`.
+- `beta`: Multiplier for the negative output term, defaults to `1.0`.
+
+# References
+[1]: S. Bach et al., On Pixel-Wise Explanations for Non-Linear Classifier Decisions by
+    Layer-Wise Relevance Propagation
+[2]: G. Montavon et al., Layer-Wise Relevance Propagation: An Overview
+"""
+struct AlphaBetaRule{T} <: AbstractLRPRule
+    α::T
+    β::T
+    function AlphaBetaRule(alpha=2.0f0, beta=1.0f0)
+        alpha < 0 && throw(ArgumentError("Parameter `alpha` must be ≥0."))
+        beta < 0 && throw(ArgumentError("Parameter `beta` must be ≥0."))
+        !isone(alpha - beta) && throw(ArgumentError("`alpha - beta` must be equal one."))
+        return new{Float32}(alpha, beta)
+    end
+end
+
+# The AlphaBetaRule requires its own implementation of relevance propagation.
+function lrp!(Rₖ, rule::AlphaBetaRule, 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!()
+    modify_layer!(Val(:keep_negative), layer)
+    out_3, pullback_3 = Zygote.pullback(layer, aₖ⁺)
+    reset!()
+    modify_layer!(Val(:keep_positive_zero_bias), layer)
+    out_4, pullback_4 = Zygote.pullback(layer, aₖ⁻)
+    reset!()
 
+    y_α = Rₖ₊₁ ./ modify_denominator(rule, out_1 + out_2)
+    y_β = Rₖ₊₁ ./ modify_denominator(rule, out_3 + out_4)
+    Rₖ .=
+        rule.α .* (aₖ⁺ .* only(pullback_1(y_α)) + aₖ⁻ .* only(pullback_2(y_α))) .-
+        rule.β .* (aₖ⁺ .* only(pullback_3(y_β)) + aₖ⁻ .* only(pullback_4(y_β)))
+    return nothing
+end
 
 # Special cases for rules that don't modify params for extra performance:
 for R in (ZeroRule, EpsilonRule)