feat(pt): add only_angle_gated_mlp

Author: iProzd

deepmd/dpmodel/descriptor/dpa3.py

@@ -68,6 +68,7 @@ def __init__(
         force_embedding_on_edge: bool = False,
         use_gated_mlp: bool = False,
         gated_mlp_norm: str = "none",
+        only_angle_gated_mlp: bool = False,
         use_res_gnn: bool = False,
         res_gnn_layer: int = 6,
         node_use_rmsnorm: bool = False,
@@ -199,6 +200,7 @@ def __init__(
         self.rk_order = rk_order
         self.rk_update_diff_layer = rk_update_diff_layer
         self.angle_use_node = angle_use_node
+        self.only_angle_gated_mlp = only_angle_gated_mlp
 
     def __getitem__(self, key):
         if hasattr(self, key):

deepmd/pt/model/descriptor/dpa3.py

@@ -205,6 +205,7 @@ def init_subclass_params(sub_data, sub_class):
             force_embedding_on_edge=self.repflow_args.force_embedding_on_edge,
             use_gated_mlp=self.repflow_args.use_gated_mlp,
             gated_mlp_norm=self.repflow_args.gated_mlp_norm,
+            only_angle_gated_mlp=self.repflow_args.only_angle_gated_mlp,
             node_use_rmsnorm=self.repflow_args.node_use_rmsnorm,
             use_res_gnn=self.repflow_args.use_res_gnn,
             res_gnn_layer=self.repflow_args.res_gnn_layer,

deepmd/pt/model/descriptor/repflow_layer.py

@@ -87,6 +87,7 @@ def __init__(
         use_slim_message: bool = False,
         use_gated_mlp: bool = False,
         gated_mlp_norm: str = "none",
+        only_angle_gated_mlp: bool = False,
         node_use_rmsnorm: bool = False,
         angle_use_node: bool = True,
         activation_function: str = "silu",
@@ -177,6 +178,7 @@ def __init__(
 
         self.use_gated_mlp = use_gated_mlp
         self.gated_mlp_norm = gated_mlp_norm
+        self.only_angle_gated_mlp = only_angle_gated_mlp
         if self.use_gated_mlp:
             assert not self.optim_update, "Gated MLP does not support optim update!"
         self.node_use_rmsnorm = node_use_rmsnorm
@@ -265,7 +267,7 @@ def __init__(
             residual_idx += 1
 
         # node edge message
-        if not self.use_gated_mlp:
+        if not self.use_gated_mlp or self.only_angle_gated_mlp:
             self.node_edge_linear = MLPLayer(
                 self.edge_info_dim
                 if not self.use_ffn_node_edge_message
@@ -309,7 +311,7 @@ def __init__(
                 residual_idx += 1
 
         # edge self message
-        if not self.use_gated_mlp:
+        if not self.use_gated_mlp or self.only_angle_gated_mlp:
             self.edge_self_linear = MLPLayer(
                 self.edge_info_dim
                 if not self.use_ffn_edge_edge_message
@@ -1248,7 +1250,7 @@ def forward(
         if not self.optim_update:
             assert edge_info is not None
             if not self.use_ffn_node_edge_message:
-                if not self.use_gated_mlp:
+                if not self.use_gated_mlp or self.only_angle_gated_mlp:
                     node_edge_update = self.act(
                         self.node_edge_linear(edge_info)
                     ) * sw.unsqueeze(-1)
@@ -1320,7 +1322,7 @@ def forward(
         if not self.optim_update:
             assert edge_info is not None
             if not self.use_ffn_edge_edge_message:
-                if not self.use_gated_mlp:
+                if not self.use_gated_mlp or self.only_angle_gated_mlp:
                     edge_self_update = self.act(self.edge_self_linear(edge_info))
                 else:
                     edge_self_update = self.edge_self_linear(edge_info)

deepmd/pt/model/descriptor/repflows.py

@@ -148,6 +148,7 @@ def __init__(
         force_embedding_on_edge: bool = False,
         use_gated_mlp: bool = False,
         gated_mlp_norm: str = "none",
+        only_angle_gated_mlp: bool = False,
         use_res_gnn: bool = False,
         res_gnn_layer: int = 6,
         node_use_rmsnorm: bool = False,
@@ -356,6 +357,7 @@ def __init__(
             assert self.rk_order == 4, "rk_order must be 4 for now"
         self.use_gated_mlp = use_gated_mlp
         self.gated_mlp_norm = gated_mlp_norm
+        self.only_angle_gated_mlp = only_angle_gated_mlp
         self.use_res_gnn = use_res_gnn
         self.res_gnn_layer = res_gnn_layer
         if self.use_res_gnn:
@@ -496,6 +498,7 @@ def __init__(
                     use_slim_message=self.use_slim_message,
                     use_gated_mlp=self.use_gated_mlp,
                     gated_mlp_norm=self.gated_mlp_norm,
+                    only_angle_gated_mlp=self.only_angle_gated_mlp,
                     node_use_rmsnorm=self.node_use_rmsnorm,
                     angle_use_node=self.angle_use_node,
                     seed=child_seed(child_seed(seed, 1), ii),

deepmd/utils/argcheck.py

@@ -1821,6 +1821,12 @@ def dpa3_repflow_args():
             optional=True,
             default="none",
         ),
+        Argument(
+            "only_angle_gated_mlp",
+            bool,
+            optional=True,
+            default=False,
+        ),
         Argument(
             "use_res_gnn",
             bool,