Masking in Molecular Graphs Leveraging Reaction Context

Token masking has been proven useful for self-supervised learning in various modalities, including the sequential SMILES representation of molecules. Yet, research for masking over molecular graph structures is scarce. We propose to leverage reaction knowledge to provide critical context outside of the molecular structures themselves to guide the graph masking. We show that graph transformers are able to exploit the additional knowledge by applying a unified masking scheme, within and across molecules inside a reaction. Our experiments cover probing and transfer learning, compare to various baselines, and give insights into the intricate nature of the task. Overall, the results demonstrate the effectiveness of our approach and the usefulness of reaction context in graph pre-training more generally.