zero init the attention and feedforward branch outputs by default, but turn it off for testing

Author: lucidrains

equiformer_pytorch/equiformer_pytorch.py

@@ -168,6 +168,10 @@ def __init__(
             self.weights.append(nn.Parameter(torch.randn(dim_in, dim_out) / sqrt(dim_in)))
             self.degrees.append(degree)
 
+    def init_zero_(self):
+        for weight in self.weights:
+            weight.data.zero_()
+
     def forward(self, x):
         out = {}
 
@@ -453,7 +457,8 @@ def __init__(
         fiber: Tuple[int, ...],
         fiber_out: Optional[Tuple[int, ...]] = None,
         mult = 4,
-        include_htype_norms = True
+        include_htype_norms = True,
+        init_out_zero = True
     ):
         super().__init__()
         self.fiber = fiber
@@ -474,6 +479,9 @@ def __init__(
         self.gate        = Gate(project_in_fiber_hidden)
         self.project_out = Linear(fiber_hidden, fiber_out)
 
+        if init_out_zero:
+            self.project_out.init_zero_()
+
     def forward(self, features):
         outputs = self.prenorm(features)
 
@@ -542,7 +550,8 @@ def __init__(
         single_headed_kv = False,
         radial_hidden_dim = 64,
         splits = 4,
-        num_linear_attn_heads = 0
+        num_linear_attn_heads = 0,
+        init_out_zero = True
     ):
         super().__init__()
         num_degrees = len(fiber)
@@ -580,6 +589,9 @@ def __init__(
 
         self.to_out = Linear(hidden_fiber, fiber)
 
+        if init_out_zero:
+            self.to_out.init_zero_()
+
     @beartype
     def forward(
         self,
@@ -669,6 +681,7 @@ def __init__(
         attn_hidden_dim_mult = 4,
         radial_hidden_dim = 16,
         num_linear_attn_heads = 0,
+        init_out_zero = True,
         **kwargs
     ):
         super().__init__()
@@ -738,6 +751,9 @@ def __init__(
 
         self.to_out = Linear(hidden_fiber, fiber)
 
+        if init_out_zero:
+            self.to_out.init_zero_()
+
     @beartype
     def forward(
         self,

equiformer_pytorch/version.py

@@ -1 +1 @@
-__version__ = '0.3.0'
+__version__ = '0.3.1'

tests/test_equivariance.py

@@ -12,7 +12,8 @@ def test_transformer(dim):
     model = Equiformer(
         dim = dim,
         depth = 2,
-        num_degrees = 3
+        num_degrees = 3,
+        init_out_zero = False
     )
 
     feats = torch.randn(1, 32, dim)
@@ -39,7 +40,8 @@ def test_equivariance(
         l2_dist_attention = l2_dist_attention,
         reversible = reversible,
         num_degrees = 3,
-        reduce_dim_out = True
+        reduce_dim_out = True,
+        init_out_zero = False
     )
 
     feat_dim = dim if not isinstance(dim, tuple) else dim[0]