Revisiting the Form of Attention with Positional Encoding for Molecular Structures
Abstract
Molecular property prediction is a core task in machine learning for physics, with applications to interatomic potential learning and materials discovery. Recent transformer-based models achieve strong performance, but their gains often conflate two coupled choices: the design of SE(3)-aware positional encodings and how these encodings are incorporated into attention. Consequently, it remains unclear whether the empirical improvements stem from the positional encoding itself or from the attention formulation used to exploit it. Here, we present a controlled study of attention design in molecular transformers that fixes the positional encoding and only varies the attention formulation. Our results show that attention design alone has a significant effect on performance. Under a shared positional encoding, vector-valued positional encodings are most effective when used to modulate attention logits, while the best overall results are obtained when geometric information is introduced through a separate pathway that is adaptively gated by the main attention message. Our study isolates the effect of attention design, identifies key ingredients for accurate prediction of molecular properties, and offers practical guidelines for future research.