Transformers Discover Molecular Structure Without Graph Priors

TL;DR

This study demonstrates that a standard Transformer, trained directly on Cartesian coordinates without graph priors, can achieve energy and force prediction accuracy comparable to state-of-the-art equivariant GNNs on OMol25, with faster inference.

cs.LG 🔴 Advanced 2025-10-03 42 views
Tobias Kreiman Yutong Bai Fadi Atieh Elizabeth Weaver Eric Qu Aditi S. Krishnapriyan
molecular ML Transformer graph-free energy prediction scalability

Key Findings

Methodology

Using the LLaMA2 architecture, the authors remove positional embeddings and input both discretized and continuous atomic features directly. The model undergoes a pre-training stage with autoregressive causal masking to learn joint distributions of atomic positions, forces, and energies, followed by fine-tuning with bidirectional attention to predict continuous outputs. Discretization employs quantile binning, preserving real-valued information, and special tokens mark sequence boundaries. Training on OMol25 with up to 1B parameters demonstrates that the Transformer can match or surpass GNN performance in energy and force MAE, while benefiting from mature software and hardware optimizations.

Key Results

  • The 1B-parameter Transformer achieves an energy MAE of 117.99 meV and force MAE of 18.35 meV/Å on OMol25, comparable to the equivariant GNN eSEN (129.77 meV, 13.01 meV/Å), with faster training and inference.
  • Attention maps reveal physically consistent patterns, such as inverse distance decay, and adaptive receptive fields across different molecular environments.
  • Scaling to larger models follows empirical laws, indicating strong potential for further size expansion without performance saturation.

Significance

This work challenges the necessity of hard-coded graph and physical biases in molecular modeling, showing that Transformers can learn physically meaningful relationships directly from data. It opens avenues for scalable, flexible architectures capable of handling vast chemical spaces, reducing reliance on domain-specific priors, and accelerating scientific discovery in chemistry and materials science.

Technical Contribution

The paper introduces a pure Transformer architecture that processes raw Cartesian coordinates, combining discretized and continuous features without positional embeddings or graph biases. It demonstrates how large-scale training can yield models that learn physical patterns, such as distance-dependent attention, and scale predictably with data and compute. This approach broadens the scope of Transformer applications, providing a new paradigm for molecular modeling that leverages the model's inherent adaptability.

Novelty

First to employ an unmodified, standard Transformer for direct molecular energy and force prediction without any graph or physics-inspired biases, achieving competitive results at large scale. This work emphasizes the self-attention mechanism's ability to learn relational structures adaptively, marking a significant departure from traditional GNN reliance on fixed graph priors.

Limitations

  • Performance on highly complex or unusual molecules remains untested, and the model's physical interpretability needs further validation.
  • Training large models requires substantial computational resources, limiting accessibility.
  • Generalization to unseen chemical spaces and extreme conditions still poses challenges, necessitating further research.

Future Work

Future directions include integrating explicit physical constraints to improve interpretability, exploring multi-modal inputs for richer representations, and applying these models to real-world tasks like drug discovery and materials design. Additionally, scaling models beyond 1B parameters and enhancing their robustness and explainability are key goals.

AI Executive Summary

The field of molecular machine learning has traditionally relied on graph neural networks (GNNs), which encode molecules as fixed graphs with physical biases like rotational equivariance. While effective, these models face limitations in scalability, expressivity, and inference speed, especially as chemical spaces grow vast. Inspired by the success of Transformers in NLP and vision, this study investigates whether a standard Transformer architecture, trained directly on atomic Cartesian coordinates without any predefined graph or physical priors, can achieve comparable or superior performance.

Using the OMol25 dataset, which encompasses diverse molecular structures, the authors trained a 1-billion-parameter Transformer model. They employed a two-stage training process: autoregressive pre-training with causal masking to learn joint distributions of positions, forces, and energies, followed by fine-tuning with bidirectional attention for direct regression of continuous outputs. To handle real-valued features, discretization via quantile binning was combined with continuous embeddings, ensuring precise representation.

Results show that the Transformer matches the accuracy of state-of-the-art equivariant GNNs in predicting molecular energies and forces, with MAEs of approximately 118 meV and 18.35 meV/Å, respectively. Notably, the model learns physically consistent attention patterns, such as inverse-distance decay, and adapts its receptive field based on molecular environment density. Scaling experiments confirm that larger models follow predictable empirical laws, indicating strong potential for further expansion.

This research demonstrates that the core advantages of GNNs—relational learning and physical consistency—can emerge naturally in Transformers without explicit biases. It paves the way for more scalable, flexible, and data-driven approaches to molecular modeling, reducing reliance on handcrafted graph structures. Despite current limitations in complex systems and interpretability, the findings suggest a paradigm shift toward universal architectures capable of handling the chemical universe with minimal domain-specific assumptions.

