alphagenome-pytorch 0.0.2__tar.gz → 0.0.4__tar.gz

{alphagenome_pytorch-0.0.2 → alphagenome_pytorch-0.0.4}/PKG-INFO

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

{alphagenome_pytorch-0.0.2 → alphagenome_pytorch-0.0.4}/alphagenome_pytorch/alphagenome.py

@@ -2,11 +2,11 @@ from __future__ import annotations
 from functools import partial
 
 import torch
-from torch import nn, cat, stack, arange
+from torch import nn, cat, stack, arange, logspace
 import torch.nn.functional as F
 from torch.nn import Linear, Sequential, Module, ModuleList
 
-import einx
+from einx import add, multiply, greater
 from einops.layers.torch import Rearrange, Reduce
 from einops import rearrange, repeat, einsum
 
@@ -41,6 +41,12 @@ def default(v, d):
 def softclamp(t, value = 5.):
     return (t / value).tanh() * value
 
+def relative_shift(t):
+    *leading_dims, seq_len, dim = t.shape
+    t = F.pad(t, (1, 0), value = 0.)
+    t = t.reshape(*leading_dims, dim + 1, seq_len)
+    return t[..., 1:, :].reshape(*leading_dims, seq_len, dim)
+
 # rotary, but with attenuation of short relative distance frequencies
 
 class RotaryEmbedding(Module):
@@ -51,7 +57,7 @@ class RotaryEmbedding(Module):
     ):
         super().__init__()
         num_freqs = dim // 2
-        inv_freq = 1. / (arange(num_freqs).float() + torch.logspace(1, max_positions - num_freqs + 1, num_freqs))
+        inv_freq = 1. / (arange(num_freqs).float() + logspace(1, max_positions - num_freqs + 1, num_freqs))
         self.register_buffer('inv_freq', inv_freq)
 
     def forward(
@@ -70,6 +76,26 @@ def rotate_half(x):
 def apply_rotary_pos_emb(pos, t):
     return t * pos.cos() + rotate_half(t) * pos.sin()
 
+# 'central mask features' - relative positions for constituting pairwise rep
+
+class RelativePosFeatures(Module):
+    def forward(self, single):
+
+        _, seq_len, dim = single.shape
+        half_dim = dim // 2
+
+        rel_pos = arange(2 * seq_len - 1) - (seq_len - 1)
+
+        center_widths = (
+            arange(half_dim) +
+            logspace(1, seq_len - half_dim + 1, half_dim + 1)[:-1] # endpoint = False
+        )
+
+        abs_rel_pos, rel_pos_sign = rel_pos.abs(), rel_pos.sign()
+        embeds = greater('j, i -> i j', center_widths, abs_rel_pos).float()
+
+        return cat((embeds, multiply('i, i j', rel_pos_sign, embeds)), dim = -1)
+
 # prenorm and sandwich norm - they use sandwich norm for single rep, prenorm for pairwise rep
 
 class NormWrapper(Module):
@@ -90,10 +116,11 @@ class NormWrapper(Module):
     def forward(
         self,
         x,
+        *args,
         **kwargs
     ):
         x = self.pre_rmsnorm(x)
-        out = self.block(x, **kwargs)
+        out = self.block(x, *args, **kwargs)
         out = self.post_block_dropout(out)
         return self.post_rmsnorm(out)
 
@@ -210,17 +237,26 @@ class SingleToPairwise(Module):
 
         dim_inner = heads * dim_pairwise
 
-        self.split_heads = Rearrange('b n (h d) -> b n h d', h = heads)
+        self.split_heads = Rearrange('... (h d) -> ... h d', h = heads)
 
         self.to_outer_sum = Sequential(
+            nn.GELU(),
             LinearNoBias(dim, dim_pairwise * 2),
-            nn.GELU()
         )
 
         self.to_qk = LinearNoBias(dim, dim_inner * 2)
         self.qk_to_pairwise = Linear(heads, dim_pairwise)
 
-    def forward(self, single):
+        # relative position related
+
+        self.to_rel_pos_encoding = Linear(dim, heads * dim_pairwise)
+        self.qk_rel_pos_bias = nn.Parameter(torch.zeros(2, 1, 1, heads, dim_pairwise))
+
+    def forward(
+        self,
+        single,
+        rel_pos_feats = None
+    ):
 
         single = self.avg_pool(single)
 
@@ -229,11 +265,24 @@ class SingleToPairwise(Module):
 
         sim = einsum(q, k, 'b i h d, b j h d -> b i j h')
 
+        if exists(rel_pos_feats):
+            rel_pos_encoding = self.to_rel_pos_encoding(rel_pos_feats)
+            rel_pos_encoding = self.split_heads(rel_pos_encoding)
+
+            q_rel_bias, k_rel_bias = self.qk_rel_pos_bias
+
+            rel_q = relative_shift(einsum(q + q_rel_bias, rel_pos_encoding, 'b i h d, j h d -> b i j h'))
+            rel_k = relative_shift(einsum(k + k_rel_bias, rel_pos_encoding, 'b i h d, j h d -> b i j h'))
+
+            rel_sim = einsum(rel_q, rel_k, 'b i p h, b j p h -> b i j h')
+
+            sim = (sim + rel_sim) * 0.5
+
         pairwise_from_sim = self.qk_to_pairwise(sim)
 
         outer_q, outer_k = self.to_outer_sum(single).chunk(2, dim = -1)
 
-        outer_sum = einx.add('b i d, b j d -> b i j d', outer_q, outer_k)
+        outer_sum = add('b i d, b j d -> b i j d', outer_q, outer_k)
 
         return outer_sum
 
@@ -314,6 +363,8 @@ class TransformerTower(Module):
 
         self.pairwise_every = pairwise_every_num_single_blocks
 
+        self.rel_pos_features = RelativePosFeatures()
+
         self.rotary_emb = RotaryEmbedding(dim_head_qk, max_positions = max_positions)
 
         for layer_index in range(depth):
@@ -360,6 +411,8 @@ class TransformerTower(Module):
 
         pairwise = None
 
+        rel_pos_feats = self.rel_pos_features(single)
+
         rotary_emb = self.rotary_emb(seq_len)
 
         for (
@@ -374,7 +427,7 @@ class TransformerTower(Module):
             single = ff(single) + single
 
             if exists(maybe_single_to_pair):
-                pairwise = maybe_single_to_pair(single) + default(pairwise, 0.)
+                pairwise = maybe_single_to_pair(single, rel_pos_feats) + default(pairwise, 0.)
                 pairwise = maybe_pairwise_attn(pairwise) + pairwise
                 pairwise = maybe_pairwise_ff(pairwise) + pairwise
 

{alphagenome_pytorch-0.0.2 → alphagenome_pytorch-0.0.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphagenome-pytorch"
-version = "0.0.2"
+version = "0.0.4"
 description = "AlphaGenome"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }