Memahami Vision Transformer (ViT): Implementasi Arsitektur Attention untuk Klasifikasi Gambar dengan PyTorch

Dari CNN ke Vision Transformer

Selama bertahun-tahun, Convolutional Neural Network (CNN) menjadi standar untuk semua tugas computer vision. Arsitektur ini bekerja dengan konvolusi yang menggeser filter kecil di seluruh gambar, menangkap pola lokal seperti tepi, tekstur, dan bentuk. Pendekatan ini sangat efektif karena gambar memang memiliki struktur spasial lokal yang kuat. Namun CNN memiliki keterbatasan mendasar: receptive field yang sempit. Setiap neuron hanya melihat sebagian kecil gambar pada satu waktu, sehingga menangkap hubungan antar-region yang berjauhan membutuhkan banyak layer konvolusi bertumpuk.

Di sisi lain, dunia Natural Language Processing (NLP) mengalami revolusi dengan kemunculan transformer. Arsitektur ini mengandalkan mekanisme attention yang memungkinkan setiap token berinteraksi langsung dengan semua token lain dalam satu langkah. Hubungan global antar kata tertangkap tanpa harus melewati lapisan bertingkat. Keberhasilan besar model seperti BERT dan GPT membuktikan bahwa attention mampu menangkap dependensi jarak jauh dengan sangat baik.

Vision Transformer (ViT) lahir dari pertanyaan sederhana: apakah kita bisa menerapkan arsitektur transformer langsung pada gambar? Jawabannya adalah bisa, dengan satu langkah transformasi penting. Gambar tidak bisa dimasukkan ke transformer secara langsung seperti teks. Gambar perlu dipecah menjadi potongan-potongan kecil yang disebut patch, lalu setiap patch diperlakukan seperti token kata dalam kalimat. Inilah ide inti dari ViT yang pertama kali dipublikasikan oleh tim Google Research pada tahun 2020.

Hasilnya mengejutkan komunitas. Ketika dilatih pada dataset yang cukup besar, ViT mampu mengungguli CNN tercanggih sekaligus membutuhkan komputasi yang lebih efisien untuk training. Kunci utamanya terletak pada kemampuan attention untuk menangkap representasi global gambar sejak layer pertama, bukan bertahap seperti konvolusi.

Memahami Mekanisme Patch Embedding

Langkah pertama dalam ViT adalah mengubah gambar menjadi urutan patch. Kita memecah gambar berukuran H×W menjadi grid patch berukuran tetap, misalnya 4×4 atau 16×16 piksel. Untuk gambar 32×32 dengan patch 4×4, kita mendapatkan 64 patch. Setiap patch kemudian diproyeksikan secara linear menjadi vektor embedding berdimensi D menggunakan sebuah konvolusi 2D dengan kernel dan stride sebesar ukuran patch.

Setiap vektor embedding mewakili satu patch sebagai "token". Posisi patch dalam gambar sangat penting — patch di pojok kiri atas jelas berbeda maknanya dengan patch di pojok kanan bawah. Transformer sendiri tidak memiliki konsep urutan bawaan, sehingga kita menambahkan positional embedding yang dijumlahkan dengan setiap token. Ada dua hal yang juga perlu ditambahkan: sebuah class token [CLS] yang ditempatkan di awal urutan dan menjadi representasi global dari seluruh gambar untuk tugas klasifikasi.

Implementasi patch embedding dengan PyTorch cukup ringkas menggunakan nn.Conv2d. Konvolusi dengan kernel dan stride sama dengan ukuran patch melakukan linear projection sekaligus membagi gambar menjadi grid.

!pip install torch torchvision

import torch
import torch.nn as nn

