allow for attending to nothing

Author: lucidrains

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "transfusion-pytorch"
-version = "0.1.5"
+version = "0.1.6"
 description = "Transfusion in Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

transfusion_pytorch/transfusion.py

@@ -600,7 +600,8 @@ def __init__(
         heads = 8,
         dropout = 0.,
         softcap_value = 50.,
-        use_flex_attn = False
+        use_flex_attn = False,
+        gate_values = True
     ):
         super().__init__()
         self.scale = dim_head ** -0.5
@@ -616,6 +617,11 @@ def __init__(
             Rearrange('b n (qkv h d) -> qkv b h n d', qkv = 3, h = heads)
         )
 
+        self.to_gates = nn.Sequential(
+            nn.Linear(dim, heads, bias = False),
+            Rearrange('b n h -> b h n 1', h = heads)
+        ) if gate_values else None
+
         self.softcap_value = softcap_value
 
         self.dropout = nn.Dropout(dropout)
@@ -705,6 +711,11 @@ def forward(
 
             out = einsum(attn, v, 'b h i j, b h j d -> b h i d')
 
+        # maybe gate values
+
+        if exists(self.to_gates):
+            out = out * self.to_gates(x).sigmoid()
+
         # combine heads and out
 
         out = self.to_out(out)