Interpretable Geometric Deep Learning via Learnable Randomness Injection
Abstract
Point cloud data is ubiquitous in scientific fields.
Recently, geometric deep learning (GDL) has been widely applied to solve prediction tasks with such data.
However,
GDL models are often complicated and hardly interpretable, which poses concerns to scientists when deploying these models in scientific analysis and experiments.
This work proposes a general mechanism named learnable randomness injection (LRI), which allows building inherently interpretable models based on general GDL backbones.
LRI-induced models, once being trained, can detect the points in the point cloud data that carry information indicative of the prediction label.
We also propose four datasets from real scientific applications that cover the domains of high energy physics and biochemistry to evaluate the LRI mechanism.
Compared with previous post-hoc interpretation methods, the points detected by LRI align much better and stabler with the ground-truth patterns that have actual scientific meanings.
LRI is grounded by the information bottleneck principle.
LRI-induced models also show more robustness to the distribution shifts between training and test scenarios.
Our code and datasets are available at https://github.com/Graph-COM/LRI .
Introduction
The measurement of many scientific research objects can be represented as a point cloud, i e., a set of featured points in some geometric space.
For example, in high energy physics (HEP), particles generated from collision experiments leave spacial signals on the detectors they pass through [1];
In biology, a protein is often measured and represented as a collection of amino acids with spacial locations [2, 3].
Geometric quantities of such point cloud data often encode important properties of the research object, analyzing which researchers may expect to achieve new scientific discoveries [4, 5].
Recently, machine learning techniques have been employed to accelerate the procedure of scientific discovery [6, 7].
For geometric data as above, geometric deep learning (GDL) [8, 9] has shown great promise and has been applied to the fields such as HEP [10, 11], biochemistry [12, 13] and so on.
However, geometric data in practice is often irregular and high-dimensional.
Think about a collision event in HEP that generates hundreds to thousands of particles, or a protein that consists of tens to hundreds of amino acids.
Although each particle or each amino acid is located in a low-dimensional space, the whole set of points eventually is extremely irregular and high-dimensional.
So, current research on GDL primarily focuses on designing neural network (NN) architectures for GDL models to deal with the above data challenge.
GDL models have to preserve some symmetries of the system and incorporate the inductive biases reflected by geometric principles to guarantee their prediction quality [14, 15], and therefore often involve dedicated-designed complex NN architectures.
[Figure]
Protein
BindingSite
Ligand
[Figure]
Albeit with outstanding prediction performance, the complication behind GDL models makes them hardly interpretable.
However, in many scientific applications, interpretable models are in need [16]:
For example, in drug discovery, compared with just predicting the binding affinity of a protein-ligand pair, it is more useful to know which groups of amino acids determine the affinity and where can be the binding site, as the obtained knowledge may guide future research directions [17–19].
Moreover, scientists tend to only trust interpretable models in many scenarios, e.g., most applications in HEP, where data from real experiments lack labels and models have to be trained on simulation data [20].
Here, model interpretation is used to verify if a model indeed captures the patterns that match scientific principles instead of some spurious correlation between the simulation environment and labels.
Unfortunately, to the best of our knowledge, there have been no studies on interpretable GDL models let alone their applications in scientific problems.
Some previous post-hoc methods may be extended to interpret a pre-trained GDL model while they suffer from some limitations as to be reviewed in Appendix A. Moreover, recent works [21–24] have shown that the data patterns detected by post-hoc methods are often inconsistent across interpretation methods and pre-trained models, and may hardly offer reliable scientific insights.
To fill the gap, this work proposes to study interpretable GDL models.
Inspired by the recent work [24], we first propose a general mechanism named Learnable Randomness Injection (LRI) that allows building inherently interpretable GDL models based on a broad range of GDL backbones.
We then propose four datasets from real-world scientific applications in HEP and biochemistry and provide an extensive comparison between LRI-induced GDL models and previous post-hoc interpretation approaches (after being adapted to GDL models) over these datasets.
Our LRI mechanism provides model interpretation by detecting a subset of points from the point cloud that is most likely to determine the label of interest.
The idea of LRI is to inject learnable randomness to each point, where, along with training the model for label prediction, injected randomness on the points that are important to prediction gets reduced.
The convergent amounts of randomness on points essentially reveal the importance of the corresponding points for prediction.
Specifically in
GDL, as the importance of a point may be indicated by either the existence of this point in the system or its geometric location, we propose to inject two types of randomness, Bernoulli randomness , with the framework name LRI-Bernoulli to test existence importance of points and Gaussian randomness on geometric features, with the framework name LRI-Gaussian to test location importance of points.
Moreover, by properly parameterized such Gaussian randomness, we may tell for a point, how in different directions perturbing its location affects the prediction result more.
With such fine-grained geometric information, we may estimate the direction of the particle velocity when analyzing particle collision data in HEP.
LRI is theoretically sound as it essentially uses a variational objective derived from the information bottleneck principle [25].
LRI-induced models also show better robustness to the distribution shifts between training and test scenarios, which gives scientists more confidence in applying them in practice.
We note that one obstacle to studying interpretable GDL models is the lack of valid datasets that consist of both classification labels and scientifically meaningful patterns to verify the quality of interpretation.
Therefore, another significant contribution of this work is to prepare four benchmark datasets grounded on real-world scientific applications to facilitate interpretable GDL research.
These datasets cover important applications in HEP and biochemistry.
We briefly illustrate these four datasets in Fig.
1 and introduce them as follows.
• ActsTrack is a particle tracking dataset in HEP that is used to reconstruct the properties, such as the kinematics of a charged particle given a set of position measurements from a tracking detector.
Tracking is an indispensable step in analyzing HEP experimental data as well as particle tracking used in medical applications such as proton therapy [26–28].
Our task is formulated differently from traditional track reconstruction tasks: We predict the existence of a z → µµ decay and use the set of points from the µ ’s to verify model interpretation, which can be used to reconstruct µ tracks.
ActsTrack also provides a controllable environment (e.g., magnetic field strength) to study fine-grained geometric patterns. • Tau3Mu , another application in HEP, is to detect a challenging signature of charged lepton flavor violating decays, i.e., the τ → µµµ decay, given simulated muon detector hits in protonproton collisions.
Such decays are greatly suppressed in the Standard Model (SM) of particle physics [29, 30], therefore, any detection of them is a clear signal of new physics beyond the Standard Model [31, 32].
Unfortunately, τ → µµµ contains particles of extremely low momentum, thus technologically impossible to trigger with traditional human-engineered algorithms.
Hence, online detection with advanced models that explores the correlations between signal hits on top of background hits is required to capture such decays at the Large Hadron Collider.
Our task is to predict the existence of τ → µµµ and use the detector hits left by the µ ’s to verify model interpretation.
• SynMol is a molecular property prediction task.
Although some works have studied model interpretability in such tasks [33, 34], they limit their focus on the chemical-bond-graph representations of molecules, and largely ignore their geometric features.
In this work, we put focus on 3D molecular representations.
Our task is to predict the property given by two functional groups carbonyl and unbranched alkane [33] and use atoms in these functional groups to verify model interpretation. • PLBind is to predict protein-ligand binding affinities given the 3D structures of proteins and ligands, which is a crucial step in drug discovery, because a high affinity is one of the major drug selecting criteria [35, 18].
Accurately predicting their affinities with interpretable models is useful for rational drug design and may help the understanding of the underlying biophysical mechanism that enables protein-ligand binding [36–38].
Our task is to predict whether the affinity is above a given threshold and use amino acids in the binding site of the test protein to verify model interpretation.
We evaluate LRI with three popular GDL backbone models DGCNN [39], Point Transformer [40], and EGNN [41] over the above datasets.
We also extend five baseline interpretation methods to GDL for comparison.
We find that interpretation results given by LRI align much better with the scientific facts than those of the baselines.
Also, we observe over some datasets, LRI-Gaussian outperforms LRI-Bernoulli while on others vice versa.
This implies different GDL applications may have different interpretation requirements.
Effective data patterns may vary regarding how the task depends on the geometric features of the points.
Interestingly, we find LRI-Gaussian can discover some fine-grained geometric patterns, such as providing high-quality estimations of the directions of particle velocities in ActsTrack , and a high-quality estimation of the strength of the used magnetic field.
Moreover, neither of LRI mechanisms degrades the prediction performance of the used backbone models.
LRI mechanisms even improve model generalization when there exist some distribution shifts from the training to test scenarios.
Preliminaries and Problem Formulation
In this section, we define some useful concepts and notations, while detailed review of related works can be found in Appendix A.
GDL Tasks.
We consider a data sample is a point cloud C
=
( V , X , r ) , where V = { v 1 , v 2 , ..., v n } is a set of n points, X ∈ R n × d includes d -dimensional features for all points, and r ∈ R n × 3 denotes 3D spacial coordinates of points.
In this work, we introduce our notations by assuming the points are in 3D euclidean space while our method can be generalized.
We focus on building a classification model ˆ y = f ( C ) to predict the class label y of C .
Regression tasks are left for future studies.
GDL Models.
The first class of DGL models view each sample of points as an unordered set.
It learns a dense representation z v for each v ∈ V , and then applies a permutation invariant function, e.g., sum/mean/max pooling, to aggregate all point representations so that they can handle irregular data [42, 43].
The second class of methods can better utilize geometric features and local information.
These methods first construct a k -nn graph G over the points in each sample based on their distances, e.g., ∥ r v − r u ∥ , and iteratively update the representation of point v via aggregation AGG ( { z u | u ∈ N ( v ) } ) , where N ( v ) is the neighbours of point v in graph G and AGG is a permutation invariant function.
Then, another function is used to aggregate all point representations to make predictions.
[Figure] Compared with graph neural networks (GNNs) [44–46] that encode graph-structured data without geometric features, GDL models often process geometric features carefully:
Typically, these features are transformed into some scalars such as distances, angles and used as features so that some group (e.g., E (3) , SO (3) ) invariances of the prediction can be kept [47–49, 41, 50].
Some models that perform 3D convolution over 3D data also belong to the second class because the convolution kernels can be viewed as one way to define the distance scalars for graph construction and neighborhood aggregation [51–53].
The third class will dynamically construct the k -nn graphs based on the hidden representations of points [39, 11, 40].
In this work, we focus on using the second class of models’ architectures as the backbones because most scientific applications adopt this class of models.
Interpretable Patterns in GDL.
Given a sample C =
( V , X , r ) , our goal of building an inherently interpretable model is that the model by itself can identify a subset of points C s =
( V s , X s , r s ) that best indicates the label y .
Mostly, C s will have a scientific meaning.
For example, in the task to detect τ → µµµ decay, our model should identify the detector hits left by the three µ ’s but not the hits from other particles.
We consider two types of indications of the label given by a point: existence importance , i.e., whether the point exists in the cloud is important to determine the label, and location importance , i.e. whether the geometric location of the point is important to determine the label.
In the above example, the existence of the detector hits left by the three µ ’s is of course important.
On the other hand, the locations of these hits are also crucial because location features reflect the momentum of these particles when they pass through detectors, which should satisfy equations regarding certain invariant mass if they are indeed generated from a τ → µµµ decay.
Methodology
In this section, we introduce our method Learnable Randomness Injection (LRI).
In the high-level framework of LRI, we have an interpreter g and a classifier f .
g is used to encode the original data and generate randomness to perturb the data.
f is used to encode the perturbed data and make predictions.
g and f are trained together to make accurate predictions while providing interpretability.
LRI can be applied to a large class of GDL models to make them interpretable.
We may choose a GDL model architecture as the backbone to build the data encoders in g and f .
These encoders could share or not share parameters.
Below, we will introduce specifics about this procedure.
We first describe LRI-Bernoulli, where Bernoulli randomness is injected to measure the existence important of points.
Then, we introduce LRI-Gaussian, which injects Gaussian randomness into geometric features to test the location importance of points.
Finally, we connect our objectives with the information bottleneck principle [25].
3.1 LRI-Bernoulli to Test Existence Importance
Pipeline.
Given a sample C =
( V , X , r ) , we first construct a k -nn graph G based on the euclidean distance ∥ r v − r u ∥ between every pair of points v, u ∈ V .
As shown in the top of Fig.
2, the interpreter g encodes C , generates a representation z v for each point v and uses the last component h to map z v to p v ∈ [0 , 1] .
Here, h consists of an MLP plus a sigmoid layer, and samples a Bernoulli mask for each point via m v ∼ Bern ( p v ) .
The sampling is based on a reparameterization trick [54, 55] to make dm v dp v computable.
The perturbed data ˜ C is yielded by removing the points with m v = 0 in
C .
The edges in G connected to these points are also masked and removed, which gives a graph ˜ G .
Finally, the classifier f takes as inputs ˜ C and ˜ G to make predictions.
Objective.
Eq.
1 shows our objective for each sample C , where the first term is a cross-entropy loss for classification and the second term is a KL divergence regularizer.
β is the regularization [Table] coefficient and Bern ( α ) is a predefined Bernoulli distribution with hyperparameter α < 1 .
[Equation]
Here, f is optimized via the first term.
The interpreter g is optimized via the gradients that pass through the masks { m v } v ∈V contained in ˜ C and ˜ G in the first term, and { p v } v ∈V in the second term.
Interpretation Rationale.
The interpretation is given by the competition between the two terms in
Eq. 1.
The first term is to achieve good classification performance so it tends to denoise the data ˜ C by reducing the randomness generated by g , i.e., p v → 1 .
The second term, on the other hand, tends to keep the level of randomness, i.e., p v → α .
After training, p v will converge to some value.
If the existence of a point v is important to the label y , then p v should be close to 1 .
Otherwise, p v is close to α .
We use p v ’s to rank the points v ∈ V according to their existence importance.
3.2 LRI-Gaussian to Test Location Importance
Pipeline.
We start from the same graph G as LRI-Bernoulli.
As shown in the bottom of Fig.
2, here, the interpreter g will encode the data and map it to a covariance matrix Σ v ∈ R 3 × 3 for each point v .
Gaussian randomness ϵ v ∼ N ( 0 , Σ v ) is then sampled to perturb the geometric features ˜ r v = r v
+ ϵ v of v .
Note that, to test location importance, a new k -nn graph ˜ G is constructed based on perturbed distances ∥
˜
r v − ˜ r u ∥ .
Reconstructing ˜ G is necessary because using the original graph G will leak information from the original geometric features r .
Finally, the classifier f takes as inputs the location-perturbed data ˜ C
=
( V , X , ˜ r ) and ˜ G to make predictions.
Objective.
Eq.
2 shows the objective of LRI-Gaussian for each sample C .
Different from Eq.
1, here the regularization is given by a predefined Gaussian distribution N ( 0 , σ I ) , where I is an identity matrix and σ is a hyperparameter.
[Equation]
Again, the classifier f will be optimized via the first term.
The interpreter g will be optimized via the gradients that pass through the perturbation { ϵ v } v ∈V implicitly contained in ˜ C and ˜ G in the first term, and { Σ v } v ∈V in the second term.
However, there are two technical difficulties to be addressed.
First, how to parameterize Σ v as it should be positive definite, and then how to make dϵ v d Σ v computable?
Our solution is to let the last component h in g map representation z v not to Σ v directly but to a dense matrix U v ∈ R 3 × 3 via an MLP and two scalars a 1 , a 2 ∈ R + via a softplus layer.
Then, the covariance matrix is computed by Σ v = a 1 U v U Tv + a 2 I .
Moreover, we find using Σ v and the reparameterization trick for multivariate Gaussian implemented by PyTorch is numerically unstable as it includes Cholesky decomposition.
So, instead, we use the reparameterization trick ϵ v = √ a 1 U v s 1 + √ a 2 Is 2 , where s 1 , s 2 ∼ N ( 0 , I ) .
It is not hard to show that E [ ϵ v ϵ Tv ]
= Σ v .
Second, the construction of the k -nn graph ˜ G based on ˜ r v is not differentiable, which makes the gradients of { ϵ v } v ∈V that pass through the structure of ˜ G not computable.
We address this issue by associating each edge v, u in ˜ G with a weight w vu ∈ (0 , 1) that monotonically decreases w.r.t.
the distance, w vu = ϕ ( ∥ ˜ r v − ˜ r u ∥ ) .
These weights are used in the neighborhood aggregation procedure in f .
Specifically, for the central point v , f adopts aggregation AGG ( { w vu z u | u ∈ N ( v ) } ) , where z u is the representation of the neighbor point u in the current layer.
This design makes the structure of ˜ G differentiable.
Moreover, because we set w uv < 1 , the number of used nearest neighbors to construct ˜ G is “conceptually” smaller than k .
So, in practice, we choose a slightly larger number (say 1.5 k ) of nearest neighbors to construct ˜ G and adopt the above strategy.
Interpretation Rationale.
The interpretation rationale is similar to that of LRI-Bernoulli, i.e., given by the competition between the two terms in Eq.
The first term is to achieve good classification [Table] performance by reducing the randomness generated by g .
The second term, on the other hand, tends to keep the level of randomness, i [Table] e., Σ v → σ I .
After training, the convergent determinant | Σ v | which characterizes the entropy of injected Gaussian randomness, indicates the location importance of point v .
We use | Σ v | ’s to rank the points v ∈ V according to their location importance.
[Table]
B Fine-grained Interpretation on Location Importance.
Interestingly, the convergent Σ v implies more fine-grained geometric information, i.e., how different directions of perturbations on point v affect the prediction.
This can be analyzed by checking the eigenvectors of Σ v .
As illustrated in the figure on the right, Σ v of point v at A is represented by the ellipses { x : x T Σ v x < θ } for different θ ’s.
It tells perturbing v towards the direction B affects the prediction less than perturbing v towards the orthogonal direction.
As a showcase, later, we use such fine-grained information to conduct an in-depth analysis of HEP data.
3.3 Connecting LRI and the Information Bottleneck Principle.
Our objectives Eq.
1 and Eq.
2 are essentially variational upper bounds of the information bottleneck (IB) principle [25, 56] whose goal is to reduce the mutual information between C and ˜ C while keeping the mutual information between ˜ C and the label, i.e., min − I ( ˜ C ; Y ) + βI ( ˜ C ; C ) .
We provide derivations in Appendix B. Grounded on the IB principle, LRI tends to extract minimal sufficient information to make predictions and can be more robust to distribution shifts between training and test datasets [57, 58, 24].
Benchmarking Interpretable GDL
In this section, we will evaluate LRI-Bernoulli, LRI-Gaussian and some baseline methods extended to
GDL over the proposed four datasets.
Here, we briefly describe our setup and put more details on the datasets, hyperparameter tuning, and method implementations, in Appendix C, D, and E, respectively. [Table]
We also report generalization performance and ablation studies of LRI in Appendix F. A summary of dataset statistics is shown in Table 1.
Backbone Models include DGCNN [39], Point Transformer [40], and EGNN [41] which have been widely used for scientific GDL [11, 59, 60].
Baseline interpretation methods include three masking-based methods BernMask, BernMask-P and PointMask, and two gradient-based methods GradGeo and GradGAM.
Masking-based methods attempt to learn a mask ∈ [0 , 1] for each point and may help with testing existence importance of points.
Among them, BernMask and BernMask-P are post-hoc methods extended from two previous methods on graph-structured data, i.e., from [61] and [62], respectively.
BernMask and BernMask-P differ in the way they generate the masks, where BernMask-P utilizes a parameterized mask generator while BernMask optimizes a randomly initialized mask with no other parameterization.
PointMask is an inherently interpretable model adopted from [63].
Thus, they are the baselines of LRI-Bernoulli.
GradGeo is extended from [64] and checks the gradients w.r.t.
geometric features, which may help with testing location importance, and thus is a baseline of LRI-Gaussian.
GradGAM is extended from [65] and leverages the gradients w.r.t.
the learned representations of points, which also reflects the importance of points to the prediction.
Metrics.
On each dataset, the model is trained based on the prediction (binary classification) task.
And all hyperparameters are tuned based on validation prediction performance for a fair comparison, because tuning on interpretation performance is impossible in a real-world setting.
We compare interpretation labels ( ∈ { 0 , 1 } ) of points with the learned point importance scores to measure interpretation performance.
We report two metrics: interpretation ROC AUC and precision@ m .
Beyond interpretation, the model prediction performance is also measured in ROC AUC and reported in Sec.
F.1, which is to make sure that all models are pre-trained well for those post-hoc baselines and to verify that LRI-induced models also have good prediction accuracy.
4.1 ActsTrack: End-to-end Pattern Recognition and Track Reconstruction
Here, we are to predict whether a collision event contains a z → µµ decay.
Each point in a point cloud sample is a detector hit where a particle passes through.
We use the hits from the decay to test interpretation performance.
We evaluate all methods on the data generated with a magnetic field B = 2 T parallel to the z axis (see Fig. 1a).
Table 2 shows the results, where LRI-Gaussian works consistently the best on all backbones and metrics, and LRI-Bernoulli achieves the second best performance.
This indicates that both the existence and locations of those detector hits left by the µ ’s are important to predictions.
GradGAM is also a competitive baseline, which even outperforms both BernMask and BernMask-P. While BernMask-P performs decently on two backbones, it fails to provide interpretability for EGNN and its results have high variances, probably due to the unstable issue of post-hoc interpretations, as described in [24].
BernMask, PointMask, and GradGeo seem unable to provide valid interpretability for ActsTrack .
[Figure]
Fine-grained Geometric Patterns.
We find LRI-Gaussian can discover fine-grained geometric patterns.
Intuitively, slight perturbation of each point along the underlying track direction will affects the model prediction less than the same level of perturbation orthogonal to the track.
Therefore, the principal component of Σ v in x-y space, i.e., the space orthogonal to the direction of the background magnetic field, can give an estimation of the track direction at v .
Table 3 provides the evaluation of track direction (velocity direction) estimation based on analyzing Σ v .
Here, we test the background magnetic field changing from 2 T to 20 T. LRI-Gaussian is far more accurate than GradGeo.
The latter uses the gradients of different coordinates to compute the most sensitive direction.
Moreover, the ratio between the lengths of two principal components of Σ v in x-y space gives an estimation of the curvature of the track at v , which is proportion to the strength of the magnetic field B up to a constant multiplier (due to the law – Lorentz force F ∝ B ).
Therefore, we can estimate B by analyzing Σ v .
Fig.
3 shows this approach provides an almost accurate estimation of B up to a constant multiplier.
4.2 Tau3Mu: τ → µµµ Decay Detection in Proton-Proton Collisions
This task is to predict whether a proton-proton collision event contains a τ → µµµ decay, which is similar to ActsTrack , while in Tau3Mu the µ ’s are a lot softer, and only the hits from the first layer of detectors are used.
Each point in a point cloud sample is a detector hit associated with a local bending angle and a 2D coordinate in the pseudorapidity-azimuth ( η - ϕ ) space.
We use the hits generated from this decay to test interpretation performance.
As shown in Table 2, LRI-Gaussian still works the best.
LRI-Bernoulli and GradGeo are close, and both are the second best.
While GradGAM still works well on some backbones, all masking-based methods do not perform well.
4.3 SynMol: Molecular Property Prediction with Synthetic Properties
This task is to predict whether a molecule contains both functional groups branched alkanes and carbonyl , which together give certain synthetic properties [33, 34].
Each point in a point cloud sample is an atom associated with a 3D coordinate and a categorical feature indicating the atom type.
We use the atoms in these two functional groups to test interpretation performance.
As shown in Table 2,
LRI-Gaussian performs consistently the best by only perturbing geometric features in molecules, and
LRI-Bernoulli works the second best and achieves comparable performance with LRI-Gaussian on
Point Transformer.
This shows that both the existence and locations of atoms are critical and further validates the benefit of using 3D representations of molecules in the tasks like molecular property prediction.
Among other methods, GradGAM, BernMask-P and PointMask are unstable and can only provide some interpretability for one or two backbones, while GradGeo and BernMask seem to fail to perform well on SynMol .
4.4 PLBind: Protein-ligand Binding Affinity Prediction
This task is to predict whether a protein-ligand pair is of affinity K D < 10 nM. Each point in a protein is an amino acid associated with a 3D coordinate, a categorical feature indicating the amino acid type, and two scalar features.
Each point in a ligand is an atom associated with a 3D coordinate, a categorical feature indicating the atom type, and a scalar feature.
Different from other datasets, each sample in PLBind contains two sets of points.
So, for each sample, two encoders will be used to encode the ligand and the protein separately, and the obtained two embeddings will be added to make a prediction.
As shown in Table 2, LRI-Bernoulli outperforms all other methods, while LRI-Gaussian achieves comparable performance on EGNN.
This might indicate that to make good predictions on
PLBind , the existence of certain groups of amino acids is more important than their exact locations.
Interestingly, all other methods do not seem to perform well on PLBind .
Moreover, all methods have low ROC AUC, which suggests only a part of but not the entire binding site is important to decide the binding affinity.
Conclusion
This work systematically studies interpretable GDL models by proposing a framework Learnable
Randomness Injection (LRI) and four datasets with ground-truth interpretation labels from realworld scientific applications.
We have studied interpretability in GDL from the perspectives of existence importance and location importance of points, and instantiated LRI with LRI-Bernoulli and LRI-Gaussian to test the two types of importance, respectively.
We observe LRI-induced models provide interpretation best aligning with scientific facts, especially LRI-Gaussian that tests location importance.
Grounded on the IB principle, LRI never degrades model prediction performance, and may often improve it when there exist distribution shifts between the training and test scenarios.
A Related Work
We review two categories of methods that can provide interpretability in the following.
Post-hoc Interpretation Methods.
Interpretation methods falling into this category assume a pre-trained model is given and attempts to further analyze it to provide post-hoc interpretation.
Among them gradient-based methods [66, 65, 67, 64, 68] may be extended to interpret geometric data by checking the gradients w.r t. the input features or intermediate embeddings of each point.
Some methods to interpret graph neural networks can be applied to geometric data [61, 62, 69, 70].
However, these methods need to mask graph structures pre-constructed by geometric features and cannot fully evaluate the effectiveness of geometric features.
Among other methods, [71, 72] study pattern selection for regular data, [73–75] utilize a local surrogate model, and [74, 76–78] leverage the shapley value to evaluate feature importance.
These methods either cannot utilize geometric features or cannot be easily applied to irregular geometric data.
Inherently Interpretable Models.
Although vanilla attention mechanisms [79, 80] were widely used for inherent interpretability, multiple recent studies show that they cannot provide reliable interpretation, especially for data with irregular structures [81, 82, 61, 62].
So, some works focusing on improving the attention mechanism for better interpretability [83, 24], some propose to identify representative prototypes during training [84, 85], and some methods [63, 86, 87] adopt the information bottleneck principle [25].
However, all these methods cannot analyze geometric features in
GDL.
Along another line, although invariant learning methods [88–91] based on causality analysis may also provide some interpretability, these methods are typically of great complexity and seem to provide subpar interpretability even on graph-structured data without geometric features [24].
B Variational Bounds of the IB Principle
Let’s assume the sample with its label ( C , Y ) ∼ P C × Y .
We ignore the ˜ G in our objectives to keep the notation simple, then, the IB objective is:
[Equation]
where ˜ C is the perturbed sample, and I ( . ; . )
denotes the mutual information between two random variables.
For the first term − I ( ˜ C ; Y ) , by definition we have:
[Equation]
We introduce a variational approximation P θ ( Y | ˜ C ) for P ( Y | ˜ C ) as it is intractable.
Then, we yield a variational upper bound:
[Equation]
where H ( Y ) is the entropy of Y which is a constant. We use the prediction model f paired with the cross-entropy loss L CE ( f ( ˜ C ) , Y ) to represent − E ˜ C ,Y (cid:104) log P θ ( Y | ˜ C ) (cid:105) , minimizing which is thus equivalent to minimizing a variational upper bound of − I ( ˜ C ; Y ) .
For the second term I ( ˜ C ; C ) , because ˜ C = g ( C ) . Suppose ϕ is the parameter of g . By definition, we have:
[Equation]
As P ( ˜ C ) is intractable, we introduce a variational approximation Q ( ˜ C ) .
Then, we yield a variational upper bound:
[Equation]
For LRI-Bernoulli, g ϕ takes as input C =
( V , X , r ) and first outputs p v ∈ [0 , 1] for each point v ∈ V .
Then, it samples m v ∼ Bern ( p v ) and yields ˜ C by removing all points with m v = 0 in C .
This procedure gives P ϕ ( ˜ C | C )
= (cid:81) v ∈V P ( m v | p v ) .
In this case, we define Q ( ˜ C ) as follows.
For every point cloud C ∼ P C , we sample m ′ v ∼ Bern ( α ) , where α ∈ [0 , 1] is a hyperparameter.
We remove all points in C and add points when their m ′ v = 1 .
This procedure gives Q ( ˜ C ) =
(cid:80) C P ( m ′ | C ) P ( C ) .
As m ′ is independent from C given its size n , Q ( ˜ C )
= (cid:80) n P ( m ′ | [Equation] [Equation] where P ( n ) is a constant.
Then, we yield:
[Equation]
where c ( n, α ) does not contain parameters to be optimized. Therefore, minimizing the second term of LRI-Bernoulli is equivalent to minimizing a variational upper bound of I ( ˜ C ; C ) .
Now, we can conclude that the objective of LRI-Bernoulli is a variational upper bound of the IB principle.
For LRI-Gaussian, g ϕ now takes as input C
=
( V , X , r ) and first outputs a covariance matrix Σ v ∈
R 3 × 3 for each point v ∈ V .
Then, it samples ϵ v ∼ N ( 0 , Σ v ) and yields ˜ C
=
( V , X , r + ϵ ) , where ϵ is a vector containing ϵ v for each point.
This procedure gives P ϕ ( ˜ C | C )
= (cid:81) v ∈V P ( ϵ v | Σ v ) .
In this case, we define Q ( ˜ C ) as follows.
For every point cloud C ∼ P C , we sample ϵ ′ v ∼ N ( 0 , σ I ) and yield ˜
C =
( V , X , r + ϵ ′ ) , where σ is a hyperparameter.
Similarly, we obtain Q ( ˜ C ) =
P ( n ) (cid:81) nv =1 P ( ϵ ′ v ) , where P ( n ) is a constant.
Finally, for this case, we have:
[Equation]
where c ( n, α ) is a constant, therefore, minimizing the second term of LRI-Gaussian is equivalent to minimizing a variational upper bound of I ( ˜ C ; C ) .
Hence, the objective of LRI-Gaussian is also a variational upper bound of the IB principle.
C Dataset Collection
We describe how we collect the four datasets in this section.
The statistics of the four datasets are shown in Table 1.
Basic Settings.
For each sample in the four datasets, all points are centered at the origin.
Only positive samples in our datasets have ground-truth interpretation labels, so we only evaluate interpretation performance on positive samples.
For any pair of points v and u in the four datasets, they have an edge feature of ( ∥ r v − r u ∥ , r v − r u ∥ r v − r u ∥ ) if they are connected in the constructed k -nn graph.
ActsTrack.
z → µµ events are simulated with PYTHIA generator [92] overlaid with soft QCD pileup events, and particle tracks are simulated using Acts Common Tracking Software [28].
For all samples, ten pileup interactions are generated with a center-of-mass energy of 14 TeV, and the additional hard scatter interaction is only generated for positive samples.
Particle tracks are simulated with a magnetic field parallel to the z axis, and the default ActsTrack is simulated with B = 2 T .
To make sure both µ ’s from the decay are properly measured, we calculate the invariant mass m ij for every pair of µ ’s in the generated data using Eq.
10 so that every positive sample in our dataset has at least a pair of µ ’ with an invariant mass close to 91 .
19 GeV, i.e., the mass of the z bosons.
[Equation]
Then, those detector hits left by the µ ’s from the z → µµ decay is labeled as ground-truth interpretation.
As our focus is on model interpretation performance, we only keep 10 tracks in each sample to train and test models to reduce classification difficulty, while raw files with all tracks are also provided.
In the future works, we will study a more extensive evaluation of different approaches over the samples with all tracks.
Each point in a sample is a detector hit left by some particle, and it is associated with a 3D coordinate.
As points in ActsTrack do not have any features in X , we use a dummy X with all ones when training models.
Momenta of particles measured by the detectors are also provided, which can help evaluate fine-grained geometric patterns in the data, but it is not used for model training.
Because each particle on average leaves 12 hits and a model may classify samples well if it captures the track of any one of the µ ’s from the decay, we report both precision@ 12 and precision@ 24 .
Finally, we randomly split the dataset into training/validation/test sets with a ratio of 70 : 15 : 15 .
Tau3Mu.
The τ leptons produced in decays of D and B mesons simulated by the PYTHIA generator [92] are used to generate the signal samples.
The background events are generated with multiple soft QCD interactions modeled by the PYTHIA generator with a setting that resembles the collider environment at the High Luminosity LHC.
The generated muons’ interactions with the material in the endcap muon chambers are simulated including multiple scattering effects that resemble the CMS detector.
The signal sample is mixed with the background samples at per-event level (per point cloud), with the ground-truth labels preserved for the interpretation studies.
The hits left by the µ ’s from the τ → µµµ decay are labeled as ground-truth interpretation.
We only use hits on the first layer of detectors to train models and make sure every sample in the dataset has at least three detector hits.
Each point in the sample contains measurements of a local bending angle and a 2D coordinate in the pseudorapidity-azimuth ( η - ϕ ) space.
Because in the best case the model only needs to capture hits from each µ , we report precision@ 3 .
And because 80% of positive samples have less than 7 hits labeled as ground-truth interpretation, we also report precision@ 7 .
Finally, we randomly split the dataset into training/validation/test sets with a ratio of 70 : 15 : 15 .
SynMol.
We utilize the molecules in ZINC [93], and follow [33] and [34] to create synthetic properties based on the existence of certain functional groups.
Specifically, if a molecule contains both the unbranched alkane and carbonyl, then we label it as a positive sample; otherwise it is labeled as a negative sample.
So, the atoms in branched alkanes and carbonyl are labeled as ground-truth interpretation.
Instead of 2D representations of molecules, we associate each atom a 3D coordinate by generating a conformer for each molecule.
To do this, we first add hydrogens to the molecule and apply the ETKDG method [94].
After that, the generated structures are cleaned up using the MMFF force field [95] with a maximum iteration of 1000 , and the added hydrogens are removed once it is finished.
Both ETKDG and MMFF94 are implemented using RDKit.
Besides a 3D coordinate, each point in a sample also has a categorical feature indicating the atom type.
Even though the two functional groups may only have five atoms in total, some molecules may contain multiple such functional groups.
So, we report both precision@ 5 and precision@ 8 ( 80% of positive samples have less than 8 atoms labeled as ground-truth interpretation).
Finally, we split the dataset into training/validation/test sets in a way that the number of molecules with or without either of these functional groups is uniformly distributed following [33] so that the dataset bias is minimized.
PLBind.
We utilize protein-ligand complexes from PDBBind [96], which annotates binding affinities for a subset of complexes in the Protein Data Bank (PDB) [97].
In PDBBind, each protein-ligand pair is annotated with a dissociation constant K d , which measures the binding affinity between a pair of protein and ligand.
We use a threshold of 10 nM on K d to obtain a binary classification task, and the model interpretability is studied on the protein-ligand pairs with high affinities.
To augment negative data, during training, there is a 10% change of switching the ligand of a complex to a random ligand, and the new protein-ligand pair will be labeled as a negative sample, i.e., low affinity.
The ground-truth interpretation labels consist of two parts.
First, as shown in previous studies [98], using the part of the protein that is within 15 Å of the ligand is enough to even learn to generate ligands that bind to a certain protein, so, we define the amino acids that are within 15 Å of the ligand to be the binding site and label them as ground-truth interpretation.
Second, we retrieve all atomic contacts (hydrogen bond and hydrophobic contact) for every protein-ligand pair from PDBsum [99] and label the corresponding amino acids in the protein as the ground-truth interpretation.
Each amino acid in a protein is associated with a 3D coordinate, the amono acid type, solvent accessible surface area (SASA), and the B-factor.
Each atom in a ligand is associated with a 3D coordinate, the atom type, and Gasteiger charges.
Finally, we split the dataset into training/validation/test sets according to the year the complexes are discovered following [100].
D Details on Hyperparameter Tuning
All hyperparameters are tuned based on validation classification AUC for a fair comparison.
All settings are trained with 5 different seeds and the average performance on the 5 seeds are reported.
Basic Settings.
We use a batch size of 128 on all datasets, except on Tau3Mu
we use a batch size of 256 due to its large dataset size.
The Adam [101] optimizer with a learning rate of 1 .
× 10 − 3 and a weight decay of 1 .
× 10 − 5 are used.
ActsTrack , SynMol , and PLBind construct k -nn graphs with k being 5 ; Tau3Mu constructs the graph by drawing edges for any pair of points with a distance less than 1 .
For a fair comparison, all models will be trained with 300 epochs on ActsTrack and SynMol and will be trained with 100 epochs on Tau3Mu and PLBind , so that all models are converged.
For post-hoc methods, a classifier will be first pretrained with the epochs mentioned above, and those methods will further work on the model from the epoch with the best validation classification AUC during pretraining.
LRI-Bernoulli.
β is tuned from { 1 . 0 , 0 . 1 , 0 . 01 } after normalizing the summation of the KL divergence by the total number of points.
α is tuned from { 0 . 5 , 0 . 7 } .
Note that α should not be less than
[Table] [Table] 5 , because Bern (0 . 5) injects the most randomness from an information-theoretic perspective.
For
ActsTrack and SynMol , we either train LRI-Bernoulli 300 epochs from scratch or first pre-train f with 200 epochs and then train both f and g with 100 epochs (also based on the best validation prediction performance).
Similarly, for Tau3Mu and PLBind , we either train it 100 epochs from scratch or first pre-train f by 50 epochs and then train both f and g by 50 epochs.
The temperature in the Gumbel-softmax trick is not tuned and is set to 1 .
LRI-Gaussian.
β and the training time is tuned in the same way as for LRI-Bernoulli.
Instead of directly tuning σ , we tune it by rescaling the magnitude of the geometric features to make sure the KL divergence is stable for optimization, because the magnitude of r may vary significantly across different datasets, e.g., extremely small or large.
Specifically, we keep σ = 1 and rescale r by multiplying a constant c on it so that 1 would be roughly the 5 th or 10 th percentile of the entries in abs ( c · r ) .
In this way, c is tuned from { 200 , 300 } for ActsTrack , { 10 , 15 } for SynMol , { 7 , 11 } for Tau3Mu , and { 0 . 9 , 1 . 2 } for PLBind .
BernMask.
This approach generalizes GNNExplainer [61] and learns a node mask m ∈ [0 , 1] n for each sample C . Similar to GNNExplainer, there are two regularizers in its loss function, that is, mask size and entropy constraints. Specifically, the coefficient of the ℓ 1 penalty on the normalized mask size, i.e., ∥ m ∥ 1 n , is tuned from { 1 . 0 , 0 . 1 , 0 . 01 } . Element-wise entropy of m is applied, i.e.,n (cid:80) i − m i log m
i − (1 − m i ) log(1 − m i ) and m i is the i th entry of m , to encourage generating discrete masks, and the coefficient of this entropy constraint is tuned from { 1 . 0 , 0 . 1 , 0 .
After pretraining the classifier, we set the number of iterations to learn masks per sample as 500 and tune the learning rate for each iteration from { 0 . 1 , 0 . 001 } .
BernMask-P. This approach generalizes PGExplainer [62] and learns a node mask m ∈ [0 , 1] n for each sample C . Similar to BernMask, it also has mask size and entropy constraints in its loss function, and their coefficients are both tuned from { 1 . 0 , 0 . 1 , 0 . 01 } as well.
After pretraining the classifier, the explainer module in BernMask-P will be trained with extra 100 epochs on ActsTrack and SynMol and with extra 50 epochs on Tau3Mu and PLBind .
The temperature in the Gumbel-softmax trick is set to 1 .
PointMask [63] .
It is trained with 300 epochs on ActsTrack and SynMol , and with 100 epochs on
Tau3Mu and PLBind .
It has two hyperparameters as specified in their paper, i.e., α and t , where α is tuned from { 1 . 0 , 0 . 1 , 0 . 01 } and t is tuned from { 0 . 2 , 0 . 5 , 0 . 8 } .
E Implementation Details
Backbone Models.
All backbone models have 4 layers with a hidden size of 64 .
Batch normalization [102] and ReLU activation [103] are used.
All MLP layers use a dropout [104] ratio of 0 . 2 .
SUM pooling is used to generate point cloud embeddings to make predictions.
All backbone models utilize implementations available in Pytorch-Geometric (PyG) [105].
LRI-Bernoulli.
Directly removing points from C to yield ˜ C makes it non-differentiable.
We provide two differentiable ways to approximate this step.
The first way is to use another MLP to map the raw point features X to a latent feature space H , and yield ˜ C
= ( V , ( m1 T ) ⊙ H , r ) .
Here m is a vector containing m v for each point v ∈ V , 1 is a vector of all ones, and ⊙ denotes element-wise product.
This is because masking on X or r is inappropriate as values in them have specific physical [Table] meanings.
When the backbone model is implemented by using a message passing scheme, e.g., using PyG, another possible way is to use m to mask the message sent by the points to be removed.
We find both ways can work well and we adopt the second way in our experiments [Table]
LRI-Gaussian.
We use a softplus layer to output a 1 and a 2 to parameterize the Gaussian distribution.
To make it numerically stable we clip the results of the softplus layer to [1 .
× 10 − 6 , 1 .
× 10 6 ] .
For ϕ in the differentiable graph reconstruction module, we find it empirically a simple MLP can work well enough and thus we adopt it to implement ϕ .
Baseline Methods.
BernMask is extended from [61] based the authors’ code and the implementation available in PyG. BernMask-P is extended from [62] based on the authors’ code and a recent PR in
PyG. PointMask is reproduced based on the authors’ code. GradGAM and GradGeo are extended based on the code from [106].
Discussions on LRI.
We note that both f and g in LRI needs a permutation equivariant encoder to learn point representations, and we find for simple tasks these two encoders can share parameters to reduce model size without degrading interpretation performance, while for challenging tasks using two different encoders is beneficial.
In our experiments we use the same encoder for ActsTrack ,
Tau3Mu , and SynMol , and use two different encoders for PLBind .
If the model size is not a concern, using two encoders is generally recommended.
The other thing is that one of the key components in LRI is the perturbation function h , as shown in Fig. 2, and we have shown two ways to design h , i.e., Bernoulli perturbation and Gaussian perturbation.
Nonetheless, h can be generalized in many different ways.
For example, h is where one can incorporate domain knowledge and provide contextualized interpretation results, i.e., human understandable results.
For instance, instead of perturbing molecules in the atom level, it is possible to perturb molecules in the functional group level, e.g., by averaging the learned perturbation in each functional group, so that the interpretation results can be more contextualized and more human understandable.
In this work, we experiment this feature on PLBind by replacing the learned { p u } u ∈N ( v ) (cid:83) v or { Σ u } u
[Table] a [Table] with the minimum p or with the Σ having the maximum determinant in that neighbourhood, which encourages either aggressively perturbing all amino acids in a neighborhood or not perturbing any amino acids in the neighborhood, and we find this can better help discover binding sites for PLBind .
F Supplementary Experiments F.1 Generalization Performance of LRI
LRI-induced models can generalize better while being interpretable.
As shown in Table 4, both LRIinduced models never degrade prediction performance and sometimes may even boost it compared with models trained without LRI, i.e., using empirical risk minimization (ERM).
Moreover, LRIinduced models are more robust to distribution shifts as LRI is grounded on the IB principle.
Tableshows a study with shifts on the numbers of particle tracks, where all models are trained on samples with 10 particle tracks, and tested on samples with a different number of (from 10 to 70 ) particle tracks.
We observe LRI-induced models work consistently better than models naively trained without
LRI.
F.2 Ablation Studies
We also conduct ablation studies on the differentiable graph reconstruction module proposed specifically for geometric data in LRI-Gaussian, we find that without this module the interpretation performance of LRI-Gaussian may be reduced up to 49% on ActsTrack and up to 23% on SynMol , as shown in Table 6, which matches our expectations and validates the necessity of the proposed module.
In this study, c is set to 200 on ActsTrack and to 5 on SynMol so that large perturbations are encouraged to better show the significance of this module, β is set to 0 .
01 , and f is first trained by 200 epochs and then f and g are trained together by 100 epochs.
F.3 Fine-grained Geometric Pattern Learning
To discover fine-grained geometric information in ActsTrack using LRI-Gaussian as shown in
Table 3, we conduct experiments based on DGCNN, where f is first trained with 100 epochs, and then g is further trained with 500 epochs while f is finetuned with a learning rate of 1 e - 8 .
f and g use two different encoders, β is set to 10 , c is 100 , all dropout ratio is set to 0 , and r v ∥ r v ∥ is used as point features.
Finally, we let g directly output covariance matrix Σ v ∈ R 2 × 2 , i.e., in x-y space, to only perturb the first two dimensions of r .