class PatchEmbedding(nn.Module):
    def __init__(self, in_channels=1, patch_size=4, embed_dim=64):
        super().__init__()
        self.proj = nn.Conv2d(in_channels, embed_dim,
                              kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        # x: (B, C, H, W)
        x = self.proj(x)          # (B, embed_dim, H/p, W/p)
        x = x.flatten(2)          # (B, embed_dim, num_patches)
        x = x.transpose(1, 2)     # (B, num_patches, embed_dim)
        return x

x = torch.randn(1, 1, 28, 28)
pe = PatchEmbedding(patch_size=4, embed_dim=64)
out = pe(x)
print("Patch embedding output shape:", tuple(out.shape))

Output:

Patch embedding output shape: (1, 64, 64)

Perhatikan alurnya: konvolusi menghasilkan tensor berukuran (1, 64, 7, 7) untuk input 28×28 dengan patch 4×4. Flatten menggabungkan dua dimensi spasial menjadi 49 patch, lalu transpose menyusunnya menjadi (batch, 49 patch, 64 dimensi). Output (1, 64, 64) ini berarti 64 token (49 patch + 1 class token akan ditambahkan nanti), masing-masing dengan embedding 64 dimensi. Bentuk ini persis seperti input transformer di NLP, sehingga kita bisa langsung memakai mekanisme attention standar.

Membangun Encoder Transformer dengan Multi-Head Self-Attention

Inti dari ViT adalah encoder transformer. Setiap blok encoder terdiri dari tiga komponen utama: Multi-Head Self-Attention, MLP (Multi-Layer Perceptron), dan LayerNorm, semuanya dihubungkan dengan residual connection.

Self-attention bekerja dengan tiga representasi yang diturunkan dari input: query, key, dan value. Setiap token menghasilkan query yang menanyakan "informasi apa yang saya butuhkan?", key yang menjawab "informasi apa yang saya miliki?", dan value yang berisi informasi aktual. Skor attention dihitung dengan dot product antara query dan key, dinormalisasi dengan softmax, lalu digunakan untuk membobot nilai value. Dengan cara ini, setiap patch dapat "melihat" semua patch lain dan mengumpulkan informasi yang relevan secara langsung.

Gambar: Alur perhitungan self-attention — input X diproyeksikan melalui matriks bobot W_q, W_k, dan W_v menjadi query, key, serta value, lalu skor attention dihitung untuk membentuk konteks — Sumber: Wikimedia Commons (Numiri, CC BY-SA 4.0)

Multi-head attention menjalankan proses ini beberapa kali secara paralel dengan proyeksi berbeda. Setiap head mempelajari jenis relasi yang berbeda — satu head mungkin fokus pada tekstur lokal, head lain pada struktur global. Hasil semua head digabungkan dan diproyeksikan kembali. LayerNorm menstabilkan training dengan menormalkan aktivasi, sementara residual connection menambahkan input asli ke output blok, mencegah informasi hilang saat jaringan semakin dalam.

class TransformerEncoderBlock(nn.Module):
    def __init__(self, embed_dim=64, num_heads=4, mlp_ratio=2.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(embed_dim)
        self.attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
        self.norm2 = nn.LayerNorm(embed_dim)
        hidden_dim = int(embed_dim * mlp_ratio)
        self.mlp = nn.Sequential(
            nn.Linear(embed_dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, embed_dim),
        )

    def forward(self, x):
        x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]
        x = x + self.mlp(self.norm2(x))
        return x

enc = TransformerEncoderBlock(embed_dim=64, num_heads=4)
out_enc = enc(out)
print("Encoder block output shape:", tuple(out_enc.shape))

Output:

Encoder block output shape: (1, 64, 64)

nn.MultiheadAttention dari PyTorch membungkus seluruh mekanisme scaled dot-product attention, query-key-value projection, dan penggabungan multi-head. Kita meneruskan output LayerNorm sebagai query, key, dan value secara bersamaan — inilah makna "self" dalam self-attention. Fungsi GELU pada MLP dipilih karena memberikan kurva aktivasi yang lebih halus dibanding ReLU. Bentuk output tetap (1, 64, 64): dimensi token dan embedding dipertahankan, hanya representasinya yang diperkaya dengan informasi kontekstual global.

Merakit Model ViT Lengkap untuk Klasifikasi Gambar

Sekarang kita merakit semua komponen menjadi model VisionTransformer yang utuh. Urutan forward pass-nya: gambar diproses oleh patch embedding, class token disisipkan di awal, positional embedding dijumlahkan, urutan token melewati stack encoder transformer, lalu output class token diambil dan dipetakan ke logits kelas melalui classification head.

Class token berperan penting. Setelah melewati encoder, representasinya telah "menyerap" informasi dari semua patch melalui attention. Head klasifikasi cukup membaca token pertama ini tanpa perlu agregasi tambahan seperti average pooling. Pendekatan ini mengikuti desain asli ViT dan terbukti efektif.

Gambar: Arsitektur Vision Transformer — gambar dipecah menjadi grid patch, setiap patch diproyeksikan secara linear, dan urutan embedding diproses oleh stack encoder transformer — Sumber: Dive into Deep Learning (CC BY-SA 4.0)

class VisionTransformer(nn.Module):
    def __init__(self, in_channels=1, patch_size=4, embed_dim=64,
                 depth=4, num_heads=4, num_classes=10, img_size=28):
        super().__init__()
        self.patch_embed = PatchEmbedding(in_channels, patch_size, embed_dim)
        num_patches = (img_size // patch_size) ** 2
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
        self.blocks = nn.Sequential(*[
            TransformerEncoderBlock(embed_dim, num_heads) for _ in range(depth)
        ])
        self.norm = nn.LayerNorm(embed_dim)
        self.head = nn.Linear(embed_dim, num_classes)

    def forward(self, x):
        B = x.shape[0]
        x = self.patch_embed(x)
        cls_token = self.cls_token.expand(B, -1, -1)
        x = torch.cat([cls_token, x], dim=1)
        x = x + self.pos_embed
        x = self.blocks(x)
        x = self.norm(x)
        return self.head(x[:, 0])

model = VisionTransformer(patch_size=4, embed_dim=64, depth=4, num_heads=4)
logits = model(x)
print("ViT logits shape:", tuple(logits.shape))
n_params = sum(p.numel() for p in model.parameters())
print(f"Total parameters: {n_params:,}")

Output:

ViT logits shape: (1, 10)
Total parameters: 142,026

Model dengan 4 layer encoder hanya memiliki sekitar 142 ribu parameter — sangat ringkas untuk arsitektur transformer. Ukuran pos_embed adalah (1, 50, 64) karena 49 patch ditambah 1 class token. Perhatikan bahwa jumlah token menentukan panjang konteks maksimum: model ini hanya bisa menerima gambar yang menghasilkan tepat 49 patch. Untuk gambar berukuran lain, dimensi pos_embed harus disesuaikan. Output akhir berbentuk (1, 10) sesuai jumlah kelas pada dataset.

Melatih dan Mengevaluasi ViT pada Dataset Klasifikasi

Kita akan melatih ViT pada dataset MNIST, dataset klasifikasi digit tulisan tangan yang terdiri dari 60.000 gambar training dan 10.000 gambar test berukuran 28×28. Gambar di-normalisasi menggunakan mean dan standar deviasi dataset. Training loop-nya standar: forward pass menghasilkan logits, CrossEntropyLoss menghitung error terhadap label, backward pass menghitung gradien, dan optimizer Adam memperbarui bobot.

import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,)),
])

trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
testset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)
train_loader = DataLoader(trainset, batch_size=128, shuffle=True)
test_loader = DataLoader(testset, batch_size=128, shuffle=False)

device = 'cuda' if torch.cuda.is_available() else ('mps' if torch.backends.mps.is_available() else 'cpu')
model = VisionTransformer(patch_size=4, embed_dim=64, depth=2, num_heads=4).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
num_epochs = 5
loss_history = []

for epoch in range(num_epochs):
    model.train()
    total_loss = 0
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    avg_loss = total_loss / len(train_loader)
    loss_history.append(avg_loss)

    model.eval()
    correct, total = 0, 0
    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    acc = 100 * correct / total
    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {avg_loss:.4f}, Test Acc: {acc:.2f}%')

print(f'Final Test Accuracy: {acc:.2f}%')

Output:

Epoch [1/5], Loss: 0.8686, Test Acc: 90.23%
Epoch [2/5], Loss: 0.2355, Test Acc: 94.37%
Epoch [3/5], Loss: 0.1516, Test Acc: 95.58%
Epoch [4/5], Loss: 0.1143, Test Acc: 95.59%
Epoch [5/5], Loss: 0.0918, Test Acc: 96.14%
Final Test Accuracy: 96.14%

Hasil ini menunjukkan pola training yang sehat. Pada epoch pertama, loss masih tinggi (0.8686) dan akurasi 90.23 persen — model baru mulai mempelajari pola dasar digit. Penurunan loss paling drastis terjadi antara epoch 1 dan 2, dari 0.8686 ke 0.2355. Setelah epoch 3, akurasi mulai mendatar di kisaran 95-96 persen. Dengan hanya 2 layer encoder dan 5 epoch, ViT mencapai akurasi test 96.14 persen pada MNIST. Kurva loss yang menurun secara konsisten tanpa osilasi menandakan learning rate dan arsitektur sudah sesuai.

