Implementasi Multi-Head Self-Attention Python From Scratch

Mekanisme self-attention, khususnya varian Multi-Head Self-Attention (MHSA), merupakan salah satu inovasi paling berpengaruh dalam deep learning modern, terutama sebagai tulang punggung arsitektur Transformer. Kemampuannya untuk menimbang relevansi berbagai bagian dari data input saat memproses bagian tertentu telah merevolusi pemrosesan bahasa alami (NLP) dan berbagai bidang lainnya. Meskipun banyak library seperti Hugging Face Transformers menyediakan implementasi siap pakai, memahami cara kerja MHSA dari dasar sangat penting untuk pemahaman mendalam dan potensi kustomisasi. Tutorial ini akan memandu Anda langkah demi langkah untuk mengimplementasikan Multi-Head Self-Attention from scratch menggunakan Python dan library dasar seperti PyTorch, tanpa bergantung pada modul MultiHeadAttention bawaan.

Memahami Konsep Dasar: Self-Attention dan MHSA

Apa Itu Self-Attention?

Self-attention, atau dikenal juga sebagai intra-attention, adalah mekanisme yang memungkinkan model untuk ‘melihat’ bagian lain dari urutan input saat memproses satu bagian tertentu (misalnya, sebuah kata dalam kalimat). Bayangkan Anda membaca kalimat, “Kucing itu mengejar tikus karena lapar.” Saat memproses kata “itu”, mekanisme self-attention membantu model memahami bahwa “itu” merujuk pada “kucing”. Hal ini dilakukan dengan menghitung ‘skor attention’ antara kata saat ini dan kata-kata lain dalam urutan, yang kemudian digunakan untuk menciptakan representasi kontekstual dari kata tersebut. Pendekatan ini memungkinkan model menangkap dependensi jangka panjang maupun pendek dalam data sekuensial secara efektif.

Pengenalan Multi-Head Self-Attention (MHSA)

Multi-Head Self-Attention (MHSA) merupakan pengembangan dari mekanisme self-attention dasar. Alih-alih menghitung attention hanya sekali, MHSA melakukannya beberapa kali secara paralel dalam ‘head’ yang berbeda. Setiap ‘head’ belajar untuk fokus pada aspek atau jenis hubungan yang berbeda dalam data. Sebagai contoh, satu head mungkin fokus pada hubungan sintaksis, sementara head lain fokus pada hubungan semantik. Hasil dari semua head ini kemudian digabungkan (dikombinasikan) untuk menghasilkan representasi akhir yang lebih kaya dan komprehensif.

Mengapa MHSA Penting dalam Arsitektur Transformer?

MHSA merupakan komponen fundamental dalam blok encoder dan decoder pada arsitektur Transformer. Kemampuannya untuk memodelkan dependensi tanpa memperhatikan jarak dalam urutan input (berbeda dengan RNN/LSTM yang memproses secara sekuensial) dan kemudahan paralelisasi komputasinya menjadi kunci keberhasilan model-model state-of-the-art seperti BERT, GPT, T5, dan banyak lainnya dalam berbagai tugas AI, mulai dari terjemahan mesin hingga pembuatan teks. Memahami cara kerja multi-head attention adalah kunci untuk memahami model-model canggih ini.

Baca juga: AI dalam Analisis Data Pendidikan Tingkatkan Kualitas Pembelajaran

Tujuan Tutorial Ini

Tujuan utama tutorial ini adalah memberikan panduan praktis untuk mengimplementasikan mekanisme Multi-Head Self-Attention dari awal (from scratch) menggunakan Python dan PyTorch. Kami akan memecah konsep matematis yang mendasarinya dan menerjemahkannya ke dalam kode langkah demi langkah. Fokus utamanya adalah pada pemahaman fundamental di balik MHSA, bukan sekadar menggunakan fungsi high-level yang sudah jadi. Dengan demikian, Anda akan memiliki fondasi yang kuat untuk memahami, bahkan memodifikasi, arsitektur berbasis Transformer.

Prasyarat Sebelum Memulai Implementasi

Pengetahuan Dasar yang Dibutuhkan

