Skip to content
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 26 additions & 0 deletions doc/model/change-bias.md
Original file line number Diff line number Diff line change
Expand Up @@ -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`.
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Outdated
- **`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`.
Comment thread
njzjz marked this conversation as resolved.
Outdated

The `dp change-bias` command supports the following methods for adjusting the bias:

::::{tab-set}
Expand Down
11 changes: 11 additions & 0 deletions doc/model/dpa4.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
37 changes: 37 additions & 0 deletions source/tests/common/dpmodel/test_zbl_bridging.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
50 changes: 50 additions & 0 deletions source/tests/pt/model/test_sezm_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -2201,6 +2201,56 @@ 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)
Comment thread
njzjz marked this conversation as resolved.

def test_zbl_respects_exclusions(self) -> None:
"""Excluded atoms and pairs contribute neither learned nor ZBL energy."""
coord = torch.tensor(
Expand Down
Loading