jinjinhe2001/NHMO

Neural Harmonic Measure Operator: checkpoints

Paper (arXiv) · OpenReview · Project page · Code

Trained models of Neural Harmonic Measure Operator (Jinjin He, Sinan Wang, Yuchen Sun, Bo Zhu, NeurIPS 2026). The data/ folder holds the data needed to rerun the experiments (see data/README.md).

NHMO solves Δu = f in Ω, u = h on ∂Ω on varying geometry as u(p) = ⟨h, K_θ(p, ·; Ω)⟩ + v_φ(p; Ω, f), where the kernel K_θ approximates the density of the harmonic measure (trained from Walk-on-Spheres exits, without boundary data) and the lift v_φ supplies the source contribution.

Files

jinjinhe2001/NHMO
├── 2d/                  kernel_2d_mask.pt, lift_2d_l17c_mask.pt, lift_2d_msf_mask.pt, rhead_2d_mfr_mask.pt
├── 3d/                  kernel_<category>.pt, lift_<category>.pt, init/
├── data/                5.5 GB, see data/README.md
│   ├── mnist_pde_2d_paramBC_lf_train.tar.gz.part-0{0,1,2}   2D benchmark, train split
│   ├── mnist_pde_2d_paramBC_lf_test.tar.gz                  2D benchmark, test and test_ood splits
│   ├── mnist_kernel_targets_mask.tar.gz                     Walk-on-Spheres training targets of the 2D kernel
│   ├── mcb_ood3d_problems.tar.gz                            3D coefficient-OOD problem sets
│   ├── nhmo_reference_results.tar.gz                        per-pair errors of these checkpoints
│   └── README.md
├── MANIFEST.md
└── manifest.json

file	model	params	paper item
3d/kernel_<category>.pt	3D kernel K_θ, one per MCB-B category	2.77 M	Table 4
3d/lift_<category>.pt	3D source lift v_φ, one per category	2.33 M (fitting 2.36 M)	Tables 2 and 3 (with the kernel)
3d/init/kernel3d_init_warmup.pt	kernel after sphere pretraining and MCB warm-up	2.77 M	initialization for retraining
3d/init/kernel3d_init_fitting_pre44k.pt	initialization of the fitting kernel	2.77 M	initialization for retraining
2d/kernel_2d_mask.pt	2D kernel K_θ (MNIST domains)	4.09 M	Table 1, kernel only
2d/lift_2d_l17c_mask.pt	2D field lift, inputs (mask, h, f, u_h)	6.37 M	Table 1, NHMO (kernel + lift)
2d/lift_2d_msf_mask.pt	2D source-only lift, inputs (mask, sdf, f)	11.32 M	Table 1, + residual head (paper Section 6)
2d/rhead_2d_mfr_mask.pt	2D residual head, inputs (mask, h, u_h)	6.37 M	Table 1, + residual head (paper Section 6)

Categories: nut, gear, motor, fitting, screws_and_bolts. MANIFEST.md lists byte sizes, architectures, training data and training steps.

Each file is a PyTorch dictionary with the weights (model_state_dict, lift_state_dict or r_state_dict) and the configuration needed to rebuild the model (cfg, lift_cfg or r_cfg). Optimizer states are not included. All files load with torch.load(path, weights_only=True). The 2D kernel expects the boundary geometry computed from each shape's mask by nhmo/data/mask_geometry.py, which is the default of the 2D evaluator and trainers in the code release.

Usage

hf download jinjinhe2001/NHMO --exclude "data/*" --local-dir checkpoints
git clone https://github.com/jinjinhe2001/NHMO-Neural-Harmonic-Measure-Operator.git nhmo && cd nhmo && pip install -r requirements.txt
export NHMO_CKPT_DIR=$PWD/../checkpoints MCB_ROOT=<MCB-B data> NHMO_DATA_2D=<2D benchmark>
python -m nhmo.eval.mcb_lift --category nut \
    --kernel-ckpt $NHMO_CKPT_DIR/3d/kernel_nut.pt --lift-ckpt $NHMO_CKPT_DIR/3d/lift_nut.pt
python -m nhmo.eval.mnist_pde_2d --split test --kernel-ckpt $NHMO_CKPT_DIR/2d/kernel_2d_mask.pt \
    --lift-ckpt $NHMO_CKPT_DIR/2d/lift_2d_l17c_mask.pt

Training data

• 3D: MCB-B, a subset of MCB, as released by NGF (Yoo et al., NeurIPS 2025) in the Hugging Face dataset DveloperY0115/ngf-mcb with NGF's split files; 200 training shapes per category. Kernels see only the training meshes and Walk-on-Spheres exits; lifts are trained on the FEM solutions of the training shapes (32 problems per shape).
• 2D: domains [-1, 1]² \ digit built from MNIST training images. The kernel was trained on KDE targets for 5000 training digits that exclude every test and test_ood shape of the benchmark. The lifts are trained on the 991 training shapes of the 2D benchmark (parametric boundary families poly3 / exp_mix, four source families, 5-point FD references at 256²).

Results

benchmark	metric	paper
MCB-B (Table 2), nut / gear / motor / fitting / screws	mean rel-L2 over 320 pairs	.216 / .188 / .284 / .147 / .131
2D MNIST (Table 1), NHMO (kernel + lift)	mean / median rel-L2, test; test_ood	2.1 / 2.0 %; 2.6 / 2.5 %
2D MNIST (Table 1), kernel only	mean / median rel-L2, test; test_ood	7.7 / 5.4 %; 6.8 / 5.6 %
2D MNIST (Table 1), + residual head (paper Section 6)	mean / median rel-L2, test; test_ood	1.84 / 1.74 %; 2.56 / 2.39 %

Limitations

• The 3D models are per category and were trained on MCB-B shapes only; they are not expected to transfer to other shape families without retraining.
• The 2D models expect the resolution (256² masks, 128² lift) and the boundary and source families of the benchmark; coefficient extrapolation was tested on U[1, 2] only.
• The lifts are trained on the discrete reference solutions of the training data and inherit their discretization error.

Citation

@inproceedings{he2026nhmo,
  title     = {Neural Harmonic Measure Operator},
  author    = {He, Jinjin and Wang, Sinan and Sun, Yuchen and Zhu, Bo},
  booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
  year      = {2026}
}