hippoformer 0.0.11__py3-none-any.whl → 0.0.12__py3-none-any.whl

hippoformer/hippoformer.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import torch
-from torch import nn, Tensor, cat, stack, zeros_like, einsum, tensor
+from torch import nn, Tensor, cat, stack, arange, zeros_like, einsum, tensor
 import torch.nn.functional as F
 from torch.nn import Module
 from torch.jit import ScriptModule, script_method
@@ -221,8 +221,11 @@ class PathIntegration(Module):
 
         return self.rnn(transitions, prev_structural)
 
-# custom transformer
-# with the connections proposed by James Whittington that bridges to hippocampal models
+# custom transformer proposed by James Whittington that bridges to hippocampal models with a few twists
+
+# the mmTEM can be seen as a linear attention / TTT variant of what he proposed
+# needed for the baseline as well as the parallel block to bolster local time prediction
+
 # https://arxiv.org/abs/2112.04035
 
 def FeedForward(dim, mult = 4.):
@@ -238,8 +241,9 @@ class Attention(Module):
         self,
         dim_q,
         dim_kv,
+        window_size,
         dim_head = 64,
-        heads = 8
+        heads = 8,
     ):
         super().__init__()
         dim_inner = dim_head * heads
@@ -251,6 +255,8 @@ class Attention(Module):
         self.split_heads = Rearrange('b n (h d) -> b h n d', h = heads)
         self.merge_heads = Rearrange('b h n d -> b n (h d)')
 
+        self.window_size = window_size
+
         self.to_out = nn.Linear(dim_inner, dim_q, bias = False)
         self.attn_head_sink = nn.Parameter(torch.randn(heads) * 1e-2) # needed as the diagonal is masked out, and for attention sink
 
@@ -280,9 +286,13 @@ class Attention(Module):
         # the diagonal is masked out
 
         i, j = sim.shape[-2:]
-        causal_mask_without_diagonal = torch.ones((i, j), dtype = torch.bool, device = device).triu(j - i)
 
-        sim = sim.masked_fill(causal_mask_without_diagonal, -torch.finfo(sim.dtype).max)
+        j_seq = arange(j, device = device)[:, None]
+        i_seq = arange(i, device = device)[None, :] + (j - i)
+
+        windowed_causal_mask_without_diagonal = (i_seq > j_seq) & ((i_seq - j_seq) <= self.window_size)
+
+        sim = sim.masked_fill(windowed_causal_mask_without_diagonal, -torch.finfo(sim.dtype).max)
 
         # attention sink, for token as well as for attention sinking - from gpt-oss
 
@@ -314,7 +324,7 @@ class TEMTransformerBlock(Module):
     ):
         super().__init__()
 
-        self.attn = Attention(dim_structure, dim_structure + dim_encoded_sensory, dim_head = dim_head, heads = heads)
+        self.attn = Attention(dim_structure, dim_structure + dim_encoded_sensory, window_size, dim_head = dim_head, heads = heads)
         self.ff = FeedForward(dim_structure, ff_expansion_factor)
 
         self.window_size = window_size

{hippoformer-0.0.11.dist-info → hippoformer-0.0.12.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hippoformer
-Version: 0.0.11
+Version: 0.0.12
 Summary: hippoformer
 Project-URL: Homepage, https://pypi.org/project/hippoformer/
 Project-URL: Repository, https://github.com/lucidrains/hippoformer

hippoformer-0.0.12.dist-info/RECORD

@@ -0,0 +1,6 @@
+hippoformer/__init__.py,sha256=A7N8GsRAZH4yP-L5hb7IVDnNjnhfjNyolg5MZ6vnGyE,71
+hippoformer/hippoformer.py,sha256=DjKAUfMpcoaAOqyuWauKp8n8e2YTzGVSOctNXxagkiA,18166
+hippoformer-0.0.12.dist-info/METADATA,sha256=bH3GaJniFX2zCgxFFg8v0amGEPdWlBbG401Ml3_hDCs,2773
+hippoformer-0.0.12.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+hippoformer-0.0.12.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+hippoformer-0.0.12.dist-info/RECORD,,

hippoformer-0.0.11.dist-info/RECORD

@@ -1,6 +0,0 @@
-hippoformer/__init__.py,sha256=A7N8GsRAZH4yP-L5hb7IVDnNjnhfjNyolg5MZ6vnGyE,71
-hippoformer/hippoformer.py,sha256=PLMfdype8AMwlVWrtItDBkE3gU_BCUaL42NMjB4vhAY,17795
-hippoformer-0.0.11.dist-info/METADATA,sha256=6NlqhZSEApQkUKsncBxmDIE03x_xZktHH-JCeYlYfcg,2773
-hippoformer-0.0.11.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-hippoformer-0.0.11.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-hippoformer-0.0.11.dist-info/RECORD,,