hippoformer 0.0.10__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,6 +221,132 @@ class PathIntegration(Module):
 
         return self.rnn(transitions, prev_structural)
 
+# 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.):
+    dim_inner = int(dim * mult)
+    return nn.Sequential(
+        nn.Linear(dim, dim_inner),
+        nn.GELU(),
+        nn.Linear(dim_inner, dim)
+    )
+
+class Attention(Module):
+    def __init__(
+        self,
+        dim_q,
+        dim_kv,
+        window_size,
+        dim_head = 64,
+        heads = 8,
+    ):
+        super().__init__()
+        dim_inner = dim_head * heads
+        self.scale = dim_head ** -0.5
+
+        self.to_queries = nn.Linear(dim_q, dim_inner, bias = False)
+        self.to_key_values = nn.Linear(dim_kv, dim_inner * 2, bias = False)
+
+        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
+
+    def forward(
+        self,
+        queries_input,
+        key_values_input,
+        kv_cache = None
+    ):
+        batch, seq_len, device = *queries_input.shape[:2], queries_input.device
+
+        q = self.to_queries(queries_input)
+
+        k, v = self.to_key_values(key_values_input).chunk(2, dim = -1)
+
+        q, k, v = tuple(self.split_heads(t) for t in (q, k, v))
+
+        if exists(kv_cache):
+            ck, cv = kv_cache
+            k = cat((ck, k), dim = -2)
+            v = cat((cv, v), dim = -2)
+
+        q = q * self.scale
+
+        sim = einsum('b h i d, b h j d -> b h i j', q, k)
+
+        # the diagonal is masked out
+
+        i, j = sim.shape[-2:]
+
+        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
+
+        attn_sink = repeat(self.attn_head_sink, 'h -> b h i 1', b = batch, i = seq_len)
+
+        sim = cat((attn_sink, sim), dim = -1)
+
+        attn = sim.softmax(dim = -1)
+
+        attn = attn[..., 1:] # remove sink
+
+        # aggregate
+
+        out = einsum('b h i j, b h j d -> b h i d', attn, v)
+
+        out = self.merge_heads(out)
+
+        return self.to_out(out), stack((k, v))
+
+class TEMTransformerBlock(Module):
+    def __init__(
+        self,
+        dim_structure,
+        dim_encoded_sensory,
+        dim_head = 64,
+        heads = 8,
+        ff_expansion_factor = 4.,
+        window_size = 64
+    ):
+        super().__init__()
+
+        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
+
+    def forward(
+        self,
+        structural_codes,
+        encoded_sensory,
+        kv_cache = None
+    ):
+        structure_and_sensory = cat((structural_codes, encoded_sensory), dim = -1)
+
+        retrieved, next_kv_cache = self.attn(structural_codes, structure_and_sensory, kv_cache = kv_cache)
+
+        x = retrieved + structural_codes
+
+        x = self.ff(x) + x
+
+        next_kv_cache = next_kv_cache[:, -self.window_size:]
+
+        return x, next_kv_cache
+
 # proposed mmTEM
 
 class mmTEM(Module):

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

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

@@ -1,6 +0,0 @@
-hippoformer/__init__.py,sha256=A7N8GsRAZH4yP-L5hb7IVDnNjnhfjNyolg5MZ6vnGyE,71
-hippoformer/hippoformer.py,sha256=GWFJy2idp0FWBoVFw8T_6inTXYtY4i47hfhKj88_I0A,14463
-hippoformer-0.0.10.dist-info/METADATA,sha256=IB7iybYMwOkee3Q5ji-B_dnOB62LyK_6t1FPM_UT-FM,2773
-hippoformer-0.0.10.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-hippoformer-0.0.10.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-hippoformer-0.0.10.dist-info/RECORD,,