ပရောဂျက် - Mini Transformer Block တစ်ခု တည်ဆောက်ခြင်း

နားလည်ထားရမယ့် အချက်

Transformer block တစ်ခုဆိုတာ တကယ်တော့ sub-layer နှစ်ခုကို ပုံစံတစ်ခုတည်းနဲ့ နှစ်ကြိမ် ဆက်ကပ်ထားတာပါပဲ — sub-layer တစ်ခုလုပ်၊ ပြီးရင် block ရဲ့ input ကို output ထဲ ပြန်ပေါင်း၊ ပြီးရင် normalize လုပ်တယ်။ Self-attention (nn.MultiheadAttention ကနေ) က sequence ထဲက position တိုင်းကို position တခြားတိုင်းကို ကြည့်ခွင့်ပေးပြီး၊ သက်ဆိုင်ရာ context ကို ဆွဲယူနိုင်စေပြီး context-mixed representation တစ်ခု ထုတ်ပေးတယ်။ ပြီးရင် feedforward network သေးလေးတစ်ခုက position တစ်ခုချင်းစီကို သီးခြားစီ လုပ်ဆောင်ပေးပြီး attention တစ်ခုတည်းမှာ မရှိတဲ့ representational capacity ကို ထပ်ဖြည့်ပေးတယ်။ sub-layer နှစ်ခုစလုံးကို x = LayerNorm(x + sublayer(x)) ပုံစံနဲ့ wrap ထားတယ်၊ x = sublayer(x) ရိုးရိုးလေးတော့ မဟုတ်ဘူး။ Naive အနေနဲ့ residual connection မပါဘဲ sub-layer တွေကို stack ချည်းလုပ်ရင် backpropagation အချိန် gradient တွေက sub-layer တစ်ခုစီရဲ့ transformation ထဲကနေပဲ ဖြတ်သွားရမှာဖြစ်ပြီး၊ block အများကြီးပါတဲ့ stack နက်နက်ရှိရင် gradient signal က layer တိုင်းမှာ ကျုံ့သွားတာ ဒါမှမဟုတ် ပေါက်ကွဲသွားတာ ဖြစ်တတ်တယ်။ x + sublayer(x) ထဲက ပေါင်းခြင်းက gradient တွေကို stack ဘယ်လောက်နက်နက် pre-layer တွေဆီ တိုက်ရိုက်၊ အနှောင့်အယှက်မရှိတဲ့ လမ်းကြောင်းတစ်ခု ပေးထားတယ် — ဒါကြောင့်ပဲ block ဆယ်ချက်ကျော် stack လုပ်ထားတဲ့ transformer တွေက train လုပ်လို့ရနေတာပါ။ LayerNorm ကလည်း ပေါင်းလိုက်တဲ့တန်ဖိုးတွေကို numeric range တည်ငြိမ်အောင် ထိန်းပေးတယ်၊ training drift မဖြစ်အောင်။ ဒီ attention/add-norm/feedforward/add-norm ပုံစံကို N ကြိမ်ထပ်လိုက်ရင် transformer encoder တစ်ခုလုံးရပါပြီ။

လက်တွေ့ scenario နဲ့ ချိတ်ကြည့်မယ်

ဒီ block အတိုင်းအတာနဲ့ တည်ဆောက်ထားတာက Tutorial Platform ရဲ့ တကယ်အသုံးဝင်တဲ့ feature တစ်ခု ဖြစ်တဲ့ learned lesson-recommendation model ရဲ့ building unit ပါပဲ — learner တစ်ဦးရဲ့ မကြာသေးမီက lesson history ကို sequence တစ်ခုအဖြစ် သတ်မှတ်ပြီး၊ ဘယ် lesson တွေက နောက်ဘာသင်သင့်တယ်ဆိုတာ ခန့်မှန်းရာမှာ အရေးအကြီးဆုံးလဲဆိုတာ self-attention ကနေ ချိန်ညှိပေးနိုင်တယ် — 'chapter တူတဲ့ lesson နောက်တစ်ခု' လိုမျိုး fixed rule ထက် ပိုအားကောင်းတယ်။ ဒီ block ကို အနည်းငယ် stack လုပ်ပြီး၊ နောက်ဆုံး sequence output ကို average/pool လုပ်ပြီး၊ available lesson တွေအပေါ် classifier head သေးလေးထဲ ကျွေးလိုက်ရင် toy shape-check exercise တစ်ခုတည်းမက real personalized recommendation engine ရဲ့ architectural core တစ်ခု ဖြစ်လာနိုင်ပါတယ်။

• Transformer block တစ်ခုက ပုံစံတစ်ခုတည်းကို အမြဲလိုက်နာတယ် — sub-layer, ပြီးရင် input ကို ပြန်ပေါင်း, ပြီးရင် normalize — block တစ်ခုစီအတွက် နှစ်ကြိမ်။
• Residual connection တွေက gradient တွေကို နောက်ကို တိုက်ရိုက်သွားနိုင်တဲ့ လမ်းကြောင်းပေးတယ်၊ ဒါကြောင့် block stack နက်နက်ကြီးတွေကို train လုပ်လို့ရနေတာ။
• Transformer block တစ်ခုက input shape ကို အတိအကျ ထိန်းထားတယ် — sequence length ဒါမှမဟုတ် embedding dimension မပြောင်းစေဘဲ representation ကိုသာ ပြောင်းလဲပေးတယ်။

အတူတူ စမ်းရေးကြည့်မယ်

import torch
import torch.nn as nn

class TransformerBlock(nn.Module):
    def __init__(self, embed_dim, num_heads, ff_dim):
        super().__init__()
        self.attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
        self.norm1 = nn.LayerNorm(embed_dim)
        self.ff = nn.Sequential(
            nn.Linear(embed_dim, ff_dim),
            nn.ReLU(),
            nn.Linear(ff_dim, embed_dim),
        )
        self.norm2 = nn.LayerNorm(embed_dim)

    def forward(self, x):
        attn_out, _ = self.attn(x, x, x)
        x = self.norm1(x + attn_out)
        ff_out = self.ff(x)
        x = self.norm2(x + ff_out)
        return x

batch_size, seq_len, embed_dim = 4, 10, 32
x = torch.randn(batch_size, seq_len, embed_dim)

block = TransformerBlock(embed_dim=embed_dim, num_heads=4, ff_dim=64)
output = block(x)

print("Input shape:", x.shape)
print("Output shape:", output.shape)

Input shape: torch.Size([4, 10, 32])
Output shape: torch.Size([4, 10, 32]) — output shape က input shape နဲ့ အတိအကျ တူညီနေတာကို တွေ့ရပါလိမ့်မယ်၊ transformer block တစ်ခုက sequence length နဲ့ embedding dimension ကို မပြောင်းလဲကြောင်း သက်သေပြသတယ်။

PyTorch Docs — Language Modeling with nn.Transformer — Deep Learning