alphagenome-pytorch 0.0.10__py3-none-any.whl → 0.0.12__py3-none-any.whl

alphagenome_pytorch/__init__.py

@@ -8,4 +8,5 @@ from alphagenome_pytorch.alphagenome import (
     TransformerTower,
     UpresBlock,
     DownresBlock,
+    BatchRMSNorm
 )

alphagenome_pytorch/alphagenome.py

@@ -4,7 +4,7 @@ from functools import partial
 import torch
 from torch import nn, cat, stack, arange, logspace
 import torch.nn.functional as F
-from torch.nn import Conv1d, Linear, Sequential, Module, ModuleList
+from torch.nn import Conv1d, Linear, Sequential, Module, ModuleList, LayerNorm, RMSNorm
 
 from torch.nn.utils.parametrize import register_parametrization
 
@@ -47,6 +47,49 @@ def default(v, d):
 def softclamp(t, value = 5.):
     return (t / value).tanh() * value
 
+def append_dims(t, ndims):
+    return t.shape(*t.shape, *((1,) * ndims))
+
+# batch rmsnorm
+
+class BatchRMSNorm(Module):
+    def __init__(
+        self,
+        dim_feat,
+        channel_first = False,
+        momentum = 0.9,
+        eps = 1e-5,
+    ):
+        super().__init__()
+        self.scale = dim_feat ** 0.5
+
+        self.eps = eps
+        self.momentum = 1. - momentum
+        self.gamma = nn.Parameter(torch.zeros(dim_feat))
+        self.channel_first = channel_first
+
+        self.register_buffer('running_var', torch.ones((dim_feat,)))
+
+    def forward(
+        self,
+        x,
+        update_running_var = None
+    ):
+        update_running_var = default(update_running_var, self.training)
+        running_var = self.running_var
+
+        if update_running_var:
+
+            to_reduce = rearrange(x, 'b d ... -> b ... d') if self.channel_first else x
+
+            batch_var = torch.var(to_reduce, dim = tuple(range(x.ndim - 1)))
+
+            running_var.lerp_(batch_var, self.momentum)
+
+        std = running_var.clamp(min = self.eps).sqrt()
+
+        return x * self.scale * (self.gamma + 1.) / std
+
 # convolutional unet related
 
 class WeightStandardConv(Conv1d):
@@ -60,7 +103,7 @@ class WeightStandardConv(Conv1d):
     ):
         super().__init__(dim, dim_out, width, *args, **kwargs)
 
-        register_parametrization(self, 'weight', nn.LayerNorm(self.weight.shape, elementwise_affine = False))
+        register_parametrization(self, 'weight', LayerNorm(self.weight.shape, elementwise_affine = False))
 
 class ConvBlock(Module):
     def __init__(
@@ -219,10 +262,10 @@ class NormWrapper(Module):
     ):
         super().__init__()
         self.block = block
-        self.pre_rmsnorm = nn.RMSNorm(dim) # they use an interesting variant of batchnorm, batch-rmsnorm. craft later and make sure it works distributed
+        self.pre_rmsnorm = RMSNorm(dim) # they use an interesting variant of batchnorm, batch-rmsnorm. craft later and make sure it works distributed
 
         self.post_block_dropout = nn.Dropout(dropout)
-        self.post_rmsnorm = nn.RMSNorm(dim) if sandwich else nn.Identity()
+        self.post_rmsnorm = RMSNorm(dim) if sandwich else nn.Identity()
 
     def forward(
         self,
@@ -243,10 +286,12 @@ class Attention(Module):
         dim,
         dim_head = 64,
         heads = 8,
+        kv_heads = 1,
         dim_head_qk = 128,
         dim_head_v = 192,
         dim_pairwise = None,
-        softclamp_value = 5. # they employ attention softclamping
+        softclamp_value = 5., # they employ attention softclamping
+        use_qk_rmsnorm = True
     ):
         super().__init__()
         dim_pairwise = default(dim_pairwise, dim)
@@ -257,8 +302,13 @@ class Attention(Module):
 
         # splitting and merging of attention heads
 
-        self.split_q_heads = Rearrange('b n (h d) -> b h n d', h = heads)
-        self.merge_heads = Rearrange('b h n d -> b n (h d)')
+        assert divisible_by(heads, kv_heads)
+        groups = heads // kv_heads
+
+        self.split_q_heads = Rearrange('b n (g h d) -> b g h n d', h = kv_heads, g = groups)
+        self.split_kv_heads = Rearrange('b n (h d) -> b h n d', h = kv_heads)
+
+        self.merge_heads = Rearrange('b g h n d -> b n (g h d)')
 
         # projections
 
@@ -267,17 +317,18 @@ class Attention(Module):
 
         # they add layernorms to queries, keys, and interestingly enough, values as well. first time i've seen this
 
