Faris
Published on

Yuk kita buat arsitektur transformer

Authors

Beberapa waktu lalu aku nonton video dari channel arfy slowy berjudul Yuk kita buat arsitektur transformer — video sepanjang 2 jam lebih di mana beliau membangun arsitektur Transformer dari nol (from scratch) pakai PyTorch, tanpa nn.Transformer bawaan. Catatan ini adalah rangkumannya, biar bisa dipelajari ulang kapan saja.

Video ini merupakan implementasi langsung dari paper legendaris Google tahun 2017, Attention Is All You Need. Kalau mau ikut dari awal, disarankan baca dulu paper-nya biar konsepnya kebayang sebelum lihat kodenya. Kode lengkapnya juga dibagikan gratis di GitHub WargaSlowy/transformer.

Pengenalan

Apa sih Transformer itu? Singkatnya, Transformer adalah arsitektur jaringan saraf yang jadi fondasi hampir semua model AI modern — dari BERT, GPT, sampai model-model multimodal. Daya tariknya: dia tidak pakai rekurensi sama sekali. Tidak ada loop per token seperti RNN/LSTM, hanya attention untuk menangkap hubungan antar token, sehingga prosesnya bisa diparalelkan penuh.

Inisialisasi project

Project-nya disetup pakai uv untuk manajemen dependensi Python yang cepat, dengan PyTorch sebagai library utama.

Ada satu gotcha Python klasik yang dibahas di sini: saat menyusun stack dari beberapa blok, jangan asal buat list yang isinya referensi objek yang sama — memodifikasi satu blok akan memengaruhi blok lain. Solusinya pakai copy.deepcopy supaya tiap blok jadi objek yang independen:

import copy
import torch.nn as nn

class Encoder(nn.Module):
    def __init__(self, block, num_layers):
        super().__init__()
        # deepcopy: tiap layer harus objek independen, bukan referensi yang sama
        self.layers = nn.ModuleList([copy.deepcopy(block) for _ in range(num_layers)])

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

Menariknya, pendaftaran modul ke nn.Module ini penting banget untuk operasi seperti backpropagation — kalau modul tidak terdaftar dengan benar, parameter-nya tidak ikut ter-update saat training.

Embedding & positional encoding

Langkah pertama: mengubah token menjadi vektor lewat embedding, lalu menambahkan informasi urutan lewat positional encoding. Kenapa perlu positional encoding? Karena attention sendiri tidak "tahu" urutan — tanpa itu, kalimat "aku makan nasi" dan "nasi makan aku" akan diperlakukan identik.

Implementasinya mengikuti paper: sinusoidal encoding (sin dan cos) yang ditambahkan ke hasil embedding:

import torch
import torch.nn as nn

class PositionalEncoding(nn.Module):
    def __init__(self, embed_dim, max_len=5000):
        super().__init__()
        pos = torch.arange(max_len).unsqueeze(1)
        i = torch.arange(0, embed_dim, 2)
        div = torch.exp(-torch.log(torch.tensor(10000.0)) * i / embed_dim)
        pe = torch.zeros(max_len, embed_dim)
        pe[:, 0::2] = torch.sin(pos * div)
        pe[:, 1::2] = torch.cos(pos * div)
        self.register_buffer("pe", pe)

    def forward(self, x):
        return x + self.pe[: x.size(1)]

Cara verifikasi yang dipakai sederhana tapi efektif: cetak beberapa nilai output (misalnya dua kata pertama dari batch pertama, lima dimensi pertama) dan pastikan angkanya masuk akal sebelum lanjut ke bagian berikutnya.

Attention

Ini jantungnya Transformer. Mekanisme scaled dot-product attention menghitung seberapa "perhatian" setiap token ke token lain lewat tiga vektor: Query, Key, dan Value:

import torch
import torch.nn as nn

class SelfAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.qkv = nn.Linear(embed_dim, embed_dim * 3)
        self.out = nn.Linear(embed_dim, embed_dim)

    def forward(self, x):
        B, T, C = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        # split heads: (B, T, C) -> (B, num_heads, T, head_dim)
        q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)

        attn = (q @ k.transpose(-2, -1)) / self.head_dim ** 0.5
        attn = attn.softmax(dim=-1)
        out = (attn @ v).transpose(1, 2).contiguous().view(B, T, C)
        return self.out(out)

Skor Query dan Key di-scale dengan akar dimensi head, lalu di-softmax jadi bobot, dan dipakai untuk menggabungkan nilai-nilai Value. Untuk multi-head attention, proses ini dilakukan paralel di beberapa head berukuran lebih kecil, sehingga model bisa menangkap berbagai jenis relasi sekaligus.

Encoder

Encoder terdiri dari stack Transformer block. Tiap block punya resep yang konsisten:

  1. Multi-head attention (dengan residual connection)
  2. Layer normalization
  3. Feed-forward layer dengan faktor ekspansi (biasanya menggandakan dimensi 4x lalu kembali)
  4. Normalisasi + residual lagi
class TransformerBlock(nn.Module):
    def __init__(self, embed_dim, num_heads, expansion_factor=4, dropout=0.1):
        super().__init__()
        self.attn = SelfAttention(embed_dim, num_heads)
        self.norm1 = nn.LayerNorm(embed_dim)
        self.ffwd = nn.Sequential(
            nn.Linear(embed_dim, embed_dim * expansion_factor),
            nn.GELU(),
            nn.Linear(embed_dim * expansion_factor, embed_dim),
        )
        self.norm2 = nn.LayerNorm(embed_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        x = x + self.dropout(self.attn(self.norm1(x)))
        x = x + self.dropout(self.ffwd(self.norm2(x)))
        return x

Decoder

Decoder strukturnya mirip encoder, tapi dengan tambahan penting: masked attention pada bagian self-attention, supaya tiap posisi hanya boleh "melihat" token sebelumnya (casual masking) — ini yang bikin model bisa untuk tugas generatif. Decoder juga menerima output dari encoder sebagai konteks, sehingga tiap block decoder punya cross-attention ke representasi input. Hasil akhirnya dilewatkan layer linear untuk menghasilkan probabilitas output per token.

Menerapkan Transformer

Bagian terakhir merangkai semuanya: class Transformer utama yang menggabungkan embedding, positional encoding, encoder stack, dan decoder stack — lengkap dengan parameter seperti ukuran vocab, dimensi embedding, jumlah head, jumlah block, dan dropout yang dikonfigurasi di satu tempat.

Setelah dirakit, arsitekturnya diuji dengan input acak dan memeriksa bentuk output — representasi kontekstual dengan dimensi yang sesuai harapan (misalnya batch 32, panjang sekuens 50, dimensi 512).

Penutup

Membangun Transformer dari nol ternyata isinya cuma beberapa modul kecil yang disusun berlapis: embedding → positional encoding → attention → block → stack → encoder + decoder. Tidak ada trik sihir — hanya matriks, softmax, dan residual connection yang disusun dengan rapi.

Referensi dari video: