x-transformers 1.40.11__py3-none-any.whl → 1.41.0__py3-none-any.whl

x_transformers/dpo.py

@@ -5,6 +5,7 @@ from torch.nn import Module
 import torch.nn.functional as F
 from x_transformers.x_transformers import TransformerWrapper
 
+import einx
 from einops import rearrange
 
 # helper functions
@@ -17,16 +18,18 @@ def freeze_all_layers_(module):
         param.requires_grad = False
 
 def log_prob_from_model_and_seq(model, seq):
-    logits = model(seq)
+    src_seq, tgt_seq = seq[:, :-1], seq[:, 1:]
+    logits = model(src_seq)
     log_prob = logits.log_softmax(dim = -1)
-    indices = rearrange(seq, '... -> ... 1')
-    log_probs = log_prob.gather(-1, indices)
-    return rearrange(log_probs, '... 1 -> ...')
+    return einx.get_at('b n [l], b n -> b n', log_prob, tgt_seq)
 
 def masked_mean(log_probs, mask = None):
     if not exists(mask):
         return log_probs.mean(dim = -1)
 
+    if mask.shape[-1] == (log_probs.shape[-1] + 1):
+        mask = mask[:, :-1]
+
     log_probs = log_probs.masked_fill(~mask, 0.)
     num = log_probs.sum(dim = -1)
     den = mask.sum(dim = -1)

x_transformers/x_transformers.py

@@ -16,8 +16,9 @@ from collections import namedtuple
 from contextlib import nullcontext
 from dataclasses import dataclass
 
-from einops import rearrange, repeat, reduce, pack, unpack
+import einx
 from einops.layers.torch import Rearrange
+from einops import rearrange, repeat, reduce, pack, unpack
 
 from x_transformers.attend import Attend, Intermediates
 from x_transformers.autoregressive_wrapper import AutoregressiveWrapper
@@ -100,6 +101,12 @@ def log(t, eps = 1e-20):
 def max_neg_value(tensor):
     return -torch.finfo(tensor.dtype).max
 
+def reverse_cumsum(t, dim = -1):
+    t = t.flip(dims = (dim,))
+    t = t.cumsum(dim = dim)
+    t = t.flip(dims = (dim,))
+    return t
+
 def l2norm(t, groups = 1):
     t = rearrange(t, '... (g d) -> ... g d', g = groups)
     t = F.normalize(t, p = 2, dim = -1)
@@ -324,7 +331,7 @@ class RelativePositionBias(Module):
         device = self.device
         q_pos = torch.arange(j - i, j, dtype = torch.long, device = device)
         k_pos = torch.arange(j, dtype = torch.long, device = device)
-        rel_pos = k_pos[None, :] - q_pos[:, None]
+        rel_pos = einx.subtract('j, i -> i j', k_pos, q_pos)
         rp_bucket = self._relative_position_bucket(rel_pos, causal = self.causal, num_buckets = self.num_buckets, max_distance = self.max_distance)
         values = self.relative_attention_bias(rp_bucket)
         bias = rearrange(values, 'i j h -> h i j')
@@ -351,8 +358,10 @@ class CoPE(Module):
         self.soft_onehot = soft_onehot
         self.soft_onehot_temp = soft_onehot_temp
 
-        if soft_onehot:
-            self.register_buffer('positions', torch.arange(max_pos))
+        if not soft_onehot:
+            return
+
+        self.register_buffer('positions', torch.arange(max_pos))
 
     def forward(self, query, attn_logits):
 
@@ -374,7 +383,7 @@ class CoPE(Module):
         logits_int = einsum('b h n d, p d -> b h n p', query, self.pos_emb)
 
         if self.soft_onehot:
-            diff_pos = (pos[..., None] - self.positions).abs()
+            diff_pos = einx.subtract('i, j -> i j', pos, self.positions).abs()
             soft_onehot_pos = F.softmax(-diff_pos / self.soft_onehot_temp, dim = -1)
             cope_pos_emb = einsum('b h i j p, b h i p -> b h i j', soft_onehot_pos, logits_int)
         else:
@@ -423,7 +432,7 @@ class DynamicPositionBias(Module):
         # get the (n x n) matrix of distances
         seq_arange = torch.arange(n, device = device)
         context_arange = torch.arange(n, device = device)
-        indices = rearrange(seq_arange, 'i -> i 1') - rearrange(context_arange, 'j -> 1 j')
+        indices = einx.subtract('i, j -> i j', seq_arange, context_arange)
         indices += (n - 1)
 
         # input to continuous positions MLP
@@ -453,9 +462,9 @@ class AlibiPositionalBias(Module):
         self.register_buffer('bias', None, persistent = False)
     
     def get_bias(self, i, j, device):