• Python: Pemahaman yang baik tentang sintaks dasar Python, termasuk struktur data (list, tuple, dictionary), fungsi, dan konsep pemrograman berorientasi objek (kelas).
• Aljabar Linear: Familiaritas dengan operasi matriks dasar, terutama perkalian matriks (dot product) dan transposisi matriks.
• Deep Learning: Konsep dasar seperti tensor, vektor embedding, lapisan linear (dense layer), dan fungsi aktivasi (khususnya softmax).

Library Python yang Akan Digunakan

Dalam tutorial ini, kita akan mengandalkan PyTorch sebagai library utama untuk implementasi. PyTorch menyediakan operasi tensor yang efisien dan fungsionalitas autograd yang sangat berguna untuk membangun model deep learning. Meskipun demikian, konsep dan logika yang sama sebagian besar dapat diterapkan menggunakan TensorFlow atau bahkan NumPy untuk bagian-bagian tertentu, dengan penyesuaian sintaks. Kita akan melihat contoh kode self-attention python secara rinci.

# Impor library yang diperlukan
import torch
import torch.nn as nn
import torch.nn.functional as F
import math

Konsep Inti: Membedah Cara Kerja Multi-Head Attention

Sebelum kita melangkah ke implementasi kode, mari kita pahami terlebih dahulu komponen dan mekanisme utama yang membentuk Multi-Head Self-Attention. Penjelasan multi-head self-attention ini akan menjadi panduan Anda saat mengikuti bagian implementasi.

Komponen Kunci: Matriks Query (Q), Key (K), dan Value (V)

Inti dari mekanisme attention terletak pada interaksi antara tiga representasi yang berasal dari input yang sama:

• Query (Q): Representasi dari elemen (misalnya, kata) saat ini yang sedang diproses. Ini bertindak sebagai ‘pertanyaan’ tentang bagian mana dari input yang relevan dengan elemen ini.
• Key (K): Representasi dari semua elemen dalam urutan input (termasuk elemen saat ini). Ini berfungsi sebagai ‘kunci’ atau label yang dapat dicocokkan dengan Query untuk menentukan relevansi.
• Value (V): Representasi dari konten atau informasi yang terkandung dalam setiap elemen input. Ini adalah nilai yang akan diambil dan diagregasi berdasarkan seberapa cocok Query dengan Key terkait.

Sebagai analogi sederhana, bayangkan sistem pencarian di perpustakaan. Query Anda adalah topik yang ingin Anda cari. Anda membandingkan Query ini dengan Key (judul atau kata kunci) dari setiap buku di rak. Jika Query cocok dengan Key sebuah buku, Anda mengambil Value (informasi di dalam buku) dari buku tersebut. Dalam konteks self-attention, Q, K, dan V (atau query key value attention python) biasanya dihasilkan dengan memproyeksikan vektor embedding input yang sama menggunakan tiga matriks bobot linier yang berbeda (Wq, Wk, Wv), yang parameternya dipelajari selama proses pelatihan model.

Mekanisme Dasar: Scaled Dot-Product Attention

Ini adalah blok bangunan inti dari self-attention, sering disebut sebagai scaled dot-product attention python. Perhitungannya mengikuti rumus matematis berikut:

Attention(Q, K, V) = softmax( (QK^T) / sqrt(d_k) ) V

Langkah-langkah perhitungannya adalah sebagai berikut:

1. Hitung Skor Attention: Lakukan perkalian matriks antara matriks Query (Q) dan transposisi dari matriks Key (K^T). Hasilnya adalah matriks skor yang menunjukkan tingkat kecocokan atau relevansi antara setiap query dengan setiap key.
2. Scaling (Pembagian Skala): Bagi skor attention yang diperoleh dengan akar kuadrat dari dimensi Key (d_k). Langkah scaling ini sangat penting untuk menstabilkan gradien selama pelatihan, terutama ketika dimensi d_k besar. Tanpa scaling, nilai input ke fungsi softmax bisa menjadi sangat besar, yang dapat menyebabkan gradien menjadi sangat kecil (vanishing gradient) dan menghambat proses pembelajaran.
3. Softmax: Terapkan fungsi softmax pada skor yang telah di-scaling (biasanya dilakukan sepanjang dimensi Key). Operasi ini mengubah skor mentah menjadi distribusi probabilitas, di mana setiap nilai berada di antara 0 dan 1, dan jumlah semua nilai (untuk satu query) adalah 1. Nilai-nilai ini merepresentasikan ‘bobot attention’ untuk setiap elemen Value.
4. Output Tertimbang: Kalikan matriks bobot attention (hasil dari softmax) dengan matriks Value (V). Langkah ini menghasilkan output attention, yang merupakan penjumlahan tertimbang dari vektor-vektor Value, di mana bobotnya ditentukan oleh skor attention yang telah dihitung. Elemen Value yang memiliki skor attention tinggi akan memberikan kontribusi yang lebih besar pada vektor output akhir.

Konsep Multiple Heads: Memparalelkan Proses Attention

Alih-alih melakukan satu perhitungan Scaled Dot-Product Attention dengan Q, K, V berdimensi penuh (embed_dim), MHSA melakukan langkah-langkah berikut:

1. Pertama, input embedding awal diproyeksikan secara linear menggunakan matriks bobot Wq, Wk, dan Wv untuk menghasilkan representasi Q, K, dan V awal.
2. Matriks Q, K, V ini kemudian ‘dibagi’ (secara konseptual) menjadi num_heads bagian yang lebih kecil. Secara praktis, ini dicapai dengan melakukan reshape tensor sehingga dimensi embedding (embed_dim) dipecah menjadi num_heads x head_dim, di mana head_dim = embed_dim / num_heads.
3. Perhitungan Scaled Dot-Product Attention kemudian dilakukan secara independen dan paralel untuk setiap set Q, K, V dari masing-masing ‘head’.

Keuntungannya adalah setiap head dapat belajar untuk fokus pada jenis informasi atau pola hubungan yang berbeda dalam data, beroperasi pada subspace representasi yang berbeda pula. Misalnya, satu head mungkin belajar menangkap dependensi sintaksis lokal, sementara head lain menangkap dependensi semantik jarak jauh.

Langkah Akhir: Konkatenasi dan Proyeksi Linear Output

Setelah perhitungan attention selesai untuk semua head secara paralel, langkah selanjutnya adalah:

1. Output attention dari masing-masing head digabungkan kembali (dikonkatenasi) sepanjang dimensi embedding. Ini menyatukan informasi yang dipelajari oleh semua head.
2. Hasil konkatenasi ini kemudian dilewatkan melalui satu lapisan linear akhir (dengan matriks bobot Wo) untuk menghasilkan output akhir dari blok Multi-Head Self-Attention. Proyeksi akhir ini memungkinkan model untuk mencampur dan mengintegrasikan informasi yang dipelajari dari semua head secara efektif.

Tutorial: Implementasi Multi-Head Self-Attention dari Awal dengan Python (PyTorch)

Sekarang, mari kita terjemahkan konsep-konsep yang telah dibahas ke dalam kode Python menggunakan PyTorch. Ini adalah bagian inti dari tutorial multi-head self-attention python ini, di mana kita akan membangun MHSA from scratch.

1. Setup Awal dan Persiapan Data

Pertama-tama, kita perlu mengimpor library yang diperlukan, mendefinisikan beberapa parameter penting untuk arsitektur kita, dan membuat data input dummy sebagai contoh.

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

# Parameter Model
embed_dim = 512  # Dimensi embedding (misal: output dari layer embedding kata)
num_heads = 8    # Jumlah attention head paralel
seq_length = 100 # Panjang sekuens input
batch_size = 32  # Ukuran batch data

# Validasi: embed_dim harus dapat dibagi habis oleh num_heads
assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"

# Hitung dimensi untuk setiap head
head_dim = embed_dim // num_heads

# Buat input tensor dummy (random)
# Shape: (batch_size, seq_length, embed_dim)
x = torch.randn(batch_size, seq_length, embed_dim)

print(f"Input shape: {x.shape}")
print(f"Embed dim: {embed_dim}, Num heads: {num_heads}, Head dim: {head_dim}")

