โมเดล AI ที่ทรงพลังที่สุดซึ่งฝึกบนแล็ปท็อปได้ใน 5 นาทีคืออะไร?

• ฝึกโมเดลทรานส์ฟอร์เมอร์สไตล์ GPT ขนาดประมาณ 1.8M พารามิเตอร์ บน MacBook Pro ภายใน 5 นาที โดยใช้โทเค็น TinyStories ราว 20M โทเค็น และทำได้ประมาณ 9.6 perplexity
• ข้อจำกัดหลักของการ ฝึกภายใน 5 นาที คือขนาดโมเดลและจำนวนโทเค็นที่ประมวลผลได้ โดยเมื่อโมเดลใหญ่ขึ้นจะคอนเวิร์จช้าลงและได้ประโยชน์จากข้อมูลจำนวนน้อยลดลง
• ในด้าน การเพิ่มประสิทธิภาพ การใช้ MPS ให้ผลดีที่สุด ขณะที่ การเลือกโมเดลขนาดเล็ก มีประสิทธิภาพกว่าการใช้คอมไพล์/ควอนไทซ์/gradient accumulation หรือการใช้ทางเลือกแทน PyTorch
• ชุดข้อมูลที่เรียบง่ายและสม่ำเสมออย่าง TinyStories ส่งผลเชิงบวกต่อประสิทธิภาพของโมเดลขนาดเล็กมากกว่าข้อมูลแนวสารานุกรม
• สถาปัตยกรรม ทรานส์ฟอร์เมอร์ ให้ผลดีกว่า LSTM หรือ diffusion ภายใต้เงื่อนไขของขนาดเล็กและเวลาฝึกสั้น

---

ภาพรวม

บทความนี้นำเสนอผลการทดลองหา โมเดลภาษา AI ที่ให้ประสิทธิภาพสูงสุด ซึ่งสามารถฝึกบนแล็ปท็อป (MacBook Pro) ได้ภายใน 5 นาที พร้อมอินไซต์เกี่ยวกับกลยุทธ์การเทรนที่เหมาะสม การเลือกชุดข้อมูล และสถาปัตยกรรมของโมเดล

สรุปผลการทดลอง

• ฝึกโมเดลทรานส์ฟอร์เมอร์สไตล์ GPT ขนาดประมาณ 1.8M พารามิเตอร์ ด้วยข้อมูล TinyStories ราว 20M และทำสถิติ 9.6 perplexity
• ตัวอย่างข้อความที่สร้างได้ยังสั้น แต่คงรูปแบบของเรื่องเล่าที่สอดคล้องกัน และไวยากรณ์ภาษาอังกฤษโดยรวมยังถูกต้อง
• เน้นย้ำว่าผลลัพธ์ของโมเดลที่ใช้งานได้จริงภายใน 5 นาที นั้นดีกว่าที่คาดไว้

ที่มาของการทดลองและข้อจำกัด

• การทดลองนี้เริ่มจากความสงสัยเชิงเล่น ๆ ว่าใน สภาพแวดล้อมแบบแล็ปท็อป จะฝึกโมเดลที่ทรงพลังได้เร็วแค่ไหน ซึ่งในทางปฏิบัติอาจไม่สมจริงนัก
• ตามความเป็นจริงสามารถฝึกโมเดลที่ทรงพลังกว่าได้บน GPU ประสิทธิภาพสูงบนคลาวด์ (เช่น H100) แต่เงื่อนไขจำกัดของการทดลองนี้คือเวลา 5 นาที
• เมื่อ ขนาดโมเดลใหญ่ขึ้น ความเร็วในการประมวลผลโทเค็นจะลดลง ทำให้ยากที่จะได้ผลลัพธ์ที่ดีภายใน 5 นาที
  • โมเดลที่เล็กเกินไป (เช่น 10K พารามิเตอร์) ไม่สามารถเรียนรู้ความซับซ้อนได้เพียงพอ
  • ช่วงที่ใช้งานได้จริงอยู่ที่ราว 1M~2M พารามิเตอร์

การเพิ่ม throughput

• การใช้ MPS (Metal Performance Shaders ของ Apple) ได้ผลดีที่สุด
• การเพิ่มประสิทธิภาพเชิงคำนวณหลากหลายแบบ เช่น torch.compile, float16, MLX ให้ผลดีน้อยกว่าที่คาด หรือบางกรณีกลับทำให้ประสิทธิภาพลดลง
• gradient accumulation มีไว้เพื่อจัดการหน่วยความจำ แต่ในทางปฏิบัติกลับทำให้ความเร็วลดลงอย่างมาก
• โมเดลต้องสามารถอัปเดต weight ได้อย่างรวดเร็วภายในหน่วยความจำภายใน จึงจะมีประสิทธิภาพ

