Install from PyPI. (#3)

f003c26 about 1 month ago

6.69 kB

---
license: other
license_name: fair-chemistry-license-v1
license_link: LICENSE
library_name: saddleflow
pipeline_tag: other
tags:
  - chemistry
  - materials-science
  - transition-states
  - saddle-points
  - flow-matching
  - uma
datasets:
  - SciLM/MaterialsSaddles
extra_gated_prompt: >-
  SaddleFlow's weights are a fine-tuned version of Meta's UMA-M-1.1 and are shared under the FAIR Chemistry
  License v1 (LICENSE) and the FAIR Chemistry Acceptable Use Policy (USE_POLICY.md). By requesting access you
  agree to both.
extra_gated_fields:
  I agree to the FAIR Chemistry License v1 and the FAIR Chemistry Acceptable Use Policy: checkbox
---
# SaddleFlow

SaddleFlow proposes **transition states (first-order saddle points) of periodic materials** from a single
structure. Give it a minimum (a bulk crystal or an adsorbate on a surface); it adds Gaussian noise to the
mobile atoms and integrates a learned flow to a nearby saddle. A saddle optimizer such as
[Sella](https://github.com/zadorlab/sella) then converges the proposal in a few dozen force calls instead of
a few hundred.

This is the unconditional model of the MaterialsSaddles paper, trained on all four subsets of
[MaterialsSaddles](https://huggingface.co/datasets/SciLM/MaterialsSaddles). Code:
[github.com/SciLM-ai/SaddleFlow](https://github.com/SciLM-ai/SaddleFlow).

## Use

```bash
pip install "saddleflow[search]"
hf auth login        # accept the license on this page first
python -c "import torch; print(torch.cuda.is_available())"   # should print True
```

On Linux on ARM (aarch64, e.g. NVIDIA GH200), PyPI has only CPU builds of torch, so first run
`pip install torch==2.13.0 --index-url https://download.pytorch.org/whl/cu129` (or `torch==2.8.0` from the same
index with fairchem-core 2.19–2.21). The first `from_pretrained()` downloads 5.6 GB.

**Hardware.** One GPU. On one GH200 a proposal takes about 1.8 s on its own, and 16–64 proposals batched together
run at 2–3 per second for 40–80-atom cells, using about 6 GB of GPU memory for one proposal, 12–15 GB for 16, and
26–34 GB for 64 (`batch_size`). fairchem compiles kernels on first use and caches them under `~/.triton` and
`~/.cache`; with a small home quota, point `TRITON_CACHE_DIR`, `TORCHINDUCTOR_CACHE_DIR`, and `XDG_CACHE_HOME` at
a scratch directory. A runnable example on a 61-atom battery cell is in
[`examples/quickstart/`](https://github.com/SciLM-ai/SaddleFlow/tree/main/examples/quickstart).

Propose saddles:

```python
from ase.io import read
from saddleflow import SaddleFlow
sf = SaddleFlow.from_pretrained("SciLM/SaddleFlow")          # GPU if available
atoms = read("reactant.traj")
guesses = sf.propose(atoms, n=8, seed=0, task_name="omat")   # 8 ase.Atoms, one per noise draw
```

Converge them to verified first-order saddles with Sella on UMA-S-1.2 (needs access to
[facebook/UMA](https://huggingface.co/facebook/UMA)):

```python
from fairchem.core import FAIRChemCalculator
from fairchem.core.calculate.pretrained_mlip import get_predict_unit
from saddleflow import SaddleFlowSearch, find_saddles
calc = FAIRChemCalculator(get_predict_unit("uma-s-1p2", device="cuda"), task_name="omat")
atoms.calc = calc
SaddleFlowSearch(atoms, sf, seed=0).run(fmax=0.01)           # like an ASE optimizer: atoms ends on a saddle
saddles = find_saddles(atoms, calc, n=8, model=sf, seed=0)   # distinct first-order saddles from 8 draws
```

Or from the command line:

```bash
saddleflow propose reactant.traj -o proposals.traj --n 8 --seed 0
saddleflow search  reactant.traj -o saddles.traj   --n 8 --seed 0 --task omat
```

**Tasks.** Pass `task_name="omat"` for bulk materials and `"oc20"` or `"oc22"` for adsorbates on catalyst
surfaces (or set `atoms.info["task_name"]`). These are the tasks the model was trained on; `oc22` structures
are routed through UMA-M-1.1's `oc20` expert, since UMA-M-1.1 has no `oc22` task.

**Defaults.** Starts are the input plus N(0, 0.3² Å²) on the mobile atoms, the training noise; larger
`sigma` is refused because the model degrades sharply beyond it, smaller values change little. Each
proposal spends 20 network evaluations as five passes of four Euler steps, each pass restarting the flow from
the previous endpoint (better than one pass of 20 steps). `FixAtoms` constraints are respected. Inference runs
in fp32 with TF32 disabled: TF32 moves endpoints enough to change the saddle Sella reaches.

## Performance

On 200 test reactants per subset, from element sets absent from training, Sella started from the
SaddleFlow proposal reaches a first-order saddle a median 0.123–0.179 Å (maximum per-atom displacement) from
the proposal on `lemat`, `oc20` and `oc22` (0.280 Å on `mp20bat`) and needs a median of 25–61 force calls,
versus 196–352 when started from the perturbed reactant (4.6–8.8× fewer). Most of the saddles it reaches are
not the dataset's labelled saddle: the model finds *a* nearby first-order saddle, as a Dimer or Sella search
would, not a specific one. Details in the paper.

## Model

UMA-M-1.1 (10 blocks, ℓ_max 4, 128 channels, 32 experts, 1.4 B parameters, about 50 M active per structure)
with every block trained, an equivariant time-FiLM before each block, and a depth-3 SO(3)-equivariant velocity
head (1,397 M parameters in total). Trained by flow matching for one epoch over the MaterialsSaddles training
split (67.7 M records; `oc20` ×2, `oc22` and `mp20bat` ×8) on 64 GH200 GPUs, 132,233 steps at global batch
512, fp32; these are the EMA weights. `config.json` holds the architecture and the inference defaults.
## Limitations
- Proposals are at the UMA level. Refine with DFT before using barriers quantitatively.
- Trained on periodic bulk materials and slab–adsorbate systems with the `omat`, `oc20` and `oc22` tasks;
  molecules and other tasks are out of scope.
- The 1.4 B-parameter backbone wants a GPU (see Hardware above); CPU inference works but is slow.
## License
The weights are a derivative of Meta's UMA-M-1.1 and are distributed under the **FAIR Chemistry License v1**
(`LICENSE`), including the **FAIR Chemistry Acceptable Use Policy** (`USE_POLICY.md`). The SaddleFlow code is
MIT-licensed. Publications using SaddleFlow must also acknowledge UMA (Wood et al., 2025).

## Citation

```bibtex
@inproceedings{materialssaddles2026,
  title     = {{MaterialsSaddles}: 34 Million Transition States and a Flow-Matching
               Saddle-Point Predictor for Materials},
  author    = {Baghishov, Ilgar and Jung, Sung Hoon and Henkelman, Graeme},
  booktitle = {Advances in Neural Information Processing Systems},
  volume    = {39},
  year      = {2026},
  note      = {Evaluations and Datasets Track}
}
```

	---
	license: other
	license_name: fair-chemistry-license-v1
	license_link: LICENSE
	library_name: saddleflow
	pipeline_tag: other
	tags:
	- chemistry
	- materials-science
	- transition-states
	- saddle-points
	- flow-matching
	- uma
	datasets:
	- SciLM/MaterialsSaddles
	extra_gated_prompt: >-
	SaddleFlow's weights are a fine-tuned version of Meta's UMA-M-1.1 and are shared under the FAIR Chemistry
	License v1 (LICENSE) and the FAIR Chemistry Acceptable Use Policy (USE_POLICY.md). By requesting access you
	agree to both.
	extra_gated_fields:
	I agree to the FAIR Chemistry License v1 and the FAIR Chemistry Acceptable Use Policy: checkbox
	---

	# SaddleFlow

	SaddleFlow proposes transition states (first-order saddle points) of periodic materials from a single
	structure. Give it a minimum (a bulk crystal or an adsorbate on a surface); it adds Gaussian noise to the
	mobile atoms and integrates a learned flow to a nearby saddle. A saddle optimizer such as
	[Sella](https://github.com/zadorlab/sella) then converges the proposal in a few dozen force calls instead of
	a few hundred.

	This is the unconditional model of the MaterialsSaddles paper, trained on all four subsets of
	[MaterialsSaddles](https://huggingface.co/datasets/SciLM/MaterialsSaddles). Code:
	[github.com/SciLM-ai/SaddleFlow](https://github.com/SciLM-ai/SaddleFlow).

	## Use

	```bash
	pip install "saddleflow[search]"
	hf auth login # accept the license on this page first
	python -c "import torch; print(torch.cuda.is_available())" # should print True
	```

	On Linux on ARM (aarch64, e.g. NVIDIA GH200), PyPI has only CPU builds of torch, so first run
	`pip install torch==2.13.0 --index-url https://download.pytorch.org/whl/cu129` (or `torch==2.8.0` from the same
	index with fairchem-core 2.19–2.21). The first `from_pretrained()` downloads 5.6 GB.

	Hardware. One GPU. On one GH200 a proposal takes about 1.8 s on its own, and 16–64 proposals batched together
	run at 2–3 per second for 40–80-atom cells, using about 6 GB of GPU memory for one proposal, 12–15 GB for 16, and
	26–34 GB for 64 (`batch_size`). fairchem compiles kernels on first use and caches them under `~/.triton` and
	`~/.cache`; with a small home quota, point `TRITON_CACHE_DIR`, `TORCHINDUCTOR_CACHE_DIR`, and `XDG_CACHE_HOME` at
	a scratch directory. A runnable example on a 61-atom battery cell is in
	[`examples/quickstart/`](https://github.com/SciLM-ai/SaddleFlow/tree/main/examples/quickstart).

	Propose saddles:

	```python
	from ase.io import read
	from saddleflow import SaddleFlow

	sf = SaddleFlow.from_pretrained("SciLM/SaddleFlow") # GPU if available
	atoms = read("reactant.traj")
	guesses = sf.propose(atoms, n=8, seed=0, task_name="omat") # 8 ase.Atoms, one per noise draw
	```

	Converge them to verified first-order saddles with Sella on UMA-S-1.2 (needs access to
	[facebook/UMA](https://huggingface.co/facebook/UMA)):

	```python
	from fairchem.core import FAIRChemCalculator
	from fairchem.core.calculate.pretrained_mlip import get_predict_unit
	from saddleflow import SaddleFlowSearch, find_saddles

	calc = FAIRChemCalculator(get_predict_unit("uma-s-1p2", device="cuda"), task_name="omat")

	atoms.calc = calc
	SaddleFlowSearch(atoms, sf, seed=0).run(fmax=0.01) # like an ASE optimizer: atoms ends on a saddle

	saddles = find_saddles(atoms, calc, n=8, model=sf, seed=0) # distinct first-order saddles from 8 draws
	```

	Or from the command line:

	```bash
	saddleflow propose reactant.traj -o proposals.traj --n 8 --seed 0
	saddleflow search reactant.traj -o saddles.traj --n 8 --seed 0 --task omat
	```

	Tasks. Pass `task_name="omat"` for bulk materials and `"oc20"` or `"oc22"` for adsorbates on catalyst
	surfaces (or set `atoms.info["task_name"]`). These are the tasks the model was trained on; `oc22` structures
	are routed through UMA-M-1.1's `oc20` expert, since UMA-M-1.1 has no `oc22` task.

	Defaults. Starts are the input plus N(0, 0.3² Å²) on the mobile atoms, the training noise; larger
	`sigma` is refused because the model degrades sharply beyond it, smaller values change little. Each
	proposal spends 20 network evaluations as five passes of four Euler steps, each pass restarting the flow from
	the previous endpoint (better than one pass of 20 steps). `FixAtoms` constraints are respected. Inference runs
	in fp32 with TF32 disabled: TF32 moves endpoints enough to change the saddle Sella reaches.

	## Performance

	On 200 test reactants per subset, from element sets absent from training, Sella started from the
	SaddleFlow proposal reaches a first-order saddle a median 0.123–0.179 Å (maximum per-atom displacement) from
	the proposal on `lemat`, `oc20` and `oc22` (0.280 Å on `mp20bat`) and needs a median of 25–61 force calls,
	versus 196–352 when started from the perturbed reactant (4.6–8.8× fewer). Most of the saddles it reaches are
	not the dataset's labelled saddle: the model finds a nearby first-order saddle, as a Dimer or Sella search
	would, not a specific one. Details in the paper.

	## Model

	UMA-M-1.1 (10 blocks, ℓ_max 4, 128 channels, 32 experts, 1.4 B parameters, about 50 M active per structure)
	with every block trained, an equivariant time-FiLM before each block, and a depth-3 SO(3)-equivariant velocity
	head (1,397 M parameters in total). Trained by flow matching for one epoch over the MaterialsSaddles training
	split (67.7 M records; `oc20` ×2, `oc22` and `mp20bat` ×8) on 64 GH200 GPUs, 132,233 steps at global batch
	512, fp32; these are the EMA weights. `config.json` holds the architecture and the inference defaults.

	## Limitations

	- Proposals are at the UMA level. Refine with DFT before using barriers quantitatively.
	- Trained on periodic bulk materials and slab–adsorbate systems with the `omat`, `oc20` and `oc22` tasks;
	molecules and other tasks are out of scope.
	- The 1.4 B-parameter backbone wants a GPU (see Hardware above); CPU inference works but is slow.

	## License

	The weights are a derivative of Meta's UMA-M-1.1 and are distributed under the FAIR Chemistry License v1
	(`LICENSE`), including the FAIR Chemistry Acceptable Use Policy (`USE_POLICY.md`). The SaddleFlow code is
	MIT-licensed. Publications using SaddleFlow must also acknowledge UMA (Wood et al., 2025).

	## Citation

	```bibtex
	@inproceedings{materialssaddles2026,
	title = {{MaterialsSaddles}: 34 Million Transition States and a Flow-Matching
	Saddle-Point Predictor for Materials},
	author = {Baghishov, Ilgar and Jung, Sung Hoon and Henkelman, Graeme},
	booktitle = {Advances in Neural Information Processing Systems},
	volume = {39},
	year = {2026},
	note = {Evaluations and Datasets Track}
	}
	```