I’ve always been fascinated by how ideas from one field can transform another. When transformers revolutionized natural language processing, I couldn’t help but wonder: could this same architecture work for images? That curiosity led me down the path of building Vision Transformers from scratch. Today, I want to share that journey with you, showing exactly how to implement and train a ViT model in PyTorch for image classification tasks.

Why did I choose to explore this approach? Traditional convolutional neural networks have served us well, but I wanted to understand if global attention mechanisms could offer something different. The results surprised even me. Let me walk you through building this model step by step.

Have you ever considered how an image could be treated as a sequence? That’s the fundamental insight behind Vision Transformers. Instead of processing pixels through convolutional filters, we divide the image into patches and treat each patch as a token, similar to words in a sentence.

Here’s how we start with patch embedding:

class PatchEmbedding(nn.Module):
    def __init__(self, img_size=32, patch_size=4, in_channels=3, embed_dim=768):
        super().__init__()
        self.img_size = img_size
        self.patch_size = patch_size
        self.n_patches = (img_size // patch_size) ** 2
        self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size)
    
    def forward(self, x):
        x = self.proj(x)  # Shape: (B, C, H, W) -> (B, embed_dim, H/P, W/P)
        x = x.flatten(2)  # Flatten spatial dimensions
        x = x.transpose(1, 2)  # Final shape: (B, n_patches, embed_dim)
        return x

This code converts our 32x32 CIFAR-10 images into 64 patches of 4x4 pixels each. Notice how we use a convolution layer cleverly to extract and embed patches simultaneously. It’s efficient and elegant.

What makes transformers so powerful? The multi-head self-attention mechanism allows the model to focus on different parts of the image simultaneously. Here’s my implementation:

class MultiHeadSelfAttention(nn.Module):
    def __init__(self, embed_dim=768, num_heads=12, dropout=0.1):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.scale = self.head_dim ** -0.5
        
        self.qkv = nn.Linear(embed_dim, embed_dim * 3)
        self.proj = nn.Linear(embed_dim, embed_dim)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x):
        batch_size, seq_len, _ = x.shape
        qkv = self.qkv(x).reshape(batch_size, seq_len, 3, self.num_heads, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]
        
        attn_scores = (q @ k.transpose(-2, -1)) * self.scale
        attn_weights = F.softmax(attn_scores, dim=-1)
        attn_weights = self.dropout(attn_weights)
        
        attn_output = attn_weights @ v
        attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, self.embed_dim)
        return self.proj(attn_output)

Did you notice how each head can learn to attend to different spatial relationships? This is what gives ViT its ability to capture both local and global features without convolutional inductive biases.

Now, let’s combine attention with feed-forward networks in a transformer block:

class TransformerBlock(nn.Module):
    def __init__(self, embed_dim=768, num_heads=12, mlp_ratio=4.0, dropout=0.1):
        super().__init__()
        self.norm1 = nn.LayerNorm(embed_dim)
        self.attn = MultiHeadSelfAttention(embed_dim, num_heads, dropout)
        self.norm2 = nn.LayerNorm(embed_dim)
        self.mlp = nn.Sequential(
            nn.Linear(embed_dim, int(embed_dim * mlp_ratio)),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(int(embed_dim * mlp_ratio), embed_dim),
            nn.Dropout(dropout)
        )
    
    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.mlp(self.norm2(x))
        return x

The residual connections and layer normalization are crucial for stable training. I found that without them, the model struggles to converge properly.

When I first trained this on CIFAR-10, the results were impressive. With just 12 transformer blocks and careful hyperparameter tuning, we can achieve around 75% accuracy. That’s remarkable for a model without any convolutional operations.

How do we handle the fact that transformers don’t inherently understand spatial relationships? We add learnable positional embeddings to give our model a sense of patch positions:

class VisionTransformer(nn.Module):
    def __init__(self, img_size=32, patch_size=4, in_channels=3, num_classes=10, 
                 embed_dim=768, depth=12, num_heads=12, mlp_ratio=4.0, dropout=0.1):
        super().__init__()
        self.patch_embed = PatchEmbedding(img_size, patch_size, in_channels, embed_dim)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, 1 + self.patch_embed.n_patches, embed_dim))
        self.blocks = nn.ModuleList([TransformerBlock(embed_dim, num_heads, mlp_ratio, dropout) for _ in range(depth)])
        self.norm = nn.LayerNorm(embed_dim)
        self.head = nn.Linear(embed_dim, num_classes)
    
    def forward(self, x):
        batch_size = x.shape[0]
        x = self.patch_embed(x)
        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)
        x = x + self.pos_embed
        for block in self.blocks:
            x = block(x)
        x = self.norm(x)
        return self.head(x[:, 0])

The [CLS] token serves as a special token whose final representation we use for classification. It’s fascinating how this single vector can capture the essence of the entire image.

Training requires some careful considerations. I use AdamW optimizer with cosine learning rate scheduling and gradual warmup. Data augmentation through random crops and horizontal flips significantly improves performance. The model needs more data than CNNs to shine, but when it works, the results are worth it.

What surprised me most was how quickly ViT learns global relationships. While CNNs build features hierarchically, ViT can attend to any part of the image from the first layer. This global perspective often leads to different failure modes and strengths compared to convolutional approaches.

I hope this practical guide helps you understand and implement Vision Transformers. Building this from scratch gave me deep appreciation for both the elegance and practical considerations of transformer architectures in computer vision.

If you found this walkthrough helpful, I’d love to hear about your experiences in the comments. Feel free to share this with others who might benefit from it, and let me know what other topics you’d like me to cover. Your feedback helps me create better content for our community.