Repository logo
 

Modelling local and general quantum mechanical properties with attention-based pooling.

Published version
Peer-reviewed

Repository DOI


Change log

Authors

Kiddle, Steven J 
Oglic, Dino 
Liò, Pietro 

Abstract

Atom-centred neural networks represent the state-of-the-art for approximating the quantum chemical properties of molecules, such as internal energies. While the design of machine learning architectures that respect chemical principles has continued to advance, the final atom pooling operation that is necessary to convert from atomic to molecular representations in most models remains relatively undeveloped. The most common choices, sum and average pooling, compute molecular representations that are naturally a good fit for many physical properties, while satisfying properties such as permutation invariance which are desirable from a geometric deep learning perspective. However, there are growing concerns that such simplistic functions might have limited representational power, while also being suboptimal for physical properties that are highly localised or intensive. Based on recent advances in graph representation learning, we investigate the use of a learnable pooling function that leverages an attention mechanism to model interactions between atom representations. The proposed pooling operation is a drop-in replacement requiring no changes to any of the other architectural components. Using SchNet and DimeNet++ as starting models, we demonstrate consistent uplifts in performance compared to sum and mean pooling and a recent physics-aware pooling operation designed specifically for orbital energies, on several datasets, properties, and levels of theory, with up to 85% improvements depending on the specific task.

Description

Acknowledgements: We are thankful for the Ph.D grant and access to the Scientific Computing Platform within AstraZeneca.


Funder: AstraZeneca; doi: https://doi.org/100004325

Keywords

34 Chemical Sciences, Machine Learning and Artificial Intelligence, Networking and Information Technology R&D (NITRD)

Journal Title

Commun Chem

Conference Name

Journal ISSN

2399-3669
2399-3669

Volume Title

6

Publisher

Springer Science and Business Media LLC
Sponsorship
The first and corresponding author (David Buterez) has been supported by a fully-funded PhD grant from AstraZeneca. The funding covers yearly University fees.
Relationships
Is derived from: