Nahual: A Sequence Model for Language and Atoms

Code | Results

Nahual is a decoder-only autoregressive diffusion model for any sequence of text and 3D molecules. Instead of calling tools to handle 3D structures, Nahual predicts the next word or diffuses the next atom's 3D position. Nahual uses next-token diffusion, whereby a causal transformer denoises each token while attending to the clean tokens before it, without increasing context size.

Nahual is initialized from Qwen3-0.6B and trained end-to-end on twelve 3D chemistry tasks, from transition-state search and microsolvation to conformer search and adsorption on metal surfaces. A single set of 1.6B-parameter weights performs all tasks and beats state-of-the-art specialist models on crystal structure prediction (MP-20) and property-conditional molecule generation (QM9).

This repository has the checkpoints, data and evaluation samples of the NeurIPS 2026 paper Nahual: A Sequence Model for Language and Atoms.

Each token pairs a word piece or an atomic number with a 3D position. A transformer initialized from Qwen3-0.6B predicts the next discrete token, and for an atom a second causal transformer generates its position by diffusion, conditioned on the clean sequence so far (next-token diffusion). DiffMLP, the paper's ablation, replaces the second transformer with a per-token diffusion MLP and is otherwise trained the same way.

Contents

Path	Contents
checkpoints/nahual/seed{0,1}/	Nahual, two seeds. 1.58B parameters, 25 GB per checkpoint
checkpoints/diffmlp/seed{0,1}/	DiffMLP, two seeds. 1.31B parameters, 21 GB per checkpoint
data/	training, validation and test data
samples/{nahual,diffmlp}/seed{0,1}/	evaluation predictions and metrics for every task
assets/qm9/	QM9 files used by the QM9 evaluations (Ramakrishnan et al., 2014)
assets/egnn_oracles/	EGNN property oracles that score QM9 conditional generation, trained by Bao et al. (2023)
release_manifest.json	size and SHA256 of every file

Checkpoints were trained for 115,000 steps. Each checkpoint folder holds step_115000.pt, the final training checkpoint, and its config.json. The checkpoint contains the model weights, their exponential moving average (EMA), the optimizer and scheduler state, and the configuration. Inference and evaluation use the EMA weights. Training can continue from the checkpoint or start a new run from its weights.

Usage

Install the code from GitHub (pip install -e .), then:

hf download the-matter-lab/nahual --include "checkpoints/nahual/seed0/*" --local-dir hub
python scripts/download_assets.py      # Qwen3-0.6B tokenizer and initialization weights

# Generate from a JSON prompt of text/XYZ segments
python -m nahual.infer --checkpoint hub/checkpoints/nahual/seed0/step_115000.pt \
  --prompt prompt.json --output samples.json --samples 8 --steps 50

# Evaluate one task
python -m nahual.prepare download      # the data/ folder of this repository
python -m evals run --backend checkpoint --task t1x --checkpoint hub/checkpoints/nahual/seed0/step_115000.pt

A prompt is a list of segments such as [{"type": "text", "text": "..."}, {"type": "xyz", "atoms": [8, 1, 1], "coords": [[0, 0, 0], [0.96, 0, 0], [-0.24, 0.93, 0]]}], with atomic numbers and coordinates in Å. The code repository also has a browser viewer (scripts/viewer/) for the data and samples here, which can prompt a checkpoint. Its README covers training, continuing from these checkpoints, the full evaluation, and rescoring the samples.

Training

Both models were trained from Qwen3-0.6B for 115,000 steps on 32 GPUs with 24,480 packed tokens per GPU per step (90.1B tokens): AdamW at 4e-4, 5,000 warmup steps, a constant rate to step 110,000 and 5,000 linear-decay steps, loss on response tokens only, and EMA with decay 0.999. The training mixture is configs/datamix.json in the code repository.

Results

Full test sets (microsolvation: the first 500 of 1,135 prompts) with the paper's sample budgets. Values are the mean ± standard error over the two seeds, as in the paper.

Task / metric	Nahual (next-token)	DiffMLP	Test prompts
Transition state (T1x) median RMSD@1 ↓	0.197 ± 0.013	0.247 ± 0.004	1073
Noncovalent (NCI Atlas) median RMSD@10 ↓	0.335 ± 0.039	0.442 ± 0.009	149
Metal complexes (LnQM) ConnectivityMatch@1 ↑	73.3 ± 1.1	73.5 ± 0.3	488
Microsolvation (QCG-10) EwT@10, crude xTB relax ↑	50.6 ± 0.2	43.0 ± 1.4	500
Crystal structure (MP-20) MatchRate@20 ↑	79.9 ± 0.1	76.5 ± 0.3	9046
Conformers (GEOM-DRUGS) median COV-R@0.75 Å ↑	95.4 ± 0.1	93.0 ± 0.3	932
Alloy median adsRMSD@1 ↓	0.438 ± 0.055	0.810 ± 0.111	1126
FG median adsRMSD@1 ↓	0.230 ± <0.001	0.255 ± 0.001	305
OCD median adsRMSD@1 ↓	4.384 ± 0.114	4.308 ± 0.092	98
ODAC25 median adsRMSD@1 ↓	4.15 ± 0.48	4.39 ± 0.01	270
CoRE MOF pore volume MAE (cm³/g) ↓	0.115 ± 0.002	0.114 ± 0.003	1014
CoRE MOF surface area MAE (m²/g) ↓	470 ± 7	451 ± 6	1014
QMOF bandgap MAE (eV) ↓	0.465 ± 0.003	0.668 ± 0.210	2037
QM9 property α MAE (Bohr³) ↓	0.523 ± 0.016	0.524 ± 0.013	13083
QM9 property Δε MAE (meV) ↓	165 ± 1	160 ± 4	13083
QM9 property ε_HOMO MAE (meV) ↓	126 ± <1	123 ± 4	13083
QM9 property ε_LUMO MAE (meV) ↓	136 ± 1	130 ± 1	13083
QM9 property μ MAE (D) ↓	0.459 ± 0.001	0.330 ± 0.111	13083
QM9 property C_v MAE (cal/mol·K) ↓	0.245 ± 0.014	0.216 ± 0.003	13083
QM9 conditional α EGNN-MAE self-reward@5 (Bohr³) ↓	0.578 ± 0.005	0.423 ± 0.006	50000
QM9 conditional Δε EGNN-MAE self-reward@5 (meV) ↓	123.8 ± 0.4	99.7 ± 1.3	50000
QM9 conditional ε_HOMO EGNN-MAE self-reward@5 (meV) ↓	85.0 ± 0.1	67.1 ± 0.4	50000
QM9 conditional ε_LUMO EGNN-MAE self-reward@5 (meV) ↓	97.6 ± 0.2	73.4 ± 0.9	50000
QM9 conditional μ EGNN-MAE self-reward@5 (D) ↓	0.274 ± 0.001	0.171 ± 0.003	50000
QM9 conditional C_v EGNN-MAE self-reward@5 (cal/mol·K) ↓	0.296 ± 0.004	0.214 ± 0.004	50000
OC20NEB dissociations median RMSD@3 ↓	0.242 ± 0.010	0.206 ± 0.017	10
OC20NEB desorptions median RMSD@3 ↓	0.108 ± 0.006	0.096 ± 0.004	10
OC20NEB transfers median RMSD@3 ↓	0.271 ± 0.006	0.259 ± 0.006	10

The per-prompt predictions behind every number are in samples/. QM9 conditional generation is scored by the EGNN property oracles of Bao et al. after each model selects one of five samples with its own property prediction (self-reward@5). The DiffMLP QMOF bandgap MAE includes one malformed seed-0 prediction of 822 eV. Without it, the two-seed mean is 0.467 eV.

Data

data/ holds all training, validation and test data of the paper, one folder per split under data/splits/. Each split keeps the license of the dataset it comes from. data/README.md lists every source with its reference and license. All are CC BY 4.0, CC0 1.0 or MIT, except LnQM (lnqm*), which is CC BY-NC 4.0 and may only be used non-commercially. The QCG-10 microsolvation splits (qcg_solv*) are new in the paper and are released under CC BY 4.0.

Limitations

• Generated coordinates are not relaxed. Check them, for example with DFT or xTB, before downstream use.
• The model shows catastrophic forgetting of natural language, and will likely produce malformed output for prompt formats outside of the training mixture.
• Sampling is slower than with specialized models, because every atom position takes a multi-step diffusion solve.

License and attribution

The code and these weights are released under the Apache 2.0 license. The first stack is initialized from Qwen3-0.6B (Apache 2.0). The data in data/ keeps the licenses of the original datasets (see above), and the QM9 files are CC0 1.0. The EGNN property oracles are those of Bao et al., Equivariant Energy-Guided SDE for Inverse Molecular Design (ICLR 2023).

Citation

@inproceedings{cheng2026nahual,
  title     = {Nahual: A Sequence Model for Language and Atoms},
  author    = {Cheng, Austin H. and M{\"u}ller, Marcel and Burger, Andreas and Kang, Yeonghun and Ko, Tsz Wai and Skreta, Marta and Mucko, Luka and Rajaonson, Ella Miray and Ser, Cher Tian and Gonthier, J{\'e}r{\^o}me F. and Zook, Alex and Bernales, Varinia and Aspuru-Guzik, Al{\'a}n},
  booktitle = {Advances in Neural Information Processing Systems},
  year      = {2026}
}