sinAE — A Simplified Autoencoder + Latent Diffusion Framework for Multi-Modal 3D Structure Generation

Overview

sinAE is a single, non-equivariant Transformer framework that generates 3D atomic structures across multiple modalities with one shared model:

data_type	Modality	Coordinates	Extra tokens
0	Molecules	Cartesian	—
1	Crystals / Materials	Fractional	3 lattice tokens
2	Proteins (Cα backbone)	Cartesian	—

Generation is a two-stage pipeline:

1. Autoencoder (AE) — a flow-matching VAE (FlowMatchingModel) that encodes a structure into a continuous latent and reconstructs coordinates + atom types.
2. Latent Diffusion (LDM) — a DiT denoiser (LatentDiffusionLitModule) trained on the frozen AE latents for unconditional generation, with classifier-free guidance conditioned on the modality.

All modalities share one vocabulary of 140 atom types (120 chemical symbols + 20 amino acids) and are dispatched at runtime by an integer data_type field, so a single network can be trained on any subset of modalities — including mixed batches.

Main features

• Single non-equivariant Transformer — no explicit bonds, no equivariant layers
• One model, three modalities (molecules, crystals, proteins) via data_type dispatch
• Modular flow-matching Interpolant classes (Cartesian / torus / lattice / discrete)
• Two-stage AE + latent-diffusion training
• Built on the lightning-hydra-template

Getting started

This repo follows the lightning-hydra-template; training and model behaviour are configured through Hydra under configs/.

Installation

conda env create -f environment.yaml
conda activate mol_vq
# register the `sinae` package (editable)
pip install -e .

Data

Datasets are stored as LMDB. On first run, a dataset is built once from a preprocessed .pt file into an lmdb_* directory under data/ (this can take a while); subsequent runs load the prebuilt LMDB and its *_stats.yaml. Paths are set per experiment under configs/experiment/ and configs/datamodule/.

• Molecules (QM9 / GEOM-Drugs) and crystals (MP-20) use preprocessed .pt splits placed in data/.
• Proteins use Cα-coordinate LMDBs; the joint/protein configs point at the prebuilt protein LMDB (/scratch/yuxuan.ren/protein/data/lmdb_pdb, schema coords_ca + x).

Training

Run all commands from the repository root.

Stage 1 — Autoencoder (src/train.py)

# single modality
python src/train.py experiment=hot_qm9      trainer=gpu   # molecules
python src/train.py experiment=hot_mp_20    trainer=gpu   # crystals
python src/train.py experiment=hot_protein  trainer=gpu   # proteins

# joint molecule + crystal (+ protein)
python src/train.py experiment=hot_joint    trainer=gpu

Available AE experiments: hot_qm9, hot_geom, mild_geom, spicy_geom, hot_mp_20, hot_protein, hot_joint, hot_joint_kl{4,6,8,16,32}.

Stage 2 — Latent Diffusion (src/train_diffusion.py)

Trains the DiT denoiser on top of a (frozen) autoencoder checkpoint. The latent dimension is selected by the diffusion_experiment config (ldm_kl{4,8,16,32}) and the modality mix by the experiment config.

# joint molecule + crystal + protein latent diffusion
python src/train_diffusion.py experiment=hot_joint_protein_kl8 trainer=gpu

# evaluate a trained LDM checkpoint
python src/test_diffusion.py  experiment=hot_joint_protein_kl8 ckpt_path=path/to/ldm.ckpt

The denoiser conditions on the modality via dataset_idx (0=null/CFG, 1=molecule, 2=crystal, 3=protein); set num_datasets in the diffusion_experiment config accordingly (currently 3).

Multi-GPU training uses torchrun; trainer settings live in configs/trainer:

torchrun --nproc_per_node=2 --nnodes=1 src/train.py experiment=hot_joint trainer=ddp

Note on mixed batches: items from different modalities are padded to the per-batch maximum sequence length, so co-training short molecules/crystals with long proteins (padded to ~256) is compute-inefficient. Prefer hot_protein for protein-only runs; length-bucketed sampling is a recommended follow-up for large joint runs.

Sampling (molecules)

# unconditional sampling
python src/sample.py \
    --num_mols 1000 --num_steps 100 \
    --checkpoint path/to/model.ckpt \
    --output_path path/to/output/folder

# boosted physical plausibility (UFF-bound guidance)
python src/sample_uff_bounds.py \
    --guidance 0.01 --step-switch 90 --to-center False \
    --ckpt path/to/model.ckpt --output-dir path/to/output/folder

Crystal and protein generation are evaluated inside the LDM validation loop; protein generation metrics (e.g. designability) are a work in progress.

Repository layout

configs/                 Hydra configs (experiment / model / datamodule / trainer / ...)
src/
  train.py               Stage-1 autoencoder training
  train_diffusion.py     Stage-2 latent-diffusion training
  test_diffusion.py      LDM evaluation
  sample*.py             Molecule sampling scripts
  sinae/
    data/                LMDB datasets (mol / crystal / protein) + datamodule
    flow/                Flow-matching interpolants & paths
    models/              FlowMatchingModel (AE), LatentDiffusionLitModule (LDM), DiT, Transformer
    chem/                Atom-type constants & RDKit conversion
    callbacks/ utils/    Lightning callbacks and helpers
  eval/                  Molecule / crystal / MOF generation evaluators

Architecture notes

The model treats structure generation as a sequence-modelling problem with a standard non-equivariant Transformer (see positional encodings and the Transformer blocks). Coordinates and atom types are jointly embedded with time and positional encodings; crystals additionally prepend lattice tokens and proteins reuse the molecule path on Cα coordinates. No explicit bond information is used — the model relies on producing physically sensible coordinates so standard tools can infer bonds. The flow-matching logic is unified in a base Interpolant class with four operations (noise sampling, path creation, loss, Euler step), making it easy to mix interpolation strategies per field.