Di sini, `embed_dim` adalah dimensi representasi vektor untuk setiap token dalam sekuens. `num_heads` menentukan jumlah head paralel, dan `seq_length` adalah panjang urutan input. Input `x` adalah tensor dummy yang merepresentasikan satu batch data sekuensial.

2. Inisialisasi Bobot dan Proyeksi Linear Awal Q, K, V

Kita memerlukan tiga lapisan linear terpisah (tanpa bias, sesuai implementasi standar Transformer) untuk memproyeksikan input `x` menjadi Q, K, dan V. Kita juga memerlukan satu lapisan linear lagi untuk proyeksi output akhir (Wo).

# Lapisan linear untuk menghasilkan Q, K, V dari input x
# Input dim: embed_dim, Output dim: embed_dim
Wq = nn.Linear(embed_dim, embed_dim, bias=False)
Wk = nn.Linear(embed_dim, embed_dim, bias=False)
Wv = nn.Linear(embed_dim, embed_dim, bias=False)

# Lapisan linear untuk proyeksi output akhir
Wo = nn.Linear(embed_dim, embed_dim, bias=False)

print("\nLapisan Linear untuk Q, K, V, dan Output (Wo) telah dibuat.")

Setiap lapisan `nn.Linear(embed_dim, embed_dim)` akan mengambil input berdimensi `embed_dim` dan menghasilkan output dengan dimensi yang sama. Ini merepresentasikan perkalian matriks: Q = XWq, K = XWk, V = XWv.

3. Memproyeksikan Input dan Membaginya untuk Multi-Head

Selanjutnya, kita terapkan lapisan linear yang baru dibuat ke input `x`. Hasilnya kemudian di-reshape untuk memisahkan representasi ke dalam `num_heads` yang berbeda.

# 1. Proyeksikan input x menjadi Q, K, V
# Input x: (batch_size, seq_length, embed_dim)
# Output Q, K, V: (batch_size, seq_length, embed_dim)
Q = Wq(x)
K = Wk(x)
V = Wv(x)

print(f"\nShape Q, K, V setelah proyeksi awal: {Q.shape}")

# 2. Reshape Q, K, V untuk multi-head
# Tujuannya adalah membagi dimensi embed_dim menjadi num_heads * head_dim
# dan mengatur ulang dimensi agar head menjadi dimensi terpisah untuk komputasi paralel.
# Target shape: (batch_size, num_heads, seq_length, head_dim)

# Proses reshape dan transpose untuk Q
# (batch_size, seq_length, embed_dim) -> (batch_size, seq_length, num_heads, head_dim) -> (batch_size, num_heads, seq_length, head_dim)
Q = Q.view(batch_size, seq_length, num_heads, head_dim).transpose(1, 2)

# Proses reshape dan transpose untuk K
# (batch_size, seq_length, embed_dim) -> (batch_size, seq_length, num_heads, head_dim) -> (batch_size, num_heads, seq_length, head_dim)
K = K.view(batch_size, seq_length, num_heads, head_dim).transpose(1, 2)

# Proses reshape dan transpose untuk V
# (batch_size, seq_length, embed_dim) -> (batch_size, seq_length, num_heads, head_dim) -> (batch_size, num_heads, seq_length, head_dim)
V = V.view(batch_size, seq_length, num_heads, head_dim).transpose(1, 2)

print(f"Shape Q, K, V setelah reshape & transpose untuk multi-head: {Q.shape}")

Operasi `view(batch_size, seq_length, num_heads, head_dim)` pertama-tama membagi dimensi terakhir (`embed_dim`) menjadi dua dimensi baru: `num_heads` dan `head_dim`. Kemudian, `transpose(1, 2)` menukar dimensi `seq_length` (indeks 1) dan `num_heads` (indeks 2) sehingga shape akhirnya menjadi `(batch_size, num_heads, seq_length, head_dim)`. Penataan ulang ini krusial karena memungkinkan kita melakukan perhitungan attention secara efisien dan paralel untuk semua head menggunakan operasi tensor batch.

4. Implementasi Scaled Dot-Product Attention per Head

Kini, saatnya mengimplementasikan inti mekanisme attention: rumus `Attention(Q, K, V) = softmax(QK^T / sqrt(d_k))V`. Perhitungan ini dilakukan secara simultan untuk semua head berkat penataan tensor pada langkah sebelumnya. Ini adalah jantung dari transformer attention mechanism code.

