feat(dpa4): preserve cutoff smoothness in edge feature branches

deepmd/dpmodel/descriptor/dpa4.py

@@ -176,6 +176,15 @@ class DescrptDPA4(NativeOP, BaseDescriptor):
     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.
+    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.
     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
@@ -457,6 +466,7 @@ def __init__(
         basis_type: str = "bessel",
         n_radial: int = 16,
         radial_mlp: list[int] | None = None,
+        edge_norm: bool = True,
         use_env_seed: bool = True,
         random_gamma: bool = True,
         edge_cartesian: bool = False,
@@ -551,6 +561,7 @@ 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 sandwich_norm is None:
             sandwich_norm = [False, True, True, False]
         if not isinstance(sandwich_norm, (list, tuple)) or len(sandwich_norm) != 4:
@@ -849,20 +860,28 @@ def __init__(
                     seed=seed_env_seed,
                 )
             )
-            self.film_scale_norm = ScalarRMSNorm(
-                channels=self.channels,
-                n_focus=1,
-                eps=self.eps,
-                precision=self.compute_precision,
-                trainable=self.trainable,
-            )
-            self.film_shift_norm = ScalarRMSNorm(
-                channels=self.channels,
-                n_focus=1,
-                eps=self.eps,
-                precision=self.compute_precision,
-                trainable=self.trainable,
-            )
+            # The FiLM logits derive from the env-seed matrix D = envᵀenv, which
+            # 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:
+                self.film_scale_norm = ScalarRMSNorm(
+                    channels=self.channels,
+                    n_focus=1,
+                    eps=self.eps,
+                    precision=self.compute_precision,
+                    trainable=self.trainable,
+                )
+                self.film_shift_norm = ScalarRMSNorm(
+                    channels=self.channels,
+                    n_focus=1,
+                    eps=self.eps,
+                    precision=self.compute_precision,
+                    trainable=self.trainable,
+                )
+            else:
+                self.film_scale_norm = None
+                self.film_shift_norm = None
             film_strength_init = 0.01
             # Use 1D tensor (not scalar) for FSDP2 compatibility
             self.film_scale_strength_log = np.full(
@@ -902,6 +921,7 @@ def __init__(
             activation_function=self.activation_function,
             precision=self.compute_precision,  # force fp32+
             trainable=self.trainable,
+            radial_norm=self.edge_norm,
             seed=seed_radial_embedding,
         )
 
@@ -964,6 +984,7 @@ def __init__(
                     channels=self.channels,
                     n_focus=self.n_focus,
                     focus_dim=self.focus_dim,
+                    focus_norm=self.edge_norm,
                     so2_norm=self.so2_norm,
                     mixing_layers=self.mixing_layers,
                     so2_attn_res=self.so2_attn_res_mode,
@@ -998,6 +1019,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_activation_function=self.so2_activation_function,
                     ffn_pre_norm=self.ffn_pre_norm,
                     ffn_post_norm=self.ffn_post_norm,
@@ -1267,8 +1289,12 @@ def call(
             )  # (N, 2*C)
             scale_logits = film[:, : self.channels]  # (N, C)
             shift_logits = film[:, self.channels :]  # (N, C)
-            scale_hat = self.film_scale_norm(scale_logits)  # (N, C)
-            shift_hat = self.film_shift_norm(shift_logits)  # (N, C)
+            scale_hat = (
+                self.film_scale_norm(scale_logits) if self.edge_norm else scale_logits
+            )  # (N, C)
+            shift_hat = (
+                self.film_shift_norm(shift_logits) if self.edge_norm else shift_logits
+            )  # (N, C)
             scale_strength = xp.exp(
                 xp_asarray_nodetach(
                     xp, self.film_scale_strength_log[...], device=device
@@ -1522,8 +1548,12 @@ def call_with_edges(
             )  # (N, 2*C)
             scale_logits = film[:, : self.channels]  # (N, C)
             shift_logits = film[:, self.channels :]  # (N, C)
-            scale_hat = self.film_scale_norm(scale_logits)  # (N, C)
-            shift_hat = self.film_shift_norm(shift_logits)  # (N, C)
+            scale_hat = (
+                self.film_scale_norm(scale_logits) if self.edge_norm else scale_logits
+            )  # (N, C)
+            shift_hat = (
+                self.film_shift_norm(shift_logits) if self.edge_norm else shift_logits
+            )  # (N, C)
             scale_strength = xp.exp(
                 xp_asarray_nodetach(
                     xp, self.film_scale_strength_log[...], device=device
@@ -2394,10 +2424,15 @@ 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
-            for key, value in self.film_scale_norm.serialize()["@variables"].items():
-                variables[f"film_scale_norm.{key}"] = value
-            for key, value in self.film_shift_norm.serialize()["@variables"].items():
-                variables[f"film_shift_norm.{key}"] = value
+            if self.edge_norm:
+                for key, value in self.film_scale_norm.serialize()[
+                    "@variables"
+                ].items():
+                    variables[f"film_scale_norm.{key}"] = value
+                for key, value in self.film_shift_norm.serialize()[
+                    "@variables"
+                ].items():
+                    variables[f"film_shift_norm.{key}"] = value
             variables["film_scale_strength_log"] = to_numpy_array(
                 self.film_scale_strength_log
             )
@@ -2496,8 +2531,9 @@ def load(module: Any, prefix: str) -> Any:
             self.env_seed_embedding = load(
                 self.env_seed_embedding, "env_seed_embedding."
             )
-            self.film_scale_norm = load(self.film_scale_norm, "film_scale_norm.")
-            self.film_shift_norm = load(self.film_shift_norm, "film_shift_norm.")
+            if self.edge_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(
                 variables["film_scale_strength_log"], dtype=compute_prec
             )
@@ -2547,6 +2583,7 @@ 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,
                 "use_env_seed": self.use_env_seed,
                 "random_gamma": self.random_gamma,
                 "edge_cartesian": self.edge_cartesian,

deepmd/dpmodel/descriptor/dpa4_nn/block.py

@@ -145,6 +145,10 @@ class SeZMInteractionBlock(NativeOP):
         ``focus_dim=0`` means using ``channels``.
     focus_compete
         If True, enable cross-focus softmax competition in SO(2) convolution.
+    focus_norm
+        If True, RMS-normalize the cross-focus competition scalars before the
+        softmax. The competition input is envelope-gated and vanishes at the
+        cutoff, so ``False`` drops the norm to keep the competition smooth there.
     so2_norm
         If True, apply intermediate ReducedEquivariantRMSNorm between SO(2) mixing layers.
         When False (default), no normalization is applied between layers.
@@ -192,6 +196,10 @@ class SeZMInteractionBlock(NativeOP):
         If True, apply pre-norm before SO(2) convolution.
     so2_post_norm
         If True, apply post-norm on SO(2) output before the residual add.
+    so2_post_norm_eps
+        Variance floor for the SO(2) post-norm. A value of ``1`` preserves small
+        residual messages instead of amplifying them by ``1/sqrt(eps)``. Other
+        normalization sites retain their own RMSNorm floors.
     ffn_pre_norm
         If True, apply pre-norm before each FFN subblock.
     ffn_post_norm
@@ -293,6 +301,7 @@ def __init__(
         n_focus: int = 1,
         focus_dim: int = 0,
         focus_compete: bool = True,
+        focus_norm: bool = True,
         so2_norm: bool = False,
         mixing_layers: int = 4,
         so2_attn_res: str = "none",
@@ -306,6 +315,7 @@ def __init__(
         atten_o_proj: bool = False,
         so2_pre_norm: bool = True,
         so2_post_norm: bool = False,
+        so2_post_norm_eps: float = 1e-5,
         ffn_pre_norm: bool = True,
         ffn_post_norm: bool = False,
         ffn_neurons: int = 96,
@@ -359,6 +369,7 @@ def __init__(
         if self.focus_dim < 0:
             raise ValueError("`focus_dim` must be >= 0")
         self.focus_compete = bool(focus_compete)
+        self.focus_norm = bool(focus_norm)
         self.so2_norm = bool(so2_norm)
         self.mixing_layers = int(mixing_layers)
         self.so2_attn_res_mode = str(so2_attn_res).lower()
@@ -376,6 +387,7 @@ def __init__(
         self.use_atten_o_proj = bool(atten_o_proj)
         self.so2_pre_norm = bool(so2_pre_norm)
         self.so2_post_norm = bool(so2_post_norm)
+        self.so2_post_norm_eps = float(so2_post_norm_eps)
         self.ffn_pre_norm = bool(ffn_pre_norm)
         self.ffn_post_norm = bool(ffn_post_norm)
         self.ffn_neurons = int(ffn_neurons)
@@ -456,6 +468,7 @@ def __init__(
                 self.lmax,
                 self.channels,
                 n_focus=1,
+                eps=self.so2_post_norm_eps,
                 precision=self.compute_precision,
                 trainable=trainable,
             )
@@ -470,6 +483,7 @@ def __init__(
             n_focus=self.n_focus,
             focus_dim=self.focus_dim,
             focus_compete=self.focus_compete,
+            focus_norm=self.focus_norm,
             so2_norm=self.so2_norm,
             mixing_layers=self.mixing_layers,
             so2_attn_res=self.so2_attn_res_mode,
@@ -1065,6 +1079,7 @@ def serialize(self) -> dict[str, Any]:
                 "n_focus": self.n_focus,
                 "focus_dim": self.focus_dim,
                 "focus_compete": self.focus_compete,
+                "focus_norm": self.focus_norm,
                 "so2_norm": self.so2_norm,
                 "mixing_layers": self.mixing_layers,
                 "so2_attn_res": self.so2_attn_res_mode,
@@ -1078,6 +1093,7 @@ def serialize(self) -> dict[str, Any]:
                 "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": self.so2_post_norm_eps,
                 "ffn_pre_norm": self.ffn_pre_norm,
                 "ffn_post_norm": self.ffn_post_norm,
                 "ffn_neurons": self.ffn_neurons,

deepmd/dpmodel/descriptor/dpa4_nn/radial.py

@@ -50,7 +50,7 @@
 
 class RadialMLP(NativeOP):
     """
-    Radial MLP with channel RMSNorm and configurable activation.
+    Radial MLP with optional channel RMSNorm and configurable activation.
 
     Parameters
     ----------
@@ -63,11 +63,14 @@ class RadialMLP(NativeOP):
         Floating point precision for the linear layers.
     trainable : bool
         Whether the parameters are trainable.
+    radial_norm : bool
+        Whether to insert a channel RMSNorm in each hidden layer.
 
     Architecture
     ------------
-    Linear → RMSNorm → Activation for all hidden layers,
-    with the final layer being a plain Linear (no norm, no activation).
+    ``radial_norm=True``  : Linear → RMSNorm → Activation for each hidden layer.
+    ``radial_norm=False`` : Linear → Activation for each hidden layer.
+    The final layer is always a plain Linear (no norm, no activation).
 
     Notes
     -----
@@ -76,6 +79,15 @@ class RadialMLP(NativeOP):
     pads masked edges with zero ``edge_rbf``; any non-zero bias would leak
     spurious features into GIE scatter, causing energy divergence between
     compile and non-compile paths.
+
+    The hidden RMSNorm normalizes each edge's radial features by their own RMS.
+    The input ``edge_rbf`` carries the C^3 cutoff envelope and therefore
+    vanishes at ``rcut``; the RMSNorm divides that envelope out, and its ``eps``
+    floor is crossed as the edge approaches ``rcut``. On a sparse neighborhood
+    (e.g. a dimer) this floor-crossing produces a sharp kink in the potential
+    energy surface just inside the cutoff. Setting ``radial_norm=False`` drops
+    the RMSNorm so the radial features vanish smoothly with the envelope, which
+    restores C^3 smoothness at the cutoff.
     """
 
     def __init__(
@@ -85,6 +97,7 @@ def __init__(
         activation_function: str = "silu",
         precision: str = DEFAULT_PRECISION,
         trainable: bool = True,
+        radial_norm: bool = True,
         seed: int | list[int] | None = None,
     ) -> None:
         if len(mlp_layers) < 2:
@@ -93,6 +106,7 @@ def __init__(
         self.activation_function = str(activation_function)
         self.precision = precision
         self.trainable = bool(trainable)
+        self.radial_norm = bool(radial_norm)
 
         modules: list = []
         n_layers = len(mlp_layers)
@@ -109,13 +123,14 @@ def __init__(
             modules.append(linear)
             # Last layer: no RMSNorm/activation
             if i < n_layers - 2:
-                modules.append(
-                    RMSNorm(
-                        channels=mlp_layers[i + 1],
-                        precision=self.precision,
-                        trainable=trainable,
+                if self.radial_norm:
+                    modules.append(
+                        RMSNorm(
+                            channels=mlp_layers[i + 1],
+                            precision=self.precision,
+                            trainable=trainable,
+                        )
                     )
-                )
                 modules.append(get_activation_fn(self.activation_function))
 
         self.net = modules
@@ -153,6 +168,7 @@ def serialize(self) -> dict[str, Any]:
             "activation_function": self.activation_function,
             "dtype": np.dtype(PRECISION_DICT[self.precision]).name,
             "trainable": self.trainable,
+            "radial_norm": self.radial_norm,
             "@variables": variables,
         }