hippoformer 0.0.6__tar.gz → 0.0.8__tar.gz

{hippoformer-0.0.6 → hippoformer-0.0.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hippoformer
-Version: 0.0.6
+Version: 0.0.8
 Summary: hippoformer
 Project-URL: Homepage, https://pypi.org/project/hippoformer/
 Project-URL: Repository, https://github.com/lucidrains/hippoformer
@@ -35,6 +35,7 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.9
 Requires-Dist: assoc-scan
+Requires-Dist: beartype
 Requires-Dist: einops>=0.8.1
 Requires-Dist: einx>=0.3.0
 Requires-Dist: torch>=2.4
@@ -50,8 +51,6 @@ Description-Content-Type: text/markdown
 
 Implementation of [Hippoformer](https://openreview.net/forum?id=hxwV5EubAw), Integrating Hippocampus-inspired Spatial Memory with Transformers
 
-[Temporary Discord](https://discord.gg/MkACrrkrYR)
-
 ## Citations
 
 ```bibtex

{hippoformer-0.0.6 → hippoformer-0.0.8}/README.md

@@ -4,8 +4,6 @@
 
 Implementation of [Hippoformer](https://openreview.net/forum?id=hxwV5EubAw), Integrating Hippocampus-inspired Spatial Memory with Transformers
 
-[Temporary Discord](https://discord.gg/MkACrrkrYR)
-
 ## Citations
 
 ```bibtex

{hippoformer-0.0.6 → hippoformer-0.0.8}/hippoformer/hippoformer.py

@@ -7,6 +7,8 @@ from torch.nn import Module
 from torch.jit import ScriptModule, script_method
 from torch.func import vmap, grad, functional_call
 
+from beartype import beartype
+
 from einx import multiply
 from einops import repeat, rearrange, pack, unpack
 from einops.layers.torch import Rearrange
@@ -36,6 +38,80 @@ def pack_with_inverse(t, pattern):
 def l2norm(t):
     return F.normalize(t, dim = -1)
 
+# sensory encoder decoder for 2d
+
+grid_sensory_enc_dec = (
+    create_mlp(
+        dim = 32 * 2,
+        dim_in = 9,
+        dim_out = 32,
+        depth = 3,
+    ),
+    create_mlp(
+        dim = 32 * 2,
+        dim_in = 32,
+        dim_out = 9,
+        depth = 3,
+    ),
+)
+
+# sensory encoder decoder for 3d maze
+
+class EncoderPackTime(Module):
+    def __init__(self, fn: Module):
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, x):
+        x = rearrange(x, 'b c t h w -> b t c h w')
+        x, packed_shape = pack([x], '* c h w')
+
+        x = self.fn(x)
+
+        x, = unpack(x, packed_shape, '* d')
+        print(x.shape)
+        return x
+
+class DecoderPackTime(Module):
+    def __init__(self, fn: Module):
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, x):
+        x, packed_shape = pack(x, '* d')
+
+        x = self.fn(x)
+
+        x = unpack(x, packed_shape, '* c h w')
+        x = rearrange(x, 'b t c h w -> b c t h w')
+        return x
+
+maze_sensory_enc_dec = (
+    EncoderPackTime(nn.Sequential(
+        nn.Conv2d(3, 16, 7, 2, padding = 3),
+        nn.ReLU(),
+        nn.Conv2d(16, 32, 3, 2, 1),
+        nn.ReLU(),
+        nn.Conv2d(32, 64, 3, 2, 1),
+        nn.ReLU(),
+        nn.Conv2d(64, 128, 3, 2, 1),
+        nn.ReLU(),
+        Rearrange('b ... -> b (...)'),
+        nn.Linear(2048, 32)
+    )),
+    DecoderPackTime(nn.Sequential(
+        nn.Linear(32, 2048),
+        Rearrange('b (c h w) -> b c h w', c = 128, h = 4),
+        nn.ConvTranspose2d(128, 64, 3, 2, 1, output_padding = (1, 1)),
+        nn.ReLU(),
+        nn.ConvTranspose2d(64, 32, 3, 2, 1, output_padding = (1, 1)),
+        nn.ReLU(),
+        nn.ConvTranspose2d(32, 16, 3, 2, 1, output_padding = (1, 1)),
+        nn.ReLU(),
+        nn.ConvTranspose2d(16, 3, 3, 2, 1, output_padding = (1, 1))
+    ))
+)
+
 # path integration
 
 class RNN(ScriptModule):
@@ -114,12 +190,12 @@ class PathIntegration(Module):
 # proposed mmTEM
 
 class mmTEM(Module):
+    @beartype
     def __init__(
         self,
         dim,
         *,
-        sensory_encoder: Module,
-        sensory_decoder: Module,
+        sensory_encoder_decoder: tuple[Module, Module],
         dim_sensory,
         dim_action,
         dim_encoded_sensory,
@@ -139,6 +215,8 @@ class mmTEM(Module):
 
         # sensory
 
+        sensory_encoder, sensory_decoder = sensory_encoder_decoder
+
         self.sensory_encoder = sensory_encoder
         self.sensory_decoder = sensory_decoder
 
@@ -179,7 +257,7 @@ class mmTEM(Module):
 
         grad_fn = grad(forward_with_mse_loss)
 
-        self.per_sample_grad_fn = vmap(vmap(grad_fn, in_dims = (None, 0, 0)), in_dims = (None, 0, 0))
+        self.per_sample_grad_fn = vmap(vmap(grad_fn, in_dims = (None, 0, 0)), in_dims = (0, 0, 0))
 
         # mlp decoder (from meta mlp output to joint)
 
@@ -213,6 +291,19 @@ class mmTEM(Module):
 
         self.integration_ratio = nn.Parameter(tensor(0.), requires_grad = integration_ratio_learned)
 
+    def init_params_and_momentum(
+        self,
+        batch_size
+    ):
+
+        params_dict = dict(self.meta_memory_mlp.named_parameters())
+
+        params = {name: repeat(param, '... -> b ...', b = batch_size) for name, param in params_dict.items()}
+
+        momentums = {name: zeros_like(param) for name, param in params.items()}
+
+        return params, momentums
+
     def retrieve(
         self,
         structural_codes,
@@ -230,7 +321,9 @@ class mmTEM(Module):
         self,
         sensory,
         actions,
-        return_losses = False
+        memory_mlp_params = None,
+        return_losses = False,
+        return_memory_mlp_params = False
     ):
         batch = actions.shape[0]
 
@@ -291,22 +384,28 @@ class mmTEM(Module):
 
         lr, forget, beta = self.to_learned_optim_hparams(joint_code_to_store).unbind(dim = -1)
 
-        params = dict(self.meta_memory_mlp.named_parameters())
+        if exists(memory_mlp_params):
+            params, momentums = memory_mlp_params
+        else:
+            params, momentums = self.init_params_and_momentum(batch)
+
+        # store by getting gradients of mse loss of keys and values
+
         grads = self.per_sample_grad_fn(params, keys, values)
 
-        # update the meta mlp parameters
+        # update the meta mlp parameters and momentums
 
-        init_momentums = {k: zeros_like(v) for k, v in params.items()}
         next_params = dict()
+        next_momentum = dict()
 
         for (
             (key, param),
             (_, grad),
-            (_, init_momentum)
+            (_, momentum)
         ) in zip(
             params.items(),
             grads.items(),
-            init_momentums.items()
+            momentums.items()
         ):
 
             grad, inverse_pack = pack_with_inverse(grad, 'b t *')
@@ -315,9 +414,7 @@ class mmTEM(Module):
 
             expanded_beta = repeat(beta, 'b t -> b t w', w = grad.shape[-1])
 
-            init_momentum = repeat(init_momentum, '... -> b ...', b = batch)
-
-            update = self.assoc_scan(grad, expanded_beta.sigmoid(), init_momentum)
+            update = self.assoc_scan(grad, expanded_beta.sigmoid(), momentum)
 
             expanded_forget = repeat(forget, 'b t -> b t w', w = grad.shape[-1])
 
@@ -325,7 +422,10 @@ class mmTEM(Module):
 
             acc_update = inverse_pack(acc_update)
 
+            # set the next params and momentum, which can be passed back in
+
             next_params[key] = param - acc_update[:, -1]
+            next_momentum[key] = update[:, -1]
 
         # losses
 
@@ -343,6 +443,9 @@ class mmTEM(Module):
             inference_pred_loss
         )
 
+        if return_memory_mlp_params:
+            return next_params, next_momentum
+
         if not return_losses:
             return total_loss
 

{hippoformer-0.0.6 → hippoformer-0.0.8}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "hippoformer"
-version = "0.0.6"
+version = "0.0.8"
 description = "hippoformer"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -25,6 +25,7 @@ classifiers=[
 
 dependencies = [
     "assoc-scan",
+    "beartype",
     "einx>=0.3.0",
     "einops>=0.8.1",
     "torch>=2.4",

hippoformer-0.0.8/tests/test_hippoformer.py

@@ -0,0 +1,64 @@
+import pytest
+param = pytest.mark.parametrize
+
+import torch
+
+def test_path_integrate():
+    from hippoformer.hippoformer import PathIntegration
+
+    path_integrator = PathIntegration(32, 64)
+
+    actions = torch.randn(2, 16, 32)
+
+    structure_codes = path_integrator(actions)
+    structure_codes = path_integrator(actions, structure_codes) # pass in previous structure codes, it will auto use the last one as hidden and pass it to the RNN
+
+    assert structure_codes.shape == (2, 16, 64)
+
+@param('sensory_type', ('naive', '2d', '3d'))
+def test_mm_tem(
+    sensory_type
+):
+    import torch
+    from hippoformer.hippoformer import mmTEM
+
+    from torch.nn import Linear
+
+    if sensory_type == 'naive':
+        enc_dec = (
+            Linear(11, 32),
+            Linear(32, 11)
+        )
+        sensory = torch.randn(2, 16, 11)
+
+    elif sensory_type == '2d':
+
+        from hippoformer.hippoformer import grid_sensory_enc_dec
+
+        enc_dec = grid_sensory_enc_dec
+        sensory = torch.randn(2, 16, 9)
+
+    elif sensory_type == '3d':
+
+        from hippoformer.hippoformer import maze_sensory_enc_dec
+
+        enc_dec = maze_sensory_enc_dec
+
+        sensory = torch.randn(2, 3, 16, 64, 64)
+
+    model = mmTEM(
+        dim = 32,
+        sensory_encoder_decoder = enc_dec,
+        dim_sensory = 11,
+        dim_action = 7,
+        dim_structure = 32,
+        dim_encoded_sensory = 32
+    )
+
+    actions = torch.randn(2, 16, 7)
+
+    next_params = model(sensory, actions, return_memory_mlp_params = True)
+    next_params = model(sensory, actions, memory_mlp_params = next_params, return_memory_mlp_params = True)
+
+    loss = model(sensory, actions, memory_mlp_params = next_params)
+    loss.backward()

hippoformer-0.0.6/tests/test_hippoformer.py

@@ -1,37 +0,0 @@
-import pytest
-
-import torch
-
-def test_path_integrate():
-    from hippoformer.hippoformer import PathIntegration
-
-    path_integrator = PathIntegration(32, 64)
-
-    actions = torch.randn(2, 16, 32)
-
-    structure_codes = path_integrator(actions)
-    structure_codes = path_integrator(actions, structure_codes) # pass in previous structure codes, it will auto use the last one as hidden and pass it to the RNN
-
-    assert structure_codes.shape == (2, 16, 64)
-
-def test_mm_tem():
-    import torch
-    from hippoformer.hippoformer import mmTEM
-
-    from torch.nn import Linear
-
-    model = mmTEM(
-        dim = 32,
-        sensory_encoder = Linear(11, 32),
-        sensory_decoder = Linear(32, 11),
-        dim_sensory = 11,
-        dim_action = 7,
-        dim_structure = 32,
-        dim_encoded_sensory = 32
-    )
-
-    actions = torch.randn(2, 16, 7)
-    sensory = torch.randn(2, 16, 11)
-
-    loss = model(sensory, actions)
-    loss.backward()