Vision Transformers

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

Deep Learning course က learnထားပြီးသားအတိုင်း transformer တွေဟာ token sequence တစ်ခု (token တစ်ခုစီက vector တစ်ခုစီ) ပေါ်မှာ အလုပ်လုပ်ပြီး၊ token တစ်ခုစီက token တခြားတိုင်းကို attend လုပ်ကာ contextual representation တွေ တည်ဆောက်ပါတယ်။ Text မှာဆိုရင် word ဒါမှမဟုတ် subword တွေကို အစဉ်လိုက် tokenize လုပ်ရုံပါပဲ ရှင်းရှင်းလင်းလင်းရှိပါတယ်။ Image တွေမှာတော့ inherent sequence structure မရှိပါဘူး— pixel တွေဟာ 2D grid ပေါ်မှာ စီစဉ်ထားပြီး၊ pixel တစ်ခုချင်းစီကို token အဖြစ် သတ်မှတ်လိုက်ရင် element ထောင်ချီ၊ သောင်းချီရှိတဲ့ sequence တွေ ဖြစ်လာမှာပါ၊ self-attention (sequence length ရဲ့ quadratic ဖြစ်တဲ့) ကို computationally မဖြစ်နိုင်တော့ပါဘူး။ Vision Transformer ရဲ့ အဓိက idea ကတော့ pixel level မှာ tokenize မလုပ်ဘဲ patch level မှာ tokenize လုပ်ခြင်းပါပဲ— ပုံတစ်ပုံကို 16x16 pixel square grid လိုမျိုး ခွဲပြီး square တစ်ခုချင်းစီကို token တစ်ခုအဖြစ် သတ်မှတ်ပါတယ်။ 224x224 ပုံတစ်ပုံဟာ 14x14 patch grid — token 196 ခု ဖြစ်လာပြီး၊ standard transformer encoder တစ်ခုအတွက် လုံးဝ ကိုင်တွယ်နိုင်တဲ့ sequence length ဖြစ်ပါတယ်။

ချက်ချာသော implementation trick ကတော့ 'patch တွေအဖြစ် ဖြတ်ပြီး patch တစ်ခုချင်းစီကို embedding vector တစ်ခုအဖြစ် linear project လုပ်ခြင်း' ဆိုတာဟာ kernel_size ကို patch size နဲ့ တူညီအောင်၊ stride ကို ကိုက်ညီအောင် သတ်မှတ်ထားတဲ့ convolution တစ်ခုက operation တစ်ခုတည်းနဲ့ အတိအကျ လုပ်ဆောင်နိုင်ခြင်းပါပဲ။ kernel_size=16, stride=16 ရှိတဲ့ Conv2d တစ်ခုဟာ 16x16 window ကို ပုံပေါ်မှာ overlap မရှိဘဲ လျှောသွားပြီး window position တစ်ခုချင်းစီက embed_dim length ရှိတဲ့ output vector တစ်ခု ထုတ်ပေးပါတယ်— ဒီ vector က patch ရဲ့ token embedding ဖြစ်ပြီး convolution ရဲ့ learned weight တွေက linear projection ဖြစ်ပါတယ်။ ရလာတဲ့ (batch, embed_dim, grid_h, grid_w) feature map ကို (batch, num_patches, embed_dim) အဖြစ် reshape လုပ်လိုက်ရင် transformer encoder တစ်ခု လိုချင်နေတဲ့ token sequence အတိအကျ ရရှိသွားပါပြီ၊ text model တွေကနေ သင်ခဲ့ပြီးသား multi-head self-attention နဲ့ feed-forward block တွေကို တိုက်ရိုက် အသုံးပြုနိုင်ပါတယ်— plain sequence တစ်ခုက ဆုံးရှုံးသွားနိုင်တဲ့ 2D spatial information ကို ထိန်းသိမ်းဖို့ positional embedding တစ်ခု ထပ်ထည့်ရုံပါပဲ။

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

Tutorial Platform မှာ lesson screenshot တွေကို content type အလိုက် (code editor, terminal output, architecture diagram, browser UI) auto-tag လုပ်ဖို့ ViT-style patch-embedding အဆင့်ကို သုံးနိုင်ပါတယ်— screenshot တစ်ခုစီကို 16x16 patch sequence အဖြစ် ခွဲပြီး lightweight transformer classifier တစ်ခုကို ပို့လိုက်ခြင်းအားဖြင့် terminal ရဲ့ monospace grid၊ diagram ရဲ့ ကျဲသောနေရာလွတ်နဲ့ ဆက်စပ်လိုင်းများ လိုမျိုး layout pattern တွေကို global CNN feature vector တစ်ခုတည်းက ဝေဝါးစွာ ရောစပ်ပစ်တတ်တာမျိုး မဖြစ်ဘဲ ဖမ်းယူနိုင်ပါတယ်။

• Raw pixel တွေပေါ်မှာ self-attention လုပ်ဖို့ computationally မဖြစ်နိုင်တဲ့အတွက် ViT က ပုံတွေကို patch level မှာ tokenize လုပ်ပြီး sequence length ကို သိသိသာသာ လျှော့ချပါတယ်။
• kernel_size နဲ့ stride ကို patch size အတိုင်း သတ်မှတ်ထားတဲ့ Conv2d တစ်ခုတည်းက 'patch အဖြစ်ဖြတ်ခြင်း + linear projection' ကို operation တစ်ခုတည်းနဲ့ လုပ်ဆောင်ပါတယ်။
• ရလာတဲ့ (batch, num_patches, embed_dim) tensor ဟာ standard transformer encoder တစ်ခုက text embedding တွေလိုပဲ တိုက်ရိုက် သုံးနိုင်တဲ့ token sequence တစ်ခု ဖြစ်ပါတယ်။

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

import torch
import torch.nn as nn

patch_size = 16
embed_dim = 768
image_size = 224

patch_embed = nn.Conv2d(
    in_channels=3,
    out_channels=embed_dim,
    kernel_size=patch_size,
    stride=patch_size,
)

images = torch.randn(2, 3, image_size, image_size)  # fake batch of 2 RGB images
patches = patch_embed(images)                # (batch, embed_dim, H/patch, W/patch)
print(patches.shape)

tokens = patches.flatten(2).transpose(1, 2)  # (batch, num_patches, embed_dim)
print(tokens.shape)

torch.Size([2, 768, 14, 14])
torch.Size([2, 196, 768]) — convolution ဟာ 16x16 pixel patch တစ်ခုစီကို 768-dimensional vector တစ်ခုအဖြစ် ပြောင်းလဲပေးပြီး၊ 14x14 patch grid ကို ထုတ်ပေးပါတယ် (224/16=14 ဖြစ်လို့ပါ)၊ ပြီးရင် flatten လုပ်ပြီး transpose လုပ်ကာ token 196 ခု (14*14) ရှိတဲ့ sequence တစ်ခု၊ dimension 768 စီ ရှိတဲ့ token sequence ဖြစ်လာပါတယ်— transformer encoder တစ်ခု လိုချင်နေတဲ့ (batch, num_patches, embed_dim) shape အတိအကျပါပဲ။

Wikipedia — Vision transformer — Computer Vision