# --- Langkah Scaled Dot-Product Attention --- 

# 1. Hitung skor attention: Matmul(Q, K_transpose)
# Q shape: (batch_size, num_heads, seq_length, head_dim)
# K shape: (batch_size, num_heads, seq_length, head_dim)
# K.transpose(-2, -1) shape: (batch_size, num_heads, head_dim, seq_length)
# Hasil (attention_scores) shape: (batch_size, num_heads, seq_length, seq_length)
attention_scores = torch.matmul(Q, K.transpose(-2, -1))

# 2. Scaling skor attention
# d_k dalam rumus adalah dimensi head_dim di sini
scale_factor = math.sqrt(head_dim)
scaled_attention_scores = attention_scores / scale_factor

# 3. Terapkan Softmax
# Softmax diterapkan pada dimensi terakhir (dimensi Key, yaitu seq_length ke-2)
# Ini menghasilkan bobot attention yang ternormalisasi untuk setiap query.
# attention_weights shape: (batch_size, num_heads, seq_length, seq_length)
attention_weights = F.softmax(scaled_attention_scores, dim=-1)

# 4. Kalikan bobot attention dengan Value
# attention_weights shape: (batch_size, num_heads, seq_length, seq_length)
# V shape: (batch_size, num_heads, seq_length, head_dim)
# Hasil (attention_output) shape: (batch_size, num_heads, seq_length, head_dim)
attention_output = torch.matmul(attention_weights, V)

print(f"\nShape skor attention mentah: {attention_scores.shape}")
print(f"Shape bobot attention (setelah softmax): {attention_weights.shape}")
print(f"Shape output attention per head (setelah matmul dengan V): {attention_output.shape}")

Perhatikan penggunaan `K.transpose(-2, -1)` yang menukar dua dimensi terakhir dari tensor K, mempersiapkannya untuk perkalian matriks dengan Q. Fungsi `F.softmax` diterapkan pada `dim=-1`, yang merupakan dimensi `seq_length` dari Key (kolom dalam matriks skor), memastikan bahwa bobot attention untuk setiap query (setiap baris) berjumlah 1. Hasilnya, `attention_output`, adalah representasi V yang telah ditimbang berdasarkan skor attention, dihitung secara terpisah untuk setiap head.

5. Konkatenasi Hasil Attention dari Semua Head

Setelah perhitungan attention selesai di setiap head secara paralel, langkah berikutnya adalah menggabungkan kembali (concatenate) output dari semua head menjadi satu tensor tunggal. Proses ini pada dasarnya adalah kebalikan dari langkah pemisahan head (langkah 3).

# --- Menggabungkan Output dari Semua Head --- 

# attention_output shape saat ini: (batch_size, num_heads, seq_length, head_dim)

# 1. Transpose kembali untuk menukar num_heads dan seq_length
# Target shape: (batch_size, seq_length, num_heads, head_dim)
attention_output_transposed = attention_output.transpose(1, 2)
# print(f"Shape setelah transpose: {attention_output_transposed.shape}")

# 2. Reshape untuk menggabungkan dimensi num_heads dan head_dim -> embed_dim
# Penting: Gunakan .contiguous() sebelum .view() setelah operasi transpose
# untuk memastikan tensor disimpan dalam blok memori yang berdekatan.
# Target shape: (batch_size, seq_length, embed_dim)
concatenated_output = attention_output_transposed.contiguous().view(batch_size, seq_length, embed_dim)

print(f"\nShape setelah konkatenasi semua head: {concatenated_output.shape}")

Operasi `transpose(1, 2)` mengembalikan dimensi `seq_length` dan `num_heads` ke posisi relatif semula. Penggunaan `.contiguous()` sangat penting di sini karena operasi `transpose` dapat mengubah cara tensor disimpan di memori (menjadi non-contiguous), sementara `view` memerlukan tensor yang contiguous. Akhirnya, `view(batch_size, seq_length, embed_dim)` menggabungkan dua dimensi terakhir (`num_heads` dan `head_dim`) kembali menjadi satu dimensi `embed_dim`.

