Add SmoothGrad and InputAugmentation (#50)

Author: adrhill

Project.toml

@@ -5,17 +5,20 @@ version = "0.3.1"
 
 [deps]
 ColorSchemes = "35d6a980-a343-548e-a6ea-1d62b119f2f4"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 ImageCore = "a09fc81d-aa75-5fe9-8630-4744c3626534"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 ColorSchemes = "3"
+Distributions = "0.25"
 Flux = "0.12"
 ImageCore = "0.8, 0.9"
 PrettyTables = "1"

README.md

@@ -54,6 +54,7 @@ Currently, the following analyzers are implemented:
 ```
 ├── Gradient
 ├── InputTimesGradient
+├── SmoothGrad
 └── LRP
     ├── LRPZero
     ├── LRPEpsilon
@@ -65,7 +66,6 @@ Individual LRP rules like `ZeroRule`, `EpsilonRule`, `GammaRule` and `ZBoxRule`
 
 ## Roadmap
 In the future, we would like to include:
-- [SmoothGrad](https://arxiv.org/abs/1706.03825)
 - [Integrated Gradients](https://arxiv.org/abs/1703.01365)
 - [PatternNet](https://arxiv.org/abs/1705.05598)
 - [DeepLift](https://arxiv.org/abs/1704.02685)

benchmark/benchmarks.jl

@@ -21,6 +21,7 @@ algs = Dict(
     "InputTimesGradient" => InputTimesGradient,
     "LRPZero" => LRPZero,
     "LRPCustom" => LRPCustom, #modifies weights
+    "SmoothGrad" => model -> SmoothGrad(model, 10),
 )
 
 # Define benchmark

docs/literate/example.jl

@@ -105,6 +105,7 @@ mosaic(heatmap(batch, analyzer, 1); nrow=10)
 # ```
 # ├── Gradient
 # ├── InputTimesGradient
+# ├── SmoothGrad
 # └── LRP
 #     ├── LRPZero
 #     ├── LRPEpsilon

docs/src/api.md

@@ -10,6 +10,12 @@ heatmap
 LRP
 Gradient
 InputTimesGradient
+SmoothGrad
+```
+
+`SmoothGrad` is a special case of `InputAugmentation`, which can be applied as a wrapper to any analyzer:
+```@doc
+InputAugmentation
 ```
 
 # LRP

src/ExplainableAI.jl

@@ -2,6 +2,8 @@ module ExplainableAI
 
 using Base.Iterators
 using LinearAlgebra
+using Distributions
+using Random: AbstractRNG, GLOBAL_RNG
 using Flux
 using Zygote
 using Tullio
@@ -18,6 +20,7 @@ include("neuron_selection.jl")
 include("analyze_api.jl")
 include("flux.jl")
 include("utils.jl")
+include("input_augmentation.jl")
 include("gradient.jl")
 include("lrp_checks.jl")
 include("lrp_rules.jl")
@@ -29,6 +32,7 @@ export analyze
 # Analyzers
 export AbstractXAIMethod
 export Gradient, InputTimesGradient
+export InputAugmentation, SmoothGrad
 export LRP, LRPZero, LRPEpsilon, LRPGamma
 
 # LRP rules

src/gradient.jl

@@ -1,4 +1,4 @@
-function gradient_wrt_input(model, input::T, output_indices) where {T}
+function gradient_wrt_input(model, input, output_indices)
     return only(gradient((in) -> model(in)[output_indices], input))
 end
 
@@ -8,7 +8,7 @@ function gradients_wrt_batch(model, input::AbstractArray{T,N}, output_indices) w
     return mapreduce(
         (gs...) -> cat(gs...; dims=N), zip(eachslice(input; dims=N), output_indices)
     ) do (in, idx)
-        gradient_wrt_input(model, batch_dim_view(in), drop_batch_dim(idx))
+        gradient_wrt_input(model, batch_dim_view(in), drop_batch_index(idx))
     end
 end
 
@@ -46,3 +46,16 @@ function (analyzer::InputTimesGradient)(input, ns::AbstractNeuronSelector)
     attr = input .* gradients_wrt_batch(analyzer.model, input, output_indices)
     return Explanation(attr, output, output_indices, :InputTimesGradient, Nothing)
 end
+
+"""
+    SmoothGrad(analyzer, [n=50, std=0.1, rng=GLOBAL_RNG])
+    SmoothGrad(analyzer, [n=50, distribution=Normal(0, σ²=0.01), rng=GLOBAL_RNG])
+
+Analyze model by calculating a smoothed sensitivity map.
+This is done by averaging sensitivity maps of a `Gradient` analyzer over random samples
+in a neighborhood of the input, typically by adding Gaussian noise with mean 0.
+
+# References
+[1] Smilkov et al., SmoothGrad: removing noise by adding noise
+"""
+SmoothGrad(model, n=50, args...) = InputAugmentation(Gradient(model), n, args...)

src/input_augmentation.jl

@@ -0,0 +1,111 @@
+"""
+    InputAugmentation(analyzer, n, [std=1, rng=GLOBAL_RNG])
+    InputAugmentation(analyzer, n, distribution, [rng=GLOBAL_RNG])
+
+A wrapper around analyzers that augments the input with `n` samples of additive noise sampled from `distribution`.
+This input augmentation is then averaged to return an `Explanation`.
+"""
+struct InputAugmentation{A<:AbstractXAIMethod,D<:Distribution,R<:AbstractRNG} <:
+       AbstractXAIMethod
+    analyzer::A
+    n::Integer
+    distribution::D
+    rng::R
+end
+function InputAugmentation(analyzer, n, distr, rng=GLOBAL_RNG)
+    return InputAugmentation(analyzer, n, distr, rng)
+end
+function InputAugmentation(analyzer, n, σ::Real=0.1f0, args...)
+    return InputAugmentation(analyzer, n, Normal(0.0f0, Float32(σ)^2), args...)
+end
+
+function (aug::InputAugmentation)(input, ns::AbstractNeuronSelector)
+    # Regular forward pass of model
+    output = aug.analyzer.model(input)
+    output_indices = ns(output)
+
+    # Call regular analyzer on augmented batch
+    augmented_input = add_noise(augment_batch_dim(input, aug.n), aug.distribution, aug.rng)
+    augmented_indices = augment_indices(output_indices, aug.n)
+    augmented_expl = aug.analyzer(augmented_input, AugmentationSelector(augmented_indices))
+
+    # Average explanation
+    return Explanation(
+        reduce_augmentation(augmented_expl.attribution, aug.n),
+        output,
+        output_indices,
+        augmented_expl.analyzer,
+        Nothing,
+    )
+end
+
+function add_noise(A::AbstractArray{T}, distr::Distribution, rng::AbstractRNG) where {T}
+    return A + T.(rand(rng, distr, size(A)))
+end
+
+"""
+    augment_batch_dim(input, n)
+
+Repeat each sample in input batch n-times along batch dimension.
+This turns arrays of size `(..., B)` into arrays of size `(..., B*n)`.
+
+## Example
+```julia-repl
+julia> A = [1 2; 3 4]
+2×2 Matrix{Int64}:
+ 1  2
+ 3  4
+
+julia> augment_batch_dim(A, 3)
+2×6 Matrix{Int64}:
+ 1  1  1  2  2  2
+ 3  3  3  4  4  4
+```
+"""
+function augment_batch_dim(input::AbstractArray{T,N}, n) where {T,N}
+    return repeat(input; inner=(ntuple(_ -> 1, Val(N - 1))..., n))
+end
+
+"""
+    reduce_augmentation(augmented_input, n)
+
+Reduce augmented input batch by averaging the explanation for each augmented sample.
+"""
+function reduce_augmentation(input::AbstractArray{T,N}, n) where {T<:AbstractFloat,N}
+    return cat(
+        (
+            Iterators.map(1:n:size(input, N)) do i
+                augmentation_range = ntuple(_ -> :, Val(N - 1))..., i:(i + n - 1)
+                sum(view(input, augmentation_range...); dims=N) / n
+            end
+        )...; dims=N
+    )::Array{T,N}
+end
+"""
+    augment_indices(indices, n)
+
+Strip batch indices and return inidices for batch augmented by n samples.
+
+## Example
+```julia-repl
+julia> inds = [CartesianIndex(5,1), CartesianIndex(3,2)]
+2-element Vector{CartesianIndex{2}}:
+ CartesianIndex(5, 1)
+ CartesianIndex(3, 2)
+
+julia> augment_indices(inds, 3)
+6-element Vector{CartesianIndex{2}}:
+ CartesianIndex(5, 1)
+ CartesianIndex(5, 2)
+ CartesianIndex(5, 3)
+ CartesianIndex(3, 4)
+ CartesianIndex(3, 5)
+ CartesianIndex(3, 6)
+```
+"""
+function augment_indices(inds::Vector{CartesianIndex{N}}, n) where {N}
+    indices_wo_batch = [i.I[1:(end - 1)] for i in inds]
+    return map(enumerate(repeat(indices_wo_batch; inner=n))) do (i, idx)
+        CartesianIndex{N}(idx..., i)
+    end
+end

src/neuron_selection.jl

@@ -27,3 +27,13 @@ function (s::IndexSelector{I})(out::AbstractArray{T,N}) where {I,T,N}
     N < 2 && throw(BATCHDIM_MISSING)
     return CartesianIndex{N}.(s.index..., 1:size(out, N))
 end
+
+"""
+    AugmentationSelector(index)
+
+Neuron selector that passes through an augmented neuron selection.
+"""
+struct AugmentationSelector{I} <: AbstractNeuronSelector
+    indices::I
+end
+(s::AugmentationSelector)(out) = s.indices

src/utils.jl

@@ -42,15 +42,15 @@ julia> batch_dim_view(A)
 batch_dim_view(A::AbstractArray{T,N}) where {T,N} = view(A, ntuple(_ -> :, Val(N + 1))...)
 
 """
-    drop_batch_dim(I)
+    drop_batch_index(I)
 
 Drop batch dimension index (last value) from CartesianIndex.
 
 ## Example
-julia> drop_batch_dim(CartesianIndex(5,3,2))
+julia> drop_batch_index(CartesianIndex(5,3,2))
 CartesianIndex(5, 3)
 """
-drop_batch_dim(C::CartesianIndex) = CartesianIndex(C.I[1:(end - 1)])
+drop_batch_index(C::CartesianIndex) = CartesianIndex(C.I[1:(end - 1)])
 
 # Utils for printing model check summary using PrettyTable.jl
 _print_name(layer) = "$layer"