Revisit loss function

[forklady42]

balance towards bonding regions/high density; incorporating cell size

[forklady42]

Loss Function Experiment Results

Background

Motivated by the voxel error analysis showing systematic underprediction of peak charge densities (worst structures compress peaks by up to 40%), we ran a controlled experiment comparing four alternative loss functions against the baseline NormMAE on the full MP dataset_2.

The hypothesis: since NormMAE treats all voxels equally, space with ~no electron density is overrepresented. Reweighting the loss could improve peak density reconstruction and tail performance.

What was tested

All experiments used identical setup: ResUNet3D (32 channels, depth 2, K=5), 4× A100 80GB GPUs on Della, same dataset/seed/architecture, ~15 epochs. Only the loss function varied.

Loss	Formula	Rationale	Config key
NormMAE (baseline)	|pred - target| / nelec	Current default	normmae
Density-weighted MAE	w(ρ) × |pred - target| / nelec where w = 1 + α(ρ/ρ_max)	Upweight errors at high-density voxels to fix peak saturation	density_weighted_normmae
NormMAE + Peak penalty	NormMAE + λ((1 - max_ratio)² + (1 - p99_ratio)²)/2	Directly penalize peak density compression	normmae_peak
Hellinger distance	(√pred - √target)² / nelec	sqrt transform compresses peaks, expands low-density gradient signal	hellinger

Loss selection is now config-driven via loss_fn YAML key (branch loss/all-experiments).

Training convergence

W&B project: mp-loss-experiment

All three viable losses converged to similar val_normmae by epoch 14 (~2.4-2.6%), with density-weighted briefly leading at epoch 11. The Hellinger loss diverged to NaN (see section below).

Test inference results (epoch ~15, 29 val structures)

Metric	Baseline	Peak Penalty	Density-Weighted
Mean NMAE	2.25%	2.36%	3.05%
Median NMAE	1.88%	2.02%	2.71%
Max NMAE	4.64%	5.04%	6.75%
Structures >5% NMAE	0	1	3
Structures improved vs baseline	—	6/29 (21%)	0/29 (0%)

Baseline NormMAE wins on every metric. Neither experimental loss outperforms it — not on mean, median, tail, or per-structure comparisons.

• Peak penalty is close (~0.1 pp worse) but consistently degrades 79% of structures (Wilcoxon p=0.0001)
• Density-weighted is clearly worse (0.8 pp behind), degrades every structure, and creates 3 new tail failures
• Hard structures remain hard across all losses (Spearman rank correlation 0.96-0.98)

Hellinger loss: NaN divergence

Both v1 (lr=0.01) and v2 (lr=0.001, grad_clip=1.0) diverged to NaN within epoch 0. Root cause: the gradient (√pred - √target) / √pred has a 1/√pred singularity. The model's bare Conv3d output layer produces exact zeros, and the charge conservation normalization propagates infinite gradients to all parameters.

Fix implemented (v3): epsilon floor inside sqrt (√(pred + 1e-8)) + optional softplus_output flag on ResUNet3D to guarantee strictly positive outputs before normalization. Hellinger v3 is currently training.

Interpretation

The charge conservation normalization (out = out / sum(out) * nelec) creates a zero-sum constraint: the model must allocate a fixed electron budget across voxels. Upweighting high-density errors doesn't create more electrons — it just forces the optimizer to steal accuracy from well-predicted regions to marginally improve peaks. This explains why both weighted losses degrade overall performance despite the weighted loss values themselves decreasing during training.

What's left

• Hellinger v3 results (running now)
• Consider whether architectural changes (e.g., separate head for peak regions) would be more effective than loss engineering
• The saturation problem may need to be addressed at the model level (e.g., removing the charge conservation constraint during training, or using a two-stage approach)

[forklady42]

Key plots

Training convergence — all 3 losses approach ~2.4-2.6% val NormMAE by epoch 14, but test inference reveals baseline maintains a clear edge:

See full training curves and run details on W&B: mp-loss-experiment

Per-structure scatter (epoch ~15) — nearly every point sits above y=x for both experimental losses, with density-weighted showing the largest deviations:

Analysis scripts and all plots are in test_output/ on branch loss/all-experiments:

• ep15_scatter_baseline_vs_experimental.png — paired per-structure comparison

• ep15_nmae_cdf.png — CDF showing baseline dominates at every percentile

• ep15_tail_scatter_highlighted.png — worst-10% structures highlighted

• ep15_per_structure_delta.png — sorted delta bars (experimental - baseline)

• val_normmae_vs_epoch_zoomed.png — training convergence curves

[forklady42]

Final Loss Function Experiment Results (updated)

Previous results had a bug where test_step wrote the training loss value (not NormMAE) to metrics.csv — this inflated density-weighted errors and deflated Hellinger scores. That's now fixed (commit 93da379). Here are the corrected final numbers.

Corrected 4-way comparison

Loss	Epochs	Mean NMAE	Median	Max	>5%
Baseline (NormMAE)	18	2.19%	1.92%	4.66%	0
Density-weighted	18	2.19%	1.89%	4.17%	0
Peak-penalty	18	2.27%	2.02%	4.75%	0
Hellinger v3	25	2.48%	2.34%	4.79%	0

Key changes from earlier (buggy) analysis

• Density-weighted is now tied with baseline (was previously reported as 0.8 pp worse). The earlier analysis was comparing density-weighted loss values against NormMAE values — an apples-to-oranges error. Corrected per-structure comparison: 11/29 structures improved, Wilcoxon p=0.70 (not significant). It also has the lowest max NMAE (4.17%) and improved 2/3 worst-10% tail structures.
• Hellinger v3 successfully trained (first stable Hellinger variant) after adding epsilon inside sqrt and softplus output activation. At epoch 25 it's at 2.48% but plateauing (~0.03 pp/epoch improvement). The 10× lower learning rate and softplus activation likely limit its ceiling.

Hellinger v3 stability fix

Two changes were needed together to prevent NaN divergence:

1. Loss: sqrt(pred + 1e-8) instead of sqrt(pred.clamp(min=0)) — caps the 1/sqrt(pred) gradient singularity
2. Model: softplus_output=True on ResUNet3D — guarantees strictly positive outputs before charge conservation normalization

Without both fixes, the gradient at pred=0 is infinite and propagates to all parameters via the normalization layer's coupled Jacobian. See commit fabe4a8.

Conclusions

1. NormMAE remains the best overall loss. None of the alternatives beat it on mean NMAE.
2. Density-weighted is the only promising alternative — tied on mean, better on max/tail. Worth exploring with tuned alpha/power parameters or as a curriculum (switch from NormMAE to density-weighted after N epochs).
3. The charge conservation normalization creates a zero-sum constraint that fundamentally limits what loss reweighting can achieve. Architectural changes (removing normalization during training, two-stage approaches, separate peak heads) may be more productive than further loss engineering.
4. Hellinger is viable but not competitive at the same training budget. The softplus activation and lower lr may be limiting factors.

What's on the branch

loss/all-experiments contains:

• All 4 loss implementations with config-driven selection (loss_fn YAML key)
• val_normmae always logged during validation for cross-experiment comparison
• Test inference scripts and analysis in test_output/
• 23 unit tests for the new losses