From 2370fc1b042065acf22a54fb29d4e1956964fc13 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Thu, 9 Jul 2026 01:59:37 +0800 Subject: [PATCH 1/2] feat(dpa4): add radial_norm --- deepmd/dpmodel/descriptor/dpa4.py | 85 ++++++++---- deepmd/dpmodel/descriptor/dpa4_nn/block.py | 16 +++ deepmd/dpmodel/descriptor/dpa4_nn/radial.py | 34 +++-- deepmd/dpmodel/descriptor/dpa4_nn/so2.py | 34 +++-- deepmd/pt/model/descriptor/sezm.py | 51 +++++-- deepmd/pt/model/descriptor/sezm_nn/block.py | 18 +++ deepmd/pt/model/descriptor/sezm_nn/radial.py | 34 +++-- deepmd/pt/model/descriptor/sezm_nn/so2.py | 30 ++++- deepmd/utils/argcheck.py | 8 ++ examples/water/dpa4/input.json | 1 + source/tests/pt/model/test_descriptor_sezm.py | 127 ++++++++++++++++++ .../pt/model/test_dpa4_dpmodel_parity.py | 72 +++++++++- 12 files changed, 435 insertions(+), 75 deletions(-) diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py index d7b4287199..8d59a5caba 100644 --- a/deepmd/dpmodel/descriptor/dpa4.py +++ b/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, diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/block.py b/deepmd/dpmodel/descriptor/dpa4_nn/block.py index c13acf2617..751efe3f95 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/block.py +++ b/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, diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/radial.py b/deepmd/dpmodel/descriptor/dpa4_nn/radial.py index 3766588de9..6012cc4256 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/radial.py +++ b/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, } diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/so2.py b/deepmd/dpmodel/descriptor/dpa4_nn/so2.py index c0c4c0b262..4c8cd80e9a 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/so2.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/so2.py @@ -941,6 +941,10 @@ class SO2Convolution(NativeOP): If True, apply cross-focus softmax competition in SO(2) local layout. Competition logits are constructed only from l=0 scalar channels and the resulting invariant weights are broadcast to all (l, m) components. + focus_norm + If True, RMS-normalize the competition l=0 scalars before the softmax. + Those scalars are envelope-gated and vanish at the cutoff, so ``False`` + drops the norm to keep the competition smooth there. so2_norm If True, apply intermediate ReducedEquivariantRMSNorm as pre-norm before each SO(2) mixing layer. The last SO(2) layer always uses Identity. @@ -1056,6 +1060,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", @@ -1105,6 +1110,7 @@ def __init__( self.hidden_channels = int(self.n_focus * self.so2_focus_dim) self.use_hidden_projection = self.hidden_channels != self.channels self.focus_compete = bool(focus_compete) + self.focus_norm = bool(focus_norm) self.focus_softmax_tau = 1.0 self.focus_label_smoothing = 0.02 self.so2_norm = bool(so2_norm) @@ -1367,13 +1373,17 @@ def __init__( self.adamw_focus_compete_w: np.ndarray | None = None self.focus_compete_bias: np.ndarray | None = None if self.focus_compete and self.n_focus > 1: - self.focus_compete_norm = ScalarRMSNorm( - channels=self.so2_focus_dim, - n_focus=self.n_focus, - eps=self.eps, - precision=self.compute_precision, - trainable=trainable, - ) + # ``focus_norm=False`` drops the competition-input RMSNorm (which + # would cross its eps floor as the envelope-gated scalars vanish at + # rcut); the competition weights then decay smoothly to uniform. + if self.focus_norm: + self.focus_compete_norm = ScalarRMSNorm( + channels=self.so2_focus_dim, + n_focus=self.n_focus, + eps=self.eps, + precision=self.compute_precision, + trainable=trainable, + ) self.adamw_focus_compete_w = ( np.random.default_rng(child_seed(seed_gate, 4)) .normal( @@ -2381,10 +2391,13 @@ def _focus_alpha(self, focus_gate_src: Array) -> Array: """ xp = array_api_compat.array_namespace(focus_gate_src) device = array_api_compat.device(focus_gate_src) + focus_in = xp.astype( + focus_gate_src, get_xp_precision(xp, self.compute_precision) + ) + if self.focus_norm: + focus_in = self.focus_compete_norm(focus_in) focus_logits = xp.sum( - self.focus_compete_norm( - xp.astype(focus_gate_src, get_xp_precision(xp, self.compute_precision)) - ) + focus_in * xp.permute_dims( xp_asarray_nodetach(xp, self.adamw_focus_compete_w[...], device=device), (1, 0), @@ -2706,6 +2719,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, diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py index 5ec9a1da19..d39f7a1028 100644 --- a/deepmd/pt/model/descriptor/sezm.py +++ b/deepmd/pt/model/descriptor/sezm.py @@ -159,6 +159,15 @@ class DescrptSeZM(BaseDescriptor, nn.Module): 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 @@ -441,6 +450,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, @@ -542,6 +552,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: @@ -845,20 +856,28 @@ def __init__( seed=seed_env_seed, ) ) - self.film_scale_norm = ScalarRMSNorm( - channels=self.channels, - n_focus=1, - eps=self.eps, - dtype=self.compute_dtype, - trainable=self.trainable, - ) - self.film_shift_norm = ScalarRMSNorm( - channels=self.channels, - n_focus=1, - eps=self.eps, - dtype=self.compute_dtype, - 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: nn.Module = ScalarRMSNorm( + channels=self.channels, + n_focus=1, + eps=self.eps, + dtype=self.compute_dtype, + trainable=self.trainable, + ) + self.film_shift_norm: nn.Module = ScalarRMSNorm( + channels=self.channels, + n_focus=1, + eps=self.eps, + dtype=self.compute_dtype, + trainable=self.trainable, + ) + else: + self.film_scale_norm = nn.Identity() + self.film_shift_norm = nn.Identity() film_strength_init = 0.01 # Use 1D tensor (not scalar) for FSDP2 compatibility self.film_scale_strength_log = nn.Parameter( @@ -906,6 +925,7 @@ def __init__( activation_function=self.activation_function, dtype=self.compute_dtype, # force fp32+ trainable=self.trainable, + radial_norm=self.edge_norm, seed=seed_radial_embedding, ) @@ -968,6 +988,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, @@ -1002,6 +1023,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, @@ -2439,6 +2461,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, diff --git a/deepmd/pt/model/descriptor/sezm_nn/block.py b/deepmd/pt/model/descriptor/sezm_nn/block.py index 825ff2a5e3..6b170a8935 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/block.py +++ b/deepmd/pt/model/descriptor/sezm_nn/block.py @@ -151,6 +151,12 @@ class SeZMInteractionBlock(nn.Module): ``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 (via radial modulation) + and vanishes at the cutoff, so normalizing it crosses the norm ``eps`` + floor near ``rcut``; ``False`` drops the norm so the competition decays + smoothly to uniform weights at the cutoff. so2_norm If True, apply intermediate ReducedEquivariantRMSNorm between SO(2) mixing layers. When False (default), no normalization is applied between layers. @@ -198,6 +204,10 @@ class SeZMInteractionBlock(nn.Module): 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 @@ -299,6 +309,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", @@ -312,6 +323,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, @@ -366,6 +378,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() @@ -383,6 +396,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) @@ -463,6 +477,7 @@ def __init__( self.lmax, self.channels, n_focus=1, + eps=self.so2_post_norm_eps, dtype=self.compute_dtype, trainable=trainable, ) @@ -477,6 +492,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, @@ -1037,6 +1053,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, @@ -1050,6 +1067,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, diff --git a/deepmd/pt/model/descriptor/sezm_nn/radial.py b/deepmd/pt/model/descriptor/sezm_nn/radial.py index 4f4e4d888c..cef4e402c3 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/radial.py +++ b/deepmd/pt/model/descriptor/sezm_nn/radial.py @@ -52,7 +52,7 @@ class RadialMLP(nn.Module): """ - Radial MLP with channel RMSNorm and configurable activation. + Radial MLP with optional channel RMSNorm and configurable activation. Parameters ---------- @@ -65,11 +65,14 @@ class RadialMLP(nn.Module): Floating point dtype 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 ----- @@ -78,6 +81,15 @@ class RadialMLP(nn.Module): 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__( @@ -87,6 +99,7 @@ def __init__( activation_function: str = "silu", dtype: torch.dtype = torch.float32, trainable: bool = True, + radial_norm: bool = True, seed: int | list[int] | None = None, ) -> None: super().__init__() @@ -98,6 +111,7 @@ def __init__( self.device = env.DEVICE self.precision = RESERVED_PRECISION_DICT[self.dtype] self.trainable = bool(trainable) + self.radial_norm = bool(radial_norm) modules: list[nn.Module] = [] n_layers = len(mlp_layers) @@ -114,13 +128,14 @@ def __init__( modules.append(linear) # Last layer: no RMSNorm/activation if i < n_layers - 2: - modules.append( - RMSNorm( - channels=mlp_layers[i + 1], - dtype=self.dtype, - trainable=trainable, + if self.radial_norm: + modules.append( + RMSNorm( + channels=mlp_layers[i + 1], + dtype=self.dtype, + trainable=trainable, + ) ) - ) modules.append(ActivationFn(self.activation_function)) self.net = nn.Sequential(*modules) @@ -151,6 +166,7 @@ def serialize(self) -> dict[str, Any]: "activation_function": self.activation_function, "dtype": RESERVED_PRECISION_DICT[self.dtype], "trainable": self.trainable, + "radial_norm": self.radial_norm, "@variables": {k: np_safe(v) for k, v in state.items()}, } diff --git a/deepmd/pt/model/descriptor/sezm_nn/so2.py b/deepmd/pt/model/descriptor/sezm_nn/so2.py index fd194c94d7..b7548507f3 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/so2.py +++ b/deepmd/pt/model/descriptor/sezm_nn/so2.py @@ -903,6 +903,11 @@ class SO2Convolution(nn.Module): If True, apply cross-focus softmax competition in SO(2) local layout. Competition logits are constructed only from l=0 scalar channels and the resulting invariant weights are broadcast to all (l, m) components. + focus_norm + If True, RMS-normalize the competition l=0 scalars before the softmax. + Those scalars are envelope-gated (radial modulation) and vanish at the + cutoff, so the norm crosses its eps floor near ``rcut``; ``False`` uses an + identity pass-through and lets the competition decay smoothly to uniform. so2_norm If True, apply intermediate ReducedEquivariantRMSNorm as pre-norm before each SO(2) mixing layer. The last SO(2) layer always uses Identity. @@ -1018,6 +1023,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", @@ -1068,6 +1074,7 @@ def __init__( self.hidden_channels = int(self.n_focus * self.so2_focus_dim) self.use_hidden_projection = self.hidden_channels != self.channels self.focus_compete = bool(focus_compete) + self.focus_norm = bool(focus_norm) self.focus_softmax_tau = 1.0 self.focus_label_smoothing = 0.02 self.so2_norm = bool(so2_norm) @@ -1370,16 +1377,24 @@ def __init__( ) # === Step 7.5. Optional cross-focus competition === - self.focus_compete_norm: ScalarRMSNorm | None = None + self.focus_compete_norm: nn.Module | None = None self.adamw_focus_compete_w: nn.Parameter | None = None self.focus_compete_bias: nn.Parameter | None = None if self.focus_compete and self.n_focus > 1: - self.focus_compete_norm = ScalarRMSNorm( - channels=self.so2_focus_dim, - n_focus=self.n_focus, - eps=self.eps, - dtype=self.compute_dtype, - trainable=trainable, + # The competition scalars are envelope-gated (radial modulation) and + # vanish at rcut; normalizing them crosses the eps floor near the + # cutoff, so ``focus_norm=False`` uses an identity pass-through and + # lets the softmax decay smoothly to uniform weights there. + self.focus_compete_norm = ( + ScalarRMSNorm( + channels=self.so2_focus_dim, + n_focus=self.n_focus, + eps=self.eps, + dtype=self.compute_dtype, + trainable=trainable, + ) + if self.focus_norm + else nn.Identity() ) self.adamw_focus_compete_w = nn.Parameter( torch.empty( @@ -2488,6 +2503,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, diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index fe5f542933..ab14e59bb1 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -416,6 +416,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_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 " @@ -734,6 +735,13 @@ def descrpt_se_zm_args() -> list[Argument]: default=[0], doc=doc_radial_mlp, ), + Argument( + "edge_norm", + bool, + optional=True, + default=True, + doc=doc_edge_norm, + ), Argument( "use_env_seed", bool, diff --git a/examples/water/dpa4/input.json b/examples/water/dpa4/input.json index 1126c7107a..1819a3afad 100644 --- a/examples/water/dpa4/input.json +++ b/examples/water/dpa4/input.json @@ -10,6 +10,7 @@ "rcut": 6.0, "channels": 32, "n_radial": 16, + "edge_norm": false, "use_env_seed": true, "edge_cartesian": false, "node_cartesian": "none", diff --git a/source/tests/pt/model/test_descriptor_sezm.py b/source/tests/pt/model/test_descriptor_sezm.py index 161e8a87e6..835d7c6a5d 100644 --- a/source/tests/pt/model/test_descriptor_sezm.py +++ b/source/tests/pt/model/test_descriptor_sezm.py @@ -24,6 +24,8 @@ ForceEmbedding, InnerClamp, NodeCartesianTensorProduct, + RadialBasis, + RadialMLP, SeZMDirectForceHead, SO2Linear, SpinEmbedding, @@ -1826,6 +1828,131 @@ def test_invalid_params(self) -> None: InnerClamp(1.0, 1.0) +class TestEdgeNorm(_SeZMTestCase): + """The ``edge_norm`` switch and its effect on cutoff smoothness. + + The descriptor exposes a single ``edge_norm`` flag; internally it drives the + ``RadialMLP.radial_norm`` hidden RMSNorm, the FiLM scale/shift norms, and the + cross-focus competition norm, and selects the post-SO(2) residual scaling + floor. The RadialMLP-level tests exercise the radial mechanism directly; the + descriptor test checks the umbrella propagation. + """ + + def setUp(self) -> None: + super().setUp() + self.dtype = torch.float64 + self.rcut = 6.0 + + def _radial_feature_curve(self, *, radial_norm: bool, seed: int) -> torch.Tensor: + """Radial features over a near-cutoff distance sweep. + + ``RadialBasis`` bakes in the C^3 envelope, so ``edge_rbf`` vanishes at + ``rcut``. With ``radial_norm=True`` the hidden RMSNorm divides that + envelope out and its ``eps`` floor is crossed near ``rcut``, injecting a + localized curvature spike; ``radial_norm=False`` drops the RMSNorm so the + feature stays smooth. Both variants share the same linear weights (same + ``seed``), isolating the effect of the norm. + + Parameters + ---------- + radial_norm : bool + Whether the RadialMLP keeps its hidden RMSNorm. + seed : int + Seed shared by both variants so the linear weights match. + + Returns + ------- + torch.Tensor + Radial features with shape (N, out_dim) over the distance sweep. + """ + basis = RadialBasis(rcut=self.rcut, n_radial=8, exponent=7, dtype=self.dtype) + mlp = RadialMLP( + [8, 12, 8], radial_norm=radial_norm, dtype=self.dtype, seed=seed + ) + r = torch.linspace( + 0.5 * self.rcut, + 0.9995 * self.rcut, + 4000, + dtype=self.dtype, + device=self.device, + ).view(-1, 1) + with torch.no_grad(): + return mlp(basis(r)) + + @staticmethod + def _peak_curvature(feat: torch.Tensor) -> float: + """Peak absolute second finite difference of the feature L2 norm.""" + y = feat.norm(dim=1) + return float((y[2:] - 2.0 * y[1:-1] + y[:-2]).abs().max()) + + def test_radial_norm_false_removes_cutoff_curvature_spike(self) -> None: + """``radial_norm=False`` suppresses the eps-crossing kink near ``rcut``.""" + feat_norm = self._radial_feature_curve(radial_norm=True, seed=3) + feat_smooth = self._radial_feature_curve(radial_norm=False, seed=3) + # Both vanish at rcut: edge_rbf -> 0 and RadialMLP(0) = 0 (bias=False). + self.assertLess(feat_smooth[-1].abs().max().item(), 1.0e-8) + # The normalized variant floor-crosses just inside rcut; dropping the + # RMSNorm removes that localized curvature spike by a wide margin. + self.assertLess( + self._peak_curvature(feat_smooth) * 5.0, + self._peak_curvature(feat_norm), + ) + + def test_radial_norm_structure_and_serialization(self) -> None: + """The flag toggles the hidden RMSNorm and round-trips through serialize.""" + for radial_norm in (True, False): + with self.subTest(radial_norm=radial_norm): + mlp = RadialMLP( + [8, 12, 8], radial_norm=radial_norm, dtype=self.dtype, seed=5 + ) + has_norm = any(type(m).__name__ == "RMSNorm" for m in mlp.net) + self.assertEqual(has_norm, radial_norm) + + restored = RadialMLP.deserialize(mlp.serialize()) + self.assertEqual(restored.radial_norm, radial_norm) + x = torch.rand(16, 8, dtype=self.dtype, device=self.device) + with torch.no_grad(): + 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): + with self.subTest(edge_norm=edge_norm): + desc = DescrptSeZM( + **_descriptor_kwargs( + edge_norm=edge_norm, + use_env_seed=True, + n_focus=2, + sandwich_norm=[True, True, True, True], + precision="float64", + ) + ) + # radial MLP hidden RMSNorm + radial_has_norm = any( + type(m).__name__ == "RMSNorm" for m in desc.radial_embedding.net + ) + self.assertEqual(radial_has_norm, edge_norm) + # env-seed FiLM scale/shift norms + self.assertEqual( + type(desc.film_scale_norm).__name__ == "ScalarRMSNorm", edge_norm + ) + self.assertEqual( + type(desc.film_shift_norm).__name__ == "ScalarRMSNorm", edge_norm + ) + # 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 + ) + # Only the post-SO(2) residual branch uses unit-floor scaling. + expected_eps = 1.0e-5 if edge_norm 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) + + class TestDescriptorEnergyCurveSmoothness(_SeZMTestCase): """Test PES smoothness from scaled symmetric eight-atom probes.""" diff --git a/source/tests/pt/model/test_dpa4_dpmodel_parity.py b/source/tests/pt/model/test_dpa4_dpmodel_parity.py index 345583b092..f5d97a709b 100644 --- a/source/tests/pt/model/test_dpa4_dpmodel_parity.py +++ b/source/tests/pt/model/test_dpa4_dpmodel_parity.py @@ -806,7 +806,8 @@ def _perturb(self, pt_mod: torch.nn.Module, seed: int) -> None: @pytest.mark.parametrize("lmax", [0, 2, 3]) # 0 covers the scalar-only branch @pytest.mark.parametrize("n_focus", [1, 2]) # focus streams - def test_equivariant_rmsnorm(self, lmax, n_focus) -> None: + @pytest.mark.parametrize("eps", [1.0e-5, 1.0]) + def test_equivariant_rmsnorm(self, lmax, n_focus, eps) -> None: from deepmd.dpmodel.descriptor.dpa4_nn.norm import ( EquivariantRMSNorm as DPEquivariantRMSNorm, ) @@ -815,10 +816,16 @@ def test_equivariant_rmsnorm(self, lmax, n_focus) -> None: ) pt_mod = PTEquivariantRMSNorm( - lmax, self.channels, n_focus, dtype=torch.float64, trainable=True + lmax, + self.channels, + n_focus, + eps=eps, + dtype=torch.float64, + trainable=True, ) self._perturb(pt_mod, 2040) serialized = pt_mod.serialize() + assert serialized["config"]["eps"] == eps # pt state_dict key contract: 2 parameters + 2 persistent buffers assert set(serialized["@variables"]) == { "adam_scale", @@ -832,6 +839,40 @@ def test_equivariant_rmsnorm(self, lmax, n_focus) -> None: x[0] = 0.0 # all-zeros row exercises the eps path assert_parity(dp_mod.call(x), pt_mod(to_pt(x))) + def test_equivariant_rmsnorm_eps_one_reparameterization(self) -> None: + """An epsilon of one preserves the legacy branch function class.""" + from deepmd.pt.model.descriptor.sezm_nn.norm import ( + EquivariantRMSNorm as PTEquivariantRMSNorm, + ) + + legacy = PTEquivariantRMSNorm( + 2, + self.channels, + 1, + eps=1.0e-5, + dtype=torch.float64, + trainable=True, + ) + unit_scale = PTEquivariantRMSNorm( + 2, + self.channels, + 1, + eps=1.0, + dtype=torch.float64, + trainable=True, + ) + self._perturb(legacy, 2042) + unit_scale.load_state_dict(legacy.state_dict()) + + rng = np.random.default_rng(2043) + x = to_pt(rng.normal(size=(17, 9, 1, self.channels))) + torch.testing.assert_close( + legacy(x), + unit_scale(x / np.sqrt(legacy.eps)), + rtol=1.0e-12, + atol=1.0e-12, + ) + def test_equivariant_rmsnorm_roundtrip(self) -> None: from deepmd.dpmodel.descriptor.dpa4_nn.norm import ( EquivariantRMSNorm as DPEquivariantRMSNorm, @@ -3293,6 +3334,17 @@ def test_block_sandwich_norm(self, sandwich) -> None: ) self._assert_block_parity(pt_mod, dp_mod, kwargs) + def test_block_post_so2_eps_one(self) -> None: + pt_mod, dp_mod, kwargs = self._build_block_pair( + so2_post_norm=True, + so2_post_norm_eps=1.0, + ) + assert pt_mod.post_so2_norm.eps == 1.0 + assert dp_mod.post_so2_norm.eps == 1.0 + assert pt_mod.pre_ffn_norms[0].eps == 1.0e-5 + assert dp_mod.pre_ffn_norms[0].eps == 1.0e-5 + self._assert_block_parity(pt_mod, dp_mod, kwargs) + def test_block_ffn_blocks(self) -> None: # multiple FFN subblocks exercise the per-subblock loop and seeds pt_mod, dp_mod, kwargs = self._build_block_pair(ffn_blocks=2) @@ -3521,6 +3573,22 @@ def test_descriptor(self, use_env_seed, n_blocks) -> None: ) self._assert_descr_parity(pt_mod, dp_mod) + @pytest.mark.parametrize( + "edge_norm", [False, True] + ) # cutoff-vanishing normalization modes + def test_descriptor_edge_norm(self, edge_norm) -> None: + # edge_norm=False drops the radial MLP RMSNorm, turns the FiLM scale/shift + # norms into identity pass-throughs, drops the focus-compete norm, and + # selects unit-floor post-SO(2) residual scaling in both backends. + pt_mod, dp_mod, _ = self._build_descr_pair( + edge_norm=edge_norm, use_env_seed=True, n_focus=2 + ) + assert dp_mod.edge_norm == edge_norm + expected_eps = 1.0e-5 if edge_norm else 1.0 + assert pt_mod.blocks[0].post_so2_norm.eps == expected_eps + assert dp_mod.blocks[0].post_so2_norm.eps == expected_eps + self._assert_descr_parity(pt_mod, dp_mod) + @pytest.mark.parametrize( "exclude_types", [[], [(0, 0)]] ) # pair-exclusion off vs on From c8a26467e2931f1a5cacb043c8863884d8a44250 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sun, 12 Jul 2026 19:34:48 +0800 Subject: [PATCH 2/2] test(pt): respect device-specific parity tolerances --- source/tests/pt/model/test_dpa4_dpmodel_parity.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/source/tests/pt/model/test_dpa4_dpmodel_parity.py b/source/tests/pt/model/test_dpa4_dpmodel_parity.py index f5d97a709b..c286f0fca7 100644 --- a/source/tests/pt/model/test_dpa4_dpmodel_parity.py +++ b/source/tests/pt/model/test_dpa4_dpmodel_parity.py @@ -869,8 +869,8 @@ def test_equivariant_rmsnorm_eps_one_reparameterization(self) -> None: torch.testing.assert_close( legacy(x), unit_scale(x / np.sqrt(legacy.eps)), - rtol=1.0e-12, - atol=1.0e-12, + rtol=PT_RTOL, + atol=PT_ATOL, ) def test_equivariant_rmsnorm_roundtrip(self) -> None: