x-transformers 1.23.1__py3-none-any.whl → 1.23.3__py3-none-any.whl

x_transformers/x_transformers.py

@@ -650,6 +650,7 @@ class Attention(nn.Module):
         num_mem_kv = 0,
         dropout = 0.,
         on_attn = False,
+        gate_value_heads = False,
         gate_values = False,
         zero_init_output = False,
         max_attend_past = None,
@@ -703,7 +704,14 @@ class Attention(nn.Module):
         if gate_values:
             self.to_v_gate = nn.Linear(dim, out_dim)
             nn.init.constant_(self.to_v_gate.weight, 0)
-            nn.init.constant_(self.to_v_gate.bias, 1)
+            nn.init.constant_(self.to_v_gate.bias, 10)
+
+        # add per head gating of the output values, from 'Attend to nothing' paper
+        self.to_v_head_gate = None
+        if gate_value_heads:
+            self.to_v_head_gate = nn.Linear(dim, heads)
+            nn.init.constant_(self.to_v_head_gate.weight, 0)
+            nn.init.constant_(self.to_v_head_gate.bias, 10)
 
         # cosine sim attention
         self.qk_norm = qk_norm
@@ -905,6 +913,12 @@ class Attention(nn.Module):
         if head_scale:
             out = out * self.head_scale_params
 
+        # per head gating, from https://arxiv.org/abs/2306.12929
+
+        if exists(self.to_v_head_gate):
+            head_gate = self.to_v_head_gate(x)
+            out = out * rearrange(head_gate, 'b n h -> b h n 1').sigmoid()
+
         # merge heads
 
         out = rearrange(out, 'b h n d -> b n (h d)')
@@ -1608,6 +1622,7 @@ class ContinuousTransformerWrapper(nn.Module):
         dim_out = None,
         emb_dim = None,
         max_mem_len = 0,
+        num_memory_tokens = None,
         post_emb_norm = False,
         emb_dropout = 0.,
         use_abs_pos_emb = True,
@@ -1632,10 +1647,21 @@ class ContinuousTransformerWrapper(nn.Module):
         self.post_emb_norm = nn.LayerNorm(dim) if post_emb_norm else nn.Identity()
         self.emb_dropout = nn.Dropout(emb_dropout)
 
-        self.project_in = nn.Linear(dim_in, dim) if exists(dim_in) else nn.Identity()
+        # memory tokens
+
+        num_memory_tokens = default(num_memory_tokens, 0)
+        self.has_memory_tokens = num_memory_tokens > 0
+
+        if num_memory_tokens > 0:
+            self.memory_tokens = nn.Parameter(torch.randn(num_memory_tokens, dim))
+
+        # attention layers
 
         self.attn_layers = attn_layers
 
+        # project in and out
+
+        self.project_in = nn.Linear(dim_in, dim) if exists(dim_in) else nn.Identity()
         self.project_out = nn.Linear(dim, dim_out) if exists(dim_out) else nn.Identity()
 
     def forward(
@@ -1651,11 +1677,19 @@ class ContinuousTransformerWrapper(nn.Module):
         prepend_embeds = None,
         **kwargs
     ):
+        batch = x.shape[0]
+
         x = self.project_in(x)
         x = x + self.pos_emb(x, pos = pos)
 
         x = self.post_emb_norm(x)
 
+        # memory tokens
+
+        if self.has_memory_tokens:
+            m = repeat(self.memory_tokens, 'm d -> b m d', b = batch)
+            x, mem_ps = pack([m, x], 'b * d')
+
         # whether to append embeds, as in PaLI, for image embeddings
 
         if exists(prepend_embeds):
@@ -1666,8 +1700,15 @@ class ContinuousTransformerWrapper(nn.Module):
 
         x = self.emb_dropout(x)
 
+        # attention layers
+
         x, intermediates = self.attn_layers(x, mask = mask, mems = mems, return_hiddens = True, **kwargs)
 
+        # splice out memory tokens
+
+        if self.has_memory_tokens:
+            m, x = unpack(x, mem_ps, 'b * d')
+
         out = self.project_out(x) if not return_embeddings else x
 
         if return_intermediates:

{x_transformers-1.23.1.dist-info → x_transformers-1.23.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: x-transformers
-Version: 1.23.1
+Version: 1.23.3
 Summary: X-Transformers - Pytorch
 Home-page: https://github.com/lucidrains/x-transformers
 Author: Phil Wang

{x_transformers-1.23.1.dist-info → x_transformers-1.23.3.dist-info}/RECORD

@@ -3,10 +3,10 @@ x_transformers/attend.py,sha256=hZcz_iijzbEqbXp2_BPEVL-1LoHXmYaHE6e6Oy-7hFE,1126
 x_transformers/autoregressive_wrapper.py,sha256=f2u0usjUfAlXwgTz87O8J8XjGTbsbrx2XEP6K2beSNI,8944
 x_transformers/continuous_autoregressive_wrapper.py,sha256=pTiDqu6JRUlnQJQp_xHATYHy0lgSd6ERLqyiFO3pC-4,1575
 x_transformers/nonautoregressive_wrapper.py,sha256=AQLE4rA_Kh8VNoe9OzpwyeWson34sRkhks4dn4seNjI,10414
-x_transformers/x_transformers.py,sha256=ikf99q_g1_v_wObZed972s2hHrbiDpAq_qGJDmNcVZc,60573
+x_transformers/x_transformers.py,sha256=goudsIa79mfyJtzuI0GqTSdGQ5CXG1ga5Is9h3UBC5Y,61861
 x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
-x_transformers-1.23.1.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.23.1.dist-info/METADATA,sha256=91Gu0qU9ztioZ2_oVeOrLeVkP--n6ngneDxEOMUHJe8,661
-x_transformers-1.23.1.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
-x_transformers-1.23.1.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.23.1.dist-info/RECORD,,
+x_transformers-1.23.3.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.23.3.dist-info/METADATA,sha256=SXNDjqYSGkklnbXVRg8S52VxDR6VVO62KvRH60abY_k,661
+x_transformers-1.23.3.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
+x_transformers-1.23.3.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.23.3.dist-info/RECORD,,