[Paper Review] Language models can generate molecules, materials, and protein binding sites directly in three dimensions as XYZ, CIF, and PDB files

Daniel Flam-Shepherd, Alán Aspuru‐Guzik|arXiv (Cornell University)|May 9, 2023

Machine Learning in Materials Science21 citations

TL;DR

The paper shows that transformers trained on sequences derived from XYZ, CIF, and PDB files can directly generate 3D molecular, crystal, and protein-pocket structures, matching or approaching state-of-the-art 3D generative models without architectural changes.

ABSTRACT

Language models are powerful tools for molecular design. Currently, the dominant paradigm is to parse molecular graphs into linear string representations that can easily be trained on. This approach has been very successful, however, it is limited to chemical structures that can be completely represented by a graph -- like organic molecules -- while materials and biomolecular structures like protein binding sites require a more complete representation that includes the relative positioning of their atoms in space. In this work, we show how language models, without any architecture modifications, trained using next-token prediction -- can generate novel and valid structures in three dimensions from various substantially different distributions of chemical structures. In particular, we demonstrate that language models trained directly on sequences derived directly from chemical file formats like XYZ files, Crystallographic Information files (CIFs), or Protein Data Bank files (PDBs) can directly generate molecules, crystals, and protein binding sites in three dimensions. Furthermore, despite being trained on chemical file sequences -- language models still achieve performance comparable to state-of-the-art models that use graph and graph-derived string representations, as well as other domain-specific 3D generative models. In doing so, we demonstrate that it is not necessary to use simplified molecular representations to train chemical language models -- that they are powerful generative models capable of directly exploring chemical space in three dimensions for very different structures.

Motivation & Objective

• Motivate exploring 3D structure generation beyond graph/strings by using full chemical file formats.
• Demonstrate that next-token prediction with transformers can generate valid 3D structures from XYZ, CIF, and PDB data.
• Compare 3D generation performance against graph- and 3D-based baselines across molecules, crystals, and protein pockets.
• Assess scalability to large biomolecular structures like protein binding sites.
• Discuss implications for inverse design and future extensions to broader materials and biomolecular spaces.

Proposed method

• Represent molecules as XYZ point clouds, crystals via CIF unit cells and fractional coordinates, and protein pockets via PDB residue-atom annotations.
• Tokenize either at the character level or atom+coordinate level to form sequences for language modeling.
• Train Transformer (GPT-like) models with next-token prediction on these sequences with rotation/translation augmentation to mitigate non-invariance.
• Evaluate validity, diversity, and distributional similarity to training data using domain-specific metrics (e.g., r.m.s.d. against RDKit conformers for molecules; covariances and WA for crystals).
• Compare LM-generated 3D structures to 3D baselines (G-Schnet, ENF, EDM) and graph/SMILES-based language models across datasets.

Figure 1: A) The training datasets of structures that we benchmark language models on in this work. B) The overview of the training workflow – chemical file formats are converted to sequences of tokens using either character or coordinate-level tokenization. The language model is trained to predict

Experimental results

Research questions

• RQ1 Can language models learn to generate valid and diverse 3D molecular structures directly from XYZ-derived sequences?
• RQ2 Do models trained on CIFs and PDBs extend 3D generation capabilities to crystals and protein pockets while maintaining distributional fidelity?
• RQ3 How do 3D-capable LMs compare to state-of-the-art 3D generative models and graph-based representations on chemistry and materials benchmarks?
• RQ4 What are practical limitations (e.g., token vocabulary size, coordinate precision) and potential paths for scaling to larger biomolecules and complex materials?

Key findings

• LMs trained on XYZ sequences generate high-quality 3D molecular conformations with r.m.s.d. largely in the 1.0–2.0 range relative to RDKit conformers.
• Coordinate-level tokenization (atom+coordinate) achieves performance comparable to or better than baselines using graphs or SMILES/SELFIES.
• For crystals, language models on CIFs achieve strong validity, coverage, and property- statistics, approaching or matching CDVAE and GSchNet baselines on PEROV5 and MP20 datasets.
• For protein pockets, the LM can generate pockets with correct residue-atom composition and realistic interatomic distances, with most pockets passing validity checks and demonstrating novel residue orders.
• Overall, LMs without architectural changes can rival domain-specific 3D generative models across molecules, crystals, and pockets, and scale to larger biomolecular structures.

Figure 2: A histogram of root mean squared deviations in atomic positions between 10K molecules sampled from the language model and their corresponding conformers generated by rdkit. Six example molecules and geometries with various r.m.s.d. values are visualized explicitly and compared with their r