From 62378bab19b9369fa44fdc11d2780de2bfcf118b Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 11 Aug 2026 16:43:34 +0800 Subject: [PATCH 1/3] docs+tests: pin the set/change bias semantics for bridged models Close #5927 as designed. The model energy decomposes as E = E_model + E_bias; set-by-statistic DEFINES E_bias as the per-type statistic of the raw labels (or user values), independent of E_model -- it ignores a trained network and it equally ignores the analytical ZBL term of a bridged model. change-by-statistic fits the residual against the complete model prediction (bridged predictor since #5910). The 'double count' described in the issue is therefore not a bridging bug but the uniform, defined behavior of the set mode for any model with nonzero E_model; no code change is made. - doc/model/change-bias.md: precise definitions of the two modes and the guidance that a self-consistent calibration of a bridged model needs change-by-statistic. - doc/model/dpa4.md: note in the ZBL section. - Semantic pin tests in pt (SeZMModel) and 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. Verified conformance of the linear_ener composition path in both backends: children compute no output statistics; the composition-level set fits raw labels. --- doc/model/change-bias.md | 26 ++++++++++ doc/model/dpa4.md | 11 ++++ .../tests/common/dpmodel/test_zbl_bridging.py | 37 +++++++++++++ source/tests/pt/model/test_sezm_model.py | 52 +++++++++++++++++++ 4 files changed, 126 insertions(+) diff --git a/doc/model/change-bias.md b/doc/model/change-bias.md index 310a21e83f..bc653dcd94 100644 --- a/doc/model/change-bias.md +++ b/doc/model/change-bias.md @@ -9,6 +9,32 @@ There are several scenarios where one might want to adjust the output bias after such as zero-shot testing (similar to the procedure before the first step in fine-tuning) or manually setting the output bias. +## The two statistic modes, precisely + +The model energy decomposes as `E = E_model + E_bias`, where `E_model` is +whatever the model computes (a learned network, an analytical term such as +ZBL bridging, or a `linear_ener` combination of models) and `E_bias` is the +per-type output bias. + +- **`set` (`set-by-statistic`)** assigns `E_bias` directly: either the + user-given values (`-b`), or the per-type least-squares statistic of the + **raw data labels**. It is independent of `E_model` by definition — it + ignores a trained network, and it equally ignores an analytical + contribution such as the ZBL term of a bridged model. The result is + reproducible and idempotent for a given dataset, but it contains **no + compensation for `E_model`**: after `set`, predictions on the calibration + data are offset by the data mean of `E_model`. +- **`change` (`change-by-statistic`)** assigns `E_bias` from the residual: + the per-type statistic of the labels **minus the complete model + prediction** (including any analytical bridging term), added to the + existing bias. Use this mode for a self-consistent calibration of a + trained (or bridged) model. + +For a bridged model — or any model whose `E_model` is significantly nonzero +on the calibration data — `set` therefore yields a bias that double-counts +nothing and compensates nothing; if you want the model predictions to match +the calibration labels, use `change`. + The `dp change-bias` command supports the following methods for adjusting the bias: ::::{tab-set} diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md index 8cce7b9faf..ce9bfc2cad 100644 --- a/doc/model/dpa4.md +++ b/doc/model/dpa4.md @@ -341,6 +341,17 @@ When ZBL bridging is enabled, set `training.training_data.min_pair_dist` to the same value as `bridging_r_inner` so frames with shorter atom pairs are excluded from training. See `examples/water/dpa4/input-zbl.json` for a complete example. +> [!NOTE] +> Output-bias statistics and bridging: the model energy is +> `E = E_model + E_bias`, and the ZBL term belongs to `E_model`. The +> `set-by-statistic` bias mode (initial statistics, finetune with a +> random fitting, `dp change-bias --mode set`) fits `E_bias` to the raw +> data labels and by definition ignores `E_model` — the analytical ZBL +> contribution included. For a self-consistent calibration of a bridged +> model use `change-by-statistic`, which subtracts the complete bridged +> prediction. See [change-bias](change-bias.md) for the precise +> definitions. + ## Performance and precision ### Training-time settings diff --git a/source/tests/common/dpmodel/test_zbl_bridging.py b/source/tests/common/dpmodel/test_zbl_bridging.py index 0db497dd84..b385c0cfef 100644 --- a/source/tests/common/dpmodel/test_zbl_bridging.py +++ b/source/tests/common/dpmodel/test_zbl_bridging.py @@ -629,3 +629,40 @@ def test_forwarded_from_children(self) -> None: assert bridged.atomic_model.get_compute_stats_distinguish_types() == any( c.get_compute_stats_distinguish_types() for c in children ) + + +def test_set_by_statistic_fits_raw_labels_by_definition(): + """Semantic pin (issue #5927): ``set-by-statistic`` is E_model-blind. + + ``E = E_model + E_bias``; the set mode defines ``E_bias`` as the + per-type statistic of the raw labels, independent of ``E_model`` -- + the composition-level fit ignores the learned child and equally + ignores the analytical ZBL child. Children compute no output + statistics of their own (the composition is the one owner). Use + ``change-by-statistic`` for a calibration that compensates + ``E_model``. + """ + model = get_model(copy.deepcopy(ZBL_CONFIG)) + rng = np.random.default_rng(5) + coord = rng.uniform(1.0, 2.5, size=(1, 4, 3)) + box = (np.eye(3) * 8.0).reshape(1, 9) + samples, labels, counts_rows = [], [], [] + for types in ([[0, 0, 1, 1]], [[0, 1, 1, 1]]): + counts = np.bincount(np.asarray(types[0]), minlength=2) + label = float(rng.normal()) + samples.append( + { + "coord": coord, + "atype": np.array(types), + "box": box, + "energy": np.array([[label]]), + "find_energy": np.float32(1.0), + "natoms": np.array([[4, 4, *counts]]), + } + ) + labels.append(label) + counts_rows.append(counts) + model.atomic_model.compute_or_load_out_stat(samples) + bias = np.asarray(model.atomic_model.out_bias).reshape(-1)[:2] + raw_fit = np.linalg.solve(np.array(counts_rows, dtype=np.float64), np.array(labels)) + np.testing.assert_allclose(bias, raw_fit, atol=1.0e-8) diff --git a/source/tests/pt/model/test_sezm_model.py b/source/tests/pt/model/test_sezm_model.py index 92a1cc14a6..0f4b110abf 100644 --- a/source/tests/pt/model/test_sezm_model.py +++ b/source/tests/pt/model/test_sezm_model.py @@ -2201,6 +2201,58 @@ def test_change_out_bias_is_invariant_for_self_labels(self) -> None: ) ) + def test_set_by_statistic_fits_raw_labels_by_definition(self) -> None: + """Semantic pin (issue #5927): ``set-by-statistic`` is E_model-blind. + + The model energy decomposes as ``E = E_model + E_bias``. The set + mode DEFINES ``E_bias`` as the per-type statistic of the raw data + labels (or a user-given value), independent of ``E_model`` -- it + ignores a trained network and it equally ignores the analytical + ZBL term of a bridged model. Do NOT "fix" this by subtracting the + analytical contribution: that would silently turn the mode into a + third, model-dependent behavior. For a calibration that + compensates ``E_model``, use ``change-by-statistic`` (which since + #5910 uses the complete bridged predictor). + """ + params = self._build_model_params(bridging_method="ZBL") + params["descriptor"]["precision"] = "float64" + params["fitting_net"]["precision"] = "float64" + model = get_sezm_model(params).to(self.device) + + rng = np.random.default_rng(5) + coord = torch.tensor( + rng.uniform(1.0, 2.5, size=(1, 4, 3)), + dtype=torch.float64, + device=self.device, + ) + box = ( + torch.eye(3, dtype=torch.float64, device=self.device).reshape(1, 9) * 8.0 + ) + samples, labels, counts_rows = [], [], [] + for types in ([[0, 0, 1, 1]], [[0, 1, 1, 1]]): + counts = np.bincount(np.asarray(types[0]), minlength=2) + label = float(rng.normal()) + samples.append( + { + "coord": coord, + "atype": torch.tensor(types, device=self.device), + "box": box, + "energy": torch.tensor( + [[label]], dtype=torch.float64, device=self.device + ), + "find_energy": np.float32(1.0), + "natoms": torch.tensor([[4, 4, *counts]], device=self.device), + } + ) + labels.append(label) + counts_rows.append(counts) + model.change_out_bias(samples, bias_adjust_mode="set-by-statistic") + bias = model.get_out_bias().detach().cpu().numpy().reshape(-1)[:2] + raw_fit = np.linalg.solve( + np.array(counts_rows, dtype=np.float64), np.array(labels) + ) + np.testing.assert_allclose(bias, raw_fit, atol=1.0e-8) + def test_zbl_respects_exclusions(self) -> None: """Excluded atoms and pairs contribute neither learned nor ZBL energy.""" coord = torch.tensor( From edde5843d81b157e71297a9858e3b953ff1af652 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 11 Aug 2026 16:43:52 +0800 Subject: [PATCH 2/3] style: ruff format for the set-bias semantic pin test --- source/tests/pt/model/test_sezm_model.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/source/tests/pt/model/test_sezm_model.py b/source/tests/pt/model/test_sezm_model.py index 0f4b110abf..57e447fa6d 100644 --- a/source/tests/pt/model/test_sezm_model.py +++ b/source/tests/pt/model/test_sezm_model.py @@ -2225,9 +2225,7 @@ def test_set_by_statistic_fits_raw_labels_by_definition(self) -> None: dtype=torch.float64, device=self.device, ) - box = ( - torch.eye(3, dtype=torch.float64, device=self.device).reshape(1, 9) * 8.0 - ) + box = torch.eye(3, dtype=torch.float64, device=self.device).reshape(1, 9) * 8.0 samples, labels, counts_rows = [], [], [] for types in ([[0, 0, 1, 1]], [[0, 1, 1, 1]]): counts = np.bincount(np.asarray(types[0]), minlength=2) From a6a2bf690e693144ea6f478f3037d4ef4b574e29 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 12 Aug 2026 18:12:02 +0800 Subject: [PATCH 3/3] docs,test: address set-bias review round 2 - change-bias.md: drop the incorrect 'offset by the data mean of E_model' claim (the remaining error is the configuration-dependent E_model plus the raw-label least-squares residual) and the 'double-counts nothing' claim (the raw-label fit can absorb the composition-correlated component of E_model into E_bias, which the forward pass then adds again). - pt + dpmodel set-by-statistic tests: seed a nonzero out_bias before the call (an accidental additive implementation would shift the result) and repeat the call to pin idempotence; mirrored in dpmodel because its bias storage is separate. --- doc/model/change-bias.md | 12 +++++++----- source/tests/common/dpmodel/test_zbl_bridging.py | 9 +++++++++ source/tests/pt/model/test_sezm_model.py | 8 ++++++++ 3 files changed, 24 insertions(+), 5 deletions(-) diff --git a/doc/model/change-bias.md b/doc/model/change-bias.md index bc653dcd94..2933facb13 100644 --- a/doc/model/change-bias.md +++ b/doc/model/change-bias.md @@ -22,8 +22,9 @@ per-type output bias. ignores a trained network, and it equally ignores an analytical contribution such as the ZBL term of a bridged model. The result is reproducible and idempotent for a given dataset, but it contains **no - compensation for `E_model`**: after `set`, predictions on the calibration - data are offset by the data mean of `E_model`. + compensation for `E_model`**: after `set`, the remaining error on the + calibration data is the configuration-dependent `E_model` itself, plus + any residual of the raw-label least-squares fit. - **`change` (`change-by-statistic`)** assigns `E_bias` from the residual: the per-type statistic of the labels **minus the complete model prediction** (including any analytical bridging term), added to the @@ -31,9 +32,10 @@ per-type output bias. trained (or bridged) model. For a bridged model — or any model whose `E_model` is significantly nonzero -on the calibration data — `set` therefore yields a bias that double-counts -nothing and compensates nothing; if you want the model predictions to match -the calibration labels, use `change`. +on the calibration data — `set` leaves `E_model` uncompensated and can absorb +its composition-correlated component into `E_bias`, so the forward pass may add +that component again. Use `change` to fit the residual against the complete +model prediction. The `dp change-bias` command supports the following methods for adjusting the bias: diff --git a/source/tests/common/dpmodel/test_zbl_bridging.py b/source/tests/common/dpmodel/test_zbl_bridging.py index b385c0cfef..2ac2878423 100644 --- a/source/tests/common/dpmodel/test_zbl_bridging.py +++ b/source/tests/common/dpmodel/test_zbl_bridging.py @@ -662,7 +662,16 @@ def test_set_by_statistic_fits_raw_labels_by_definition(): ) labels.append(label) counts_rows.append(counts) + # Seed a nonzero bias: `set` must DISCARD it (an accidental additive + # implementation would shift the result by the seed). The dpmodel bias + # storage is separate from pt's, so the pin is mirrored here. + model.atomic_model.out_bias = np.ones_like(model.atomic_model.out_bias) model.atomic_model.compute_or_load_out_stat(samples) bias = np.asarray(model.atomic_model.out_bias).reshape(-1)[:2] raw_fit = np.linalg.solve(np.array(counts_rows, dtype=np.float64), np.array(labels)) np.testing.assert_allclose(bias, raw_fit, atol=1.0e-8) + # Idempotence: repeating the call from the fitted state must land on + # the same raw-label fit again. + model.atomic_model.compute_or_load_out_stat(samples) + repeated_bias = np.asarray(model.atomic_model.out_bias).reshape(-1)[:2] + np.testing.assert_allclose(repeated_bias, raw_fit, atol=1.0e-8) diff --git a/source/tests/pt/model/test_sezm_model.py b/source/tests/pt/model/test_sezm_model.py index 57e447fa6d..f3027a5a8e 100644 --- a/source/tests/pt/model/test_sezm_model.py +++ b/source/tests/pt/model/test_sezm_model.py @@ -2244,12 +2244,20 @@ def test_set_by_statistic_fits_raw_labels_by_definition(self) -> None: ) labels.append(label) counts_rows.append(counts) + # Seed a nonzero bias: `set` must DISCARD it (an accidental + # additive implementation would shift the result by the seed). + model.set_out_bias(torch.ones_like(model.get_out_bias())) model.change_out_bias(samples, bias_adjust_mode="set-by-statistic") bias = model.get_out_bias().detach().cpu().numpy().reshape(-1)[:2] raw_fit = np.linalg.solve( np.array(counts_rows, dtype=np.float64), np.array(labels) ) np.testing.assert_allclose(bias, raw_fit, atol=1.0e-8) + # Idempotence: repeating the call from the fitted state must land + # on the same raw-label fit again. + model.change_out_bias(samples, bias_adjust_mode="set-by-statistic") + repeated_bias = model.get_out_bias().detach().cpu().numpy().reshape(-1)[:2] + np.testing.assert_allclose(repeated_bias, raw_fit, atol=1.0e-8) def test_zbl_respects_exclusions(self) -> None: """Excluded atoms and pairs contribute neither learned nor ZBL energy."""