From 74b62009f8629ce7802d73942875ec68edd4ee09 Mon Sep 17 00:00:00 2001 From: Phil Wang Date: Thu, 10 Oct 2024 08:09:58 -0700 Subject: [PATCH] go for multi-headed rmsnorm for the qknorm on hypersphere vit --- setup.py | 2 +- vit_pytorch/normalized_vit.py | 7 +++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/setup.py b/setup.py index 53b727d..e7e77aa 100644 --- a/setup.py +++ b/setup.py @@ -6,7 +6,7 @@ with open('README.md') as f: setup( name = 'vit-pytorch', packages = find_packages(exclude=['examples']), - version = '1.8.1', + version = '1.8.2', license='MIT', description = 'Vision Transformer (ViT) - Pytorch', long_description=long_description, diff --git a/vit_pytorch/normalized_vit.py b/vit_pytorch/normalized_vit.py index 0c47ff1..c3a5925 100644 --- a/vit_pytorch/normalized_vit.py +++ b/vit_pytorch/normalized_vit.py @@ -76,7 +76,8 @@ class Attention(Module): self.dropout = dropout - self.qk_scale = nn.Parameter(torch.ones(dim_head) * (dim_head ** 0.25)) + self.q_scale = nn.Parameter(torch.ones(dim_inner) * (dim_head ** 0.25)) + self.k_scale = nn.Parameter(torch.ones(dim_inner) * (dim_head ** 0.25)) 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)') @@ -89,12 +90,14 @@ class Attention(Module): ): q, k, v = self.to_q(x), self.to_k(x), self.to_v(x) + q = q * self.q_scale + k = k * self.k_scale + q, k, v = map(self.split_heads, (q, k, v)) # query key rmsnorm q, k = map(l2norm, (q, k)) - q, k = (q * self.qk_scale), (k * self.qk_scale) # scale is 1., as scaling factor is moved to s_qk (dk ^ 0.25) - eq. 16