Which chemical features are captured by ChemBERTa's attention?

Transformer-based models such as ChemBERTa have demonstrated strong performance across a wide range of molecular prediction tasks by learning contextual representations from SMILES strings. Despite this success, the internal organization of the chemical and structural information learned by these models remains poorly understood. In this work, we investigate which chemical concepts, functional groups, and structural features are captured by ChemBERTa's attention mechanisms across layers and heads. To address this question, we formulate the Attention Matrix Pattern Capture (AMPC) problem and introduce two complementary algorithms: AMPC-Chem-FG, which assesses whether attention patterns reflect predefined chemical concepts and functional groups, and AMPC-Struct, which evaluates the alignment between attention patterns and three-dimensional molecular structure represented by Coulomb matrices. Using a curated dataset derived from the ZINC database, our framework identifies specialized layer-head pairs that consistently emerge as statistical outliers and exhibit strong associations with chemically meaningful features. The results reveal specialized attention patterns related to ring structures, bond types, chirality, and functional groups. Furthermore, a layer-head pair achieves approximately 65% similarity to Coulomb-matrix representations, suggesting that ChemBERTa can capture information related to three-dimensional molecular organization despite being trained solely on SMILES strings. Finally, we demonstrate that embeddings constructed from AMPC-identified attention heads can improve downstream prediction performance compared with randomly selected heads and conventional ChemBERTa representations, highlighting the practical utility of the discovered specialization. Together, these findings provide new insights into how molecular transformers organize chemical information and offer a pathway toward more interpretable and task-informed molecular representations.