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 concernsto 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.
Once being trained, LRI-induced models can detect the points inthe point cloud data that carry information indicative of the prediction label.
Suchindicative information may be reflected by either the existence of these points in thedata or the geometric locations of these points.
We also propose four datasets fromreal scientific applications in the domains of high energy physics and biochemistryto evaluate LRI.
Compared with previous post-hoc interpretation methods, thepoints detected by LRI align much better and stabler with the ground-truth patternsthat have actual scientific meanings.
LRI-induced models are also more robust tothe distribution shifts between training and test scenarios.
Introduction
The measurement of many scientific research objects can be represented as a point cloud, i e., a setof featured points in some geometric space.
For example, in high energy physics (HEP), particlesgenerated 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 spaciallocations [2, 3].
Geometric quantities of such point cloud data often encode important properties ofthe research object, analyzing which researchers may expect new scientific discoveries [4, 5].
Recently, machine learning techniques have been employed to accelerate the procedure of scientificdiscovery [6, 7].
For geometric data as above, geometric deep learning (GDL) [8, 9] has shown greatpromise 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 collisionevent in HEP that generates hundreds to thousands of particles, or a protein that consists of tens tohundreds of amino acids.
Although each particle or each amino acid is located in a low-dimensionalspace, the whole set of points eventually is extremely irregular and high-dimensional.
So, currentresearch on GDL primarily focuses on designing neural network (NN) architectures for GDL modelsto deal with the above data challenge.
GDL models have to preserve the symmetries of the system and incorporate the inductive biases reflected by geometric principles to guarantee their predictionquality [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 themhardly 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-ligandpair, it is more useful to know which groups of amino acids determine the affinity and where can bethe 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 matchscientific principles instead of some spurious correlation between the simulation environment andlabels.
Unfortunately, to the best of our knowledge, there have been no studies on interpretable GDLmodels let alone their applications in scientific problems.
Some previous post-hoc methods maybe extended to interpret a pre-trained GDL model while they suffer from some limitations as to bereviewed in Appendix A. Moreover, recent works [21–24] have shown that the data patterns detectedby post-hoc methods are often inconsistent across interpretation methods and pre-trained models, andmay hardly offer reliable scientific insights.
To fill the gap, this work proposes to study interpretable GDL models.
Inspired by the recentwork [24], we first propose a general mechanism named Learnable Randomness Injection (LRI) thatallows 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 biochemistryand provide an extensive comparison between LRI-induced GDL models and previous post-hocinterpretation 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 cloudthat is most likely to determine the label of interest.
The idea of LRI is to inject learnable randomnessto each point, where, along with training the model for label prediction, injected randomness onthe points that are important to prediction gets reduced.
The convergent amounts of randomness onpoints 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 systemor its geometric location, we propose to inject two types of randomness, Bernoulli randomness , withthe framework name LRI-Bernoulli to test existence importance of points and Gaussian randomnesson 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 indifferent directions perturbing its location affects the prediction result more.
With such fine-grainedgeometric information, we may estimate the direction of the particle velocity when analyzing particlecollision data in HEP.
LRI is theoretically sound as it essentially uses a variational objective derivedfrom the information bottleneck principle [25].
LRI-induced models also show better robustness tothe distribution shifts between training and test scenarios, which gives scientists more confidence inapplying them in practice.
We note that one obstacle to studying interpretable GDL models is the lack of valid datasets thatconsist of both classification labels and scientifically meaningful patterns to verify the quality ofinterpretation.
Therefore, another significant contribution of this work is to prepare four benchmarkdatasets grounded on real-world scientific applications to facilitate interpretable GDL research.
Thesedatasets cover important applications in HEP and biochemistry.
We briefly illustrate these fourdatasets in Fig.
1 and briefly introduce them in the following.
• ActsTrack is a particle tracking dataset in HEP that is used to reconstruct the properties, such asthe 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 trackingused in medical applications such as proton therapy [26–28].
Our task is formulated differentlyfrom traditional track reconstruction tasks: We predict the existence of a z → µµ decay and usethe 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) tostudy fine-grained geometric patterns.• Tau3Mu , another application in HEP, is to detect a challenging signature of charged leptonflavor 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 particlephysics [29, 30], therefore, any detection of them is a clear signal of new physics beyond theStandard 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 is required to capture such decays at the Large HadronCollider.
Our task is to predict the existence of τ → µµµ and use the detector hits left by the µ ’sto verify model interpretation.•
SynMol is a molecular property prediction task.
Although some works have studied modelinterpretability 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 on3D molecular representations.
Our task is to predict the property given by two functional groupscarbonyl and unbranched alkane [33] and use atoms in them to verify model interpretation.• PLBind is to predict protein-ligand binding affinities given the 3D structures of proteins andligands, which is a crucial step in drug discovery,
because a high affinity is one of the majordrug selecting criteria [35, 18].
Accurately predicting their affinities with interpretable modelsis useful for rational drug design and may help the understanding of the underlying biophysicalmechanism that enables binding [36–38].
Our task is to predict if the affinity is above a thresholdand use amino acids in the binding site of the test protein to verify model interpretation.
We evaluate LRI with three popular GDL backbone models over the above datasets.
We also extendfive baseline interpretation methods to GDL for comparison.
We find that interpretation results givenby LRI align much better with the scientific facts than those of the baselines.
Also, we observe oversome datasets, LRI-Gaussian outperforms LRI-Bernoulli while on others vice versa.
This impliesdifferent GDL applications may have different interpretation requirements.
Information indicative ofthe labels may vary regarding how it depends on the geometric features of the points.
Interestingly, wefind LRI-Gaussian can discover some fine-grained geometric patterns, such as providing high-qualityestimations of the directions of particle velocities in ActsTrack , and a high-quality estimation of thestrength of the used magnetic field.
Moreover, neither of LRI mechanisms degrades the predictionperformance of the used backbone models.
LRI mechanisms even improve model generalizationwhen there exist some distribution shifts between the training and test scenarios.
Preliminaries and Problem Formulation
In this section, we define some useful concepts and notations, while detailed review of related workscan 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 denotes3D spacial coordinates of points.
In this work, we introduce our notations by assuming the points arein 3D euclidean space while our method can be generalized.
We focus on building a classificationmodel ˆ 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.
Itlearns 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
[Figure]
data [39, 40].
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 permutationinvariant function.
Then, another function is used to aggregate all point representations to makepredictions.
Compared with graph neural networks (GNNs) [41–43] that encode graph-structureddata 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 thatsome group (e.g., E (3) , SO (3) ) invariances of the prediction can be kept [44–48].
Some models thatperform 3D convolution over 3D data also belong to the second class because the convolution kernelscan be viewed as one way to define the distance scalars for graph construction and neighborhoodaggregation [49–51].
The third class will dynamically construct the k -nn graphs based on the hiddenrepresentations of points [52, 11, 53].
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 inherentlyinterpretable model is that the model by itself can identify a subset of points C s
=
( V s , X s , r s ) thatbest 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 hitsfrom other particles.
We consider two types of indications of the label given by a point: existenceimportance , i.e., whether the point exists in the cloud is important to determine the label, and locationimportance , i.e. whether the geometric location of the point is important to determine the label.
In theabove example, the existence of the detector hits left by the three µ ’s is of course important.
On theother hand, the locations of these hits are also crucial because location features reflect the momentumof these particles when they pass through detectors, which should satisfy equations regarding certaininvariant mass if they are indeed generated from a τ → µµµ decay.
Methodology
In this section, we introduce our method Learnable Randomness Injection (LRI).
In the high-levelframework of LRI, we have an interpreter g and a classifier f .
g is used to encode the originaldata and generate randomness to perturb the data.
f is used to encode the perturbed data and makepredictions.
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 choosea GDL model architecture as the backbone to build the data encoders in g and f .
These encoderscould share or not share parameters.
Below, we will introduce specifics about this procedure.
We firstdescribe LRI-Bernoulli, where Bernoulli randomness is injected to measure the existence importantof points.
Finally, we introduce LRI-Gaussian, which injects Gaussian randomness into geometricfeatures to test the location importance of points.
We also show in Appendix B that our objectives areessentially the variational upper bounds of the information bottleneck principle [25], which makes
LRI theoretically sound.
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 euclideandistance ∥
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 hto map z v to p v ∈ [0 , 1] .
Here, h consists of an MLP plus a sigmoid layer, and samples a Bernoullimask 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-entropyloss for classification and the second term is a KL divergence regularizer.
β is the regularizationcoefficient 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 passthrough 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 byreducing the randomness generated by g , i.e., p v → 1 .
The second term, on the other hand, tends tokeep the level of randomness, i.e., p v → α .
After training, p v will converge to some value.
If theexistence of a point v is important to the label y , then p v should be close to 1 .
Otherwise, p v is closeto α .
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 eachpoint 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 onperturbed distances ∥
˜
r v − ˜ r u ∥ .
Reconstructing ˜ G is necessary because using the original graph Gwill leak information from the original geometric features r .
Finally, the classifier f takes as inputsthe 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, herethe regularization is given by a predefined Gaussian distribution N ( 0 , σ I ) , where I is an identitymatrix and σ is a hyperparameter.
[Equation]
Again, the classifier f will be optimized via the first term.
The interpreter g will be optimized via thegradients 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 adense 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 andthe reparameterization trick for multivariate Gaussian implemented by PyTorch is numericallyunstable 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 thegradients of { ϵ v } v ∈V that pass through the structure of ˜ G not computable.
We address this issue byassociating each edge v, u in ˜ G with a weight w vu ∈ (0 , 1) that monotonically decreases w.r.t.
thedistance, w vu = ϕ ( ∥ ˜ r v − ˜ r u ∥ ) .
These weights are used in the neighborhood aggregation procedurein f .
Specifically, for the central point v , f adopts aggregation AGG ( { w vu z u | u ∈ N ( v ) } ) , wherez 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 toconstruct ˜ G is “conceptually” smaller than k .
So, in practice, we choose a slightly larger number (say1.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., givenby the competition between the two terms in Eq.
The first term is to achieve good classificationperformance by reducing the randomness generated by g .
The second term, on the other hand, tendsto keep the level of randomness, i.e., Σ v → σ I .
After training, the convergent determinant | Σ v |which characterizes the entropy of injected Gaussian randomness, indicates the location importanceof point v .
We use | Σ v | ’s to rank the points v ∈ V according to their location importance.
A
B Fine-grained Interpretation on Location Importance.
Interestingly, the convergent Σ v implies more fine-grained geometric information, i.e., how differentdirections of perturbations on point v affect the prediction.
This can be analyzedby checking the eigenvectors of Σ v .
As illustrated in the figure on the right, Σ v of point v at Ais represented by the ellipses { x : x T Σ v x < θ } for different θ ’s.
It tells perturbing v towards thedirection B affects the prediction less than perturbing v towards the orthogonal direction.
As ashowcase, later, we use such fine-grained information to conduct an in-depth analysis of HEP data.
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 thedatasets, hyperparameter tuning, and method implementations, in Appendix C, D, and E, respectively.
We also report generalization performance and ablation studies of LRI in Appendix F. A summary ofdataset statistics is shown in Table 3.
Backbone Models include DGCNN [52], Point Transformer [53], and EGNN [47] which have beenwidely used for scientific GDL [11, 56, 57].
Baseline interpretation methods include three masking-based methods BernMask, BernMask-Pand PointMask, and two gradient-based methods GradGeo and GradGAM.
Masking-based methodsattempt to learn a mask ∈ [0 , 1] for each point
and may help with testing existence importance ofpoints.
Among them, BernMask and BernMask-P are post-hoc methods extended from two previousmethods on graph-structured data [58] and [59], respectively.
BernMask and BernMask-P differ inthe 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 aninherently interpretable model adopted from [60].
Thus, they are the baselines of LRI-Bernoulli.
GradGeo is extended from [61] and checks the gradients w.r.t.
geometric features, which may helpwith testing location importance, and thus is a baseline of LRI-Gaussian.
GradGAM is extended from[62] and leverages the gradients w.r.t.
the learned representations of points, which also reflects theimportance 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 compareinterpretation labels ( ∈ { 0 , 1 } ) of points with the learned point importance scores to measureinterpretation performance.
We report two metrics: interpretation ROC AUC and precision@ m .
Beyond interpretation, the model prediction performance is also measured in ROC AUC and reportedin Sec.
F.1, which is to make sure that all models are pre-trained well for those post-hoc baselinesand 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 pointcloud sample is a detector hit where a particle passes through.
We use the hits from the decay totest interpretation performance.
We evaluate all methods on the data generated with a magnetic
[Table] [Table]
field B = 2 T parallel to the z axis (see Fig.
1a).
Table 1 shows the results, where LRI-Gaussianworks consistently the best on all backbones and metrics, and LRI-Bernoulli achieves the second bestperformance.
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 outperformsboth BernMask and BernMask-P. While BernMask-P performs decently on two backbones, it fails toprovide interpretability for EGNN and its results have high variances, probably due to the unstableissue of post-hoc interpretations, as described in [24].
BernMask, PointMask, and GradGeo seemunable to provide valid interpretability for ActsTrack .
[Figure]
Fine-grained Geometric Patterns.
We find LRI-Gaussian candiscover fine-grained geometric patterns.
Intuitively, slight perturbation of each point along the underlying track direction willaffects the model prediction less than the same level of perturbation orthogonal to the track.
Therefore, the principal componentof Σ v in x-y space, i.e., the space orthogonal to the directionof the background magnetic field, can give an estimation of thetrack direction at v .
Table 2 provides the evaluation of trackdirection (velocity direction) estimation based on analyzing Σ v .
Here, we test the background magnetic field changing from 2 T to20 T. LRI-Gaussian is far more accurate than GradGeo.
The latter uses the gradients of different coordinates to compute the most sensitive direction.
Moreover, theratio between the lengths of two principal components of Σ v in x-y space gives an estimation of thecurvature of the track at v , which is proportion to the strength of the magnetic field B up to a constantmultiplier (due to the Lorentz force F ∝ B ).
Therefore, we can estimate B by analyzing Σ v .
Fig.
3shows 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, whichis similar to ActsTrack , while in Tau3Mu the µ ’s are a lot softer, and only the hits from the firstlayer of detectors are used.
Each point in a point cloud sample is a detector hit associated with alocal bending angle and a 2D coordinate in the pseudorapidity-azimuth ( η - ϕ ) space.
We use the hitsgenerated from this decay to test interpretation performance.
As shown in Table 1, LRI-Gaussian stillworks the best.
LRI-Bernoulli and GradGeo are close, and both are the second best.
While GradGAMstill 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 andcarbonyl , which together give certain synthetic properties [33, 34].
Each point in a point cloud sampleis an atom associated with a 3D coordinate and a categorical feature indicating the atom type.
We usethe atoms in these two functional groups to test interpretation performance.
As shown in Table 1,
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 furthervalidates the benefit of using 3D representations of molecules in the tasks like molecular propertyprediction.
Among other methods, GradGAM, BernMask-P and PointMask are unstable and can onlyprovide some interpretability for one or two backbones, while GradGeo and BernMask seem to failto 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 aprotein is a amino acid associated with a 3D coordinate, a categorical feature indicating the aminoacid 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 toencode the ligand and the protein separately, and the obtained two embeddings will be added to makea prediction.
As shown in Table 1, LRI-Bernoulli outperforms all other methods, while LRI-Gaussianachieves 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 havelow ROC AUC, which suggests only a part of but not the entire binding site is important to decide thebinding 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 ofexistence importance and location importance of points, and instantiated LRI with LRI-Bernoulliand LRI-Gaussian to test the two types of importance, respectively.
We observe LRI-induced modelsprovide interpretation best aligning with scientific facts, especially LRI-Gaussian that tests locationimportance.
Grounded on the IB principle, LRI never degrades model prediction performance, andmay often improve it when there exist distribution shifts between the training
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 apre-trained model is given and attempts to further analyze it to provide post-hoc interpretation.
Among them, gradient-based methods [63, 62, 64, 61, 65] may be extended to interpret geometricdata 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 [58, 59, 66, 67].
However, these methods need to mask graph structures pre-constructed by geometric features andcannot fully evaluate the effectiveness of geometric features.
Among other methods, [68, 69] studypattern selection for regular data, [70–72] utilize a local surrogate model, and [71, 73–75] leveragethe shapley value to evaluate feature importance.
These methods either cannot utilize geometricfeatures or cannot be easily applied to irregular geometric data.
Inherently Interpretable Models.
Although vanilla attention mechanisms [76, 77] were widelyused for inherent interpretability, multiple recent studies show that they cannot provide reliableinterpretation, especially for data with irregular structures [78, 79, 58, 59].
So, some works focusingon improving the attention mechanism for better interpretability [80, 24], some propose to identifyrepresentative prototypes during training [81, 82], and some methods [60, 83, 84] adopt the information bottleneck principle [25].
However, all these methods cannot analyze geometric features in
GDL.
Along another line, although invariant learning methods [85–88] based on causality analysismay also provide some interpretability, these methods are typically of great complexity and seem toprovide subpar interpretability even on graph-structured data without geometric features [24].
B 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, 89] whose goal is to reduce the mutual information between C and ˜ C whilekeeping the mutual information between ˜ C and the label, i.e., min − I ( ˜ C ; y ) + βI ( ˜ C ; C ) .
We providederivations in the following.
Grounded on the IB principle, LRI tends to extract minimal sufficientinformation to make predictions and can be more robust to distribution shifts between training andtest datasets [90, 91, 24].
B.1 Variational Bounds of the IB Principle
Let’s assume the sample C ∼ P C .
We ignore the ˜ G in our objectives to keep the notation simple, then,the IB objective is:
[Equation]
where y is the label, ˜ C is the perturbed sample, and I ( . ; . )
denotes the mutual information betweentwo random variables.
For the first term − I ( ˜ C ; y ) , by definition we have:
[Equation]
[Equation]
We introduce a variational approximation P θ ( y | ˜ C ) for P ( y | ˜ C ) as it is intractable.
Then, we yield avariational upper bound:
[Equation]
where H ( y ) is a constant, therefore, minimizing the cross-entropy loss L CE ( f θ ( ˜ C ) , y ) is equivalentto minimizing a variational upper bound of − I ( ˜ C ; y ) .
For the second term I ( ˜ C ; C ) , because ˜ C = g ϕ ( C ) , i.e., ˜ C is given by a model parameterized by ϕ , bydefinition, we have:
[Equation]
As P ( ˜ C ) is intractable, we introduce a variational approximation Q ( ˜ C ) .
Then, we yield a variationalupper 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 .
Thisprocedure gives P ϕ ( ˜ C | C ) =
(cid:81) v ∈V P ( m v | p v ) .
In this case, we define Q ( ˜ C ) as follows.
For everypoint cloud C ∼ P C , we sample m ′ v ∼ Bern ( α ) , where α ∈ [0 , 1] is a hyperparameter.
We remove allpoints in C and add points when their m ′ v = 1 .
This procedure gives Q ( ˜ C ) =
(cid:80) C P ( m ′ | C ) P C ( C ) .
As m ′ is independent from C given its size n , Q ( ˜ C )
= (cid:80) n P ( m ′ | n ) P C ( C = n ) = P ( n ) (cid:81) nv =1 P ( m ′ v ) ,where P ( n ) is a constant.
Then, we yield:
[Equation]
where c ( n, α ) is a constant, therefore, minimizing the second term of LRI-Bernoulli is equivalentto 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 thiscase, 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 equivalentto minimizing a variational upper bound of I ( ˜ C ; C ) .
Hence, the objective of LRI-Gaussian is also avariational upper bound of the IB principle.
[Table]
C Dataset Collection
We describe how we collect the four datasets in this section.
Table 3 shows the statistics of the fourdatasets.
Basic Settings.
1) All coordinates in the four datasets are centered at the origin.
2) Only positivesamples in our datasets have ground-truth interpretation labels, so we only evaluate interpretationperformance on positive samples.
For any pair of points v and u in the four datasets, they have anedge 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 Pythia8 generator [92] overlaid with soft QCD pileupevents, 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 hardscatter interaction is only generated for positive samples.
Particle tracks are simulated with a magneticfield 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 µ ’sin 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 sampleto train and test models to reduce classification difficulty, while raw files with all tracks are alsoprovided.
Each point in a sample is a detector hit left by some particle, and it is associated with a 3Dcoordinate.
As points in ActsTrack do not have any features in X , we use a dummy X with all oneswhen training models.
Momentum of particles measured by the detectors is also provided, which canhelp evaluate fine-grained geometric patterns in the data, but it is not used for model training.
Becauseeach particle on average leaves 12 hits and a model may classify samples well if it captures the trackof any one of the µ ’s from the decay, we report both precision@ 12 and precision@ 24 .
Finally, werandomly 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 generatorare used to generate the signal samples.
The background events are generated with multiple soft QCDinteractions modeled by the PYTHIA generator with a setting that resembles the collider environmentat the High Luminosity LHC.
The generated muons’ interactions with the material in the endcapmuon 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), withthe 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 ofdetectors 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 inthe pseudorapidity-azimuth ( η - ϕ ) space.
Because in the best case the model only needs to capturehits from each µ , we report precision@ 3 .
And because 80% of positive samples have less than 7 hitslabeled as ground-truth interpretation, we also report precision@ 7 .
Finally, we randomly split thedataset into training/validation/test sets with a ratio of 70 : 15 : 15 .
SynMol.
We utilize the molecules in ZINC [93], and follow [33, 34] to create synthetic propertiesbased on the existence of certain functional groups.
Specifically, if a molecule contains both theunbranched 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-truthinterpretation.
Instead of 2D representations of molecules, we associate each atom a 3D coordinateby generating a conformer for each molecule.
To do this, we first add hydrogens to the molecule andapply the ETKDG method [94].
After that, the generated structures are cleaned up using the MMFF94force field [95] with a maximum iteration of 1000 , and the added hydrogens are removed once itis 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 thetwo functional groups may only have five atoms in total, some molecules may contain multiplesuch functional groups.
So, we report both precision@ 5 and precision@ 8 ( 80% of positive sampleshave less than 8 atoms labeled as ground-truth interpretation).
Finally, we split the dataset intotraining/validation/test sets in a way that the number of molecules with or without either of thesefunctional 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 affinitiesfor a subset of complexes in the Protein Data Bank (PDB) [97].
In PDBBind, each protein-ligandpair is annotated with a dissociation constant K d , which measures the binding affinity between apair 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 augmentnegative data, during training, there is a 10% change of switching the ligand of a complex to a randomligand, and the new protein-ligand pair will be labeled as a negative sample, i.e., low affinity.
Theground-truth interpretation labels consist of two parts.
First, as shown in previous studies [98], usingthe part of the protein that is within 15 Å of the ligand is enough to even learn to generate ligands thatbind to a certain protein, so, we define the amino acids that are within 15 Å of the ligand to be thebinding 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 labelthe corresponding amino acids in the protein as the ground-truth interpretation.
Each amino acid ina 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 theyear the complexes are discovered following [100].
D Details on Hyperparameter Tuning
Gradient-based methods have no hyperparameters to tune, while all other methods are tuned basedon the validation prediction performance for a fair comparison, because tuning on interpretationperformance is impossible in a real-world setting.
All settings are trained with 5 different seeds andthe 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 sizeof 256 due to its large dataset size.
Adam [101] optimizer with a learning rate of 1 .
× 10 − 3 and aweight decay of 1 .
× 10 − 5 are used.
ActsTrack , SynMol , and PLBind construct k -nn graphs withk being 5 ; Tau3Mu constructs the graph by drawing edges for any pair of points with a distance lessthan 1 .
For a fair comparison, all models will be trained with 300 epochs on ActsTrack and SynMoland will be trained with 100 epochs on Tau3Mu and PLBind .
For post-hoc methods, a classifier willbe first pretrained with the epochs mentioned above, and those methods will further work on themodel 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 than0 .
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 fwith 200 epochs and then train both f and g with 100 epochs.
For Tau3Mu and PLBind , we eithertrain it 100 epochs from scratch or first pre-train f by 50 epochs and then train both f and g by 50epochs.
The temperature in the Gumbel-softmax trick is not tuned and is set to 1 .
[Table] [Table]
LRI-Gaussian.
β and the training time is tuned in the same way as for LRI-Bernoulli.
Instead ofdirectly 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 acrossdifferent datasets, e.g., extremely small or large.
Specifically, we rescale r by multiplying a constantc 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, cis tuned from { 200 , 300 } for ActsTrack , { 10 , 15 } for SynMol , { 7 , 11 } for Tau3Mu , and { 0 . 9 , 1 . 2 }for PLBind .
BernMask.
The size loss coefficient and the mask entropy coefficient are both tuned from{ 1 . 0 , 0 . 1 , 0 . 01 } .
After the pretraining of the classifier, the number of mask update iterations persample is set to 500 and the learning rate for each iteration is tuned from { 0 . 1 , 0 . 001 } .
BernMask-P.
After the pretraining of the classifier, the explainer module in BernMask-P will betrained with extra 100 epochs on ActsTrack and SynMol and with extra 50 epochs on Tau3Mu and
PLBind .
The size loss coefficient and the mask entropy coefficient are tuned in the same way as for
BernMask.
The temperature in the Gumbel-softmax trick is set to 1 .
PointMask.
It is trained with 300 epochs on ActsTrack and SynMol , and with 100 epochs on
Tau3Mu and PLBind .
α 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 modelsutilize implementations available in Pytorch-Geometric (PyG) [105].
LRI-Bernoulli.
Directly removing points from C to yield ˜ C makes it non-differentiable.
We providetwo differentiable ways to approximate this step.
The first way is to use another MLP to map the rawpoint features X to a latent feature space H , and yield ˜ C
= ( V , m1 T ⊙ H , r ) .
Here m is a vectorcontaining 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 meanings.
When the backbone model is implemented using a message passing scheme, e.g., using PyG, anotherpossible way is to directly use m to mask the message sent by the points to be removed.
We findboth ways can work well and we adopt the second way in our experiments.
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 workwell enough and thus we adopt it to implement ϕ .
Baseline Methods.
BernMask is extended from [58] based the authors’ code and the implementationavailable in PyG. BernMask-P is extended from [
[Table]
Discussions on LRI.
We note that both f and g in LRI needs a permutation equivariant encoder to
[Table]
learn point representations, and we find for simple tasks these two encoders can share parameters to
[Table]
reduce model size without degrading interpretation performance, while for challenging tasks usingtwo different encoders is beneficial.
In out 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 componentsin LRI is the perturbation function h , as shown in Fig. 2, and we have shown two ways to designh , i.e., Bernoulli perturbation and Gaussian perturbation.
Nonetheless, h can be generalized inmany different ways.
For example, h is where one can incorporate domain knowledge and providecontextualized interpretation results, i.e., human understandable results.
For instance, instead ofperturbing molecules in the atom level, it is possible to perturb molecules in the functional grouplevel, e.g., by averaging the learned perturbation in each functional group, so that the interpretationresults can be more contextualized and more human understandable.
In this work, we experiment thisfeature on PLBind by replacing the learned { p u } u ∈N ( v ) (cid:83) v or { Σ u } u ∈N ( v ) (cid:83) v in a neighbourhoodwith the minimum p or with the Σ having the maximum determinant in that neighbourhood, whichencourages either aggressively perturbing all amino acids in a neighborhood or not perturbing anyamino acids in the neighborhood, and we find this can better help discover binding sites for PLBind .
F Supplementary ExperimentsF.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 comparedwith 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.
Table 5shows a study with shifts on the numbers of particle tracks, where all models are trained on sampleswith 10 particle tracks, and tested on samples with a different number of (from 10 to 70 ) particletracks.
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 interpretationperformance 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 proposedmodule.
In this study, c is set to 200 on ActsTrack and to 5 on SynMol so that large perturbationsare encouraged to better show the significance of this module, β is set to 0 .
01 , and f is first trainedby 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 2, 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 guse two different encoders, β is set to 10 , c is 100 , all dropout ratio is set to 0 , and r v ∥ r v ∥ is used aspoint features.
Finally, we let g directly output covariance matrix Σ v ∈ R 2 × 2 , i.e., in x-y space, toonly perturb the first two dimensions of r .