scaling of loss with number of atoms (and other general cleanup?)

[bernstei]

We should figure out how to scale the loss with number of atoms. For GAP and ACE, we've settled on

loss = (w / sqrt(N)) ^2 (E-Eref)^2

for both energy and virial (extensive). From what I can tell, MACE uses `

loss = wM ((E-Eref)/N)^2

In addition to a (trivial) square on the magnitude of the weight, these do not scale the same with number of atoms. They can be mapped to each other through

wM = N w^2 

The former assumes that the energy errors are uncorrelated. The latter assumes something else, but I'm not sure how to define it in terms of correlations. But, for example, if you double the system size you get twice as many forces, but the same scale E/N, so same contribution from the energy.

Maybe the scaling should be selectable by the user?

Also, in general, there are many possibly duplicative loss functions coded up, and it would be very nice if they were unified. I can even imagine all of them being implemented as huber, and recover the quadratic ones by setting the threshold to infinity.

[alok-108]

Old thread, but I've been going through mace/modules/loss.py for #1767, so I'll take a shot at it.

On the correlations: say the per-atom energy errors each have variance σ² and a pairwise correlation ρ. Then

Var(E - Eref) = N σ² (1 + (N - 1) ρ)

With ρ = 0 that's N σ², so your (E - Eref)² / N is O(1) per config, whatever the size. MACE's ((E - Eref) / N)² is only O(1) when the variance grows like N², and that happens for any ρ that doesn't go to zero as the cell gets bigger, like a systematic per-atom offset. So GAP/ACE assume the errors are local and decorrelate, and MACE assumes they're fully correlated. If they're actually uncorrelated, MACE's energy term for a config drops like 1/N.

The forces make the difference bigger. The force term is a mean over every force component in the batch, while the energy term is a mean over configs. So in a mixed batch a config's share of the force loss goes like N, and its share of the energy loss is 1/(number of configs). With uncorrelated errors the energy/force balance for one config goes like 1/N² in MACE, against 1/N in GAP. That's your "twice as many forces, same E/N".

Virials get the same treatment as the energy, ((V - Vref) / N)². Stress isn't divided by N since it's already intensive, but stress is the virial over the cell volume, so --loss stress has the same N-dependence as --loss virials, just divided by the volume per atom squared. That's one more reason the stress and virial weights in #521 aren't interchangeable.

There's no switch for it today, but you can get the GAP scaling for the energy through the per-config weights, because config_energy_weight multiplies the per-atom term:

• --loss weighted, stress or virials: set config_energy_weight = N * w**2 and the term becomes w² (E - Eref)² / N. I checked it through config_from_atoms with three Cu cells (4, 32 and 24 atoms) and it matches the GAP term averaged over configs exactly. config_virials_weight works the same way for --loss virials.
• --loss universal: here the property weights multiply the huber arguments, so the weight ends up squared. config_energy_weight = sqrt(N) * w gives 0.5 w² (E - Eref)² / N in the quadratic regime (exactly half in the same check, the 0.5 is huber's). The delta then applies to w (E - Eref) / sqrt(N), which is probably what you'd want anyway.
• --loss huber: the per-config weights are ignored on main right now, so it doesn't work there yet. That's huber, universal, l1l2energyforces and dipole/polar losses ignore per-config weights: missing labels trained against zeros, config_weight has no effect #1767, and there's an open PR for it, Apply the config and property weights in every loss that skipped them #1781.

On unifying the losses: that's more or less what the v1 rewrite is doing in TRN-2 (#1576). The loss terms are generated from the declared observables, the weights and things like huber_delta come from config, and custom losses go into a registry. I don't see the N-scaling mentioned in the ticket though, so if you want it as a user option it's worth bringing up there.

On doing everything as huber: torch's huber is 0.5 x² below delta, so a huge delta gives half the MSE, which the global weights can absorb. The universal forces term also isn't a single huber. Its delta goes down with |F_ref| (1.0, 0.7, 0.4 and 0.1 times delta, switching at 100, 200 and 300 eV/Å), so that would need its own option.