Uncovering Internal Prediction Mechanisms of Transformer-Based Chemical Foundation Models

Supervised deep learning has become a standard approach to
deliver competitive predictive tools that allow relating the structure
of molecules and their physicochemical features to properties such as
binding to protein targets, performance as electronic materials, and
reactivity. However, efforts to understand how these models learn and
how specific predictions can be explained are still limited in
chemistry, which hinders trust in these predictions. To overcome both
the current widespread approach of considering supervised deep learning
models as black boxes and the limitations of explainable deep learning
based on feature attribution, we introduce two general methods from the
field of mechanistic interpretability to molecular property prediction.
Specifically, we leverage ablation and adapt existing techniques into
the regression lens to inspect model predictions of the
Transformer-based ChemBERTa foundation model. Our results for 3
ChemBERTa-based models finetuned on distinct datasets allow us to
propose the internal mechanisms operable within each of the
corresponding model layers that lead to the final predictions. Our
results are a stepping stone towards more trustworthy deep learning
models in the molecular domain.