Ablations and scaling

Protocol

Every ablation compares two or three settings of one configuration field across matched seeds, with all other settings, the dataset membership (data-split-seed 42) and the 12-hour wall-clock budget held fixed. The sweeps are W&B grid sweeps over (setting, seed) executed by scripts/wandb_run.py; ablation-setting-from names the varied field so that every run is labelled <ablation>_<setting>_seed<seed>. Each run ends with one held-out test evaluation.

Ablation

Varied field and settings

Sweep files (sweeps/paper_ablations/)

SiO2 energy guidance

loss_coef_observables: 0, \(10^{-4}\), \(10^{-3}\)

paper_siox_mature_energy_guidance_12h.yaml, paper_siox_mature_energy_guidance_3em5_12h.yaml, paper_additional_seeds_siox_mature_energy_guidance_12h.yaml

ZnCu2Sn(SeS)2 spectral guidance

spectral_loss_coef: 0, \(10^{-3}\) (\(2\times2\times2\), \(W=10\) eV, \(\tau=2\) eV, Huber 0.1 eV)

paper_zncusnses_mature_full_network_spectral_12h.yaml, paper_additional_seeds_zncusnses_mature_full_network_spectral_12h.yaml

ZnCu2Sn(SeS)2 distance envelope

hamiltonian_envelope_mode: off, multiply_prediction

paper_zncusnses_envelope_12h.yaml, paper_additional_seeds_zncusnses_envelope_12h.yaml

Silicon node aggregation

node_update_message_agg: average, attention

paper_silicon_node_aggregation_12h.yaml, paper_additional_seeds_silicon_node_aggregation_12h.yaml

The energy- and spectral-guidance studies (“mature”) fine-tune the selected demonstration models with each setting (resume-from-checkpoint, compatibility: true), measuring the effect of the objective on an already accurate network. The envelope and aggregation studies train from scratch.

Running a sweep

  1. Edit the dataset, cache and checkpoint paths (data-path, snapshot-cache-dir, checkpoint-dir and, for mature studies, resume-from-checkpoint) in the sweep file, or regenerate all sweep files from the shared matrix after editing it:

    python sweeps/paper_ablations/generate_sweeps.py

  2. Register the sweep and start agents until its run_cap is reached:

    wandb sweep sweeps/paper_ablations/paper_silicon_node_aggregation_12h.yaml
wandb agent <entity>/<project>/<sweep-id>

  3. Check the paired design (every seed present for every setting):

    PYTHONPATH=src:. python scripts/report/aggregate_paper_ablation.py \
  --entity <entity> --project paper_silicon_node_aggregation_12h \
  --group paper_silicon_node_aggregation_12h --baseline average \
  --expected-seeds 5 --output silicon_node_aggregation.json

  4. Export the run summaries and draw the figures:

    python paper/download_ablation_csvs.py
python paper/generate_result_plots.py

paper/data/*.csv contains the summaries of the published runs (seed, setting, Hamiltonian MAE, observable MAE, split), so step 4 reproduces the published figures directly from them.

Results

Ablation

Effect (medians over matched seeds)

Energy guidance, SiO2 (test)

\(10^{-4}\): band-energy MAE 0.245 → 0.078 eV/atom (−68 %), Hamiltonian MAE +4.7 %; \(10^{-3}\) is worse than \(10^{-4}\) on both.

Spectral guidance, ZnCu2Sn(SeS)2 (test, 10 seeds)

Eigenvalue MAE 0.579 → 0.0391 eV, paired mean improvement 92.1 %, all seeds improve; Hamiltonian MAE unchanged within its confidence interval.

Distance envelope, ZnCu2Sn(SeS)2 (test)

Hamiltonian MAE −3.8 %, four of five seeds improve.

Attention vs. average, silicon (validation)

Hamiltonian MAE −10.1 %.

The spectral- and energy-guidance figures are generated by step 4 above. Their rendered PNG files are not bundled in this checkout.

Inference scaling

Target-free inference on diamond-silicon supercells of 8 to 4096 atoms with the silicon model:

PYTHONPATH=src:. python scripts/generate_silicon_scaling_cifs.py --output-dir cifs/
PYTHONPATH=src:. python scripts/benchmark_silicon_cif_scaling.py \
  --checkpoint /runs/publication/silicon/best_model.pt \
  --structures-dir cifs/ --output scaling.txt --device cuda \
  --warmup-runs 2 --atom-counts 8,64,512,4096
PYTHONPATH=src:. python scripts/plot_silicon_cif_scaling.py \
  --input scaling.txt --output-dir figures/

Replace the checkpoint path with your trained silicon model. The CIFs and text timing log are generated by the preceding commands; the benchmark writes a text log, not JSON. --atom-counts selects the published range explicitly (the benchmark’s default also includes 32768 atoms).

The benchmark separately times graph construction, the forward pass and block reconstruction (two warm-up calls, ten timed repetitions) and records the peak memory of each phase. On one NVIDIA B200 GPU the total time grows from 0.364 s (8 atoms) to 97.0 s (4096 atoms; 1 216 512 blocks), with the forward pass taking 3.17 s of that at the largest size.

../_images/silicon_cif_scaling_time_total.png

Total inference wall time versus number of atoms.