titans-pytorch 0.1.31__py3-none-any.whl → 0.1.33__py3-none-any.whl

titans_pytorch/__init__.py

@@ -1,4 +1,4 @@
-from titans_pytorch.titans import (
+from titans_pytorch.neural_memory import (
     NeuralMemory,
     MemoryMLP,
     MemoryAttention,

titans_pytorch/mac_transformer.py

@@ -65,7 +65,7 @@ from hyper_connections import get_init_and_expand_reduce_stream_functions
 
 # proposed neural memory
 
-from titans_pytorch.titans import NeuralMemory
+from titans_pytorch.neural_memory import NeuralMemory
 
 # constants
 
@@ -106,7 +106,11 @@ def pad_at_dim(t, pad, dim = -1, value = 0.):
     zeros = ((0, 0) * dims_from_right)
     return F.pad(t, (*zeros, *pad), value = value)
 
-def pad_and_segment_with_inverse(seq, segment_len, fold_into_batch = True):
+def pad_and_segment_with_inverse(
+    seq,
+    segment_len,
+    fold_into_batch = True,
+):
     batch, seq_len = seq.shape[:2]
     next_seq_len_mult = round_up_multiple(seq_len, segment_len)
 
@@ -119,11 +123,15 @@ def pad_and_segment_with_inverse(seq, segment_len, fold_into_batch = True):
     if fold_into_batch:
         seq = rearrange(seq, 'b (w n) d -> (b w) n d', n = segment_len)
 
-    def inverse(out, remove_pad = True):
+    shape = seq.shape
+
+    def inverse(out):
+        unchanged_shape = out.shape == shape
+
         if fold_into_batch:
             out = rearrange(out, '(b w) ... n d -> b ... (w n) d', b = batch)
 
-        if needs_pad and remove_pad:
+        if needs_pad and unchanged_shape:
             out = out[..., :-padding, :]
 
         return out
@@ -690,7 +698,7 @@ class MemoryAsContextTransformer(Module):
         mems = repeat(self.longterm_mems, 'n d -> b n d', b = x.shape[0])
         x, inverse_pack_mems = pack_with_inverse((x, mems), 'b * d')
 
-        x = inverse_segment(x, remove_pad = False)
+        x = inverse_segment(x)
 
         # splice out unneeded tokens from padding for longterm mems
 
@@ -759,14 +767,13 @@ class MemoryAsContextTransformer(Module):
                 mem_input, add_residual = mem_hyper_conn(x)
 
                 if not is_inferencing:
-                    retrieved, mem_kv_aux_loss = mem(
+                    (retrieved, next_neural_mem_cache), mem_kv_aux_loss = mem(
                         mem_input,
                         return_aux_kv_loss = True
                     )
 
                     kv_recon_losses = kv_recon_losses + mem_kv_aux_loss
 
-                    next_neural_mem_cache = (seq_len, None, None, None)
                 else:
                     retrieved, next_neural_mem_cache = mem.forward_inference(
                         mem_input,
@@ -836,7 +843,7 @@ class MemoryAsContextTransformer(Module):
 
             x, _ = inverse_pack_mems(x)
 
-            x = inverse_segment(x, remove_pad = False)
+            x = inverse_segment(x)
 
             x = x[:, :seq_len]
 

titans_pytorch/{titans.py → neural_memory.py}

@@ -3,6 +3,7 @@ from typing import Callable
 
 import math
 from functools import partial
+from collections import namedtuple
 
 import torch
 from torch import nn, cat, Tensor
@@ -33,13 +34,18 @@ w - num memory network weight parameters
 
 LinearNoBias = partial(Linear, bias = False)
 
+NeuralMemCache = namedtuple('NeuralMemCache', ['seq', 'cache_store_segment', 'states', 'updates'])
+
 # functions
 
 def exists(v):
     return v is not None
 
-def default(v, d):
-    return v if exists(v) else d
+def default(*args):
+    for arg in args:
+        if exists(arg):
+            return arg
+    return None
 
 def xnor(x, y):
     return not (x ^ y)
@@ -468,7 +474,12 @@ class NeuralMemory(Module):
             weighted_loss = loss * loss_weights
             return weighted_loss.sum(), weighted_loss.mean()
 
-        self.per_sample_grad_fn = vmap(grad(forward_and_loss, has_aux = True), in_dims = (None, 0, 0, 0))
+        # two functions
+
+        grad_fn = grad(forward_and_loss, has_aux = True)
+
+        self.per_sample_grad_fn = vmap(grad_fn, in_dims = (None, 0, 0, 0))
+        self.per_sample_grad_fn_expanded_weights = vmap(grad_fn, in_dims = (0,) * 4)
 
         # queries for retrieving from the model
 
@@ -561,6 +572,7 @@ class NeuralMemory(Module):
         seq,
         weights: dict[str, Tensor],
         past_state: tuple[dict[str, Tensor], dict[str, Tensor]] | None = None,
+        prev_layer_updates: dict[str, Tensor] | None = None,
         return_aux_kv_loss = False,
         chunk_size = None,
         value_residual = None
@@ -583,10 +595,25 @@ class NeuralMemory(Module):
 
         seq = seq[:, :round_down_seq_len]
 
+        # per sample grad function
+
+        per_sample_grad_fn = self.per_sample_grad_fn
+
         # weights of the memory network
 
         weights = TensorDict(weights)
 
+        # allow for neural memory of a previous layer and the past to produce gradients that become the weights of the current one generating the surprise
+        # think this is necessary otherwise the memory model is static (unless if paper is misunderstood)
+        # improvise (or perhaps correcting to) a solution
+
+        if exists(prev_layer_updates):
+            prev_layer_updates = TensorDict(prev_layer_updates)
+
+            weights = weights + prev_layer_updates
+
+            per_sample_grad_fn = self.per_sample_grad_fn_expanded_weights # the weights will now have a batch * chunk dimension
+
         # derive learned hparams for optimization of memory network
 
         adaptive_lr = self.to_adaptive_step(seq)
@@ -633,9 +660,14 @@ class NeuralMemory(Module):
 
         adaptive_lr = rearrange(adaptive_lr, 'b (n c) -> (b n) c', c = chunk_size)
 
+        # flatten batch and time if surprise depends on previous layer memory model
+
+        if exists(prev_layer_updates):
+            weights = weights.apply(lambda t: rearrange(t, 'b n ... -> (b n) ...'))
+
         # get grads and extra auxiliary loss (for backwarding through qkv projection in base neural memory module)
 
-        grads, aux_kv_recon_loss = self.per_sample_grad_fn(dict(weights), keys, adaptive_lr, values)
+        grads, aux_kv_recon_loss = per_sample_grad_fn(dict(weights), keys, adaptive_lr, values)
 
         grads = TensorDict(grads)
 
@@ -713,7 +745,8 @@ class NeuralMemory(Module):
         self,
         seq,
         past_weights: dict[str, Tensor],
-        chunk_size = None
+        chunk_size = None,
+        prev_layer_updates: dict[str, Tensor] | None = None
     ):
         chunk_size = default(chunk_size, self.retrieve_chunk_size)
         batch, seq_len = seq.shape[:2]
@@ -736,6 +769,9 @@ class NeuralMemory(Module):
 
         curr_weights = TensorDict(past_weights)
 
+        if exists(prev_layer_updates):
+            curr_weights = curr_weights + TensorDict(prev_layer_updates)
+
         # sequence Float['b n d'] to queries
 
         queries = self.to_queries(seq)
@@ -781,6 +817,7 @@ class NeuralMemory(Module):
 
         return values[:, :seq_len]
 
+    @torch.no_grad()
     def forward_inference(
         self,
         token: Tensor,
@@ -813,7 +850,7 @@ class NeuralMemory(Module):
         if curr_seq_len < self.chunk_size:
             empty_mem = self.init_empty_memory_embed(batch, 1)
 
-            return empty_mem, (curr_seq_len, cache_store_seq, past_states, updates)
+            return empty_mem, NeuralMemCache(curr_seq_len, cache_store_seq, past_states, updates)
 
         # store if storage sequence cache hits the chunk size
 
@@ -823,6 +860,8 @@ class NeuralMemory(Module):
         if not exists(updates):
             updates = weights.clone().zero_()
             updates = updates.apply(lambda t: repeat(t, '... -> b 1 ...', b = batch))
+        else:
+            updates = updates.apply(lambda t: t[:, -1:])
 
         if store_seq_cache_len == self.chunk_size:
 
@@ -841,7 +880,7 @@ class NeuralMemory(Module):
 
         # next state tuple
 
-        next_state = (curr_seq_len, cache_store_seq, next_states, updates)
+        next_state = NeuralMemCache(curr_seq_len, cache_store_seq, next_states, updates)
 
         return retrieved, next_state
 
@@ -854,13 +893,19 @@ class NeuralMemory(Module):
         return_aux_kv_loss = False,
         chunk_size = None,
         store_chunk_size = None,
-        return_values = False
+        return_values = False,
+        return_next_state = False,
+        prev_layer_updates: dict[str, Tensor] | None = None
     ):
         batch, seq_len = seq.shape[:2]
 
         if seq_len < self.retrieve_chunk_size:
             out = self.init_empty_memory_embed(batch, seq_len)
 
+            next_store_state = (seq_len, seq, None, None)
+
+            out = (out, next_store_state)
+
             if not return_aux_kv_loss:
                 return out
 
@@ -869,17 +914,45 @@ class NeuralMemory(Module):
         if not exists(mem_model_weights):
             mem_model_weights = self.init_weights()
 
+        # store
+
         store_seq = default(store_seq, seq)
-        store_chunk_size = default(store_chunk_size, chunk_size)
 
-        (updates, next_state, values), aux_kv_recon_loss = self.store_memories(store_seq, mem_model_weights, chunk_size = store_chunk_size, return_aux_kv_loss = True)
+        store_seq_len = store_seq.shape[-2]
+        store_chunk_size = default(store_chunk_size, chunk_size, self.store_chunk_size)
+        remainder = store_seq_len % store_chunk_size
+
+        (updates, next_state, values), aux_kv_recon_loss = self.store_memories(
+            store_seq,
+            mem_model_weights,
+            chunk_size = store_chunk_size,
+            prev_layer_updates = prev_layer_updates,
+            return_aux_kv_loss = True
+        )
+
+        # retrieve
+
+        retrieved = self.retrieve_memories(
+            seq,
+            mem_model_weights + updates,
+            chunk_size = chunk_size,
+            prev_layer_updates = prev_layer_updates
+        )
+
+        # determine state for the storing of memories
+        # for transformer-xl like training with neural memory as well as inferencing with initial prompt
+
+        cache_store_seq = None
+
+        if remainder > 0:
+            cache_store_seq = store_seq[:, -remainder:]
 
-        retrieved = self.retrieve_memories(seq, mem_model_weights + updates, chunk_size = chunk_size)
+        next_store_state = NeuralMemCache(seq_len, cache_store_seq, next_state, updates)
 
-        output = retrieved
+        output = (retrieved, next_store_state)
 
         if return_values:
-            output = (retrieved, values)
+            output = (*output, values)
 
         if not return_aux_kv_loss:
             return output

{titans_pytorch-0.1.31.dist-info → titans_pytorch-0.1.33.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.1.31
+Version: 0.1.33
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -82,7 +82,7 @@ mem = NeuralMemory(
 ).cuda()
 
 seq = torch.randn(2, 1024, 384).cuda()
-retrieved = mem(seq)
+retrieved, mem_state = mem(seq)
 
 assert seq.shape == retrieved.shape
 ```

titans_pytorch-0.1.33.dist-info/RECORD

@@ -0,0 +1,8 @@
+titans_pytorch/__init__.py,sha256=rMT99CPQFH4Gudp0FmVPWGKfhBf6xksGEaFEcOVdqjs,230
+titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
+titans_pytorch/mac_transformer.py,sha256=Cui-hCl6X4UVGmuyoKCSKWbag9Yrc-a2MrfVkHM-z0A,24828
+titans_pytorch/neural_memory.py,sha256=9dXpSaQYomc-ur-nEwej1nG9M5NqS0c3LBBP9jUIMPU,28352
+titans_pytorch-0.1.33.dist-info/METADATA,sha256=A9BBoe0Sas2kxUcUi7w_RFl8-SIF1TLzPIRGuZlauFM,6826
+titans_pytorch-0.1.33.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.1.33.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.1.33.dist-info/RECORD,,

titans_pytorch-0.1.31.dist-info/RECORD

@@ -1,8 +0,0 @@
-titans_pytorch/__init__.py,sha256=u0tta_KqhOdfzCEDWT9P4_jejJEK2q1XxhsEzB5MnQU,223
-titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/mac_transformer.py,sha256=pKFRL_ISoHEKUyfssKwfBfwFO2eQN9objJmxLrNsYrU,24838
-titans_pytorch/titans.py,sha256=6B8ioP26RTja5kVFMsorAnM9CcxIUySJS9RZBlDPI2s,25825
-titans_pytorch-0.1.31.dist-info/METADATA,sha256=9ejOFuH2B2-yCRFK4x_C1DONPxecW8VcjEUeRh9OzXg,6815
-titans_pytorch-0.1.31.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.1.31.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.1.31.dist-info/RECORD,,