Diffeomorphic Explanations with Normalizing Flows
Abstract
Normalizing ﬂows are diffeomorphisms which are parameterized by neural networks.
As a result, they can induce coordinate transformations in the tangent space of the data manifold.
In this work, we demonstrate that such transformations can be used to generate interpretable explanations for decisions of neural networks.
More speciﬁcally, we perform gradient ascent in the base space of the ﬂow to generate counterfactuals which are classiﬁed with great conﬁdence as a speciﬁed target class.
We analyze this generation process theoretically using Riemannian differential geometry and establish a rigorous theoretical connection between gradient ascent on the data manifold and in the base space of the ﬂow.
original x counterfactual x (cid:48) heatmap δx
1. Introduction
Explaining a complex system can be drastically simpliﬁed using a suitable coordinate system.
As an example, the solar system can be explained either by using a reference system for which the sun is at rest (heliocentristic) or, alternatively, for which the earth is at rest (geocentristic).
Despite wide-held belief, both reference system are physically valid.
However, the dynamics of the planets is signiﬁcantly easier to describe in heliocentristic coordinates since the planets will follow geometrically simple trajectories.
Explanation methods for neural networks have recently gained signiﬁcant attention because they promise to make black-box classiﬁers more transparent, see (Samek et al., 2019) for a detailed overview.
In this paper, we use the bijectivity of a normalizing ﬂow to consider a classiﬁer in the base space of the ﬂow.
This amounts to a coordinate transformation in the data space (or mathematically more precise: a diffeomorphism).
We will show that in this coordinate system, the classiﬁer is more easily interpretable and can be used to construct counterfactual explanations that lie on the data manifold.
Using Riemannian differential geometry, we will analyze the advantages of creating counterfactual explanations in the base space of the ﬂow and establish a process by which the tangent space of the data manifold can be estimated from the ﬂow.
We strongly expect these theoretical insights to be useful beyond explainability.
In summary, our main contributions are as follows:
• We propose a novel application domain for ﬂows: inducing a bijective transformation to a more interpretable space on which counterfactuals can be easily generated.
• We analyze the properties of this generation process theoretically using Riemannian differential geometry.
• We experimentally demonstrate superior performance compared to more traditional approaches for generating counterfactuals for classiﬁcation tasks in three different domains.
2. Counterfactual Explanations
Let f : X → R K be a classiﬁer whose component f ( x ) k is the probability for the point x ∈ X to be of class k ∈ { 1 , . .
. , K } .
We make no assumptions on the architecture of the classiﬁer f and only require that we can evaluate f ( x ) and its derivative ∂ x f ( x ) for a given input x ∈ X .
In this work, we will follow the well-established paradigm of counterfactual explanations – see (Verma et al., 2020) for the gradient, it by ﬁnite differences a recent review.
These methods aim to explain the classiﬁer f by providing an answer to the question which minimal deformations x (cid:48)
=
x + δx need to be applied to the original input x in order to change its prediction.
Often, the difference δx is then visualized by a heatmap highlighting the relevant pixels for the change in classiﬁcation, see Figurefor an example.
In the following, we will assume that the data lies on a submanifold S ⊂ X which is of (signiﬁcantly) lower dimension n than the dimension N of its embedding space X .
We stress that this is also known as the manifold assumption and is expected to hold across a wide range of machine learning tasks, see e.g. (Goodfellow et al., 2016).
In these situations, we are often interested in only the deformations x (cid:48) which lie on the data manifold S .
As an example, a customer of a bank may want to understand how their ﬁnancial data needs to change in order to receive a loan.
If the classiﬁcation changes off-manifold, for example for zip codes that do not exist, this is of no relevance since the user is obliged to enter their correct zip code.
Furthermore, it is often required that the deformation is minimal, i.e. the perturbation δx should be as small as possible.
However, the relevant norm is the one of the data manifold S and not of its embedding pixel space X .
For example, a slightly rotated number in an MNIST image may have large pixel-wise distance but should be considered an inﬁnitesimal perturbation of the original image.
More precisely, we deﬁne counterfactuals as follows: let t = argmax j f j ( x ) be the predicted class for the data point x ∈ S .
The set ∆ k,δ ⊂ S of counterfactuals x (cid:48) of the point x with respect to the target class k ∈ { 1 , . . .
, K } \ { t } and conﬁdence δ ∈ (0 , 1] is deﬁned by ∆ k,δ = { x (cid:48)
( x ) ∈ S : argmax j f j ( x (cid:48) ) =
k ∧ f k ( x (cid:48) ) > δ } , i.e. all points on the data manifold which are classiﬁed to be of the target class k with at least the conﬁdence δ .
A minimal counterfactual x (cid:48) ∈
∆ k,δ is then a counterfactual with smallest distance d γ ( x (cid:48) , x ) to the original point x , where d γ is the distance on the data manifold (induced by its Riemannian metric γ ).
Note that there may not be a unique minimal counterfactual.
3. Construction of Counterfactual
We propose to estimate the minimal counterfactual x (cid:48) of the data sample x with respect to the classiﬁer f by using a diffeomorphism g :
Z → X modelled by a normalizing ﬂow.
The ﬂow g equips the space X with a probability density
[Equation]
[Equation]
by push-forward of a simple base density q Z , such as
N (0 , 1) , deﬁned on the base space Z .
We assume that the ﬂow was successfully trained to approximate the data distribution p X by minimizing the forward KL-divergence as usual (this assumption will be made more precise in
Section 4).
We then perform gradient ascent in the base space Z to maximize the probability of the target class k , i.e.
[Equation]
where λ is the learning rate and we initialize by mapping the original point x to the base space by z (0)
=
g − 1 ( x ) .
We then take the sample x ( T ) =
g ( z ( T ) ) as an estimator for a minimal counterfactual if x ( T ) is the ﬁrst optimization step to be classiﬁed as the target k with given conﬁdence δ : argmax j f j ( x ( T ) )
=
k and f k ( x ( T ) ) > δ .
This is because generically taking further steps only increases the distance to the original sample
||
z ( T ) − z
|| >
||
z T + t − z
|| for t >
and we want to ﬁnd (an estimate of a) minimal counterfactual.
This may also be validated by continuing optimization for a certain number of steps and selecting the sample with the minimal distance.
As discussed in Section 6, generative-model-based methods to estimate (minimal) counterfactuals have previously been proposed, for example based on Generative Adversarial
Networks or Autoencoders.
However, the relevance of normalizing ﬂows in this domain has so far not be recognized.
This is unfortunate as normalizing ﬂows have important advantages in this application domain compared to other generative models: ﬁrstly, a ﬂow g is a diffeomorphism and therefore no information is lost by considering the classiﬁer f ◦ g on Z instead of the original classiﬁer f on X , i.e. there is a unique z
= g − 1 ( x ) ∈
Z for any data point x ∈ X .
Secondly, performing gradient ascent in the base space Z of a well-trained ﬂow will ensure (to good approximation) that each optimization step x ( t ) =
g ( z ( t ) ) will stay on the data manifold S ⊂ X .
Since the base space Z has the same dimension as the data space X , the latter statement is far from obvious and is substantiated with theoretical arguments in the next section.
4. Theoretical Analysis
In the following, it will be shown that performing gradient ascent in Z space and then mapping the result in X space will stay on the data manifold S .
This is in stark contrast to gradient ascent directly in X space, i.e.
[Equation]
where λ is the learning rate.
It is well-known that such an optimization procedure would very quickly leave the data manifold S , see for example (Goodfellow et al., 2015).
For gradient ascent in Z space (2), each step z ( t ) can uniquely be mapped to X space by x ( t )
=
g ( z ( t ) ) .
In the Supplement A.1, we derive the following result:
Theorem 1.
Let z ( t ) be deﬁned as in (2) and x ( t )
=
g ( z ( t ) ) .
Then, to leading order in the learning rate λ ,
[Equation]
where γ − 1 = ∂g∂z ∂g∂zT ∈ R N,N is the pull-back of the ﬂat metric on Z under the ﬂow g .
Therefore, performing gradient ascent in X or Z space is not equivalent because (3) and (4) do not agree.
In particular, the presence of the inverse metric γ − 1 in the update formula (4) effectively induces separate learning rates for each direction in tangent space.
In the following, we prove that directions orthogonal to the data manifold S ⊂ X are heavily suppressed by the inverse metric and thus gradient ascent (4) stays on the data manifold S to very good approximation.
In practice, the data manifold is only approximately of lower dimension.
Speciﬁcally, we assume that the data manifold is a product manifold, equipped with the canonical product metric, of the form
[Equation]
where D is a n -dimensional manifold and B δ is an open one-dimensional ball with radius δ (with respect to the ﬂat metric of the embedding space X ).
Since we will choose all the radii δ i to be small, the data manifold S is thus approximately n -dimensional.
We choose Gaussian normal coordinates x
[Equation]
, x 1 ⊥ , . .
, x N − n ⊥ ) on X , where the x i ⊥ are slice coordinates for B δ i and the ( x 1 (cid:107) , . .
, x n (cid:107) ) are slice coordinates of D , see Figure 2.
We furthermore require that in our coordinates x i ⊥ ( p ) ∈ ( − δ, + δ ) for p ∈ S .
We then show in the Supplement A.2:
Theorem 2.
Let p X denote the data density with supp( p X )
=
S , and the ﬂow g be well-trained such that
KL( p X , q X ) < (cid:15) , and the base density be bounded.
Let γ − 1 =
∂g∂z
∂g∂zT be the inverse of the induced metric γ in the canonical basis of coordinates x . [Figure] X
In this basis, γ − 1 is given by
[Equation]
where γ − 1 M is the inverse of the induced metric on the submanifold M ∈ {D , B δ 1 , . .
, B δ N − n } .
Furthermore, γ − 1 B δi → 0 for vanishing radius δ i → 0 .
We therefore conclude that gradient ascent in Z space corresponds to gradient ascent in X where the learning rate of all gradient components ∂ x ⊥ f orthogonal to the data manifold are effectively scaled by a vanishingly small learning rate.
As a result, the gradient ascent in Z will, to very good approximation, not leave the data manifold.
5. Experiments
Tangent Space:
A non-trivial consequence of our theoretical insights is that we can infer the tangent plane of each point on the data manifold from our ﬂow g .
Speciﬁcally, we perform a singular value decomposition of the Jacobian ∂g∂z =
U Σ V and rewrite the inverse induced metric as γ − 1 = ∂g
[Equation]
For an approximately n -dimensional data manifold S in an
N -dimensional embedding space X , Theorem 2 shows that the inverse induced metric γ − 1 has N − n small eigenvalues.
The corresponding eigenvectors of the large eigenvalues will then approximately span the tangent space of the data manifold.
In order to demonstrate this in a toy example, we train ﬂows to approximate data manifolds with the shape of a helix and torus respectively.
Figure 3 shows that we can indeed recover the tangent planes of these data manifolds to very good approximation.
We refer to the Supplement B for details about the used ﬂow and the data generation.
Diffeomorphic Explanations: We now demonstrate applications to image classiﬁcation in several domains.
The discussion is necessarily concise, see Supplement C for more details.
Datasets: We use the MNIST (Deng, 2012), CelebA (Liu et al., 2015), as well as the CheXpert datasets (Irvin et al., 2019).
The latter is a dataset of labeled chest X-rays.
Classiﬁers: We train a ten-class classiﬁer on MNIST (test accuracy of 99%).
For CelebA, we train a binary classiﬁer on the blonde attribute (test accuracy of 94%).
For
CheXpert, we train a binary classiﬁer on the cardiomegaly attribute (test accuracy of 86%).
All classiﬁers consists of a few standard convolutional, pooling and fully-connected layers with ReLU-activations and batch normalization.
Flows: We choose a ﬂow with RealNVP-type couplings (Dinh et al., 2016) for MNIST and the Glow architecture (Kingma & Dhariwal, 2018) for CelebA and CheXpert.
Estimation of Counterfactuals: We select the classes ‘nine’, ‘blonde’, and ‘cardiomegaly’ as targets k for MNIST,
CelebA, and CheXpert, respectively, and take the conﬁdence threshold to be δ = 0 .
We use Adam for optimization.
Results: Counterfactuals produced by the ﬂow indeed show semantically meaningful deformations in particular when compared to counterfactuals produced by gradient ascent in the data space X , see Figure 4.
For Figure 5, we train a linear SVM for the same classiﬁcation tasks and show that the ﬂow’s counterfactuals generalize better to such a simple model suggesting that they indeed use semantically more relevant deformations than conventional counterfactuals produced by gradient ascent in X space.
6. Related Works
An inﬂuential reference for our work is (Singla et al., 2019) which uses generative adversarial networks (GANs) to generate counterfactuals, see also (Liu et al., 2015; Samangouei et al., 2018) for similar methods.
Other approaches (Dhurandhar et al., 2018; Joshi et al., 2019) use Autoencoders instead of GANs.
While both classes of generative models can currently sample more realistic high-dimensional samples, they are not bijective.
As a result, an encoder network [Figure] has to be used which comes at the risk of mode-dropping and without any theoretical guarantees in contrast to our work.
(Sixt et al., 2021)
propose to train a linear classiﬁer in the base space of a normalizing ﬂow and show that this classiﬁer tends to use highly interpretable features.
In contrast to their approach, our method is completely model-agnostic.
In (Rombach et al., 2020; Esser et al., 2020), an invertible neural network is used to decompose latent representations of an autoencoder into semantic factors to automatically detect interpretable concepts as well as invariances of classiﬁers.
7. Conclusion
In this work, we have used the fact that a normalizing ﬂow is a diffeomorphism to map the data space to its base space.
In this space, we can then straightforwardly perform gradient ascent on the data manifold, as we have established rigorously using Riemannian differential geometry.
For future work, more high-dimensional classiﬁcation tasks will be considered as well as the dependence of the explanations on the chosen ﬂow architecture.
Furthermore, it would be interesting to evaluate the robustness of these explanations with respect to adversarial model and input manipulations (Ghorbani et al., 2019; Dombrowski et al., 2019; Anders et al., 2020; Heo et al., 2019). Acknowledgments
We want to thank the anonymous reviewers for their valuable and detailed feedback.
A.K.D. is supported by the Research Training Group “Differential Equation- and Data-driven Models in Life Sciences and Fluid Dynamics (DAEDALUS)” (GRK 2433). J.G. is supported by the
Swedish Research Council and by the Knut and Alice Wallenberg Foundation.
P.K. is supported in part by the German
Ministry for Education and Research (BMBF) under Grants 01IS14013A-E, 01GQ1115, 1GQ0850, 01IS18025A and 01IS18037A. P.K. also wants to thank Shinichi Nakajima for insightful discussions.
A. Proofs A.1. Proof of Theorem
We repeat the theorem for convenience:
Theorem.
Let z ( t ) be deﬁned as in (2) and x ( t )
=
g ( z ( t ) ) .
Then, to leading order in the learning rate λ ,
[Equation]
where γ − 1 = ∂g∂z ∂g∂zT ∈ R N,N is the pull-back of the ﬂat metric on Z under the ﬂow g .
Proof.
The step x ( t +1)
=
g ( z ( t +1) ) can be rewritten using the update formula (2) of the gradient ascent in Z as
[Equation]
We now perform a Taylor expansion to leading order in the learning rate λ using index notation as it eases notation x ( t +1) i = g ( z ( t ) ) i + λ (cid:88) j
[Equation]
l
[Equation]
The result then follows by identifying g ( z ( t ) )
=
x ( t ) and γ − 1 il = (cid:80) j ∂g i ∂z j ∂g l ∂z j which in matrix notation is given by γ −
[Equation]
∂g∂zT .
A.2. Proof of Theorem
Following the notation used throughout the main part, we denote by p X the data probability density.
In particular, it holds that the data manifold is given by S = supp( p ) .
The ﬂow g : Z → X induces the probability density q X on the target space X by push-forward of a base density q Z on the base space Z , i.e. q X
[Equation] =
[Equation]
Before giving the proof of Theorem 2, we will ﬁrst derive the following result:
Theorem 3.
Let the ﬂow be well-trained such that
[Equation]
for some small (cid:15) ∈ R .
Then, we have for the data manifold
[Equation]
Proof.
By assumption,
[Equation]
Using the deﬁnition of the KL-divergence and the inequality ln( a ) ≤ a − 1 , it the follows that
[Equation]
and thus
[Equation]
We repeat Theorem 2 for convenience:
Theorem.
Let p X denote the data density with supp( p X )
=
S , and the ﬂow g be well-trained such that
[Equation]
and the base density be bounded.
Let γ − 1 = ∂g∂z ∂g∂z
T be the inverse of the induced metric γ in the canonical basis of coordinates x .
In this basis, γ − 1 is given by
[Equation]
where γ − 1 M is the inverse of the induced metric on the submanifold M .
Furthermore, γ − 1 B δi → 0 for vanishing radius δ i → 0 .
Proof.
In the chosen coordinates, the metric γ takes the block-diagonal form (in the canonical basis)
[Equation]
see e.g. Example 13.2 of (Lee, 2013) for a proof.
In these coordinates, we can then perform the integral (10) of Theorem 3:
[Equation]
where in the second step, we have used the deﬁnition of the induced metric γ = ∂z∂x ∂z∂xT which implies that det | γ | = det | ∂z∂x | 2 .
Using the Gaussian normal coordinates, we can rewrite the integral as
[Equation]
Using the assumption that the base density q Z is bounded, i.e. q Z ( z ) ≤ C , we arrive at the inequality
[Equation]
The integral however vanishes in the limit of vanishing radius δ i since (cid:90) δ i − δ i (cid:113) | γ
[Equation]
| d x i
[Equation]
→for δ i → 0 , unless (cid:113) | γ B δi | → ∞ .
Thus for the inequality (12) to hold the metric γ B δi has to diverge in the limit of vanishing δ i .
Since the induced metric γ is block-diagonal, its inverse is given by
[Equation]
Because γ B δi ∈ R diverges for vanishing radius, it follows that γ − 1 B δi → 0 for δ i → 0 .
B. Toy Example for Tangent Space
Flow The ﬂow used for the toy example is composed of
RealNVP-type coupling layer blocks.
Each of these blocks includes a three-layer fully-connected neural network with leaky ReLU activations for the scale and translation functions.
For training, we sample from the target distribution.
We train for 5000 epochs using a batch of 500 samples per epoch.
We use the Adam optimizer with standard parameters and learning rate λ
= 1 × 10 − 4 .
This takes aroundminutes on a standard CPU.
Latent distribution We use a 3D standard Gaussian distribution as the latent distribution.
Helix To get a data sample from the helix we sample from a uniform distribution x 3 ∼ U ( − 4 , 4) and deﬁne x 1 = sin( x 3 ) and x 2 = cos( x 3 ) .
Torus We deﬁne a torus with outer radius R = 3 and unit inner radius.
To get a data sample from the Helix we sample from a uniform distribution φ, θ ∼ U (0 , 2 π ) and deﬁne x 0 = cos( θ )( R + cos( φ )) , x 1 = sin( θ )( R + cos( φ )) , and x 3 = sin( φ ) .
C. Details on Experiments C.1. Flows
Architecture: We use the RealNVP architecture 2 for
MNIST and the Glow architecture 3 for CelebA and CheXpert.
Training: We use the Adam optimizer with a learning rate of 1 × 10 − 4 and weight decay of 5 × 10 − 4 for all ﬂows.
MNIST: we train for 30 epochs on all available training images.
Bits per dimension on the test set average to 1.21.
CelebA: we train for 8 epochs on all available training images.
We use 5 bit images.
Bits per dimension on the test set average to 1.32.
CheXpert: we train for 4 epochs on all available training images.
Bits per dimension on the test set average to 3.59.
C.2. Classiﬁer
Architecture: All classiﬁers have a similar structure consisting of convolutional, pooling and fully connected layers.
We use ReLU activations and batch normalization.
For
MNIST we use four convolutional layers and three fully connected layers.
For CelebA and CheXpert we use six convolutional layers and four fully connected layers.
Training: We use the Adam optimizer with a weight decay of 5 × 10 − 4 for all classiﬁers.
MNIST: we use training and test data as speciﬁed in torchvision.
We use 10% of the training data for validation.
We train for 4 epochs using a learning rate of 1 × 10 − 3 .
We get a test accuracy of 0.99.
CelebA: we take training and test data set as speciﬁed in torchvision.
We use 10% of the training images for validation.
We scale and crop the images to 64 × 64 pixels.
We partition the data sets into all images for which the blonde attribute is positive and the rest of the images.
We treat the imbalance by undersampling the class with more examples.
We train for 10 epochs using a learning rate of 5 × 10 − 3 .
We get a balanced test accuracy of 93.63% by averaging over true positive rate (93.95%) and true negative rate (93.31%).
CheXpert: we choose the ﬁrst 6500 patients from the training set for testing.
The remaining patients are used for training.
We select the model based on performance on the original validation set.
We only consider frontal images and scale and crop the images to 128 × 128 pixels.
For the training data the cardiomegaly attibute can have four different values: blanks, 0, 1, and -1.
We label images with the blank attribute as 0 if the no ﬁnding attribute is 1, otherwise we ignore images with blank attributes.
We also ignore images where the cardiomegaly attribute is labeled as uncertain.
Using this technique, we obtain 25717 training images labelled as healthy and 20603 training images labelled as cardiomegaly.
We do not treat the imbalance but train on the data as is.
We train for 9 epochs using a learning rate of 1 × 10 − 4 .
We test on the test set, that was produced in the same way as the training set.
We get a balanced test accuracy of 86.07% by averaging over true positive rate (84.83%) and true negative rate (87.27%)
[Figure]
C.3. Optimization Counterfactuals
Counterfactuals are found using the Adam optimizer with standard parameters.
We vary only the learning rate λ .
For MNIST we use λ = 5 × 10 − 4 for conventional counterfactuals and λ
=
×
10 − 2 for counterfactuals found via the ﬂow.
We do a maximum of 2000 steps stopping early when we reach the target conﬁdence of 0.99.
We perform attacks on 500 images of the true class ‘four’.
All conventional attacks and 498 of the attacks via the ﬂow reached the target conﬁdence of 0.99 for the target class ‘nine’.
For CelebA we use λ = 7 × 10 − 4 for conventional counterfactuals and λ =
5 × 10 − 3 for counterfactuals found via the ﬂow.
We do a maximum of 1000 steps stopping early when we reach the target conﬁdence of 0.99.
We perform attacks on 500 images of the true class ‘non-blonde’.
492 conventional attacks and 496 of the attacks via the ﬂow reached the target conﬁdence of 0.99 for the target class ‘blonde’.
For CheXpert we use λ = 5 × 10 − 4 for conventional counterfactuals and λ
= 5 × 10 − 3 for counterfactuals found via the ﬂow.
We do a maximum of 1000 steps stopping early when we reach the target conﬁdence of 0.99.
We perform attacks on 1000 images of the true class ‘healthy’.
All conventional attacks and 990 of the attacks via the ﬂow reached the target conﬁdence of 0.99 for the target class ‘cardiomegaly’.
D. Examples for Counterfactuals
In this supplement, we present results on randomly selected images from the three datasets for which we produce counterfactuals via the ﬂow.
For the heatmaps, we visualize both the sum over the absolute values of color channels as well as the sum over the color channnels.