some better defaults for scalable vit

README.md

@@ -18,12 +18,14 @@
 - [Twins SVT](#twins-svt)
 - [CrossFormer](#crossformer)
 - [RegionViT](#regionvit)
+- [ScalableViT](#scalablevit)
 - [NesT](#nest)
 - [MobileViT](#mobilevit)
 - [Masked Autoencoder](#masked-autoencoder)
 - [Simple Masked Image Modeling](#simple-masked-image-modeling)
 - [Masked Patch Prediction](#masked-patch-prediction)
 - [Adaptive Token Sampling](#adaptive-token-sampling)
+- [Patch Merger](#patch-merger)
 - [Vision Transformer for Small Datasets](#vision-transformer-for-small-datasets)
 - [Dino](#dino)
 - [Accessing Attention](#accessing-attention)
@@ -524,6 +526,38 @@ img = torch.randn(1, 3, 224, 224)
 pred = model(img) # (1, 1000)
 ```
 
+## ScalableViT
+
+<img src="./images/scalable-vit-1.png" width="400px"></img>
+
+<img src="./images/scalable-vit-2.png" width="400px"></img>
+
+This Bytedance AI <a href="https://arxiv.org/abs/2203.10790">paper</a> proposes the Scalable Self Attention (SSA) and the Interactive Windowed Self Attention (IWSA) modules. The SSA alleviates the computation needed at earlier stages by reducing the key / value feature map by some factor (`reduction_factor`), while modulating the dimension of the queries and keys (`ssa_dim_key`). The IWSA performs self attention within local windows, similar to other vision transformer papers. However, they add a residual of the values, passed through a convolution of kernel size 3, which they named Local Interactive Module (LIM).
+
+They make the claim in this paper that this scheme outperforms Swin Transformer, and also demonstrate competitive performance against Crossformer.
+
+You can use it as follows (ex. ScalableViT-S)
+
+```python
+import torch
+from vit_pytorch.scalable_vit import ScalableViT
+
+model = ScalableViT(
+    num_classes = 1000,
+    dim = 64,                               # starting model dimension. at every stage, dimension is doubled
+    heads = (2, 4, 8, 16),                  # number of attention heads at each stage
+    depth = (2, 2, 20, 2),                  # number of transformer blocks at each stage
+    ssa_dim_key = (40, 40, 40, 32),         # the dimension of the attention keys (and queries) for SSA. in the paper, they represented this as a scale factor on the base dimension per key (ssa_dim_key / dim_key)
+    reduction_factor = (8, 4, 2, 1),        # downsampling of the key / values in SSA. in the paper, this was represented as (reduction_factor ** -2)
+    window_size = (64, 32, None, None),     # window size of the IWSA at each stage. None means no windowing needed
+    dropout = 0.1,                          # attention and feedforward dropout
+).cuda()
+
+img = torch.randn(1, 3, 256, 256).cuda()
+
+preds = model(img) # (1, 1000)
+```
+
 ## NesT
 
 <img src="./images/nest.png" width="400px"></img>
@@ -542,7 +576,7 @@ nest = NesT(
     dim = 96,
     heads = 3,
     num_hierarchies = 3,        # number of hierarchies
-    block_repeats = (8, 4, 1),  # the number of transformer blocks at each heirarchy, starting from the bottom
+    block_repeats = (2, 2, 8),  # the number of transformer blocks at each heirarchy, starting from the bottom
     num_classes = 1000
 )
 
@@ -732,12 +766,58 @@ v = ViT(
 
 img = torch.randn(4, 3, 256, 256)
 
-preds = v(img) # (1, 1000)
+preds = v(img) # (4, 1000)
 
 # you can also get a list of the final sampled patch ids
 # a value of -1 denotes padding
 
-preds, token_ids = v(img, return_sampled_token_ids = True) # (1, 1000), (1, <=8)
+preds, token_ids = v(img, return_sampled_token_ids = True) # (4, 1000), (4, <=8)
+```
+
+## Patch Merger
+
+
+<img src="./images/patch_merger.png" width="400px"></img>
+
+This <a href="https://arxiv.org/abs/2202.12015">paper</a> proposes a simple module (Patch Merger) for reducing the number of tokens at any layer of a vision transformer without sacrificing performance.
+
+```python
+import torch
+from vit_pytorch.vit_with_patch_merger import ViT
+
+v = ViT(
+    image_size = 256,
+    patch_size = 16,
+    num_classes = 1000,
+    dim = 1024,
+    depth = 12,
+    heads = 8,
+    patch_merge_layer = 6,        # at which transformer layer to do patch merging
+    patch_merge_num_tokens = 8,   # the output number of tokens from the patch merge
+    mlp_dim = 2048,
+    dropout = 0.1,
+    emb_dropout = 0.1
+)
+
+img = torch.randn(4, 3, 256, 256)
+
+preds = v(img) # (4, 1000)
+```
+
+One can also use the `PatchMerger` module by itself
+
+```python
+import torch
+from vit_pytorch.vit_with_patch_merger import PatchMerger
+
+merger = PatchMerger(
+    dim = 1024,
+    num_tokens_out = 8   # output number of tokens
+)
+
+features = torch.randn(4, 256, 1024) # (batch, num tokens, dimension)
+
+out = merger(features) # (4, 8, 1024)
 ```
 
 ## Vision Transformer for Small Datasets
@@ -1294,6 +1374,28 @@ Coming from computer vision and new to transformers? Here are some resources tha
 }
 ```
 
+```bibtex
+@misc{renggli2022learning,
+    title   = {Learning to Merge Tokens in Vision Transformers},
+    author  = {Cedric Renggli and André Susano Pinto and Neil Houlsby and Basil Mustafa and Joan Puigcerver and Carlos Riquelme},
+    year    = {2022},
+    eprint  = {2202.12015},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.CV}
+}
+```
+
+```bibtex
+@misc{yang2022scalablevit,
+    title   = {ScalableViT: Rethinking the Context-oriented Generalization of Vision Transformer}, 
+    author  = {Rui Yang and Hailong Ma and Jie Wu and Yansong Tang and Xuefeng Xiao and Min Zheng and Xiu Li},
+    year    = {2022},
+    eprint  = {2203.10790},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.CV}
+}
+```
+
 ```bibtex
 @misc{vaswani2017attention,
     title   = {Attention Is All You Need},

setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'vit-pytorch',
   packages = find_packages(exclude=['examples']),
-  version = '0.26.2',
+  version = '0.28.1',
   license='MIT',
   description = 'Vision Transformer (ViT) - Pytorch',
   author = 'Phil Wang',
@@ -15,7 +15,7 @@ setup(
     'image recognition'
   ],
   install_requires=[
-    'einops>=0.3',
+    'einops>=0.4.1',
     'torch>=1.6',
     'torchvision'
   ],

vit_pytorch/crossformer.py

@@ -62,9 +62,9 @@ class LayerNorm(nn.Module):
         self.b = nn.Parameter(torch.zeros(1, dim, 1, 1))
 
     def forward(self, x):
-        std = torch.var(x, dim = 1, unbiased = False, keepdim = True).sqrt()
+        var = torch.var(x, dim = 1, unbiased = False, keepdim = True)
         mean = torch.mean(x, dim = 1, keepdim = True)
-        return (x - mean) / (std + self.eps) * self.g + self.b
+        return (x - mean) / (var + self.eps).sqrt() * self.g + self.b
 
 def FeedForward(dim, mult = 4, dropout = 0.):
     return nn.Sequential(

vit_pytorch/cvt.py

@@ -30,9 +30,9 @@ class LayerNorm(nn.Module): # layernorm, but done in the channel dimension #1
         self.b = nn.Parameter(torch.zeros(1, dim, 1, 1))
 
     def forward(self, x):
-        std = torch.var(x, dim = 1, unbiased = False, keepdim = True).sqrt()
+        var = torch.var(x, dim = 1, unbiased = False, keepdim = True)
         mean = torch.mean(x, dim = 1, keepdim = True)
-        return (x - mean) / (std + self.eps) * self.g + self.b
+        return (x - mean) / (var + self.eps).sqrt() * self.g + self.b
 
 class PreNorm(nn.Module):
     def __init__(self, dim, fn):

vit_pytorch/efficient.py

@@ -3,12 +3,16 @@ from torch import nn
 from einops import rearrange, repeat
 from einops.layers.torch import Rearrange
 
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
 class ViT(nn.Module):
     def __init__(self, *, image_size, patch_size, num_classes, dim, transformer, pool = 'cls', channels = 3):
         super().__init__()
-        assert image_size % patch_size == 0, 'image dimensions must be divisible by the patch size'
+        image_size_h, image_size_w = pair(image_size)
+        assert image_size_h % patch_size == 0 and image_size_w % patch_size == 0, 'image dimensions must be divisible by the patch size'
         assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'
-        num_patches = (image_size // patch_size) ** 2
+        num_patches = (image_size_h // patch_size) * (image_size_w // patch_size)
         patch_dim = channels * patch_size ** 2
 
         self.to_patch_embedding = nn.Sequential(

vit_pytorch/mae.py

@@ -14,13 +14,11 @@ class MAE(nn.Module):
         masking_ratio = 0.75,
         decoder_depth = 1,
         decoder_heads = 8,
-        decoder_dim_head = 64,
-        apply_decoder_pos_emb_all = False # whether to (re)apply decoder positional embedding to encoder unmasked tokens
+        decoder_dim_head = 64
     ):
         super().__init__()
         assert masking_ratio > 0 and masking_ratio < 1, 'masking ratio must be kept between 0 and 1'
         self.masking_ratio = masking_ratio
-        self.apply_decoder_pos_emb_all = apply_decoder_pos_emb_all
 
         # extract some hyperparameters and functions from encoder (vision transformer to be trained)
 
@@ -73,10 +71,9 @@ class MAE(nn.Module):
 
         decoder_tokens = self.enc_to_dec(encoded_tokens)
 
-        # reapply decoder position embedding to unmasked tokens, if desired
+        # reapply decoder position embedding to unmasked tokens
 
-        if self.apply_decoder_pos_emb_all:
-            decoder_tokens = decoder_tokens + self.decoder_pos_emb(unmasked_indices)
+        decoder_tokens = decoder_tokens + self.decoder_pos_emb(unmasked_indices)
 
         # repeat mask tokens for number of masked, and add the positions using the masked indices derived above
 

vit_pytorch/nest.py

@@ -20,9 +20,9 @@ class LayerNorm(nn.Module):
         self.b = nn.Parameter(torch.zeros(1, dim, 1, 1))
 
     def forward(self, x):
-        std = torch.var(x, dim = 1, unbiased = False, keepdim = True).sqrt()
+        var = torch.var(x, dim = 1, unbiased = False, keepdim = True)
         mean = torch.mean(x, dim = 1, keepdim = True)
-        return (x - mean) / (std + self.eps) * self.g + self.b
+        return (x - mean) / (var + self.eps).sqrt() * self.g + self.b
 
 class PreNorm(nn.Module):
     def __init__(self, dim, fn):
@@ -131,10 +131,11 @@ class NesT(nn.Module):
 
         seq_len = (fmap_size // blocks) ** 2   # sequence length is held constant across heirarchy
         hierarchies = list(reversed(range(num_hierarchies)))
-        mults = [2 ** i for i in hierarchies]
+        mults = [2 ** i for i in reversed(hierarchies)]
 
         layer_heads = list(map(lambda t: t * heads, mults))
         layer_dims = list(map(lambda t: t * dim, mults))
+        last_dim = layer_dims[-1]
 
         layer_dims = [*layer_dims, layer_dims[-1]]
         dim_pairs = zip(layer_dims[:-1], layer_dims[1:])
@@ -157,10 +158,11 @@ class NesT(nn.Module):
                 Aggregate(dim_in, dim_out) if not is_last else nn.Identity()
             ]))
 
+
         self.mlp_head = nn.Sequential(
-            LayerNorm(dim),
+            LayerNorm(last_dim),
             Reduce('b c h w -> b c', 'mean'),
-            nn.Linear(dim, num_classes)
+            nn.Linear(last_dim, num_classes)
         )
 
     def forward(self, img):

vit_pytorch/recorder.py

@@ -55,5 +55,5 @@ class Recorder(nn.Module):
         target_device = self.device if self.device is not None else img.device
         recordings = tuple(map(lambda t: t.to(target_device), self.recordings))
 
-        attns = torch.stack(recordings, dim = 1)
+        attns = torch.stack(recordings, dim = 1) if len(recordings) > 0 else None
         return pred, attns

vit_pytorch/scalable_vit.py

@@ -0,0 +1,302 @@
+from functools import partial
+import torch
+from torch import nn
+
+from einops import rearrange, repeat
+from einops.layers.torch import Rearrange, Reduce
+
+# helpers
+
+def exists(val):
+    return val is not None
+
+def default(val, d):
+    return val if exists(val) else d
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+def cast_tuple(val, length = 1):
+    return val if isinstance(val, tuple) else ((val,) * length)
+
+# helper classes
+
+class ChanLayerNorm(nn.Module):
+    def __init__(self, dim, eps = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.g = nn.Parameter(torch.ones(1, dim, 1, 1))
+        self.b = nn.Parameter(torch.zeros(1, dim, 1, 1))
+
+    def forward(self, x):
+        var = torch.var(x, dim = 1, unbiased = False, keepdim = True)
+        mean = torch.mean(x, dim = 1, keepdim = True)
+        return (x - mean) / (var + self.eps).sqrt() * self.g + self.b
+
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.norm = ChanLayerNorm(dim)
+        self.fn = fn
+
+    def forward(self, x):
+        return self.fn(self.norm(x))
+
+class Downsample(nn.Module):
+    def __init__(self, dim_in, dim_out):
+        super().__init__()
+        self.conv = nn.Conv2d(dim_in, dim_out, 3, stride = 2, padding = 1)
+
+    def forward(self, x):
+        return self.conv(x)
+
+class PEG(nn.Module):
+    def __init__(self, dim, kernel_size = 3):
+        super().__init__()
+        self.proj = nn.Conv2d(dim, dim, kernel_size = kernel_size, padding = kernel_size // 2, groups = dim, stride = 1)
+
+    def forward(self, x):
+        return self.proj(x) + x
+
+# feedforward
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, expansion_factor = 4, dropout = 0.):
+        super().__init__()
+        inner_dim = dim * expansion_factor
+        self.net = nn.Sequential(
+            nn.Conv2d(dim, inner_dim, 1),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Conv2d(inner_dim, dim, 1),
+            nn.Dropout(dropout)
+        )
+    def forward(self, x):
+        return self.net(x)
+
+# attention
+
+class ScalableSelfAttention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        heads = 8,
+        dim_key = 32,
+        dim_value = 32,
+        dropout = 0.,
+        reduction_factor = 1
+    ):
+        super().__init__()
+        self.heads = heads
+        self.scale = dim_key ** -0.5
+        self.attend = nn.Softmax(dim = -1)
+
+        self.to_q = nn.Conv2d(dim, dim_key * heads, 1, bias = False)
+        self.to_k = nn.Conv2d(dim, dim_key * heads, reduction_factor, stride = reduction_factor, bias = False)
+        self.to_v = nn.Conv2d(dim, dim_value * heads, reduction_factor, stride = reduction_factor, bias = False)
+
+        self.to_out = nn.Sequential(
+            nn.Conv2d(dim_value * heads, dim, 1),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        height, width, heads = *x.shape[-2:], self.heads
+
+        q, k, v = self.to_q(x), self.to_k(x), self.to_v(x)
+
+        # split out heads
+
+        q, k, v = map(lambda t: rearrange(t, 'b (h d) ... -> b h (...) d', h = heads), (q, k, v))
+
+        # similarity
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        # attention
+
+        attn = self.attend(dots)
+
+        # aggregate values
+
+        out = torch.matmul(attn, v)
+
+        # merge back heads
+
+        out = rearrange(out, 'b h (x y) d -> b (h d) x y', x = height, y = width)
+        return self.to_out(out)
+
+class InteractiveWindowedSelfAttention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        window_size,
+        heads = 8,
+        dim_key = 32,
+        dim_value = 32,
+        dropout = 0.
+    ):
+        super().__init__()
+        self.heads = heads
+        self.scale = dim_key ** -0.5
+        self.window_size = window_size
+        self.attend = nn.Softmax(dim = -1)
+
+        self.local_interactive_module = nn.Conv2d(dim_value * heads, dim_value * heads, 3, padding = 1)
+
+        self.to_q = nn.Conv2d(dim, dim_key * heads, 1, bias = False)
+        self.to_k = nn.Conv2d(dim, dim_key * heads, 1, bias = False)
+        self.to_v = nn.Conv2d(dim, dim_value * heads, 1, bias = False)
+
+        self.to_out = nn.Sequential(
+            nn.Conv2d(dim_value * heads, dim, 1),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        height, width, heads, wsz = *x.shape[-2:], self.heads, self.window_size
+
+        wsz = default(wsz, height) # take height as window size if not given
+        assert (height % wsz) == 0 and (width % wsz) == 0, f'height ({height}) or width ({width}) of feature map is not divisible by the window size ({wsz})'
+
+        q, k, v = self.to_q(x), self.to_k(x), self.to_v(x)
+
+        # get output of LIM
+
+        local_out = self.local_interactive_module(v)
+
+        # divide into window (and split out heads) for efficient self attention
+
+        q, k, v = map(lambda t: rearrange(t, 'b (h d) (x w1) (y w2) -> (b x y) h (w1 w2) d', h = heads, w1 = wsz, w2 = wsz), (q, k, v))
+
+        # similarity
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        # attention
+
+        attn = self.attend(dots)
+
+        # aggregate values
+
+        out = torch.matmul(attn, v)
+
+        # reshape the windows back to full feature map (and merge heads)
+
+        out = rearrange(out, '(b x y) h (w1 w2) d -> b (h d) (x w1) (y w2)', x = height // wsz, y = width // wsz, w1 = wsz, w2 = wsz)
+
+        # add LIM output 
+
+        out = out + local_out
+
+        return self.to_out(out)
+
+class Transformer(nn.Module):
+    def __init__(
+        self,
+        dim,
+        depth,
+        heads = 8,
+        ff_expansion_factor = 4,
+        dropout = 0.,
+        ssa_dim_key = 32,
+        ssa_dim_value = 32,
+        ssa_reduction_factor = 1,
+        iwsa_dim_key = 32,
+        iwsa_dim_value = 32,
+        iwsa_window_size = None,
+        norm_output = True
+    ):
+        super().__init__()
+        self.layers = nn.ModuleList([])
+        for ind in range(depth):
+            is_first = ind == 0
+
+            self.layers.append(nn.ModuleList([
+                PreNorm(dim, ScalableSelfAttention(dim, heads = heads, dim_key = ssa_dim_key, dim_value = ssa_dim_value, reduction_factor = ssa_reduction_factor, dropout = dropout)),
+                PreNorm(dim, FeedForward(dim, expansion_factor = ff_expansion_factor, dropout = dropout)),
+                PEG(dim) if is_first else None,
+                PreNorm(dim, FeedForward(dim, expansion_factor = ff_expansion_factor, dropout = dropout)),
+                PreNorm(dim, InteractiveWindowedSelfAttention(dim, heads = heads, dim_key = iwsa_dim_key, dim_value = iwsa_dim_value, window_size = iwsa_window_size, dropout = dropout))
+            ]))
+
+        self.norm = ChanLayerNorm(dim) if norm_output else nn.Identity()
+
+    def forward(self, x):
+        for ssa, ff1, peg, iwsa, ff2 in self.layers:
+            x = ssa(x) + x
+            x = ff1(x) + x
+
+            if exists(peg):
+                x = peg(x)
+
+            x = iwsa(x) + x
+            x = ff2(x) + x
+
+        return self.norm(x)
+
+class ScalableViT(nn.Module):
+    def __init__(
+        self,
+        *,
+        num_classes,
+        dim,
+        depth,
+        heads,
+        reduction_factor,
+        window_size = None,
+        iwsa_dim_key = 32,
+        iwsa_dim_value = 32,
+        ssa_dim_key = 32,
+        ssa_dim_value = 32,
+        ff_expansion_factor = 4,
+        channels = 3,
+        dropout = 0.
+    ):
+        super().__init__()
+        self.to_patches = nn.Conv2d(channels, dim, 7, stride = 4, padding = 3)
+
+        assert isinstance(depth, tuple), 'depth needs to be tuple if integers indicating number of transformer blocks at that stage'
+
+        num_stages = len(depth)
+        dims = tuple(map(lambda i: (2 ** i) * dim, range(num_stages)))
+
+        hyperparams_per_stage = [
+            heads,
+            ssa_dim_key,
+            ssa_dim_value,
+            reduction_factor,
+            iwsa_dim_key,
+            iwsa_dim_value,
+            window_size,
+        ]
+
+        hyperparams_per_stage = list(map(partial(cast_tuple, length = num_stages), hyperparams_per_stage))
+        assert all(tuple(map(lambda arr: len(arr) == num_stages, hyperparams_per_stage)))
+
+        self.layers = nn.ModuleList([])
+
+        for ind, (layer_dim, layer_depth, layer_heads, layer_ssa_dim_key, layer_ssa_dim_value, layer_ssa_reduction_factor, layer_iwsa_dim_key, layer_iwsa_dim_value, layer_window_size) in enumerate(zip(dims, depth, *hyperparams_per_stage)):
+            is_last = ind == (num_stages - 1)
+
+            self.layers.append(nn.ModuleList([
+                Transformer(dim = layer_dim, depth = layer_depth, heads = layer_heads, ff_expansion_factor = ff_expansion_factor, dropout = dropout, ssa_dim_key = layer_ssa_dim_key, ssa_dim_value = layer_ssa_dim_value, ssa_reduction_factor = layer_ssa_reduction_factor, iwsa_dim_key = layer_iwsa_dim_key, iwsa_dim_value = layer_iwsa_dim_value, iwsa_window_size = layer_window_size),
+                Downsample(layer_dim, layer_dim * 2) if not is_last else None
+            ]))
+
+        self.mlp_head = nn.Sequential(
+            Reduce('b d h w -> b d', 'mean'),
+            nn.LayerNorm(dims[-1]),
+            nn.Linear(dims[-1], num_classes)
+        )
+
+    def forward(self, img):
+        x = self.to_patches(img)
+
+        for transformer, downsample in self.layers:
+            x = transformer(x)
+
+            if exists(downsample):
+                x = downsample(x)
+
+        return self.mlp_head(x)

vit_pytorch/twins_svt.py

@@ -38,9 +38,9 @@ class LayerNorm(nn.Module):
         self.b = nn.Parameter(torch.zeros(1, dim, 1, 1))
 
     def forward(self, x):
-        std = torch.var(x, dim = 1, unbiased = False, keepdim = True).sqrt()
+        var = torch.var(x, dim = 1, unbiased = False, keepdim = True)
         mean = torch.mean(x, dim = 1, keepdim = True)
-        return (x - mean) / (std + self.eps) * self.g + self.b
+        return (x - mean) / (var + self.eps).sqrt() * self.g + self.b
 
 class PreNorm(nn.Module):
     def __init__(self, dim, fn):

vit_pytorch/vit_with_patch_merger.py

@@ -0,0 +1,144 @@
+import torch
+from torch import nn
+
+from einops import rearrange, repeat
+from einops.layers.torch import Rearrange, Reduce
+
+# helpers
+
+def exists(val):
+    return val is not None
+
+def default(val ,d):
+    return val if exists(val) else d
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+# patch merger class
+
+class PatchMerger(nn.Module):
+    def __init__(self, dim, num_tokens_out):
+        super().__init__()
+        self.scale = dim ** -0.5
+        self.norm = nn.LayerNorm(dim)
+        self.queries = nn.Parameter(torch.randn(num_tokens_out, dim))
+
+    def forward(self, x):
+        x = self.norm(x)
+        sim = torch.matmul(self.queries, x.transpose(-1, -2)) * self.scale
+        attn = sim.softmax(dim = -1)
+        return torch.matmul(attn, x)
+
+# classes
+
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.fn = fn
+    def forward(self, x, **kwargs):
+        return self.fn(self.norm(x), **kwargs)
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, hidden_dim, dropout = 0.):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(dropout)
+        )
+    def forward(self, x):
+        return self.net(x)
+
+class Attention(nn.Module):
+    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
+        super().__init__()
+        inner_dim = dim_head *  heads
+        project_out = not (heads == 1 and dim_head == dim)
+
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+
+        self.attend = nn.Softmax(dim = -1)
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim),
+            nn.Dropout(dropout)
+        ) if project_out else nn.Identity()
+
+    def forward(self, x):
+        qkv = self.to_qkv(x).chunk(3, dim = -1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        attn = self.attend(dots)
+
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+class Transformer(nn.Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0., patch_merge_layer = None, patch_merge_num_tokens = 8):
+        super().__init__()
+        self.layers = nn.ModuleList([])
+
+        self.patch_merge_layer_index = default(patch_merge_layer, depth // 2) - 1 # default to mid-way through transformer, as shown in paper
+        self.patch_merger = PatchMerger(dim = dim, num_tokens_out = patch_merge_num_tokens)
+
+        for _ in range(depth):
+            self.layers.append(nn.ModuleList([
+                PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout)),
+                PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))
+            ]))
+    def forward(self, x):
+        for index, (attn, ff) in enumerate(self.layers):
+            x = attn(x) + x
+            x = ff(x) + x
+
+            if index == self.patch_merge_layer_index:
+                x = self.patch_merger(x)
+
+        return x
+
+class ViT(nn.Module):
+    def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, patch_merge_layer = None, patch_merge_num_tokens = 8, channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0.):
+        super().__init__()
+        image_height, image_width = pair(image_size)
+        patch_height, patch_width = pair(patch_size)
+
+        assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
+
+        num_patches = (image_height // patch_height) * (image_width // patch_width)
+        patch_dim = channels * patch_height * patch_width
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),
+            nn.Linear(patch_dim, dim),
+        )
+
+        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
+        self.dropout = nn.Dropout(emb_dropout)
+
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout, patch_merge_layer, patch_merge_num_tokens)
+
+        self.mlp_head = nn.Sequential(
+            Reduce('b n d -> b d', 'mean'),
+            nn.LayerNorm(dim),
+            nn.Linear(dim, num_classes)
+        )
+
+    def forward(self, img):
+        x = self.to_patch_embedding(img)
+        b, n, _ = x.shape
+
+        x += self.pos_embedding[:, :n]
+        x = self.dropout(x)
+
+        x = self.transformer(x)
+
+        return self.mlp_head(x)