เทคนิคการเทรนโมเดล Text-to-Image ให้เร็ว แรง และแม่นยำ: บทเรียนจากการทดลอง

เทคนิคการเทรนโมเดล Text-to-Image ให้เร็ว แรง และแม่นยำ: บทเรียนจากการทดลองการสร้างโมเดล Text-to-Image ที่มีประสิทธิภาพสูงจากศูนย์นั้นเป็นเรื่องท้าทาย ในบทความตอนที่แล้ว เราได้พูดถึงการออกแบบสถาปัตยกรรมโมเดล PRX กันไปแล้ว วันนี้เราจะมาเจาะลึกเรื่อง "การเทรน" ซึ่งเป็นหัวใจสำคัญที่จะทำให้โมเดลของเราเรียนรู้ได้เร็วขึ้น มีเสถียรภาพ และสร้างภาพได้ตรงตามที่ต้องการมากขึ้นวงการ AI พัฒนาไปอย่างรวดเร็ว เทคนิคการเทรนใหม่ๆ เกิดขึ้นตลอดเวลา เพื่อให้เห็นภาพชัดเจน เราจึงขอสรุปเป็นเหมือน "สมุดบันทึกการทดลอง" โดยนำไอเดียที่น่าสนใจมาทดลองในสภาพแวดล้อมเดียวกัน เพื่อดูว่าเทคนิคไหนช่วยให้การ Optimization และ Convergence ดีขึ้นจริงบ้าง และที่สำคัญ เราไม่ได้ทดลองแค่ทีละเทคนิค แต่ยังดูด้วยว่าเมื่อนำมารวมกันแล้ว ยังมีประสิทธิภาพเหมือนเดิมหรือไม่จุดเริ่มต้น: การตั้งค่าพื้นฐาน (Baseline) ที่มั่นคงก่อนจะลองเทคนิคใหม่ๆ เราต้องมี "จุดอ้างอิง" ที่ชัดเจนเสียก่อน การตั้งค่าพื้นฐานของเราจึงเรียบง่าย เน้นใช้ส่วนประกอบมาตรฐาน ไม่มี Objective เสริมพิเศษ และไม่ใช้เทคนิคทางสถาปัตยกรรมที่ซับซ้อนเพื่อประหยัดพลังประมวลผลการตั้งค่าพื้นฐานนี้คือ Flow Matching แบบเพียวๆ (ตามที่อธิบายใน Part 1) โดยใช้โมเดล PRX-1.2B ขนาดเล็กที่เราเคยนำเสนอ โดยเทรนใน Flux VAE latent space และจะคงค่า Configuration นี้ไว้ในการทดลองเปรียบเทียบทั้งหมด (เว้นแต่จะระบุเป็นอย่างอื่น)การตั้งค่าพื้นฐานประกอบด้วย:โมเดล PRX-1.2B สถาปัตยกรรม Single Streamใช้ Global Attention สำหรับ Image Tokens และ Text Tokensเทรนใน Flux VAE latent spaceFlow Matching Loss เป็นหลักการมี "Baseline" ที่โปร่งใสและทำซ้ำได้เช่นนี้ ช่วยให้เราสามารถระบุได้ว่า การปรับปรุงหรือถดถอยที่เกิดขึ้นนั้น มาจากการเปลี่ยนแปลงเทคนิคใดเทคนิคหนึ่งโดยเฉพาะ ไม่ใช่เพราะการปรับ Hyperparameter หรือการเปลี่ยนแปลงการตั้งค่าที่ซ่อนอยู่ตัวชี้วัดประสิทธิภาพ: สังเกตการณ์ความก้าวหน้าเพื่อให้การทดลองมีหลักการ เราใช้ชุดตัวชี้วัดพื้นฐานไม่กี่อย่างในการติดตาม Checkpoint ตลอดการเทรน แม้จะไม่มีตัวชี้วัดใดวัดคุณภาพของภาพที่ได้สมบูรณ์แบบ แต่เมื่อรวมกันแล้วก็ช่วยให้เราเห็นภาพรวมของการพัฒนาได้เป็นอย่างดีFréchet Inception Distance (FID): วัดความใกล้เคียงของการกระจายตัวของภาพที่สร้างขึ้นกับภาพจริง โดยใช้สถิติจาก Inception-v3 ค่าที่ต่ำลงมักจะบ่งบอกถึงคุณภาพของภาพที่สูงขึ้นCLIP Maximum Mean Discrepancy (CMMD): วัดระยะห่างระหว่างการกระจายตัวของภาพจริงและภาพที่สร้างขึ้น โดยใช้ Image Embeddings จาก CLIP และ Maximum Mean Discrepancy (MMD) แตกต่างจาก FID ตรงที่ CMMD ไม่ได้ตั้งสมมติฐานว่าการกระจายตัวของ Feature เป็นแบบ Gaussian และอาจใช้ข้อมูลน้อยกว่าในการประเมิน ในทางปฏิบัติ มักจะสะท้อนคุณภาพที่รับรู้ได้ดีกว่า FID แม้จะยังคงเป็นเพียงการประมาณDINOv2 Maximum Mean Discrepancy (DINO-MMD): เป็นการวัดระยะทางแบบ MMD เช่นเดียวกับ CMMD แต่คำนวณจาก DINOv2 Image Embeddings แทน CLIP วิธีนี้ให้มุมมองที่แตกต่างในการประเมินการกระจายตัวโดยใช้ Vision Backbone แบบ Self-supervisedNetwork Throughput: วัดจำนวนตัวอย่างที่ประมวลผลได้ต่อวินาที (samples/s) เพื่อวัดประสิทธิภาพการเทรนโดยรวมกลุ่มเทคนิคที่ทดลอง: ปรับปรุงการเรียนรู้ให้ดีขึ้นเราแบ่งเทคนิคที่สำรวจออกเป็น 4 กลุ่มหลักๆ ได้แก่: Representation Alignment, Training Objectives, Token Routing and Sparsification, และ Data1. Representation Alignment: จัดแนวการเรียนรู้ให้ตรงจุดโมเดล Diffusion และ Flow มักจะเทรนด้วย Objective เดียว คือ การทำนายเป้าหมายที่เหมือน Noise จาก Input ที่ถูกรบกวน ในช่วงแรกของการเทรน Objective เดียวนี้ต้องทำหน้าที่สองอย่างพร้อมกัน คือ สร้าง Representation ภายในโมเดลให้มีประโยชน์ และเรียนรู้การ Denoise บนพื้นฐานนั้นRepresentation Alignment ทำให้เรื่องนี้ชัดเจนขึ้น โดยการรักษา Objective การ Denoise ไว้ และเพิ่ม Auxiliary Loss ที่ควบคุม Feature ระหว่างทาง (Intermediate Features) โดยตรง โดยใช้ Vision Encoder ที่แข็งแกร่งและถูก Freeze ไว้ เทคนิคนี้มักจะช่วยเร่งการเรียนรู้ในช่วงต้น และทำให้ Feature ของโมเดลใกล้เคียงกับ Vision Encoder สมัยใหม่มากขึ้น ส่งผลให้ใช้ Compute น้อยลงแต่ได้คุณภาพเท่าเดิมลองนึกภาพว่าเราสามารถแบ่ง Denoiser ออกเป็น Encoder แบบ Implicit ที่สร้าง Hidden States ระหว่างทาง และ Decoder ที่แปลง Hidden States เหล่านั้นไปสู่เป้าหมายการ Denoise ได้ ข้ออ้างคือ Representation Learning คือคอขวด: แม้ Diffusion และ Flow Transformers จะเรียนรู้ Feature ที่แยกแยะได้ แต่ก็ยังตามหลัง Vision Encoder ระดับ Foundation เมื่อการเทรนถูกจำกัดด้วย Compute การยืมใช้ Representation Space ที่ทรงพลังจึงช่วยให้ปัญหาการ Denoise ง่ายขึ้นREPA (Representation Learning via Patch Alignment)REPA เพิ่มเทอมการจับคู่ Representation (Representation Matching) เข้าไปบน Objective หลักของ Flow Matchingโดยพื้นฐานแล้ว โมเดลจะถูกเทรนบน State ที่ถูก Interpolate ระหว่างภาพจริง (x₀) และภาพ Noise (x₁) ที่เวลา t และทำนาย Vector Field (vθ(xt, t)) ใน REPA เราใช้ Pretrained Vision Encoder (f) ประมวลผลภาพจริง (x₀) เพื่อสร้าง Patch Embeddings (y₀) จากนั้น Denoiser จะประมวลผล xt และสร้าง Hidden Tokens (ht) (หนึ่ง Token ต่อหนึ่ง Patch)หัวใจของ REPA คือการใช้ Projection Head ขนาดเล็ก (hφ) เพื่อแปลง Student Hidden Tokens (ht) ให้อยู่ใน Teacher Embedding Space และใช้ Auxiliary Loss เพื่อเพิ่มความคล้ายคลึงกันระหว่าง Teacher และ Student Tokens ในแต่ละ Patch:LREPA(θ,φ) = -Ex₀,x₁,t [1/N * Σ sim(y₀,[n], hφ(ht,[n]))]โดย sim คือ Cosine Similarity และ N คือจำนวน Patch Tokensเทอมนี้จะถูกรวมเข้ากับ Flow Matching Loss หลัก:L = LFM + λ * LREPAโดย λ คือค่าที่ควบคุม Trade-offในทางปฏิบัติ Student จะถูกเทรนให้สร้าง Patch Representation ที่ทนต่อ Noise และสอดคล้องกับข้อมูลจาก xt ทำให้ Layer ที่อยู่ถัดไปสามารถโฟกัสกับการทำนาย Vector Field และสร้างรายละเอียดได้ โดยไม่ต้องเสียเวลาเรียนรู้โครงสร้างพื้นฐานใหม่ทั้งหมดผลการทดลอง:เราทดลองใช้ REPA ร่วมกับ Baseline PRX โดยใช้ DINOv2 และ DINOv3 ที่ Freeze เป็น Teacher พบว่ารูปแบบที่ได้ค่อนข้างสอดคล้องกัน: การเพิ่ม Alignment ช่วยเพิ่ม Quality Metrics และ Teacher ที่แข็งแกร่งกว่า (DINOv3) ให้ผลดีที่สุด แต่ก็แลกมาด้วยความเร็วที่ลดลงเล็กน้อยQuality Metrics: ทั้งสอง Teacher ให้ผลดีกว่า Baseline โดย DINOv3 ทำคะแนนได้ดีที่สุดThroughput: REPA มีค่าใช้จ่ายเพิ่มขึ้นจากการ Forward Pass ของ Teacher และ Loss ระดับ Patch ทำให้ Throughput ลดลงจาก 3.95 batches/s เป็น 3.66 (DINOv2) หรือ 3.46 (DINOv3) กล่าวคือ DINOv3 เลือกที่จะให้ความสำคัญกับคุณภาพ Representation สูงสุด โดยยอมแลกกับความเร็วที่ช้าลง ส่วน DINOv2 ให้ Trade-off ที่มีประสิทธิภาพมากกว่า โดยยังคงให้ประโยชน์ที่สำคัญโดยที่ความเร็วลดลงไม่มากนักข้อสรุปจากการใช้งานจริง:REPA เป็นเครื่องมือที่ทรงพลังสำหรับการเทรน Text-to-Image ในการตั้งค่าของเรา Trade-off ด้าน Throughput นั้นมีอยู่จริง และความเร็วสุทธิ (เวลาที่ใช้เพื่อให้ได้คุณภาพภาพระดับหนึ่ง) อาจไม่น่าทึ่งเท่าที่รายงานในเปเปอร์สำหรับ ImageNet-style generation อย่างไรก็ตาม การเพิ่มขึ้นของคุณภาพนั้นชัดเจนมาก และเมื่อดูจากภาพตัวอย่างที่สร้างขึ้น โมเดลที่เทรนด้วย Alignment มักจะสร้างโครงสร้างโดยรวมที่ชัดเจนและ Layout ที่สอดคล้องกันได้ดีตั้งแต่ช่วงต้นของการเทรน ทำให้เข้าใจได้ไม่ยากว่าทำไม REPA (และเทคนิค Alignment อื่นๆ) ถึงกลายเป็นส่วนประกอบสำคัญในสูตรการเทรน T2I สมัยใหม่iREPA (Spatial Structure Alignment)ต่อยอดจาก REPA คำถามคือ "เราควรจะ Align อะไร?" iREPA เสนอว่าคำตอบคือ โครงสร้างเชิงพื้นที่ (Spatial Structure) ไม่ใช่ความหมายโดยรวม (Global Semantics) จากการสำรวจ Vision Encoder จำนวนมาก ผู้เขียนพบว่าคุณภาพ "โดยรวม" แบบ ImageNet (เช่น Linear-probe accuracy บน Patch Tokens) มีความสัมพันธ์เพียงเล็กน้อยกับคุณภาพการสร้างภาพในภายหลังภายใต้ REPA แต่มาตรวัดความคล้ายคลึงกันเชิงพื้นที่ (Spatial Self-similarity) ของ Patch Tokens กลับสัมพันธ์กับ FID อย่างมากจากข้อสังเกตนี้ iREPA ได้ทำการปรับเปลี่ยนเล็กน้อยแต่ตรงจุด 2 อย่างในสูตร REPA เพื่อรักษาและถ่ายทอดข้อมูลเชิงพื้นที่ได้ดียิ่งขึ้น:เปลี่ยน Projection Head: แทนที่จะใช้ MLP ทั่วไป เปลี่ยนไปใช้ Convolution 3x3 ขนาดเล็กที่ทำงานบน Grid ของ PatchSpatial Normalization: ใช้การ Normalization เชิงพื้นที่กับ Teacher Patch Tokens เพื่อลบ Overlay โดยรวม (Mean ของแต่ละ Spatial Location) ซึ่งช่วยเพิ่ม Local Contrastแม้จะใช้ "โค้ดน้อยกว่า 4 บรรทัด" การปรับเปลี่ยนเหล่านี้ช่วยเร่ง Convergence และปรับปรุงคุณภาพได้อย่างสม่ำเสมอใน Vision Encoder, ขนาดโมเดล และแม้กระทั่งสูตรการเทรนที่ใกล้เคียงกับ REPAผลการทดลอง:ในการตั้งค่าของเรา เราสังเกตเห็นการ Boost ที่คล้ายกันเมื่อใช้ iREPA ร่วมกับ DINOv2 Convergence ดูราบรื่นขึ้น และ Metrics ก็ปรับปรุงอย่างต่อเนื่องในช่วง 100K Steps แรกน่าสนใจว่าการเปลี่ยนแปลงเดียวกันนี้ไม่ได้ส่งผลดีเท่าที่ควรเมื่อใช้กับ DINOv3 และกลับทำให้ประสิทธิภาพแย่ลง เราไม่ต้องการตีความผลนี้มากเกินไป เพราะอาจเกิดจากปฏิสัมพันธ์กับสถาปัตยกรรมเฉพาะของเรา, Resolution/Patching, Loss Weighting หรือแม้แต่รายละเอียดการ Implement เล็กๆ น้อยๆ แต่ถึงกระนั้น จากความไม่สอดคล้องกันนี้ เราอาจจะยังไม่รวมเทคนิคเหล่านี้ในสูตรเริ่มต้นของเรา แม้ว่าจะเป็นตัวเลือกที่น่าสนใจสำหรับการปรับแต่งเฉพาะทางก็ตามข้อควรระวังในการใช้ REPA ตลอดการเทรนเปเปอร์ "REPA Works Until It Doesn't: Early-Stopped, Holistic Alignment Supercharges Diffusion Training" ชี้ให้เห็นข้อจำกัดที่สำคัญ: REPA เป็นตัวเร่งที่มีประสิทธิภาพในช่วงต้น แต่ก็อาจจะถึงจุดที่ประสิทธิภาพคงที่ (Plateau) หรือแม้กระทั่งกลายเป็นตัวถ่วงในช่วงท้ายของการเทรน ผู้เขียนอธิบายถึง "ความไม่สอดคล้องกันของ Capacity" (Capacity Mismatch) เมื่อโมเดลสร้างภาพเริ่มปรับตัวเข้ากับการกระจายตัวของข้อมูลทั้งหมด (โดยเฉพาะรายละเอียดความถี่สูง) การบังคับให้โมเดลอยู่ใกล้กับ Manifold ของ Recognition Encoder ที่ Freeze ไว้ซึ่งมีมิติที่ต่ำกว่า อาจกลายเป็นข้อจำกัดได้ข้อสรุปเชิงปฏิบัติของผู้เขียนคือ: ใช้ Alignment ในช่วง "Burn-in" (ช่วงเริ่มต้น) จากนั้นจึงค่อยๆ ลดหรือหยุดการใช้งานลงการทำความเข้าใจเทคนิคเหล่านี้ จะช่วยให้เราสามารถปรับปรุงกระบวนการเทรนโมเดล Text-to-Image ให้มีประสิทธิภาพสูงสุดได้ ในบทความตอนต่อไป เราจะนำเสนอสูตรการเทรนฉบับเต็มพร้อมโค้ด และรายงานผลการทดสอบ "Speedrun" สาธารณะ ที่จะรวมเอาเทคนิคที่ดีที่สุดเข้าด้วยกันเพื่อทดสอบขีดจำกัดของ Pipeline การเทรนของเรา#AI #TextToImage #MachineLearning #DeepLearning #FlowMatchinghttps://huggingface.co/blog/Photoroom/prx-part2