I’ve been thinking a lot about how Vision Transformers are reshaping computer vision. Unlike traditional convolutional networks, they treat images as sequences of patches, applying attention mechanisms to understand global context. This approach often delivers superior performance on complex visual tasks. I want to share my journey of building one from the ground up.

Let’s start with the core component: patch embedding. This process converts image patches into token embeddings that the transformer can process. Here’s how we implement it in PyTorch:

class PatchEmbedding(nn.Module):
    def __init__(self, image_size=224, patch_size=16, dim=768, channels=3):
        super().__init__()
        num_patches = (image_size // patch_size) ** 2
        patch_dim = channels * patch_size * patch_size
        
        self.projection = nn.Sequential(
            Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=patch_size, p2=patch_size),
            nn.Linear(patch_dim, dim)
        )
        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))

    def forward(self, x):
        b, _, _, _ = x.shape
        x = self.projection(x)
        cls_tokens = repeat(self.cls_token, '1 1 d -> b 1 d', b=b)
        x = torch.cat([cls_tokens, x], dim=1)
        x += self.pos_embedding
        return x

But why do we need positional embeddings when working with images? The answer lies in how transformers process information. Without positional cues, the model would treat all patches equally, losing crucial spatial relationships.

The multi-head attention mechanism forms the heart of the transformer. It allows the model to focus on different parts of the image simultaneously:

class MultiHeadAttention(nn.Module):
    def __init__(self, dim, heads=8, dim_head=64, dropout=0.):
        super().__init__()
        inner_dim = dim_head * heads
        self.heads = heads
        self.scale = dim_head ** -0.5
        
        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)
        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, dim),
            nn.Dropout(dropout)
        )

    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 = dots.softmax(dim=-1)
        
        out = torch.matmul(attn, v)
        out = rearrange(out, 'b h n d -> b n (h d)')
        return self.to_out(out)

Have you ever wondered how these attention heads actually learn to focus on different aspects of an image? Each head develops unique patterns, some focusing on edges, others on textures or specific objects.

Let’s put everything together into a complete Vision Transformer:

class ViT(nn.Module):
    def __init__(self, image_size=224, patch_size=16, num_classes=1000, dim=768, 
                 depth=12, heads=12, mlp_dim=3072, dropout=0.1, emb_dropout=0.1):
        super().__init__()
        
        self.patch_embedding = PatchEmbedding(image_size, patch_size, dim)
        self.dropout = nn.Dropout(emb_dropout)
        
        self.transformer = nn.ModuleList([
            nn.ModuleDict({
                'attn': MultiHeadAttention(dim, heads, dim//heads, dropout),
                'ff': nn.Sequential(
                    nn.Linear(dim, mlp_dim),
                    nn.GELU(),
                    nn.Dropout(dropout),
                    nn.Linear(mlp_dim, dim),
                    nn.Dropout(dropout)
                ),
                'norm1': nn.LayerNorm(dim),
                'norm2': nn.LayerNorm(dim)
            }) for _ in range(depth)
        ])
        
        self.mlp_head = nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, num_classes)
        )

    def forward(self, img):
        x = self.patch_embedding(img)
        x = self.dropout(x)
        
        for block in self.transformer:
            x = block['norm1'](x)
            x = block['attn'](x) + x
            x = block['norm2'](x)
            x = block['ff'](x) + x
            
        cls_output = x[:, 0]
        return self.mlp_head(cls_output)

Training these models requires careful consideration of learning rates and optimization strategies. I’ve found that using cosine annealing with warm restarts works particularly well for ViTs:

def get_cosine_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps):
    def lr_lambda(current_step):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress)))
    
    return LambdaLR(optimizer, lr_lambda)

What happens when we scale these models to larger datasets? The performance improvements can be dramatic, but so are the computational requirements. Modern techniques like gradient checkpointing and mixed precision training make this more manageable.

The true power of Vision Transformers emerges when we combine them with modern data augmentation techniques. Mixup, CutMix, and RandAugment can significantly boost performance while improving generalization.

Building custom Vision Transformers has transformed how I approach computer vision problems. The flexibility to adapt the architecture to specific needs, combined with the power of attention mechanisms, opens up new possibilities. I encourage you to experiment with different configurations and see what works best for your specific use case.

If you found this guide helpful or have questions about implementing your own Vision Transformer, I’d love to hear your thoughts in the comments. Feel free to share this with others who might benefit from building custom vision models!