6. Proyeksi Linear Akhir (Output Projection)

Langkah terakhir dalam blok MHSA adalah melewatkan hasil konkatenasi melalui lapisan linear output (Wo) yang telah kita definisikan sebelumnya.

# Terapkan lapisan linear akhir Wo
# concatenated_output shape: (batch_size, seq_length, embed_dim)
# final_output shape: (batch_size, seq_length, embed_dim)
final_output = Wo(concatenated_output)

print(f"\nShape output akhir MHSA setelah proyeksi Wo: {final_output.shape}")

Lapisan `Wo` memproyeksikan representasi gabungan dari semua head, memungkinkan model untuk belajar kombinasi optimal dari informasi yang ditangkap oleh setiap head. Output akhir ini (`final_output`) memiliki dimensi yang sama dengan input awal (`batch_size, seq_length, embed_dim`), sehingga dapat dengan mulus diintegrasikan ke dalam lapisan Transformer berikutnya atau digunakan sebagai output akhir jika ini adalah lapisan terakhir.

7. (Opsional) Membungkus Implementasi dalam Kelas `MultiHeadSelfAttention`

Untuk meningkatkan reusabilitas dan keterbacaan kode, sangat disarankan untuk membungkus seluruh logika MHSA ke dalam sebuah kelas `nn.Module` PyTorch. Ini memberikan contoh kode multi-head self-attention pytorch yang terstruktur dan modular.

class MultiHeadSelfAttention(nn.Module):
    """Implementasi Multi-Head Self-Attention dari awal."""
    def __init__(self, embed_dim, num_heads):
        super(MultiHeadSelfAttention, self).__init__()
        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"

        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        # Lapisan linear untuk Q, K, V dan output projection (Wo)
        self.Wq = nn.Linear(embed_dim, embed_dim, bias=False)
        self.Wk = nn.Linear(embed_dim, embed_dim, bias=False)
        self.Wv = nn.Linear(embed_dim, embed_dim, bias=False)
        self.Wo = nn.Linear(embed_dim, embed_dim, bias=False)

    def split_heads(self, x, batch_size):
        """Membagi dimensi embedding menjadi beberapa head."""
        # Reshape: (batch_size, seq_length, embed_dim) -> (batch_size, seq_length, num_heads, head_dim)
        # Transpose: -> (batch_size, num_heads, seq_length, head_dim)
        x = x.view(batch_size, -1, self.num_heads, self.head_dim)
        return x.transpose(1, 2)

    def scaled_dot_product_attention(self, Q, K, V, mask=None):
        """Menghitung Scaled Dot-Product Attention."""
        # Matmul Q dan K transpose
        # Q, K, V shapes: (batch_size, num_heads, seq_length, head_dim)
        attention_scores = torch.matmul(Q, K.transpose(-2, -1))

        # Scaling
        scale_factor = math.sqrt(self.head_dim)
        scaled_attention_scores = attention_scores / scale_factor

        # (Opsional) Terapkan mask jika ada (misalnya, untuk padding atau look-ahead)
        if mask is not None:
            # Pastikan mask memiliki shape yang sesuai atau dapat di-broadcast
            # Mask biasanya berisi -inf di posisi yang ingin di-mask
            scaled_attention_scores = scaled_attention_scores.masked_fill(mask == 0, -1e9)

        # Softmax
        # attention_weights shape: (batch_size, num_heads, seq_length, seq_length)
        attention_weights = F.softmax(scaled_attention_scores, dim=-1)

        # Matmul dengan V
        # attention_output shape: (batch_size, num_heads, seq_length, head_dim)
        attention_output = torch.matmul(attention_weights, V)

        return attention_output, attention_weights # Seringkali bobot juga dikembalikan untuk analisis

    def combine_heads(self, x, batch_size):
        """Menggabungkan kembali output dari semua head."""
        # Input x shape: (batch_size, num_heads, seq_length, head_dim)
        # Transpose: -> (batch_size, seq_length, num_heads, head_dim)
        # Reshape (contiguous needed): -> (batch_size, seq_length, embed_dim)
        x = x.transpose(1, 2).contiguous()
        return x.view(batch_size, -1, self.embed_dim)

    def forward(self, x_q, x_k, x_v, mask=None):
        """Forward pass untuk Multi-Head Self-Attention.
           Memungkinkan input Q, K, V yang berbeda (misalnya untuk cross-attention).
           Untuk self-attention, x_q = x_k = x_v = input.
        """
        batch_size = x_q.shape[0] # Dapatkan batch_size dari input Query

        # 1. Proyeksi Q, K, V
        Q = self.Wq(x_q)
        K = self.Wk(x_k)
        V = self.Wv(x_v)

        # 2. Split heads
        Q = self.split_heads(Q, batch_size)
        K = self.split_heads(K, batch_size)
        V = self.split_heads(V, batch_size)

        # 3. Scaled Dot-Product Attention per head
        # Di sini kita juga meneruskan mask opsional
        attention_output, attention_weights = self.scaled_dot_product_attention(Q, K, V, mask)

        # 4. Combine heads
        concatenated_output = self.combine_heads(attention_output, batch_size)

        # 5. Proyeksi output akhir
        final_output = self.Wo(concatenated_output)

        # Mengembalikan output akhir dan bobot attention (opsional)
        return final_output #, attention_weights