-        i_arange = torch.arange(j - i, j, device = device)
-        j_arange = torch.arange(j, device = device)
-        bias = -torch.abs(rearrange(j_arange, 'j -> 1 1 j') - rearrange(i_arange, 'i -> 1 i 1'))
+        seq_arange = torch.arange(j - i, j, device = device)
+        context_arange = torch.arange(j, device = device)
+        bias = -torch.abs(einx.subtract('j, i -> 1 i j', context_arange, seq_arange))
         return bias
 
     @staticmethod
@@ -490,6 +499,44 @@ class AlibiPositionalBias(Module):
 
         return self.bias
 
+class DataDependentAlibi(Module):
+    """ https://openreview.net/forum?id=q2Lnyegkr8 """
+
+    def __init__(
+        self,
+        dim,
+        heads
+    ):
+        super().__init__()
+
+        linear = nn.Linear(dim, heads)
+
+        self.to_forget_gates = nn.Sequential(
+            linear,
+            Rearrange('b n h -> b h n'),
+            nn.Sigmoid()
+        )
+
+        nn.init.constant_(linear.bias, 5.)
+
+    def forward(self, x):
+        seq = x.shape[-2]
+
+        forget_gates = self.to_forget_gates(x).log()
+        forget_gates = repeat(forget_gates, 'b h j -> b h i j', i = seq)
+
+        # causal mask out, including diagonal (so token to itself attention is never masked out)
+
+        causal_mask = torch.ones((seq, seq), dtype = torch.bool, device = x.device).triu()
+
+        forget_gates = forget_gates.masked_fill(causal_mask, 0.)
+
+        # reverse cumulative sum in log space (equivalent to cumprod)
+
+        forget_gates = reverse_cumsum(forget_gates)
+
+        return forget_gates
+
 class RotaryEmbedding(Module):
     def __init__(
         self,
@@ -938,6 +985,7 @@ class Attention(Module):
         tensor_product = False,      # https://arxiv.org/abs/2208.06061
         add_zero_kv = False,         # same as add_zero_attn in pytorch
         rotary_embed_values = False,
+        data_dependent_alibi = False,
         use_cope = False,
         cope_max_pos = 16,
         cope_soft_onehot_pos = False,
@@ -1041,6 +1089,19 @@ class Attention(Module):
                 soft_onehot = cope_soft_onehot_pos
             )
 
