feat(dpa4): split edge_norm into per-site [radial, film, focus] switches

The edge_norm option now accepts, besides the original bool, a list of
three bools [radial, film, focus] that gates the radial-MLP hidden
RMSNorms, the environment-seed FiLM scale/shift norms and the cross-focus
competition norms individually; the post-SO(2) residual scaling is bound
to the radial entry (1e-5 if on, unit floor if off), so all-true and
all-false lists reproduce the bool behaviour bit for bit. Serialization
writes the canonical three-bool list; configs and serialized data from
models that predate the option (no key) or store a bool keep loading
unchanged. The argument doc records the recommended setting
[false, true, true]: the radial-site norms amplify noise where the radial
features vanish at the cutoff and produce a spurious long-range force
step, while the FiLM and focus norms are safe to keep.

Unit tests cover the per-site gating, the SO(2) eps coupling, the
canonical serialization, legacy no-key and bool deserialization, and
pt/dpmodel parity for bool, list and legacy forms.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: OutisLi

deepmd/dpmodel/descriptor/dpa4.py

@@ -321,14 +321,16 @@ class DescrptDPA4(NativeOP, BaseDescriptor):
         Hidden layer sizes for radial networks. An output layer of size
         `(l_schedule[0]+extra_node_l+1)*channels` will be automatically appended.
     edge_norm
-        Whether to apply channel RMSNorm on the descriptor's cutoff-vanishing
-        branches: the radial network hidden layers, the environment-seed FiLM
-        scale/shift logits, the cross-focus competition scalars, and the
-        post-SO(2) residual messages. ``False`` replaces the first three norms
-        with identity and changes only the post-SO(2) norm to unit-floor residual
-        scaling. The unit floor uses ``sqrt(1 + variance)`` so small messages
-        retain their cutoff envelope instead of receiving the standard
-        ``1/sqrt(eps)`` small-signal gain.
+        Channel RMSNorm on the descriptor's cutoff-vanishing branches: the
+        radial network hidden layers, the environment-seed FiLM scale/shift
+        logits, and the cross-focus competition scalars. A bool switches all
+        three together; a list of three bools ``[radial, film, focus]``
+        switches them individually. Disabled norms are identity
+        pass-throughs. The post-SO(2) residual scaling follows the ``radial``
+        entry: with it disabled the norm uses the unit floor
+        ``sqrt(1 + variance)``, so small messages retain their cutoff
+        envelope instead of receiving the standard ``1/sqrt(eps)``
+        small-signal gain.
     use_env_seed
         If True, seed the initial node state with local-environment information:
         apply environment matrix FiLM conditioning on l=0 features using 4D
