Improving Molecular Graph Neural Network Explainability with Orthonormalization and Induced Sparsity

Abstract

Rationalizing which parts of a molecule drive the predictions of a molecular graph convolutional neural network (GCNN) can be difficult. To help, we propose two simple regularization techniques to apply during the training of GCNNs: Batch Representation Orthonormalization (BRO) and Gini regularization. BRO, inspired by molecular orbital theory, encourages graph convolution operations to generate orthonormal node embeddings. Gini regularization is applied to the weights of the output layer and constrains the number of dimensions the model can use to make predictions. We show that Gini and BRO regularization can improve the accuracy of state-of-the-art GCNN attribution methods on artificial benchmark datasets. In a real-world setting, we demonstrate that medicinal chemists significantly prefer explanations extracted from regularized models. While we only study these regularizers in the context of GCNNs, both can be applied to other types of neural networks

arXiv commentFixed typo and reference

References in corpus (12)

• Adam: A Method for Stochastic Optimization
• Axiomatic Attribution for Deep Networks
• Learning Important Features Through Propagating Activation Differences
• Fast Graph Representation Learning with PyTorch Geometric
• Convolutional Networks on Graphs for Learning Molecular Fingerprints
• SmoothGrad: removing noise by adding noise
• On orthogonality and learning recurrent networks with long term dependencies
• Orthogonal Weight Normalization: Solution to Optimization over Multiple Dependent Stiefel Manifolds in Deep Neural Networks
• All You Need is Beyond a Good Init: Exploring Better Solution for Training Extremely Deep Convolutional Neural Networks with Orthonormality and Modulation
• Meta-Learning GNN Initializations for Low-Resource Molecular Property Prediction
• A comprehensive study on the prediction reliability of graph neural networks for virtual screening
• Gini in a Bottleneck: Sparse Molecular Representations for Graph Convolutional Neural Networks