# --- Akhir Definisi Kelas ---

# Contoh inisialisasi (jika ingin menjalankan langsung)
# mhsa_layer_class = MultiHeadSelfAttention(embed_dim, num_heads)
# output_class = mhsa_layer_class(x, x, x) # Self-attention: Q, K, V dari input yang sama
# print(f"\nShape output dari kelas MHSA: {output_class.shape}")

Kelas ini mengorganisir kode dengan rapi ke dalam metode-metode logis (`split_heads`, `scaled_dot_product_attention`, `combine_heads`) dan metode `forward` utama yang mendefinisikan alur komputasi dari input ke output. Kelas ini juga dimodifikasi sedikit untuk menerima input Q, K, V yang terpisah (`x_q`, `x_k`, `x_v`) pada metode `forward`, membuatnya lebih fleksibel dan dapat digunakan untuk *cross-attention* selain *self-attention* (di mana `x_q=x_k=x_v`). Parameter `mask` juga ditambahkan untuk menangani masking, yang umum diperlukan dalam aplikasi Transformer. Ini membuat kode lebih mudah dibaca, diuji, dan digunakan kembali sebagai bagian dari implementasi multi-head self-attention python from scratch yang lebih besar.

Contoh Penggunaan Kode Kelas Multi-Head Self-Attention

Mari kita lihat bagaimana menggunakan kelas `MultiHeadSelfAttention` yang baru saja kita buat dengan input dummy. Ini menunjukkan penerapan praktis dari contoh kode self-attention python yang telah kita bangun dalam bentuk modular.

# Inisialisasi model MHSA menggunakan kelas yang telah dibuat
mhsa_layer = MultiHeadSelfAttention(embed_dim, num_heads)

# Buat input tensor dummy baru (atau gunakan 'x' yang sudah ada)
input_tensor = torch.randn(batch_size, seq_length, embed_dim)
print(f"\nShape input tensor untuk kelas MHSA: {input_tensor.shape}")

# Lewatkan input melalui layer MHSA
# Untuk self-attention, kita gunakan input_tensor yang sama untuk Q, K, dan V
output_tensor = mhsa_layer(input_tensor, input_tensor, input_tensor)

# Cetak shape output
print(f"Shape output tensor dari kelas MHSA: {output_tensor.shape}")

# Verifikasi bahwa shape output sama dengan shape input
assert output_tensor.shape == input_tensor.shape, "Output shape mismatch!"

Seperti yang diharapkan, shape output (`batch_size, seq_length, embed_dim`) sama persis dengan shape input. Tensor `output_tensor` ini sekarang berisi representasi kontekstual dari setiap elemen dalam sekuens input, yang dihasilkan oleh mekanisme multi-head self-attention yang baru saja kita implementasikan dari awal.

Diskusi: Keuntungan dan Keterbatasan Multi-Head Self-Attention

Setelah memahami dan mengimplementasikan MHSA, penting juga untuk mengetahui keunggulan dan keterbatasannya.