การเลือกชุดข้อมูล

• เมื่อใช้จำนวนโทเค็นที่จำกัด (ราว 10~20M) แล้วเริ่มจากข้อมูลวิกิภาษาอังกฤษแบบเรียบง่าย เช่น Simple English Wikipedia พบว่าสามารถรักษาความสม่ำเสมอทางไวยากรณ์ได้ แต่ขาดความสอดคล้องด้านความหมาย
  • เน้นคำนามเฉพาะและการเรียงข้อเท็จจริงที่ดูเหมือนฝืน ทำให้มีข้อจำกัดในการสร้างเนื้อหาที่มีความหมายจริง
• เมื่อใช้ชุดข้อมูล TinyStories ผลลัพธ์มีความสอดคล้องและมีความหมายมากขึ้นอย่างชัดเจน เพราะโครงสร้างเรื่องชัดเจนและภาษาง่าย
  • เป็นเรื่องเล่าระดับเด็ก 4 ขวบ จึงเหมาะให้โมเดลขนาดเล็กเรียนรู้ได้ดี

โทเค็นไนเซอร์และการทำโทเค็น

• การเทรนโทเค็นไนเซอร์ไม่ถูกนับรวมในเวลา 5 นาที และด้วยขนาดข้อมูลที่เล็ก จึงแทบไม่มีความจำเป็นต้องปรับแต่งมาก
• การเรียนรู้โทเค็นแบบหลายไบต์ทำให้โมเดลเรียนรู้ง่ายกว่า

การทดลองสถาปัตยกรรมโมเดล

• ใช้สถาปัตยกรรม ทรานส์ฟอร์เมอร์ (สไตล์ GPT-2)

  • ปรับไฮเปอร์พารามิเตอร์ เช่น 2~3 เลเยอร์, ฟังก์ชันกระตุ้นแบบ SwiGLU, positional embedding
  • LSTM ให้ประสิทธิภาพใกล้เคียงกัน แต่ทรานส์ฟอร์เมอร์ดีกว่าในแง่ perplexity
  • Dropout, mixture-of-experts เป็นต้น ไม่มีประสิทธิภาพนักเมื่อโมเดลมีขนาดเล็ก
  • Curriculum learning แทบไม่เกิดผล เพราะเวลาฝึกสั้นเกินไป

• ทดลองใช้ โมเดล diffusion (D3PM)

  • เนื่องจากภาษาธรรมชาติเป็นโทเค็นแบบไม่ต่อเนื่อง กระบวนการ diffusion จึงล้มเหลวและสร้างได้เพียงโทเค็นสุ่มที่ไร้ความหมาย
  • การสร้างโครงสร้างประโยคอย่างรวดเร็วทำได้ยากกว่าเมื่อเทียบกับทรานส์ฟอร์เมอร์หรือ LSTM

ความสัมพันธ์ระหว่างขนาดโมเดลกับ throughput โทเค็น/วินาที

• โมเดลขนาด 1M~2M พารามิเตอร์คือ sweet spot ที่เหมาะที่สุด
  • ถ้าใหญ่เกินไปจะคอนเวิร์จไม่ทันภายใน 5 นาที ถ้าเล็กเกินไปก็จะชนเพดานประสิทธิภาพทันทีหลังเริ่มฝึก
• Chinchilla scaling law สอดคล้องกับผลการทดลองโดยรวมพอสมควร
  • ขนาดโมเดลที่เหมาะสมคือจำนวนโทเค็นฝึกทั้งหมด/20 และการทดลองนี้ก็ยืนยันแนวโน้มนั้น

บทสรุปและนัยสำคัญ

• แม้จะมีเวลาสั้นมากและใช้ฮาร์ดแวร์ขนาดเล็ก ก็ยัง สามารถ ฝึกโมเดลสำหรับการเล่าเรื่องที่มีความสอดคล้องได้
• การฝึก 5 นาทีไม่เหมาะกับการพัฒนาโมเดลทรงพลัง แต่มีความหมายต่อการออกแบบโมเดลขนาดเล็กมาก น้ำหนักเบามาก และการทดลองด้านการปรับแต่งฮาร์ดแวร์กับสถาปัตยกรรม
• ในอนาคต หาก GPU บนแล็ปท็อปและโครงสร้างโมเดลพัฒนาขึ้น ก็มีความเป็นไปได้ที่ประสิทธิภาพของโมเดลที่ฝึกได้ภายในเวลาเพียงไม่กี่นาทีจะดีขึ้นอีก