From 867b6dd3b21dce397f67bf914195a6d14670dcb3 Mon Sep 17 00:00:00 2001 From: "Lee, Kin Long Kelvin" Date: Mon, 12 Aug 2024 14:58:39 -0700 Subject: [PATCH 1/4] refactor: normalize target noise by the scale --- matsciml/datasets/transforms/pretraining/noisy_positions.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/matsciml/datasets/transforms/pretraining/noisy_positions.py b/matsciml/datasets/transforms/pretraining/noisy_positions.py index 6c9b67e5..e7cb4fed 100644 --- a/matsciml/datasets/transforms/pretraining/noisy_positions.py +++ b/matsciml/datasets/transforms/pretraining/noisy_positions.py @@ -58,7 +58,8 @@ def __call__(self, data: DataDict) -> DataDict: else: data["noisy_pos"] = noisy_pos # set targets so that tasks know what to do - data["targets"]["denoise"] = noise + # noise targets are normalized by the scale + data["targets"]["denoise"] = noise / self.scale if "pretraining" in data["target_types"]: data["target_types"]["pretraining"].append("denoise") else: From 991b77873498d479a3962f546addb00b57843354 Mon Sep 17 00:00:00 2001 From: "Lee, Kin Long Kelvin" Date: Mon, 12 Aug 2024 15:02:55 -0700 Subject: [PATCH 2/4] refactor: adding normalize flag to pretraining transform --- matsciml/datasets/transforms/pretraining/noisy_positions.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/matsciml/datasets/transforms/pretraining/noisy_positions.py b/matsciml/datasets/transforms/pretraining/noisy_positions.py index e7cb4fed..cec86500 100644 --- a/matsciml/datasets/transforms/pretraining/noisy_positions.py +++ b/matsciml/datasets/transforms/pretraining/noisy_positions.py @@ -11,7 +11,7 @@ class NoisyPositions(AbstractDataTransform): - def __init__(self, scale: float = 1e-3) -> None: + def __init__(self, scale: float = 1e-3, normalize: bool = False) -> None: """ Initializes a NoisyPositions transform. @@ -31,9 +31,13 @@ def __init__(self, scale: float = 1e-3) -> None: ---------- scale : float Scale used to multiply N~(0, I_3) Gaussian noise + normalize : bool + If True, rescales the noise tensor by the scale + for loss computation. This is done in Liao, 2024. """ super().__init__() self.scale = scale + self.normalize = normalize def __call__(self, data: DataDict) -> DataDict: if "graph" in data: From 17b3c13e93cc254d098b7bd24ef7009682cc652e Mon Sep 17 00:00:00 2001 From: "Lee, Kin Long Kelvin" Date: Mon, 12 Aug 2024 15:04:49 -0700 Subject: [PATCH 3/4] refactor: making noise rescaling conditional on flag --- matsciml/datasets/transforms/pretraining/noisy_positions.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/matsciml/datasets/transforms/pretraining/noisy_positions.py b/matsciml/datasets/transforms/pretraining/noisy_positions.py index cec86500..d1181517 100644 --- a/matsciml/datasets/transforms/pretraining/noisy_positions.py +++ b/matsciml/datasets/transforms/pretraining/noisy_positions.py @@ -63,7 +63,9 @@ def __call__(self, data: DataDict) -> DataDict: data["noisy_pos"] = noisy_pos # set targets so that tasks know what to do # noise targets are normalized by the scale - data["targets"]["denoise"] = noise / self.scale + if self.normalize: + noise /= self.scale + data["targets"]["denoise"] = noise if "pretraining" in data["target_types"]: data["target_types"]["pretraining"].append("denoise") else: From 5793e54048c1aa7b1434abb2b21b98ff2b706c2a Mon Sep 17 00:00:00 2001 From: "Lee, Kin Long Kelvin" Date: Mon, 12 Aug 2024 15:05:47 -0700 Subject: [PATCH 4/4] refactor: making default noise scale less conservative --- matsciml/datasets/transforms/pretraining/noisy_positions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/matsciml/datasets/transforms/pretraining/noisy_positions.py b/matsciml/datasets/transforms/pretraining/noisy_positions.py index d1181517..2ee31de4 100644 --- a/matsciml/datasets/transforms/pretraining/noisy_positions.py +++ b/matsciml/datasets/transforms/pretraining/noisy_positions.py @@ -11,7 +11,7 @@ class NoisyPositions(AbstractDataTransform): - def __init__(self, scale: float = 1e-3, normalize: bool = False) -> None: + def __init__(self, scale: float = 1e-1, normalize: bool = False) -> None: """ Initializes a NoisyPositions transform.