use axial positional embedding, and make sure to always show gratitude

Author: lucidrains

README.md

@@ -10,6 +10,10 @@ The big surprise is that the generations can reach this level of fidelity. Will
 
 Additionally, we will try adding an extra linear attention on the main branch, in addition to the full self attention on the latents. Self conditioning will also be applied to the non-latent images in pixel-space. Let us see how far we can push this approach.
 
+## Appreciation
+
+- <a href="https://stability.ai/">Stability.ai</a> for the generous sponsorship to work on cutting edge artificial intelligence research
+
 ## Install
 
 ```bash

rin_pytorch/rin_pytorch.py

@@ -293,7 +293,7 @@ def __init__(
             nn.Linear(pixel_patch_dim * 2, dim)
         )
 
-        self.pos_emb = nn.Parameter(torch.randn(num_patches, dim))
+        self.axial_pos_emb = nn.Parameter(torch.randn(2, patch_height_width, dim) * 0.02)
 
         self.to_pixels = nn.Sequential(
             LayerNorm(dim),
@@ -302,6 +302,7 @@ def __init__(
         )
 
         self.latents = nn.Parameter(torch.randn(num_latents, dim))
+        nn.init.normal_(self.latents, std = 0.02)
 
         self.init_self_cond_latents = nn.Sequential(
             FeedForward(dim),
@@ -362,7 +363,9 @@ def forward(
 
         patches = self.to_patches(x)
 
-        patches = patches + self.pos_emb
+        pos_emb_h, pos_emb_w = self.axial_pos_emb
+        pos_emb = rearrange(pos_emb_h, 'i d -> i 1 d') + rearrange(pos_emb_w, 'j d -> 1 j d')
+        patches = patches + rearrange(pos_emb, 'i j d -> (i j) d')
 
         # the recurrent interface network body
 

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'RIN-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.2',
+  version = '0.0.3',
   license='MIT',
   description = 'RIN - Recurrent Interface Network - Pytorch',
   author = 'Phil Wang',