-        self.q_norm = nn.LayerNorm(dim_head_qk, bias = False)
-        self.k_norm = nn.LayerNorm(dim_head_qk, bias = False)
-        self.v_norm = nn.LayerNorm(dim_head_v, bias = False)
+        norm_klass = RMSNorm if use_qk_rmsnorm else partial(LayerNorm, bias = False)
+        self.q_norm = norm_klass(dim_head_qk)
+        self.k_norm = norm_klass(dim_head_qk)
+        self.v_norm = norm_klass(dim_head_v)
 
         # to attention bias
 
         self.to_attn_bias = Sequential(
-            nn.RMSNorm(dim_pairwise), # replace with BatchRMSNorm once crafted
+            RMSNorm(dim_pairwise), # replace with BatchRMSNorm once crafted
             nn.GELU(),
             LinearNoBias(dim_pairwise, heads),
-            Rearrange('b i j h -> b h i j')
+            Rearrange('b i j (g h) -> b g h i j', g = groups)
         )
         # variables
 
@@ -296,6 +347,7 @@ class Attention(Module):
         # they use multi-query attention, with only 1 key / value head - pretty unconventional, but maybe enough for genomic modeling
 
         q = self.split_q_heads(q)
+        k, v = tuple(self.split_kv_heads(t) for t in (k, v))
 
         q, k, v = self.q_norm(q), self.k_norm(k), self.v_norm(v)
 
@@ -308,7 +360,7 @@ class Attention(Module):
 
         # similarities
 
-        sim = einsum(q, k, 'b h i d, b j d -> b h i j')
+        sim = einsum(q, k, 'b g h i d, b h j d -> b g h i j')
 
         # add attention bias + softclamping
 
@@ -318,7 +370,7 @@ class Attention(Module):
             assert divisible_by(sim.shape[-1], attn_bias.shape[-1])
             expand_factor = sim.shape[-1] // attn_bias.shape[-1]
 
-            attn_bias = repeat(attn_bias, 'b h i j -> b h (i r1) (j r2)', r1 = expand_factor, r2 = expand_factor)
+            attn_bias = repeat(attn_bias, 'b g h i j -> b g h (i r1) (j r2)', r1 = expand_factor, r2 = expand_factor)
 
             sim = softclamp(sim + attn_bias, value = self.softclamp_value)
 
@@ -328,7 +380,7 @@ class Attention(Module):
 
         # aggregate
 
-        out = einsum(attn, v, 'b h i j, b j d -> b h i d')
+        out = einsum(attn, v, 'b g h i j, b h j d -> b g h i d')
 
         out = self.merge_heads(out)
         return self.to_out(out)
@@ -542,7 +594,7 @@ class TransformerTower(Module):
                 pairwise = maybe_pairwise_attn(pairwise) + pairwise
                 pairwise = maybe_pairwise_ff(pairwise) + pairwise
 
-            single = attn(single, rotary_emb = rotary_emb, pairwise = None) + single
+            single = attn(single, rotary_emb = rotary_emb, pairwise = pairwise) + single
             single = ff(single) + single
 
         return single, pairwise

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

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

alphagenome_pytorch-0.0.12.dist-info/RECORD

@@ -0,0 +1,6 @@
+alphagenome_pytorch/__init__.py,sha256=O7Mq4mhJwDUEKAJhFTLZRgSJf2uJaccX-Saef1d5nyg,242
+alphagenome_pytorch/alphagenome.py,sha256=21HFdrp5q55BRz4xt2h8nw4EgIJgUUHuqoxfgVfXjWE,19394
+alphagenome_pytorch-0.0.12.dist-info/METADATA,sha256=eODwGA_Ig6C4cbTI1nZDHCgWL4lRgAxY_a31EjyfYH4,3382
+alphagenome_pytorch-0.0.12.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+alphagenome_pytorch-0.0.12.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+alphagenome_pytorch-0.0.12.dist-info/RECORD,,

alphagenome_pytorch-0.0.10.dist-info/RECORD

@@ -1,6 +0,0 @@
-alphagenome_pytorch/__init__.py,sha256=XPNDv0q_c3nkiQo-4ROb_RQsbbKNV8KXmD6X5VnErKI,225
-alphagenome_pytorch/alphagenome.py,sha256=gfctPlDfzjl7CUkWJBX-9omT1MrvQDIgsvKSR0lLCQc,17902
-alphagenome_pytorch-0.0.10.dist-info/METADATA,sha256=SWZoRmpq6HVnddFPVMt108RFL-PywE-mniDLVRYkI3o,3382
-alphagenome_pytorch-0.0.10.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-alphagenome_pytorch-0.0.10.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-alphagenome_pytorch-0.0.10.dist-info/RECORD,,