ml-committee-uncertainty

Files (4)

• SKILL.md
• README.md · examples/mace-lipo4-committee/
• uncertainty_summary.json · examples/mace-lipo4-committee/
• run_committee_inference.py · scripts/

SKILL.md (first lines)

# MACE Committee Model Uncertainty Quantification

<!-- mcp-tools-note -->
> [!NOTE]
> Steps written `server.tool` are MCP tool calls: `base.search_model_registry` is the `search_model_registry`
> tool of the `base` server (`mcp__base__search_model_registry`, or
> `mcp__plugin_atomistic-skills_base__search_model_registry` when installed as a plugin).
> Without a connected server, run the same tools from the shell. Tools named in
> one command share a process, so a model loaded by `load_model` stays loaded:
>
> ```bash
> ${CLAUDE_SKILL_DIR}/../../venv/run cpu python -m src.mcp_server.cli base search_model_registry key=value register_model key=value
> ```

## Goal
To estimate the epistemic uncertainty of a MACE MLIP by running inference with a committee (ensemble) of independently trained models. Structures where the committee disagrees strongly (high energy or force variance) are flagged as candidates for DFT labelling, supporting active learning workflows and validating MLIP reliability in under-sampled regions of configuration space.

The uncertainty estimate is:
- **Energy uncertainty**: standard deviation of predicted energies across committee members (meV/atom)
- **Force uncertainty**: component-wise force RMSE of the across-committee
  standard deviations (meV/Å): take the sample standard deviation for every
  atom and Cartesian component, then take one root-mean-square over all `3N`
  components. This is the conventional MLIP force-RMSE reduction used by MACE
  reporting.

> [!WARNING]
> **Energy std is only a valid disagreement signal when every committee member shares the same energy reference** — the same training dataset, level of theory, and atomic reference energies (`E0`). If the committee instead pools several *independently pretrained foundation models* (e.g. different MACE-MP / MACE-OMAT / MACE-MATPES releases) rather than same-data/different-seed checkpoints, their absolute energies are not on a common scale: an energy-std ranking will mostly reflect per-model reference-energy offsets, not genuine epistemic disagreement. For this heterogeneous committee flavor, rank structures by **force disagreement** instead — forces are invariant to each model's arbitrary atomic reference energy, so they remain a reliable cross-model signal. See the heterogeneous-committee note under Constraints.

## Instructions

### 1. Obtain Committee Models

A committee requires **N ≥ 3** independently trained MACE checkpoints covering the same chemical system. There are two ways to obtain them:

**Option A — Train with different random seeds (recommended for fine-tuned models)**

Run [ml-mace-finetune](../ml-mace-finetune/SKILL.md) N times, varying only the random seed via the `--seed` flag in `generate_mace_config.py`. Save each checkpoint to a separate directory:

```bash
# Run for seed=0, seed=1, seed=2 (at minimum)
for SEED in 0 1 2; do
    ${CLAUDE_SKILL_DIR}/../../venv/run mlip python ${CLAUDE_SKILL_DIR}/../ml-mace-finetune/scripts/generate_mace_config.py \
        --train-file ./mace_data/train.xyz \
        --valid-file ./mace_data/valid.xyz \
        --model MACE-MH-1 \
        --epochs 200 \
        --lr 1e-4 \
        --batch-size 4 \
        --freeze-backbone \
        --seed ${SEED} \
        --output-dir ./committee_models/seed_${SEED}
    ${CLAUDE_SKILL_DIR}/../../venv/run mlip mace_run_train \
        --config ./committee_models/seed_${SEED}/finetune_config.yaml
done
```

**Option B — Use pre-existing models from the registry**

```bash
base.search_model_registry(
    chemical_system="Li-Fe-P-O",
    backend="mace",
)
# If multiple models exist for the same system, they can form a committee.
```

### 2. Run Committee Inference

Pass all checkpoint paths to the inference script. It will run each model independently and compute mean ± std across the committee.

```bash
${CLAUDE_SKILL_DIR}/../../venv/run mlip python ${CLAUDE_SKILL_DIR}/scripts/run_committee_inference.py \
    --structures /path/to/structures_dir_or_file.cif \
    --models ./committee_models/seed_0/mace_finetuned.model \
             ./committee_models/seed_1/mace_finetuned.model \
             ./committee_models/seed_2/mace_finetuned.model \
    --output-dir ./uncertainty_results \
    --energy-threshold 10.0 \
    --force-threshold 200.0
# For MACE-MH foundation models, add:
#   --head omat_pbe
```

**Key parameters:**

| Parameter | Description | Typical value |
|:----------|:------------|:--------------|
| `--structures` | Path to structure file, directory, or `.xyz` trajectory | — |
| `--models` | Space-separated list of MACE checkpoint paths | ≥ 3 models |
| `--energy-threshold` | Flag if energy std > this value (meV/atom) | 5–20 meV/atom |
| `--force-threshold` | Flag if component-wise force RMSE > this value (meV/Å) | 100–300 meV/Å |
| `--output-dir` | Directory for results and plots | — |
| `--device` | `cuda` or `cpu` | `cuda` |
| `--head` | Head name for multi-head models (e.g. `omat_pbe` for MACE-MH-1) | `None` |

### 3. Interpret Results

The script produces:
- `uncertainty_summary.json` — per-structure energy/force mean ± std
- `high_uncertainty_structures/` — `.cif` files of flagged structures requiring DFT
- `uncertainty_distribution.png` — histogram of energy uncertainty across all structures

**Thresholds for DFT flagging:** There is no universal threshold. Empirically:
- Energy std > **10 meV/atom** is a conservative threshold suitable for phonon/stability work
- Energy std > **5 meV/atom** for high-accuracy MD (e.g., melting point, ionic conductivity)
- Force RMSE > **200 meV/Å** generally indicates the configuration is poorly represented in training data

> [!TIP]
> Run this skill on your MD trajectory after a few nanoseconds to check whether the MLIP remains in-distribution throughout the simulation. High uncertainty at late simulation times suggests the trajectory has drifted into unexplored configuration space.

### 4. Send High-Uncertainty Structures to DFT

Pass the flagged structures to the DFT labelling pipeline. See the [mat-sample-pes-by-md](../mat-sample-pes-by-md/SKILL.md) skill for context on when and how to label new structures.

```bash
# High-uncertainty structures needing DFT labels are in:
ls ./uncertainty_results/high_uncertainty_structures/
# → Pass these to atomate2 DFT workflow for labelling
```

From learningmatter-mit/AtomisticSkills, MIT.

More from learningmatter-mit/AtomisticSkills