residualize a section of the mlp within the radial function

Author: lucidrains

equiformer_pytorch/equiformer_pytorch.py

@@ -20,17 +20,6 @@
 
 EdgeInfo = namedtuple('EdgeInfo', ['neighbor_indices', 'neighbor_mask', 'edges'])
 
-# biasless layernorm
-
-class LayerNorm(nn.Module):
-    def __init__(self, dim):
-        super().__init__()
-        self.gamma = nn.Parameter(torch.ones(dim))
-        self.register_buffer("beta", torch.zeros(dim))
-
-    def forward(self, x):
-        return F.layer_norm(x, x.shape[-1:], self.gamma, self.beta)
-
 # fiber functions
 
 @beartype
@@ -89,6 +78,23 @@ def feature_fiber(feature):
 
 # classes
 
+class Residual(nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, x, **kwargs):
+        return self.fn(x, **kwargs) + x
+
+class LayerNorm(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.ones(dim))
+        self.register_buffer("beta", torch.zeros(dim))
+
+    def forward(self, x):
+        return F.layer_norm(x, x.shape[-1:], self.gamma, self.beta)
+
 @beartype
 class Linear(nn.Module):
     def __init__(
@@ -328,9 +334,11 @@ def __init__(
         self.rp = nn.Sequential(
             nn.Linear(edge_dim + mid_dim, mid_dim),
             nn.SiLU(),
-            LayerNorm(mid_dim),
-            nn.Linear(mid_dim, mid_dim),
-            nn.SiLU(),
+            Residual(nn.Sequential(
+                LayerNorm(mid_dim),
+                nn.Linear(mid_dim, mid_dim),
+                nn.SiLU()
+            )),
             LayerNorm(mid_dim),
             nn.Linear(mid_dim, self.num_freq * nc_in * nc_out)
         )

equiformer_pytorch/version.py

@@ -1,3 +1,3 @@
-__version__ = '0.0.25'
+__version__ = '0.0.26'
 
 __cuda_pkg_name__ = f'equiformer_pytorch_cuda_{__version__.replace(".", "_")}'