model_name: Nequix MP
model_key: nequix-mp-1
model_version: '0.1.0'
date_added: '2025-08-17'
date_published: '2025-08-17'
authors:
  - name: Teddy Koker
    affiliation: Massachusetts Institute of Technology
    email: tekoker@mit.edu
    url: https://teddykoker.com
    orcid: https://orcid.org/0000-0001-8861-9788
  - name: Tess Smidt
    affiliation: Massachusetts Institute of Technology
    email: tsmidt@mit.edu
    url: https://blondegeek.github.io/
    orcid: https://orcid.org/0000-0001-5581-5344

trained_by:
  - name: Teddy Koker
    affiliation: Massachusetts Institute of Technology
    email: tekoker@mit.edu
    url: https://teddykoker.com
    orcid: https://orcid.org/0000-0001-8861-9788

repo: https://github.com/atomicarchitects/nequix
url: https://github.com/atomicarchitects/nequix
doi: https://doi.org/10.48550/arXiv.2508.16067
paper: https://arxiv.org/abs/2508.16067
pr_url: https://github.com/janosh/matbench-discovery/pull/276
checkpoint_url: https://figshare.com/files/57245573

openness: OSOD
train_task: S2EFS
test_task: IS2RE-SR
targets: EFS_G
model_type: UIP
model_params: 707_658
trained_for_benchmark: true
n_estimators: 1

license:
  code: MIT
  code_url: https://github.com/atomicarchitects/nequix/blob/main/LICENSE
  checkpoint: CC-BY-4.0
  checkpoint_url: https://creativecommons.org/licenses/by/4.0/

hyperparams:
  max_force: 0.02
  max_steps: 500
  ase_optimizer: FIRE
  cell_filter: FrechetCellFilter
  hidden_irreps: 128x0e+64x1o+32x2e+32x3o
  graph_construction_radius: 6.0
  max_neighbors: .inf
  lmax: 3
  radial_basis_size: 8
  radial_mlp_size: 64
  radial_mlp_layers: 2
  radial_polynomial_p: 6.0
  layer_norm: true

  optimizer: muon
  weight_decay: 0.001
  warmup_epochs: 0.1
  warmup_factor: 0.2
  grad_clip_norm: 100
  batch_size: 256
  n_epochs: 100
  loss:
    energy: mae
    force: l2
    stress: mae
  loss_weights:
    energy: 20.0
    force: 20.0
    stress: 5.0
  ema_decay: 0.999

training_cost:
  Nvidia A100 GPUs: { amount: 4, hours: 125 }

training_set: [MPtrj]

notes:
  Description: Nequix is a simplified version of NequIP with added RMSnorm, implemented in JAX, and trained with Muon optimizer.

requirements:
  jax: 0.6.2
  e3nn-jax: 0.20.7
  jraph: 0.0.6.dev0
  equinox: 0.11.11
  optax: 0.2.5
  ase: 3.24.0
  tqdm: 4.67.1
  wandb: 0.19.11
  pyyaml: 6.0.2
  matscipy: 1.1.1
  h5py: 3.14.0
  wandb-osh: 1.2.2
  cloudpickle: 3.1.1

metrics:
  phonons:
    kappa_103:
      pred_file: models/nequix/nequix-mp-1/2025-08-17-kappa-103-FIRE-dist=0.03-fmax=1e-4-symprec=1e-5.json.gz
      pred_file_url: https://figshare.com/files/57262724
      κ_SRME: 0.4461
      κ_SRE: 0.2901
  geo_opt:
    pred_file: models/nequix/nequix-mp-1/2025-08-17-wbm-IS2RE-FIRE.jsonl.gz
    pred_file_url: https://figshare.com/files/57262715
    struct_col: nequix_structure
    symprec=1e-5:
      rmsd: 0.0853 # unitless
      n_sym_ops_mae: 2.0462 # unitless
      symmetry_decrease: 0.0607 # fraction
      symmetry_match: 0.6999 # fraction
      symmetry_increase: 0.2335 # fraction
      n_structures: 256963 # count
      analysis_file: models/nequix/nequix-mp-1/2025-08-17-wbm-IS2RE-FIRE-symprec=1e-5-moyo=0.4.4.csv.gz
      analysis_file_url: https://figshare.com/files/57262718
    symprec=1e-2:
      rmsd: 0.0853 # unitless
      n_sym_ops_mae: 1.7969 # unitless
      symmetry_decrease: 0.0554 # fraction
      symmetry_match: 0.8111 # fraction
      symmetry_increase: 0.1261 # fraction
      n_structures: 256963 # count
      analysis_file: models/nequix/nequix-mp-1/2025-08-17-wbm-IS2RE-FIRE-symprec=1e-2-moyo=0.4.4.csv.gz
      analysis_file_url: https://figshare.com/files/57262721
  discovery:
    pred_file: models/nequix/nequix-mp-1/2025-08-17-wbm-IS2RE.csv.gz
    pred_file_url: https://figshare.com/files/57262712
    pred_col: e_form_per_atom_nequix
    full_test_set:
      F1: 0.744 # fraction
      DAF: 3.95 # dimensionless
      Precision: 0.678 # fraction
      Recall: 0.825 # fraction
      Accuracy: 0.903 # fraction
      TPR: 0.825 # fraction
      FPR: 0.081 # fraction
      TNR: 0.919 # fraction
      FNR: 0.175 # fraction
      TP: 36363.0 # count
      FP: 17282.0 # count
      TN: 195589.0 # count
      FN: 7729.0 # count
      MAE: 0.042 # eV/atom
      RMSE: 0.084 # eV/atom
      R2: 0.78 # dimensionless
      missing_preds: 2 # count
    unique_prototypes:
      F1: 0.751 # fraction
      DAF: 4.455 # dimensionless
      Precision: 0.681 # fraction
      Recall: 0.836 # fraction
      Accuracy: 0.914 # fraction
      TPR: 0.836 # fraction
      FPR: 0.072 # fraction
      TNR: 0.928 # fraction
      FNR: 0.164 # fraction
      TP: 27898.0 # count
      FP: 13064.0 # count
      TN: 169050.0 # count
      FN: 5476.0 # count
      MAE: 0.044 # eV/atom
      RMSE: 0.086 # eV/atom
      R2: 0.782 # dimensionless
      missing_preds: 0 # count
    most_stable_10k:
      F1: 0.953 # fraction
      DAF: 5.949 # dimensionless
      Precision: 0.91 # fraction
      Recall: 1.0 # fraction
      Accuracy: 0.91 # fraction
      TPR: 1.0 # fraction
      FPR: 1.0 # fraction
      TNR: 0.0 # fraction
      FNR: 0.0 # fraction
      TP: 9095.0 # count
      FP: 905.0 # count
      TN: 0.0 # count
      FN: 0.0 # count
      MAE: 0.051 # eV/atom
      RMSE: 0.117 # eV/atom
      R2: 0.734 # dimensionless
      missing_preds: 0 # count