Gambar: Kurva training loss ViT yang menurun konsisten dari 0.87 ke 0.09 selama 5 epoch pada dataset MNIST — dihasilkan dari eksekusi kode training di atas.

Kita juga bisa memplot kurva loss untuk memverifikasi konvergensi. Kurva yang menurun mulus tanpa lonjakan menunjukkan tidak ada masalah dengan learning rate. Jika loss berosilasi, learning rate terlalu besar. Jika penurunan sangat lambat, arsitektur terlalu kecil atau learning rate terlalu rendah. Memonitor kurva ini adalah kebiasaan yang wajib dalam setiap eksperimen deep learning.

Praktik Terbaik dan Kapan Menggunakan ViT

ViT bukan pengganti CNN dalam semua situasi. Keunggulan utamanya muncul saat data sangat besar. ViT tidak memiliki inductive bias bawaan seperti CNN — model harus mempelajari sendiri bahwa piksel yang berdekatan cenderung berkorelasi. Akibatnya, ViT yang dilatih dari nol pada dataset kecil sering kalah dari CNN. Pada ImageNet-21k atau JFT-300M, ViT justru unggul karena attention memanfaatkan data besar secara lebih efisien.

Di praktik nyata, sebagian besar pengguna tidak melatih ViT dari nol. Pendekatan yang umum adalah transfer learning: menggunakan ViT yang sudah dilatih pada dataset raksasa, lalu fine-tuning pada dataset target. Varian seperti DeiT (Data-efficient Image Transformers) memperkenalkan teknik distilasi untuk membuat ViT bekerja dengan dataset yang lebih kecil. Swin Transformer menggabungkan keunggulan CNN dan ViT dengan attention yang dihitung dalam window lokal.

Pertimbangan komputasi juga penting. Self-attention memiliki kompleksitas kuadratik terhadap jumlah token. Untuk gambar resolusi tinggi, jumlah patch membengkak dan biaya komputasi melonjak. Teknik seperti patching dengan ukuran lebih besar, window attention, atau hierarchical structure membantu mengendalikan biaya ini. Sebagai pedoman: untuk dataset kecil dan resolusi tinggi, CNN atau hybrid masih pilihan yang bijak. Untuk dataset besar dengan kebutuhan representasi global, ViT adalah pilihan yang tepat.

Memahami ViT membuka jalan menuju arsitektur vision modern yang menguasai leaderboard computer vision saat ini. Konsep yang Anda pelajari — patch embedding, multi-head self-attention, dan class token — menjadi fondasi untuk model seperti DeiT, Swin, dan berbagai varian lainnya. Siap menguasai arsitektur transformer untuk computer vision? Lanjutkan perjalanan belajar Anda di program Deep Learning Rumah Coding.