receiving reports of instability when scaling up with the SiLU based …

…value gates. use the sigmoid gating from alphafold2 instead

iTransformer/iTransformer.py

@@ -47,9 +47,9 @@ def __init__(
         )
 
         self.to_v_gates = nn.Sequential(
-            nn.Linear(dim, dim_inner, bias = False),
-            nn.SiLU(),
-            Rearrange('b n (h d) -> b h n d', h = heads)
+            nn.Linear(dim, heads, bias = False),
+            nn.Sigmoid(),
+            Rearrange('b n h -> b h n 1', h = heads)
         )
 
         self.attend = Attend(flash = flash, dropout = dropout)

iTransformer/iTransformer2D.py

@@ -66,9 +66,9 @@ def __init__(
         )
 
         self.to_v_gates = nn.Sequential(
-            nn.Linear(dim, dim_inner, bias = False),
-            nn.SiLU(),
-            Rearrange('b n (h d) -> b h n d', h = heads)
+            nn.Linear(dim, heads, bias = False),
+            nn.Sigmoid(),
+            Rearrange('b n h -> b h n 1', h = heads)
         )
 
         self.attend = Attend(flash = flash, dropout = dropout, causal = causal)

iTransformer/iTransformerFFT.py

@@ -48,9 +48,9 @@ def __init__(
         )
 
         self.to_v_gates = nn.Sequential(
-            nn.Linear(dim, dim_inner, bias = False),
-            nn.SiLU(),
-            Rearrange('b n (h d) -> b h n d', h = heads)
+            nn.Linear(dim, heads, bias = False),
+            nn.Sigmoid(),
+            Rearrange('b n h -> b h n 1', h = heads)
         )
 
         self.attend = Attend(flash = flash, dropout = dropout)

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'iTransformer',
   packages = find_packages(exclude=[]),
-  version = '0.5.5',
+  version = '0.6.0',
   license='MIT',
   description = 'iTransformer - Inverted Transformer Are Effective for Time Series Forecasting',
   author = 'Phil Wang',