Keunggulan MHSA

• Menangkap Dependensi Kontekstual Jarak Jauh: Sangat efektif dalam memodelkan hubungan kompleks antara elemen-elemen dalam sebuah sekuens, bahkan jika elemen-elemen tersebut terpisah jauh satu sama lain, tidak seperti RNN/LSTM yang rentan terhadap masalah vanishing gradient pada dependensi jarak jauh.
• Pembelajaran Representasi Paralel: Pendekatan ‘multi-head’ memungkinkan model untuk secara bersamaan mempelajari berbagai jenis dependensi atau aspek informasi dari subspace representasi yang berbeda, menghasilkan pemahaman yang lebih kaya.
• Paralelisasi Komputasi: Sebagian besar operasi dalam MHSA, terutama perkalian matriks yang dominan, dapat diparalelkan secara efisien pada perangkat keras modern seperti GPU dan TPU. Hal ini secara signifikan mempercepat proses pelatihan dan inferensi model besar.

Keterbatasan MHSA

• Kompleksitas Kuadratik terhadap Panjang Sekuens: Kompleksitas komputasi dan kebutuhan memori MHSA adalah O(n² * d), di mana n adalah panjang sekuens dan d adalah dimensi embedding. Hal ini membuatnya sangat mahal dan terkadang tidak praktis untuk diterapkan pada sekuens yang sangat panjang (misalnya, dokumen teks yang panjang, genom, atau data deret waktu yang sangat panjang). Berbagai varian attention yang lebih efisien (misalnya, Sparse Attention, Longformer, Linformer, Performer) telah diusulkan untuk mengatasi keterbatasan ini.
• Kurangnya Informasi Posisi Inherent: Mekanisme self-attention murni bersifat *permutation invariant*, artinya ia memperlakukan input sebagai ‘kantong kata’ (bag of words) dan tidak secara inheren memperhitungkan urutan atau posisi elemen dalam sekuens. Oleh karena itu, arsitektur Transformer memerlukan mekanisme tambahan seperti Positional Encoding (yang ditambahkan ke input embedding) untuk menyuntikkan informasi posisi ke dalam model.

Kesimpulan: Rangkuman Implementasi MHSA dari Awal

Multi-Head Self-Attention adalah mekanisme revolusioner yang menjadi inti dari keberhasilan arsitektur Transformer dan banyak model AI modern lainnya. Dalam tutorial ini, kita telah membedah cara kerjanya secara konseptual dan mengimplementasikannya langkah demi langkah from scratch menggunakan PyTorch. Kita telah melihat bagaimana memproyeksikan input menjadi Query, Key, dan Value, membaginya menjadi beberapa head paralel, menghitung scaled dot-product attention untuk setiap head, menggabungkan kembali hasilnya, dan akhirnya melakukan proyeksi linear akhir untuk mendapatkan output kontekstual.

Meskipun implementasi from scratch ini mungkin terasa lebih kompleks daripada sekadar memanggil modul bawaan dari library seperti PyTorch atau TensorFlow, proses ini memberikan pemahaman fundamental yang tak ternilai tentang salah satu komponen terpenting dalam lanskap AI saat ini. Menyelesaikan implementasi Multi-Head Self-Attention dari awal menggunakan Python seperti ini akan memberikan Anda pemahaman yang lebih mendalam tentang cara kerja model-model canggih.

Pemahaman mendalam tentang mekanisme inti seperti Multi-Head Self-Attention adalah langkah krusial dalam perjalanan menguasai dan menerapkan AI modern secara efektif. Jika Anda tertarik bagaimana konsep-konsep AI canggih ini dapat diaplikasikan untuk mendorong pertumbuhan bisnis Anda melalui solusi digital inovatif, platform SaaS berbasis AI dengan berbagai alat (teks, audio, gambar, video), AI Agent untuk SEO otomatis, atau pengembangan aplikasi khusus, pelajari lebih lanjut tentang solusi komprehensif yang ditawarkan Kirim.ai. Kami membantu bisnis memanfaatkan kekuatan AI untuk otomatisasi, optimasi, dan penciptaan pengalaman digital yang unggul.