Microsoft BEiT: BERT Pre-Training of Image Transformers (ไทย)
BEiT, what for?
ในเปเปอร์นี้ Microsoft เสนอ BEiT (Bidirectional Encoder representation from Image Transformers) ซึ่งเป็นการนำหลักการของ BERT pre-training ที่โด่งดังจากสาย NLP (Natural Language Processing) มาลองปรับใช้กับงานด้าน image/vision ดูบ้าง
จุดประสงค์ของงาน ก็เพื่อให้ได้ pre-trained BEiT model (เหมือนอย่างที่ NLP มี pre-trained BERT) ที่เราสามารถนำไป finetune ต่อยอดกับ image tasks อื่น ๆ บน image dataset อื่น ๆ ที่เราต้องการได้ง่ายขึ้น เรียกว่าเป็นตัวช่วย jump start ทำให้ภาระการ train ของเราน้อยลง ไม่ต้องไปเริ่มทุกอย่างจาก train-from-scratch นั่นเอง … พูดง่าย ๆ มันก็อารมณ์เดียวกับที่ตอนนี้เรามี pre-trained CNN (Convolutional Neural Network) ให้ใช้เป็นตัวเริ่มต้น (jump start) เยอะแยะมากมาย อาทิ VGG, ResNet, Inception, Xception, EfficientNet แต่งานนี้ BEiT มันไม่ใช่ CNN แต่เป็น transformer-based model ก็แค่นั้น
ก่อนอ่านต่อ แนะนำว่าถ้ามีพื้นจากเปเปอร์ต่อไปนี้มาบ้าง น่าจะช่วยให้เข้าใจง่ายขึ้น (จะว่าไปทั้ง 3 เปเปอร์นี้ของ Google หมดเลยแฮะ)
Google Transformer (NIPS2017)
… Google BERT (arXiv2018): อันนี้ค่อนข้างเกี่ยวข้องเยอะ โดยเฉพาะเรื่องของ Masked Language Modeling (MLM) ซึ่งเดี๋ยว BEiT จะขอยืมมาปรับใช้และเรียกชื่อใหม่ว่า Masked Image Modeling (MIM)
… Google ViT หรือที่รู้จักกันในชื่อ Vision Transformer (arXiv2020, ICLR2021 oral): เดี๋ยวมีไอเดียจาก ViT มาโผล่ใน BEiT ด้วยเช่นกัน เช่น เอาภาพมาตัดแบ่งเป็น non-overlapping patches > ทำ linear projection ให้ทุก patch
Challenges of transformer-based models for image/vision tasks
1 st
อันนี้เป็น challenge ที่นำเสนอมาตั้งแต่เปเปอร์ ViT เลยว่า
… ถึงแม้ pure transformer จะทำงานได้ผลในระดับ state-of-the-art (SOTA) บน image recognition benchmarks แต่กว่าจะทำได้ขนาดนี้ก็ต้องใช้ training images จำนวนมหาศาลกว่า 330 ล้านภาพและ train นานเว่อร์วังระดับ 2,500 TPUv3-core-days กันเลยทีเดียว
… หลังจากเปเปอร์ ViT ก็มีหลายเปเปอร์ที่พยายามนำเสนอวิธีบรรเทาปัญหา data hungry ของพวก vision-based transformer ซึ่งหนึ่งในเทคนิคเหล่านั้นก็คือการนำไอเดียของ self-supervised pre-training มาช่วย
… สรุปว่าในเปเปอร์ BEiT นี้ เดี๋ยวเค้าจะทำ self-supervised pre-training นี่ล่ะ
==
2nd
ความยากถัดมาคือจะเอา BERT-style pre-training ซึ่งถูกนำเสนอมาสำหรับ text data มาใช้กับ image data ได้ยังไงล่ะ?
… ในขณะที่ทาง NLP มีคนทำพวก language vocabulary (text) ดี ๆ ไว้ให้หยิบยืมมาใช้ได้อยู่แล้ว แต่ทาง image/vision มันไม่เคยมีใครทำ vocabulary for vision Transformer’s input unit (i.e., image patches) อะไรอย่างนี้เอาไว้ให้เลยนี่สิ
… ครั้นจะแก้ปัญหาทื่อ ๆ โดยมองเป็นปัญหา regression และให้โมเดลพยากรณ์ raw pixels ของแต่ละ masked image patches ออกมาเลย (pixel-level recovery task) เขาก็ว่ามันดูเป็นการ waste modeling capability ไปหน่อยนะ อุตส่าห์เสียเวลา pre-training เพื่อจะให้โมเดลมาเรียนรู้ short-range dependencies and high-frequency details ในระดับ pixel-level แบบนี้
… ในเปเปอร์ BEiT เขาแก้ปัญหานี้ด้วยการเสนอ Image Tokenizer ขึ้นมา ซึ่งตัวนี้มีหน้าที่คือ >> represent the input image as a sequence of discrete visual tokens … หรือก็คือ BEiT เลี่ยงการให้โมเดลต้องมาเรียนรู้ที่จะพยากรณ์ raw pixels ของแต่ละ masked image patch โดยเปลี่ยน representation ของแต่ละ image patch จากเดิมที่เป็น raw pixels ให้กลายมาเป็นตัวเลขหนึ่งตัวที่เรียกว่า discrete visual token แทน
BEiT: Before start pre-training
ก่อนจะไปเข้าเนื้อหาว่า pre-training BEiT เขาทำยังไง หัวข้อนี้ขอเคลียร์ประเด็นบางอย่างที่ควรรู้ก่อนเริ่ม pre-training … ทั้งนี้ จะอธิบายโดยอ้างอิงจากภาพใต้โพสต์ (figure 1) เป็นหลัก
1 Two views of image
งานนี้จะแบ่งภาพอินพุตออกเป็น 2 views (2 representations) ได้แก่
VIEW1: “image patch” ซึ่งจะทำหน้าที่เป็น input representation ในระหว่าง pre-training
… เช่น จาก figure 1 ด้านซ้ายสุด จะเห็นว่าภาพน้องหมา ถูกแบ่งออกเป็น 16 image patches (4x4 non-overlapping patches) ซึ่งเดี๋ยว patches เหล่านี้เมื่อผ่านขั้นตอนต่าง ๆ แล้วก็จะถูกนำไปเป็นอินพุตของ BEiT encoder ต่อไป
VIEW2: “visual tokens” ซึ่งจะทำหน้าที่เป็น output representation ในระหว่าง pre-training
… เช่น จาก figure 1 ด้านบน ๆ จะเห็นว่าภาพน้องหมาถูกนำไปผ่าน Image Tokenizer แปลงให้กลายมาเป็น 4x4 grid ที่แต่ละช่องมี 1 ตัวเลข (one discrete visual token)
… ทั้งนี้ สังเกตว่าขนาดของ a grid of discrete visual tokens ที่ว่าจะเท่ากันกับจำนวน image patches ใน view1 … หรือก็คือ ถ้ามี 4x4 patches ก็จะได้ a 4x4 grid of discrete visual tokens นั่นเอง
2 Train the image tokenizer
ในส่วนของ Image Tokenizer ที่ทำหน้าที่แปลงจากภาพอินพุตให้เป็น a grid of discrete visual tokens นั้น มันจะต้องถูก train รอเอาไว้ให้เรียบร้อยก่อนเริ่มกระบวนการ pre-training ของ BEiT
… โดยงาน BEiT เลือกใช้ Discrete Variational AutoEncoder (dVAE) มาเป็นตัว train ส่วน Image Tokenizer นี้ (ส่วนตัวยังไม่ได้อ่านเปเปอร์ dVAE เลย แต่ด้วยความเป็น VAE ก็คาดว่าการ train ตรงนี้ *น่าจะ* ทำได้แบบ unsupervised ตามสไตล์ของ AutoEncoder มั้งนะ)
BEiT: Pre-training (self-supervised learning)
เมื่อมี Image Tokenizer ที่ถูก train ไว้พร้อมใช้แล้ว ก็ได้เวลาเริ่มงานหลักของการใช้เทคนิค Masked Image Modeling (MIM) สำหรับ pre-train vision Transformers หรือ BEiT ในที่นี้ … ส่วนนี้จะอธิบายโดยอ้างอิงจากภาพใต้โพสต์ (figure 1) เป็นหลัก
1 จากภาพอินพุตของจริงขนาด 224x224 pixels ให้นำมาแบ่งเป็น patch ขนาด 16x16 pixels จะได้ทั้งหมด 14x14 patches
2 ทำ MIM โดยแทนที่จะสุ่มมั่วเพื่อเลือกว่าจะ mask ตัว image patch อันไหนดี งานนี้ใช้อัลกอริทึมเล็ก ๆ ชื่อ Blockwise Masking (ซึ่งก็ inspire จาก blockwise masking หรือ n-gram masking ใน NLP) ในการเลือกตำแหน่งของ patch ที่จะถูก mask … ทั้งนี้ งานนี้จะทำการ mask ประมาณ 40% ของจำนวน image patches ทั้งหมด
3 สำหรับ patch ที่รอดมาได้ (ไม่โดน mask) จะถูก flatten ให้กลายเป็น 1D vector และนำไปผ่านส่วนของ linear projection เพื่อให้ได้ patch embedding ออกมา … ส่วนพวก masked patches ก็กำหนดเป็นค่า [M] หรือค่า special mask embedding ไป
4 ในส่วน backbone transformer’s encoder งานนี้ใช้ตัวเดียวกับ ViT (หรือก็คือใช้เหมือน transformer 2017 ตัวต้นฉบับนั่นเอง) เพื่อให้สามารถเปรียบเทียบผลลัพธ์กันได้
… โดย patch embedding จะถูกนำไปผ่านกระบวนการตามปกติของ transformer’s encoder (เช่น เพิ่ม special token ด้านหน้า, ทำ positional embedding, …)
… ผลลัพธ์จากส่วน BEiT encoder ในภาพ figure 1 ก็คือ the encoded representations for all image patches
5 ผลลัพธ์จาก BEiT encoder จะถูกนำไปผ่าน softmax classifier เพื่อให้ทายค่า discrete visual token ของ patch นั้น ๆ ออกมา (BEiT กำหนดค่า vocabulary size ไว้ที่ 8,192) … โดยการดูว่าค่าที่ทายถูกต้องหรือไม่ ก็แค่เอาไปเทียบกับค่า discrete visual token ที่ได้จาก Image Tokenizer
ขั้นตอน pre-training BEiT นี้ใช้เวลา train บน ImageNet-1K (1.2 ล้านภาพ) dataset นานประมาณ 5 วัน ด้วย 16 Nvidia Tesla V100 32GB GPU cards … ทั้งนี้ ระหว่าง train ไม่มีการนำ ImageNet class labels มาใช้เลย เพราะงานทุกอย่างสามารถทำเสร็จได้หมดด้วยกลไกภายในของ BEiT อยู่แล้วตามหลักของ self-supervised learning
Fine-tuning the pre-trained BEiT on downstream tasks
เมื่อได้ pre-trained BEiT มาแล้ว ถัดไปคือการทดลองนำไป finetune กับ downstream tasks อื่น ๆ โดยในงานนี้เลือกทำ 2 งานคือ image classification task และ image (semantic) segmentation task
วิธีการ คือ เอา pre-trained BEiT มา เก็บเฉพาะส่วน BEiT encoder ไว้ ส่วน softmax classifier เอาทิ้งไปไม่ใช้แต่จะเสียบแทนด้วย task-specific layers ตาม task ที่เราต้องการทำแทน … เรียกง่าย ๆ ว่าส่วน pre-trained BEiT encoder ก็เหมือนส่วน feature extractor ใน pre-trained CNN นั่นล่ะ
… สำหรับ image classification task: ส่วนที่เสียบเพิ่มเป็น linear classifier ธรรมดาที่มีแค่ average pooling layer (สำหรับรวม representation ทั้งหมดที่ได้จาก BEiT encoder) ตามด้วย softmax classifier
… สำหรับ image (semantic) segmentation task: ใช้ layers ตามงานชื่อ SETR- PUP (เห็นว่าข้างในมี deconvolutional layers หลาย ๆ อันเพื่อขยายขนาดให้ได้ segmentation result ภาพใหญ่เท่าอินพุต)
ผลการทดลองบางส่วน
… Base-size BEiT achieves 83.2% top-1 accuracy on ImageNet-1K, significantly outperforming from-scratch DeiT training (81.8%; Touvron et al., 2020) with the same setup.
… Large-size BEIT obtains 86.3% only using ImageNet-1K, even outperforming ViT-L with supervised pre-training on ImageNet-22K (85.2%; Dosovitskiy et al., 2020).
อ้างอิงข้อมูลจาก
… เปเปอร์ arXiv 15JUN2021 (16 pages): https://arxiv.org/abs/2106.08254
… Github (PyTorch): https://github.com/microsoft/unilm/tree/master/beit