@@ -610,7 +612,7 @@ def __init__(
         basis_type: str = "bessel",
         n_radial: int = 16,
         radial_mlp: list[int] | None = None,
-        edge_norm: bool = True,
+        edge_norm: bool | list[bool] = True,
         use_env_seed: bool = True,
         random_gamma: bool = True,
         edge_cartesian: bool = False,
@@ -706,7 +708,22 @@ def __init__(
         if radial_mlp is None:
             radial_mlp = [0]
         self.radial_mlp = [self.channels if x == 0 else int(x) for x in radial_mlp]
-        self.edge_norm = bool(edge_norm)
+        if isinstance(edge_norm, bool):
+            self.radial_norm = edge_norm
+            self.film_norm = edge_norm
+            self.focus_norm = edge_norm
+        elif (
+            isinstance(edge_norm, (list, tuple))
+            and len(edge_norm) == 3
+            and all(isinstance(v, bool) for v in edge_norm)
+        ):
+            self.radial_norm = bool(edge_norm[0])
+            self.film_norm = bool(edge_norm[1])
+            self.focus_norm = bool(edge_norm[2])
+        else:
+            raise ValueError(
+                "edge_norm must be a bool or a list[bool] of length 3: [radial, film, focus]"
+            )
         if sandwich_norm is None:
             sandwich_norm = [False, True, True, False]
         if not isinstance(sandwich_norm, (list, tuple)) or len(sandwich_norm) != 4:
@@ -1009,7 +1026,7 @@ def __init__(
             # vanishes at rcut; normalizing them shares the radial network's
             # cutoff-smoothness issue, so ``edge_norm=False`` also drops these
             # norms (identity pass-through) to keep the FiLM scale/shift smooth.
-            if self.edge_norm:
+            if self.film_norm:
                 self.film_scale_norm = ScalarRMSNorm(
                     channels=self.channels,
                     n_focus=1,
@@ -1066,7 +1083,7 @@ def __init__(
             activation_function=self.activation_function,
             precision=self.compute_precision,  # force fp32+
             trainable=self.trainable,
-            radial_norm=self.edge_norm,
+            radial_norm=self.radial_norm,
             seed=seed_radial_embedding,
         )
 
@@ -1129,7 +1146,7 @@ def __init__(
                     channels=self.channels,
                     n_focus=self.n_focus,
                     focus_dim=self.focus_dim,
-                    focus_norm=self.edge_norm,
+                    focus_norm=self.focus_norm,
                     so2_norm=self.so2_norm,
                     mixing_layers=self.mixing_layers,
                     so2_attn_res=self.so2_attn_res_mode,
@@ -1164,7 +1181,7 @@ def __init__(
                     atten_o_proj=self.use_atten_o_proj,
                     so2_pre_norm=self.so2_pre_norm,
                     so2_post_norm=self.so2_post_norm,
-                    so2_post_norm_eps=1.0e-5 if self.edge_norm else 1.0,
+                    so2_post_norm_eps=1.0e-5 if self.radial_norm else 1.0,
                     so2_activation_function=self.so2_activation_function,
                     ffn_pre_norm=self.ffn_pre_norm,
                     ffn_post_norm=self.ffn_post_norm,
@@ -1618,10 +1635,10 @@ def _run_graph(
             scale_logits = film[:, : self.channels]  # (N, C)
             shift_logits = film[:, self.channels :]  # (N, C)
             scale_hat = (
-                self.film_scale_norm(scale_logits) if self.edge_norm else scale_logits
+                self.film_scale_norm(scale_logits) if self.film_norm else scale_logits
             )  # (N, C)
             shift_hat = (
-                self.film_shift_norm(shift_logits) if self.edge_norm else shift_logits
+                self.film_shift_norm(shift_logits) if self.film_norm else shift_logits
             )  # (N, C)
             scale_strength = xp.exp(
                 xp_asarray_nodetach(
@@ -2605,7 +2622,7 @@ def _variables(self) -> dict[str, np.ndarray]:
         if self.use_env_seed:
             for key, value in self.env_seed_embedding.serialize()["@variables"].items():
                 variables[f"env_seed_embedding.{key}"] = value
-            if self.edge_norm:
+            if self.film_norm:
                 for key, value in self.film_scale_norm.serialize()[
                     "@variables"
                 ].items():
@@ -2712,7 +2729,7 @@ def load(module: Any, prefix: str) -> Any:
             self.env_seed_embedding = load(
                 self.env_seed_embedding, "env_seed_embedding."
             )
-            if self.edge_norm:
+            if self.film_norm:
                 self.film_scale_norm = load(self.film_scale_norm, "film_scale_norm.")
                 self.film_shift_norm = load(self.film_shift_norm, "film_shift_norm.")
             self.film_scale_strength_log = np.asarray(
@@ -2826,7 +2843,11 @@ def serialize(self) -> dict[str, Any]:
                 "basis_type": self.basis_type,
                 "n_radial": self.n_radial,
                 "radial_mlp": self.radial_mlp,
-                "edge_norm": self.edge_norm,
+                "edge_norm": [
+                    self.radial_norm,
+                    self.film_norm,
+                    self.focus_norm,
+                ],
                 "use_env_seed": self.use_env_seed,
                 "random_gamma": self.random_gamma,
                 "edge_cartesian": self.edge_cartesian,

deepmd/pt/model/descriptor/sezm.py

@@ -162,14 +162,16 @@ class DescrptSeZM(BaseDescriptor, nn.Module):
         Hidden layer sizes for radial networks. An output layer of size
         `(l_schedule[0]+extra_node_l+1)*channels` will be automatically appended.
     edge_norm
-        Whether to apply channel RMSNorm on the descriptor's cutoff-vanishing
-        branches: the radial network hidden layers, the environment-seed FiLM
-        scale/shift logits, the cross-focus competition scalars, and the
-        post-SO(2) residual messages. ``False`` replaces the first three norms
-        with identity and changes only the post-SO(2) norm to unit-floor residual
-        scaling. The unit floor uses ``sqrt(1 + variance)`` so small messages
-        retain their cutoff envelope instead of receiving the standard
-        ``1/sqrt(eps)`` small-signal gain.
+        Channel RMSNorm on the descriptor's cutoff-vanishing branches: the
+        radial network hidden layers, the environment-seed FiLM scale/shift
+        logits, and the cross-focus competition scalars. A bool switches all
+        three together; a list of three bools ``[radial, film, focus]``
+        switches them individually. Disabled norms are identity
+        pass-throughs. The post-SO(2) residual scaling follows the ``radial``
+        entry: with it disabled the norm uses the unit floor
+        ``sqrt(1 + variance)``, so small messages retain their cutoff
+        envelope instead of receiving the standard ``1/sqrt(eps)``
+        small-signal gain.
     use_env_seed
         If True, seed the initial node state with local-environment information:
         apply environment matrix FiLM conditioning on l=0 features using 4D
@@ -452,7 +454,7 @@ def __init__(
         basis_type: str = "bessel",
         n_radial: int = 16,
         radial_mlp: list[int] | None = None,
-        edge_norm: bool = True,
+        edge_norm: bool | list[bool] = True,
         use_env_seed: bool = True,
         random_gamma: bool = True,
         edge_cartesian: bool = False,
@@ -554,7 +556,22 @@ def __init__(
         if radial_mlp is None:
             radial_mlp = [0]
         self.radial_mlp = [self.channels if x == 0 else int(x) for x in radial_mlp]
-        self.edge_norm = bool(edge_norm)
+        if isinstance(edge_norm, bool):
+            self.radial_norm = edge_norm
+            self.film_norm = edge_norm
+            self.focus_norm = edge_norm
+        elif (
+            isinstance(edge_norm, (list, tuple))
+            and len(edge_norm) == 3
+            and all(isinstance(v, bool) for v in edge_norm)
+        ):
+            self.radial_norm = bool(edge_norm[0])
+            self.film_norm = bool(edge_norm[1])
+            self.focus_norm = bool(edge_norm[2])
+        else:
+            raise ValueError(
+                "edge_norm must be a bool or a list[bool] of length 3: [radial, film, focus]"
+            )
         if sandwich_norm is None:
             sandwich_norm = [False, True, True, False]
         if not isinstance(sandwich_norm, (list, tuple)) or len(sandwich_norm) != 4:
@@ -862,7 +879,7 @@ def __init__(
             # vanishes at rcut; normalizing them shares the radial network's
             # cutoff-smoothness issue, so ``edge_norm=False`` also drops these
             # norms (identity pass-through) to keep the FiLM scale/shift smooth.
-            if self.edge_norm:
+            if self.film_norm:
                 self.film_scale_norm: nn.Module = ScalarRMSNorm(
                     channels=self.channels,
                     n_focus=1,
@@ -928,7 +945,7 @@ def __init__(
             activation_function=self.activation_function,
             dtype=self.compute_dtype,  # force fp32+
             trainable=self.trainable,
-            radial_norm=self.edge_norm,
+            radial_norm=self.radial_norm,
             seed=seed_radial_embedding,
         )
 
@@ -991,7 +1008,7 @@ def __init__(
                     channels=self.channels,
                     n_focus=self.n_focus,
                     focus_dim=self.focus_dim,
-                    focus_norm=self.edge_norm,
+                    focus_norm=self.focus_norm,
                     so2_norm=self.so2_norm,
                     mixing_layers=self.mixing_layers,
                     so2_attn_res=self.so2_attn_res_mode,
@@ -1026,7 +1043,7 @@ def __init__(
                     atten_o_proj=self.use_atten_o_proj,
                     so2_pre_norm=self.so2_pre_norm,
                     so2_post_norm=self.so2_post_norm,
-                    so2_post_norm_eps=1.0e-5 if self.edge_norm else 1.0,
+                    so2_post_norm_eps=1.0e-5 if self.radial_norm else 1.0,
                     so2_activation_function=self.so2_activation_function,
                     ffn_pre_norm=self.ffn_pre_norm,
                     ffn_post_norm=self.ffn_post_norm,
@@ -2617,7 +2634,11 @@ def serialize(self) -> dict[str, Any]:
                 "basis_type": self.basis_type,
                 "n_radial": self.n_radial,
                 "radial_mlp": self.radial_mlp,
-                "edge_norm": self.edge_norm,
+                "edge_norm": [
+                    self.radial_norm,
+                    self.film_norm,
+                    self.focus_norm,
+                ],
                 "use_env_seed": self.use_env_seed,
                 "random_gamma": self.random_gamma,
                 "edge_cartesian": self.edge_cartesian,

deepmd/utils/argcheck.py

@@ -617,7 +617,7 @@ def descrpt_se_zm_args() -> list[Argument]:
     doc_basis_type = "Radial basis type. Supported values are `bessel` and `gaussian`."
     doc_n_radial = "Number of radial basis functions."
     doc_radial_mlp = "Hidden layer sizes for radial networks. An output layer of size (l_schedule[0]+extra_node_l+1)*channels will be automatically appended. Use 0 as a placeholder to be replaced by channels."
-    doc_edge_norm = "Whether to apply standard channel RMSNorm on cutoff-vanishing feature branches. Setting to `false` removes RMSNorm from the radial network, environment-seed FiLM, and cross-focus competition, and uses unit-floor residual scaling for post-SO(2) messages. Setting to `false` is recommended."
+    doc_edge_norm = "Channel RMSNorm on the cutoff-vanishing feature branches. A bool switches every site together: `false` removes the RMSNorm from the radial-network hidden layers, the environment-seed FiLM scale/shift logits and the cross-focus competition scalars, and uses unit-floor residual scaling for post-SO(2) messages. A list of three bools `[radial, film, focus]` switches the sites individually; the post-SO(2) treatment follows the first (radial) entry. Recommended: `[false, true, false]` — the radial-site norms amplify noise where the radial features vanish at the cutoff and produce a spurious long-range force step, while the FiLM and focus norms are safe to keep."
     doc_use_env_seed = (
         "If True, seed the initial node state with local-environment information: "
         "apply environment matrix FiLM conditioning on l=0 features using 4D "
@@ -938,7 +938,7 @@ def descrpt_se_zm_args() -> list[Argument]:
         ),
         Argument(
             "edge_norm",
-            bool,
+            [bool, list],
             optional=True,
             default=True,
             doc=doc_edge_norm,

source/tests/pt/model/test_descriptor_sezm.py

@@ -2346,8 +2346,19 @@ def test_radial_norm_structure_and_serialization(self) -> None:
                     torch.testing.assert_close(mlp(x), restored(x))
 
     def test_edge_norm_gates_all_cutoff_vanishing_norms(self) -> None:
-        """``edge_norm`` controls every cutoff-vanishing normalization path."""
-        for edge_norm in (True, False):
+        """``edge_norm`` controls every cutoff-vanishing normalization path.
+
+        A bool switches the radial, FiLM and focus norms together; a list of
+        three bools ``[radial, film, focus]`` switches them individually. The
+        post-SO(2) residual scaling follows the radial switch.
+        """
+        cases = [
+            (True, (True, True, True)),
+            (False, (False, False, False)),
+            ([False, True, False], (False, True, False)),
+            ([True, False, True], (True, False, True)),
+        ]
+        for edge_norm, (radial_on, film_on, focus_on) in cases:
             with self.subTest(edge_norm=edge_norm):
                 desc = DescrptSeZM(
                     **_descriptor_kwargs(
@@ -2362,26 +2373,84 @@ def test_edge_norm_gates_all_cutoff_vanishing_norms(self) -> None:
                 radial_has_norm = any(
                     type(m).__name__ == "RMSNorm" for m in desc.radial_embedding.net
                 )
-                self.assertEqual(radial_has_norm, edge_norm)
+                self.assertEqual(radial_has_norm, radial_on)
                 # env-seed FiLM scale/shift norms
                 self.assertEqual(
-                    type(desc.film_scale_norm).__name__ == "ScalarRMSNorm", edge_norm
+                    type(desc.film_scale_norm).__name__ == "ScalarRMSNorm", film_on
                 )
                 self.assertEqual(
-                    type(desc.film_shift_norm).__name__ == "ScalarRMSNorm", edge_norm
+                    type(desc.film_shift_norm).__name__ == "ScalarRMSNorm", film_on
                 )
                 # cross-focus competition norm (n_focus>1 -> competition active)
                 focus_norm_mod = desc.blocks[0].so2_conv.focus_compete_norm
                 self.assertEqual(
-                    type(focus_norm_mod).__name__ == "ScalarRMSNorm", edge_norm
+                    type(focus_norm_mod).__name__ == "ScalarRMSNorm", focus_on
                 )
-                # Only the post-SO(2) residual branch uses unit-floor scaling.
-                expected_eps = 1.0e-5 if edge_norm else 1.0
+                # Only the post-SO(2) residual branch uses unit-floor scaling,
+                # bound to the radial switch.
+                expected_eps = 1.0e-5 if radial_on else 1.0
                 self.assertEqual(desc.blocks[0].post_so2_norm.eps, expected_eps)
                 self.assertEqual(desc.blocks[0].pre_so2_norm.eps, 1.0e-5)
                 self.assertEqual(desc.blocks[0].pre_ffn_norms[0].eps, 1.0e-5)
                 self.assertEqual(desc.blocks[0].post_ffn_norms[0].eps, 1.0e-5)
-                self.assertEqual(desc.serialize()["config"]["edge_norm"], edge_norm)
+                canonical = [radial_on, film_on, focus_on]
+                self.assertEqual(desc.serialize()["config"]["edge_norm"], canonical)
+                restored = DescrptSeZM.deserialize(desc.serialize())
+                self.assertEqual(
+                    [
+                        restored.radial_norm,
+                        restored.film_norm,
+                        restored.focus_norm,
+                    ],
+                    canonical,
+                )
+
+    def test_edge_norm_legacy_checkpoint_formats(self) -> None:
+        """Serialized data from older checkpoints loads unchanged.
+
+        The oldest checkpoints predate the ``edge_norm`` option and carry no
+        such config key (the norms were always built, matching the default
+        ``True``); intermediate checkpoints store a plain bool. Both must
+        deserialize into the same module structure and reproduce the source
+        model's output exactly.
+        """
+        dtype = PRECISION_DICT["float64"]
+        cases = [
+            ("missing_key", True, None, (True, True, True)),
+            ("bool_true", True, True, (True, True, True)),
+            ("bool_false", False, False, (False, False, False)),
+        ]
+        for case_name, build_edge_norm, stored_edge_norm, expected in cases:
+            with self.subTest(case=case_name):
+                model = DescrptSeZM(
+                    **_descriptor_kwargs(
+                        edge_norm=build_edge_norm,
+                        use_env_seed=True,
+                        n_focus=2,
+                        precision="float64",
+                    )
+                )
+                data = model.serialize()
+                if stored_edge_norm is None:
+                    data["config"].pop("edge_norm")
+                else:
+                    data["config"]["edge_norm"] = stored_edge_norm
+                restored = DescrptSeZM.deserialize(data)
+                self.assertEqual(
+                    (
+                        restored.radial_norm,
+                        restored.film_norm,
+                        restored.focus_norm,
+                    ),
+                    expected,
+                )
+                coord, atype, nlist = _tiny_two_atom_system(self.device, dtype=dtype)
+                extended_coord = coord.reshape(1, -1)
+                desc1, _, _, _, sw1 = model(extended_coord, atype, nlist)
+                desc2, _, _, _, sw2 = restored(extended_coord, atype, nlist)
+                atol, rtol = _forward_tols(dtype)
+                torch.testing.assert_close(desc1, desc2, atol=atol, rtol=rtol)
+                torch.testing.assert_close(sw1, sw2, atol=atol, rtol=rtol)
 
 
 class TestDescriptorEnergyCurveSmoothness(_SeZMTestCase):