Transformers Discover Molecular Structure Without Graph Priors

This is the implementation for "Transformers Discover Molecular Structure Without Graph Priors". We plan to continue rolling out improvements and updates to the code.

Environment Setup

mamba env create -f env_simple.yml
mamba activate graph-free
pip install -e .

Example Train Command

The configs expect data to be in a data folder in the directory of the repo (data/Omol/ for example). The path to data can be modified in the configs. Logs will be written to exp_logs. OMol data can be downloaded from here.

python -m mmlm.train +models=llama_57M_ch +omol_scaling_experiments=model_scaling wandb.group_name=omol_model_scaling wandb.run_name=57M training.batch_size=32 training.gradient_accumulation_steps=8

Note that metadata files (energy and force mean/std) can be found here.

Bibtex

If you find this useful, please consider citing:

@article{kreiman2025transformers,
title={Transformers Discover Molecular Structure Without Graph Priors},
author={Kreiman, Tobias and Bai, Yutong and Atieh, Fadi and Weaver, Elizabeth and Qu, Eric and Krishnapriyan, Aditi S},
journal={arXiv preprint arXiv:2510.02259},
year={2025}
}