Deep Analysis

Background

分子机器学习从早期的手工特征工程逐步演变到深度学习方法。早期模型如Behler-Parrinello MLIPs利用物理特征,强调局部性和对称性。近年来,GNN成为主流,利用图结构偏置实现局部关系和物理对称性,提升预测性能。然而,GNN在模型扩展、长距离依赖和硬编码偏置方面存在局限,限制了其在大规模复杂体系中的应用。Transformer在自然语言处理和视觉任务中的成功激发了其在分子建模中的潜力,但尚未充分探索其在无偏架构中的表现。

Core Problem

核心问题是,能否用标准Transformer在没有预定义图结构和物理偏置的情况下,学习到与GNN相媲美的分子能量和力的表达。传统GNN依赖局部性和对称性偏置,限制了模型的表达能力和扩展性。如何在保持高预测精度的同时,提升训练和推理效率,成为关键挑战。此外,模型的物理可解释性和泛化能力也亟待验证。

Innovation

创新点包括:1)采用无偏的标准Transformer架构,直接处理原子笛卡尔坐标,省略位置嵌入和图结构偏置;2)结合离散与连续特征编码,确保数值信息完整性;3)在大规模数据集上训练,验证模型的可扩展性和物理一致性。这一方法打破了传统GNN对硬编码偏置的依赖,展示了Transformer在分子建模中的潜力,推动了无偏架构的研究方向。

Methodology

  • �� 输入:离散和连续的原子位置、能量、力信息。• 特征编码:采用量子化分箱处理连续特征,加入特殊标记。• 模型架构:基于LLaMA2,去除位置嵌入,保持多头自注意机制。• 训练策略:预训练阶段采用自回归目标学习联合分布,微调阶段预测连续能量与力。• 损失函数:结合交叉熵和回归误差。• 训练细节:在OMol25数据集上,参数规模从5M到1B,采用相同超参数,训练10轮,处理超过2B tokens。

Experiments

  • �� 数据:OMol25,涵盖多种化学结构。• 评估指标:能量和力的MAE,与GNN模型对比。• 训练设置:在单节点H100硬件上,比较训练速度和推理延迟。• 额外测试:模型在分子动力学模拟中表现稳定,能量守恒,适合作为力场使用。

Results

  • �� 在能量预测上,Transformer达到117.99 meV的MAE,优于部分早期模型,且训练推理速度明显优于GNN。• 学习到距离反比的注意力关系,能自适应不同环境。• 模型扩展到10亿参数时,性能符合经验扩展定律,显示良好可扩展性。

Applications

  • �� 直接用于分子能量与力的快速预测,支持大规模分子模拟和药物设计。• 可作为高效的力场,用于分子动力学模拟,提升模拟速度与精度。• 未来可结合物理约束,拓展到材料科学、催化等领域。

Limitations & Outlook

  • �� 在极端复杂体系或特殊条件下的泛化能力尚未验证。•训练成本高,硬件资源需求大。•模型的物理可解释性和在未知空间的表现仍需深入研究。

Plain Language Accessible to non-experts

想象你在厨房里做菜,传统的方法像是提前准备好所有食材的配比和步骤,像是用硬性规则告诉你怎么做菜。而这项研究就像是让厨师只看食材的照片(原子坐标),不用提前设定菜谱,直接用一台智能厨师(Transformer)根据照片自己学习,逐渐掌握了不同食材之间的关系,比如距离越远的食材影响越小。结果显示,这样的智能厨师可以做出和传统厨师一样好甚至更快的菜,而且还能根据不同的厨房环境灵活调整。这种方法打破了以往依赖硬性规则的限制,让机器能更自主、更灵活地理解和处理复杂的菜肴(分子结构),未来在药物开发和新材料设计中都可能大有用处。

ELI14 Explained like you're 14

想象你在学校的科学实验室里做实验,传统的方法就像是按照老师提前给的配方一步步操作,严格按照步骤来。而这次的研究就像是让你自己观察各种化学物质(原子),不用老师告诉你具体怎么做,只用一台超级聪明的机器人(Transformer)观察这些原子,自己学习它们之间的关系。这个机器人没有提前告诉它“原子A和原子B要靠得近一点”,而是通过大量的观察,自己发现了距离越远的原子影响越小的规律。结果,这个机器人可以用来预测分子里的能量和力量,和传统的方法差不多好,还更快。它还能根据不同的分子环境,自由调整自己的“注意力”,就像你在不同的实验中学会了不同的技巧。这意味着未来我们可以用这种方法,快速设计新药或新材料,不再依赖繁琐的规则,让科学变得更智能、更灵活!

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.

cs.LG cond-mat.mtrl-sci physics.chem-ph q-bio.BM