hippoformer 0.0.7__py3-none-any.whl → 0.0.9__py3-none-any.whl

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,114 @@ def pack_with_inverse(t, pattern):
 def l2norm(t):
     return F.normalize(t, dim = -1)
 
+# Muon - Jordan et al from oss community - applied to the latest version of titans
+
+def newtonschulz5(
+    t,
+    steps = 5,
+    eps = 1e-7,
+    coefs = (3.4445, -4.7750, 2.0315)
+):
+    not_weights = t.ndim <= 3
+
+    if not_weights:
+        return t
+
+    shape = t.shape
+    should_transpose = shape[-2] > shape[-1]
+
+    if should_transpose:
+        t = t.transpose(-1, -2)
+
+    t, inv_pack = pack_with_inverse(t, '* i j')
+    t = t / t.norm(dim = (-1, -2), keepdim = True).clamp(min = eps)
+
+    a, b, c = coefs
+
+    for _ in range(steps):
+        A = t @ t.transpose(-1, -2)
+        B = b * A + c * A @ A
+        t = a * t + B @ t
+
+    if should_transpose:
+        t = t.transpose(-1, -2)
+
+    return inv_pack(t)
+
+# 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 +224,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,
@@ -133,12 +243,15 @@ class mmTEM(Module):
         loss_weight_consistency = 1.,
         loss_weight_relational = 1.,
         integration_ratio_learned = True,
+        muon_update = False,
         assoc_scan_kwargs: dict = dict()
     ):
         super().__init__()
 
         # sensory
 
+        sensory_encoder, sensory_decoder = sensory_encoder_decoder
+
         self.sensory_encoder = sensory_encoder
         self.sensory_decoder = sensory_decoder
 
@@ -209,6 +322,10 @@ class mmTEM(Module):
         self.loss_weight_consistency = loss_weight_consistency
         self.register_buffer('zero', tensor(0.), persistent = False)
 
+        # update with muon
+
+        self.muon_update = muon_update
+
         # there is an integration ratio for error correction, but unclear what value this is fixed to or whether it is learned
 
         self.integration_ratio = nn.Parameter(tensor(0.), requires_grad = integration_ratio_learned)
@@ -338,6 +455,13 @@ class mmTEM(Module):
 
             update = self.assoc_scan(grad, expanded_beta.sigmoid(), momentum)
 
+            # maybe muon
+
+            if self.muon_update:
+                update = newtonschulz5(update)
+
+            # with forget gating
+
             expanded_forget = repeat(forget, 'b t -> b t w', w = grad.shape[-1])
 
             acc_update = self.assoc_scan(update, expanded_forget.sigmoid())

{hippoformer-0.0.7.dist-info → hippoformer-0.0.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hippoformer
-Version: 0.0.7
+Version: 0.0.9
 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.9.dist-info/RECORD

@@ -0,0 +1,6 @@
+hippoformer/__init__.py,sha256=A7N8GsRAZH4yP-L5hb7IVDnNjnhfjNyolg5MZ6vnGyE,71
+hippoformer/hippoformer.py,sha256=m7luQGFdMWOkZUorjd5v34hx_vjOQqpJOAGCL0njHUE,14426
+hippoformer-0.0.9.dist-info/METADATA,sha256=owgkDcdTf0_N5IbUr3e_yt7u5sIWfOMha-hA5LQWnus,2772
+hippoformer-0.0.9.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+hippoformer-0.0.9.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+hippoformer-0.0.9.dist-info/RECORD,,

hippoformer-0.0.7.dist-info/RECORD

@@ -1,6 +0,0 @@
-hippoformer/__init__.py,sha256=A7N8GsRAZH4yP-L5hb7IVDnNjnhfjNyolg5MZ6vnGyE,71
-hippoformer/hippoformer.py,sha256=yYoJ5XO0YVAyp3LcRxpunU-0HA97mpCBeQFyi-NSkF0,11549
-hippoformer-0.0.7.dist-info/METADATA,sha256=Xg6NZ6VAQGmuiOo8mMwIAM39Gf6TpVOpyn7o4PMq7JE,2800
-hippoformer-0.0.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-hippoformer-0.0.7.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-hippoformer-0.0.7.dist-info/RECORD,,