Skip to content

Commit fec4497

Browse files
wanghan-iapcmHan Wang
andauthored
docs+tests(bias): pin set/change-by-statistic semantics for bridged models (#5966)
Close #5927 — resolved **as designed**, with documentation and semantic pin tests instead of a behavior change. ## The definitions (verified against the code) The model energy decomposes as `E = E_model + E_bias`, where `E_model` is everything the model computes (learned network, analytical ZBL term, or a `linear_ener` combination) and `E_bias` is the per-type output bias. - **`set-by-statistic`** assigns `E_bias` directly: user values, or the per-type least-squares statistic of the **raw labels**. It never calls a model forward (`compute_output_stats` with `model_forward=None`, stored with `add=False`) — it is independent of `E_model` **by definition**, ignoring a trained network and the analytical ZBL term alike. - **`change-by-statistic`** fits the residual of the labels against the **complete model prediction** (the bridged predictor since #5910) and adds the delta to the existing bias. Under these definitions the "double count" described in the issue is not a bridging bug: `set` uniformly ignores all of `E_model` for every model kind. A bridged model after `set` carries no compensation for the mean ZBL contribution — exactly as a trained plain model after `set` carries no compensation for its network output. The two modes "disagree" because they are defined to answer different questions; a self-consistent calibration is `change`'s job. Verified conformance of the `linear_ener` composition path in both backends: children compute no output statistics (`compute_or_load_out_stat=False`); the composition level performs one `set` fit on the raw labels. Also verified: the ZBL term contributes exactly zero for isolated atoms, so its own statistics are trivially `bias = 0`. ## Changes - `doc/model/change-bias.md`: precise definitions of the two modes, and the guidance that calibrating a bridged (or any nonzero-`E_model`) model self-consistently requires `change-by-statistic`. - `doc/model/dpa4.md`: note in the ZBL section. - Semantic pin tests (pt `SeZMModel` + dpmodel `LinearEnergyAtomicModel` composition): `set-by-statistic` equals the raw-label least-squares fit exactly — guarding against a future "fix" that would subtract the analytical term and silently create a third, model-dependent mode. ## Known limitations - No behavior change anywhere; the pin tests cover `model.change_out_bias` (the `dp change-bias --mode set` and finetune routes) and the dpmodel composition out-stat; the pt training-init chain was verified in-session and funnels into the same pinned branch. - The spin variants share the machinery but have no dedicated pin. - Fact worth knowing when reading the docs: `InnerPotential` adds the full ZBL over the whole cutoff (not only below `bridging_r_outer`), so the label-side ZBL at equilibrium geometries is small but not strictly zero; the docs state the offset plainly. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Documentation** * Clarified how `set` and `change` statistic modes handle model energy, analytical contributions, labels, residuals, and existing bias. * Documented the interaction between output-bias calibration and ZBL bridging, including guidance for self-consistent calibration. * **Tests** * Added regression coverage confirming that `set-by-statistic` calibration uses raw energy labels independently of learned and ZBL model contributions. * Verified calibration correctly replaces seeded bias and remains consistent when repeated. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>
1 parent 7141514 commit fec4497

4 files changed

Lines changed: 143 additions & 0 deletions

File tree

doc/model/change-bias.md

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,34 @@ There are several scenarios where one might want to adjust the output bias after
99
such as zero-shot testing (similar to the procedure before the first step in fine-tuning)
1010
or manually setting the output bias.
1111

12+
## The two statistic modes, precisely
13+
14+
The model energy decomposes as `E = E_model + E_bias`, where `E_model` is
15+
whatever the model computes (a learned network, an analytical term such as
16+
ZBL bridging, or a `linear_ener` combination of models) and `E_bias` is the
17+
per-type output bias.
18+
19+
- **`set` (`set-by-statistic`)** assigns `E_bias` directly: either the
20+
user-given values (`-b`), or the per-type least-squares statistic of the
21+
**raw data labels**. It is independent of `E_model` by definition — it
22+
ignores a trained network, and it equally ignores an analytical
23+
contribution such as the ZBL term of a bridged model. The result is
24+
reproducible and idempotent for a given dataset, but it contains **no
25+
compensation for `E_model`**: after `set`, the remaining error on the
26+
calibration data is the configuration-dependent `E_model` itself, plus
27+
any residual of the raw-label least-squares fit.
28+
- **`change` (`change-by-statistic`)** assigns `E_bias` from the residual:
29+
the per-type statistic of the labels **minus the complete model
30+
prediction** (including any analytical bridging term), added to the
31+
existing bias. Use this mode for a self-consistent calibration of a
32+
trained (or bridged) model.
33+
34+
For a bridged model — or any model whose `E_model` is significantly nonzero
35+
on the calibration data — `set` leaves `E_model` uncompensated and can absorb
36+
its composition-correlated component into `E_bias`, so the forward pass may add
37+
that component again. Use `change` to fit the residual against the complete
38+
model prediction.
39+
1240
The `dp change-bias` command supports the following methods for adjusting the bias:
1341

1442
::::{tab-set}

doc/model/dpa4.md

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -341,6 +341,17 @@ When ZBL bridging is enabled, set `training.training_data.min_pair_dist` to the
341341
same value as `bridging_r_inner` so frames with shorter atom pairs are excluded
342342
from training. See `examples/water/dpa4/input-zbl.json` for a complete example.
343343

344+
> [!NOTE]
345+
> Output-bias statistics and bridging: the model energy is
346+
> `E = E_model + E_bias`, and the ZBL term belongs to `E_model`. The
347+
> `set-by-statistic` bias mode (initial statistics, finetune with a
348+
> random fitting, `dp change-bias --mode set`) fits `E_bias` to the raw
349+
> data labels and by definition ignores `E_model` — the analytical ZBL
350+
> contribution included. For a self-consistent calibration of a bridged
351+
> model use `change-by-statistic`, which subtracts the complete bridged
352+
> prediction. See [change-bias](change-bias.md) for the precise
353+
> definitions.
354+
344355
## Performance and precision
345356

346357
### Training-time settings

source/tests/common/dpmodel/test_zbl_bridging.py

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -629,3 +629,49 @@ def test_forwarded_from_children(self) -> None:
629629
assert bridged.atomic_model.get_compute_stats_distinguish_types() == any(
630630
c.get_compute_stats_distinguish_types() for c in children
631631
)
632+
633+
634+
def test_set_by_statistic_fits_raw_labels_by_definition():
635+
"""Semantic pin (issue #5927): ``set-by-statistic`` is E_model-blind.
636+
637+
``E = E_model + E_bias``; the set mode defines ``E_bias`` as the
638+
per-type statistic of the raw labels, independent of ``E_model`` --
639+
the composition-level fit ignores the learned child and equally
640+
ignores the analytical ZBL child. Children compute no output
641+
statistics of their own (the composition is the one owner). Use
642+
``change-by-statistic`` for a calibration that compensates
643+
``E_model``.
644+
"""
645+
model = get_model(copy.deepcopy(ZBL_CONFIG))
646+
rng = np.random.default_rng(5)
647+
coord = rng.uniform(1.0, 2.5, size=(1, 4, 3))
648+
box = (np.eye(3) * 8.0).reshape(1, 9)
649+
samples, labels, counts_rows = [], [], []
650+
for types in ([[0, 0, 1, 1]], [[0, 1, 1, 1]]):
651+
counts = np.bincount(np.asarray(types[0]), minlength=2)
652+
label = float(rng.normal())
653+
samples.append(
654+
{
655+
"coord": coord,
656+
"atype": np.array(types),
657+
"box": box,
658+
"energy": np.array([[label]]),
659+
"find_energy": np.float32(1.0),
660+
"natoms": np.array([[4, 4, *counts]]),
661+
}
662+
)
663+
labels.append(label)
664+
counts_rows.append(counts)
665+
# Seed a nonzero bias: `set` must DISCARD it (an accidental additive
666+
# implementation would shift the result by the seed). The dpmodel bias
667+
# storage is separate from pt's, so the pin is mirrored here.
668+
model.atomic_model.out_bias = np.ones_like(model.atomic_model.out_bias)
669+
model.atomic_model.compute_or_load_out_stat(samples)
670+
bias = np.asarray(model.atomic_model.out_bias).reshape(-1)[:2]
671+
raw_fit = np.linalg.solve(np.array(counts_rows, dtype=np.float64), np.array(labels))
672+
np.testing.assert_allclose(bias, raw_fit, atol=1.0e-8)
673+
# Idempotence: repeating the call from the fitted state must land on
674+
# the same raw-label fit again.
675+
model.atomic_model.compute_or_load_out_stat(samples)
676+
repeated_bias = np.asarray(model.atomic_model.out_bias).reshape(-1)[:2]
677+
np.testing.assert_allclose(repeated_bias, raw_fit, atol=1.0e-8)

source/tests/pt/model/test_sezm_model.py

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2259,6 +2259,64 @@ def test_change_out_bias_is_invariant_for_self_labels(self) -> None:
22592259
)
22602260
)
22612261

