Transformers Discover Molecular Structure Without Graph Priors (2510.02259v1)

Abstract: Graph Neural Networks (GNNs) are the dominant architecture for molecular machine learning, particularly for molecular property prediction and machine learning interatomic potentials (MLIPs). GNNs perform message passing on predefined graphs often induced by a fixed radius cutoff or k-nearest neighbor scheme. While this design aligns with the locality present in many molecular tasks, a hard-coded graph can limit expressivity due to the fixed receptive field and slows down inference with sparse graph operations. In this work, we investigate whether pure, unmodified Transformers trained directly on Cartesian coordinates$\unicode{x2013}$without predefined graphs or physical priors$\unicode{x2013}$can approximate molecular energies and forces. As a starting point for our analysis, we demonstrate how to train a Transformer to competitive energy and force mean absolute errors under a matched training compute budget, relative to a state-of-the-art equivariant GNN on the OMol25 dataset. We discover that the Transformer learns physically consistent patterns$\unicode{x2013}$such as attention weights that decay inversely with interatomic distance$\unicode{x2013}$and flexibly adapts them across different molecular environments due to the absence of hard-coded biases. The use of a standard Transformer also unlocks predictable improvements with respect to scaling training resources, consistent with empirical scaling laws observed in other domains. Our results demonstrate that many favorable properties of GNNs can emerge adaptively in Transformers, challenging the necessity of hard-coded graph inductive biases and pointing toward standardized, scalable architectures for molecular modeling.

• The paper introduces a Transformer model trained solely on Cartesian coordinates to predict molecular energies and forces without graph priors.
• It employs a dual input scheme with quantile-binned tokens and continuous values to accurately capture diverse molecular features.
• Empirical results show that the Transformer scales efficiently, matching equivariant GNN accuracy while providing effective uncertainty quantification.

Transformers for Molecular Modeling Without Graph Priors

Introduction

The paper "Transformers Discover Molecular Structure Without Graph Priors" (2510.02259) presents a systematic investigation into the capacity of standard Transformer architectures to learn molecular energies and forces directly from Cartesian coordinates, without relying on explicit graph-based inductive biases or physics-inspired architectural modifications. This work challenges the prevailing paradigm in molecular machine learning, where Graph Neural Networks (GNNs) with hard-coded graph construction and geometric priors have dominated the development of machine learning interatomic potentials (MLIPs). The authors demonstrate that pure Transformers, when trained on large, diverse datasets, can match the accuracy of state-of-the-art equivariant GNNs, while offering superior scalability and efficiency.

Model Architecture and Input Representation

The authors employ an unmodified LLaMA2 Transformer backbone, removing positional embeddings and introducing a dual input scheme: each token is represented by both a discretized and a continuous embedding. Discretization is performed via quantile binning, ensuring uniform bin populations for atomic positions, forces, and energies. Special tokens demarcate sequence boundaries and prediction targets. Critically, the continuous values are provided alongside the discretized tokens, enabling the model to circumvent discretization errors, especially in the heavy-tailed regions of molecular feature distributions.

Figure 1: Discretization scheme for molecular input, combining quantile-binned tokens and continuous values to mitigate discretization artifacts.

The training protocol consists of two stages: (1) autoregressive pre-training with a causal attention mask to learn the joint distribution over molecular features, and (2) fine-tuning with a bi-directional attention mask and regression heads for energy and force prediction. This approach leverages established recipes from LLM training, facilitating transferability and scalability.

Empirical Performance and Scaling Laws

On the OMol25 dataset, the 1B parameter Transformer achieves energy and force mean absolute errors (MAEs) competitive with the eSEN equivariant GNN, under a matched training compute budget. Notably, the Transformer exhibits faster training and inference in wall-clock time, attributable to the maturity of dense matrix operations and hardware/software optimizations for Transformer architectures.

The authors conduct extensive scaling experiments, demonstrating that both pre-training and fine-tuning losses follow predictable power-law relationships with model size and compute, consistent with scaling laws observed in NLP and vision domains.

Figure 2: Pre-training model scaling laws, showing power-law improvement in loss with increasing parameter count.

These results indicate no saturation in performance up to 1B parameters, suggesting that further gains are achievable with larger models and datasets, in contrast to the early saturation observed in GNN-based MLIPs.

Learned Representations: Attention Dynamics and Adaptivity

A central contribution of the paper is the analysis of learned attention patterns. In early layers, attention scores between position tokens decay steeply with interatomic distance, mirroring the locality encoded in GNN message passing. In later layers, attention shifts toward global tokens (e.g., charge, spin), enabling aggregation of global molecular properties.

Figure 3: Attention distribution by token type, illustrating the transition from local (interatomic) to global attention across layers.

Figure 4: Attention is strongly inversely correlated with interatomic distance in layers 1-9, indicating emergent locality.

The effective attention radius, defined as the minimal distance containing 90% of an atom's attention mass, is shown to adapt to the local atomic environment: atoms in sparse regions attend to more distant neighbors, while those in dense regions focus locally.

Figure 5: Effective attention radius adapts to atomic packing density, ranging from 2Å to 15Å depending on context.

This adaptivity is unattainable in GNNs with fixed cutoff radii, highlighting the flexibility of the Transformer architecture. Furthermore, individual attention heads exhibit diverse behaviors, including $k$-nearest-neighbor patterns and non-monotonic distance dependencies, which are difficult to encode a priori.

Physical Consistency and Downstream Applications

The Transformer model demonstrates approximate rotational equivariance, with force predictions in rotated frames exhibiting cosine similarity exceeding 0.99. When fine-tuned to predict conservative forces ($\mathbf{F} = -\nabla U$), the model conserves energy in NVE molecular dynamics simulations.

Figure 6: Distribution of interatomic distances $h(r)$ from MD simulations, showing close agreement between Transformer and reference models.

Figure 7: Energy conservation in NVE simulations with fine-tuned conservative force prediction.

The model's ability to learn physically consistent behaviors from data alone, without explicit symmetry constraints, is a strong empirical result. Additionally, the joint modeling of positions, forces, and energies enables uncertainty quantification: the log-probability assigned to test structures correlates with force errors in other MLIPs, providing a mechanism for out-of-distribution detection.

Figure 8: Transformer log-probabilities correlate with force errors of other MLIPs, supporting uncertainty quantification.

Implications and Future Directions

The findings challenge the necessity of hard-coded graph priors and geometric inductive biases in molecular modeling. The demonstrated scalability, efficiency, and adaptivity of Transformers suggest a path toward standardized, general-purpose architectures for chemical and materials modeling, analogous to the role of Transformers in NLP and vision.

Practically, the flexibility of the input format enables straightforward integration of multimodal data and conditioning on computational chemistry methods. The capacity for uncertainty quantification and generative modeling further broadens the applicability of Transformer-based MLIPs.

Theoretically, the results support the hypothesis that sufficiently expressive models trained on large, diverse datasets can discover domain-relevant inductive biases directly from data, obviating the need for manual architectural constraints. However, strict adherence to physical laws may still require post-hoc fine-tuning or distillation into specialized models for certain scientific applications.

Conclusion

This work provides compelling evidence that pure Transformers, trained without graph priors or physics-inspired modifications, can learn molecular structure and properties with accuracy and efficiency matching state-of-the-art GNNs. The emergent locality, adaptivity, and scaling behavior observed in the learned representations underscore the potential of standardized Transformer architectures for molecular machine learning. Future research should explore further scaling, integration of physical constraints, and extension to generative and multimodal tasks in chemistry and materials science.