+        # data dependent alibi
+        # https://openreview.net/forum?id=q2Lnyegkr8
+
+        self.data_dependent_alibi = None
+
+        if data_dependent_alibi:
+            assert causal, 'data dependent alibi only works for autoregressive for now until further research'
+
+            self.data_dependent_alibi = DataDependentAlibi(
+                dim,
+                heads = heads
+            )
+
         # attend class - includes core attention algorithm + talking heads
 
         self.attend = Attend(
@@ -1236,7 +1297,7 @@ class Attention(Module):
         if exists(self.max_attend_past):
             range_q = torch.arange(j - i, j, device = device)
             range_k = torch.arange(j, device = device)
-            dist = rearrange(range_q, 'i -> 1 1 i 1') - rearrange(range_k, 'j -> 1 1 1 j')
+            dist = einx.subtract('i, j -> 1 1 i j', range_q, range_k)
             max_attend_past_mask = dist > self.max_attend_past
             max_attend_past_mask = pad_at_dim(max_attend_past_mask, (num_mem_kv, 0), value = False, dim = -1) # handle memory key / values
             masks.append(max_attend_past_mask)
@@ -1251,6 +1312,11 @@ class Attention(Module):
             attn_bias = rel_pos(i, j)
             attn_bias = pad_at_dim(attn_bias, (num_mem_kv, 0), value = 0.) # handle memory key / values
 
+        # prepare data dependent alibi from forgetting transformers paper, if needed
+
+        if exists(self.data_dependent_alibi):
+            attn_bias = self.data_dependent_alibi(x)
+
         # if previous values passed in for residual, either invoke resformer or neutreno
 
         if exists(value_residual):
@@ -1291,7 +1357,7 @@ class Attention(Module):
 
         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()
+            out = einx.multiply('b n h, b h n d ->b h n d', head_gate.sigmoid(), out)
 
         # merge heads
 
@@ -1308,8 +1374,7 @@ class Attention(Module):
         out = self.to_out(out)
 
         if exists(mask):
-            mask = rearrange(mask, 'b n -> b n 1')
-            out = out.masked_fill(~mask, 0.)
+            out = einx.where('b n, b n d, -> b n d', mask, out, 0.)
 
         if not return_intermediates:
             return out
@@ -1389,10 +1454,13 @@ class AttentionLayers(Module):
         attn_kwargs, kwargs = groupby_prefix_and_trim('attn_', kwargs)
         cross_attn_kwargs, kwargs = groupby_prefix_and_trim('cross_attn_', kwargs)
 
+        dim_head = attn_kwargs.get('dim_head', DEFAULT_DIM_HEAD)
+        data_dependent_alibi = attn_kwargs.get('data_dependent_alibi', False)
+        neutreno_value_residual = attn_kwargs.get('neutreno_value_residual', False)
+
         assert len(kwargs) == 0, f'unrecognized kwargs passed in {kwargs.keys()}'
 
-        dim_head = attn_kwargs.get('dim_head', DEFAULT_DIM_HEAD)
-        add_value_residual |= attn_kwargs.get('neutreno_value_residual', False)
+        add_value_residual |= neutreno_value_residual
 
         self.dim = dim
         self.causal = causal
@@ -1405,7 +1473,7 @@ class AttentionLayers(Module):
         assert not (rotary_xpos and not causal), 'rotary xpos is not compatible with bidirectional attention'
         self.rotary_pos_emb = RotaryEmbedding(rotary_emb_dim, use_xpos = rotary_xpos, scale_base = rotary_xpos_scale_base, interpolation_factor = rotary_interpolation_factor, base_rescale_factor = rotary_base_rescale_factor) if rotary_pos_emb else None
 
-        assert not (alibi_pos_bias and rel_pos_bias), 'you can only choose Alibi positional bias or T5 relative positional bias, not both'
+        assert at_most_one_of(alibi_pos_bias, rel_pos_bias, data_dependent_alibi), 'you can only choose one of Alibi positional bias, data dependent Alibi (forgetting transformers), or T5 relative positional bias'
         assert rel_pos_num_buckets <= rel_pos_max_distance, 'number of relative position buckets must be less than the relative position max distance'
 
         # relative positional bias

{x_transformers-1.40.11.dist-info → x_transformers-1.41.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: x-transformers
-Version: 1.40.11
+Version: 1.41.0
 Summary: X-Transformers - Pytorch
 Home-page: https://github.com/lucidrains/x-transformers
 Author: Phil Wang
@@ -15,5 +15,6 @@ Classifier: Programming Language :: Python :: 3.6
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: torch >=2.0
+Requires-Dist: einx >=0.3.0
 Requires-Dist: einops >=0.8.0
 

{x_transformers-1.40.11.dist-info → x_transformers-1.41.0.dist-info}/RECORD

@@ -2,14 +2,14 @@ x_transformers/__init__.py,sha256=-MkQrSc37cTVDX7AOykxunYnqVtFlQ7lb0Cse5dsGWU,79
 x_transformers/attend.py,sha256=SdWlV8Vp5DtpsOzAd0LRhm4VGrJf0lJCGiV2_j_CtoA,17284
 x_transformers/autoregressive_wrapper.py,sha256=DOJJCMMDOqDYKWy_IaG5IyKsXD3AW6amzfUgdAADOLY,10500
 x_transformers/continuous.py,sha256=cIVEdhfei258__ziV7kQBrJMxCel54bExBTDrO9rfCI,6450
-x_transformers/dpo.py,sha256=LjvWgCkqTl-UuehrzQ8nkX5guLr4whYwsmm7SKSwdls,3450
+x_transformers/dpo.py,sha256=xt4OuOWhU8pN3OKN2LZAaC2NC8iiEnchqqcrPWVqf0o,3521
 x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg,9252
 x_transformers/nonautoregressive_wrapper.py,sha256=2NU58hYMgn-4Jzg3mie-mXb0XH_dCN7fjlzd3K1rLUY,10510
-x_transformers/x_transformers.py,sha256=RfpihlGygZz4ICq4IGOgGNOipInXUiYWYNs1tej2Orw,88290
+x_transformers/x_transformers.py,sha256=fz71zW2IQ3NQU_csHbzCwFzGHNwrdIF9rZTLhUjmM_Q,90260
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=7S00kCuab4tWQa-vf-z-XfzADjVj48MoFIr7VSIvttg,8575
-x_transformers-1.40.11.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.40.11.dist-info/METADATA,sha256=D97orsPC5EYEtJN6EN75bLOfOY-FBmodr2eaFIovwu8,662
-x_transformers-1.40.11.dist-info/WHEEL,sha256=OVMc5UfuAQiSplgO0_WdW7vXVGAt9Hdd6qtN4HotdyA,91
-x_transformers-1.40.11.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.40.11.dist-info/RECORD,,
+x_transformers-1.41.0.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.41.0.dist-info/METADATA,sha256=832V1ChJ77viLmQO7d9RU7R9SV_bspLgVl9vtdRdq5Q,689
+x_transformers-1.41.0.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+x_transformers-1.41.0.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.41.0.dist-info/RECORD,,

{x_transformers-1.40.11.dist-info → x_transformers-1.41.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.2.0)
+Generator: setuptools (75.3.0)
 Root-Is-Purelib: true
 Tag: py3-none-any