2262+
def test_set_by_statistic_fits_raw_labels_by_definition(self) -> None:
2263+
"""Semantic pin (issue #5927): ``set-by-statistic`` is E_model-blind.
2264+
2265+
The model energy decomposes as ``E = E_model + E_bias``. The set
2266+
mode DEFINES ``E_bias`` as the per-type statistic of the raw data
2267+
labels (or a user-given value), independent of ``E_model`` -- it
2268+
ignores a trained network and it equally ignores the analytical
2269+
ZBL term of a bridged model. Do NOT "fix" this by subtracting the
2270+
analytical contribution: that would silently turn the mode into a
2271+
third, model-dependent behavior. For a calibration that
2272+
compensates ``E_model``, use ``change-by-statistic`` (which since
2273+
#5910 uses the complete bridged predictor).
2274+
"""
2275+
params = self._build_model_params(bridging_method="ZBL")
2276+
params["descriptor"]["precision"] = "float64"
2277+
params["fitting_net"]["precision"] = "float64"
2278+
model = get_sezm_model(params).to(self.device)
2279+
2280+
rng = np.random.default_rng(5)
2281+
coord = torch.tensor(
2282+
rng.uniform(1.0, 2.5, size=(1, 4, 3)),
2283+
dtype=torch.float64,
2284+
device=self.device,
2285+
)
2286+
box = torch.eye(3, dtype=torch.float64, device=self.device).reshape(1, 9) * 8.0
2287+
samples, labels, counts_rows = [], [], []
2288+
for types in ([[0, 0, 1, 1]], [[0, 1, 1, 1]]):
2289+
counts = np.bincount(np.asarray(types[0]), minlength=2)
2290+
label = float(rng.normal())
2291+
samples.append(
2292+
{
2293+
"coord": coord,
2294+
"atype": torch.tensor(types, device=self.device),
2295+
"box": box,
2296+
"energy": torch.tensor(
2297+
[[label]], dtype=torch.float64, device=self.device
2298+
),
2299+
"find_energy": np.float32(1.0),
2300+
"natoms": torch.tensor([[4, 4, *counts]], device=self.device),
2301+
}
2302+
)
2303+
labels.append(label)
2304+
counts_rows.append(counts)
2305+
# Seed a nonzero bias: `set` must DISCARD it (an accidental
2306+
# additive implementation would shift the result by the seed).
2307+
model.set_out_bias(torch.ones_like(model.get_out_bias()))
2308+
model.change_out_bias(samples, bias_adjust_mode="set-by-statistic")
2309+
bias = model.get_out_bias().detach().cpu().numpy().reshape(-1)[:2]
2310+
raw_fit = np.linalg.solve(
2311+
np.array(counts_rows, dtype=np.float64), np.array(labels)
2312+
)
2313+
np.testing.assert_allclose(bias, raw_fit, atol=1.0e-8)
2314+
# Idempotence: repeating the call from the fitted state must land
2315+
# on the same raw-label fit again.
2316+
model.change_out_bias(samples, bias_adjust_mode="set-by-statistic")
2317+
repeated_bias = model.get_out_bias().detach().cpu().numpy().reshape(-1)[:2]
2318+
np.testing.assert_allclose(repeated_bias, raw_fit, atol=1.0e-8)
2319+
22622320
def test_zbl_respects_exclusions(self) -> None:
22632321
"""Excluded atoms and pairs contribute neither learned nor ZBL energy."""
22642322
coord = torch.tensor(

0 commit comments

Comments
 (0)