Mixture-of-Depths: เทคนิคจัดสรรทรัพยากรคอมพิวต์แบบไดนามิกใน Transformer

 (arxiv.org)
2 คะแนน โดย GN⁺ 2024-04-08 | 1 ความคิดเห็น | แชร์ทาง WhatsApp
  • Mixture-of-Depths (MoD) ของ Google DeepMind เป็นวิธีที่ทำให้โมเดลภาษาแบบ Transformer ไม่ใช้ FLOPs เท่ากันกับทุกโทเคน แต่ให้เฉพาะโทเคนที่จำเป็นในแต่ละเลเยอร์เข้าร่วมการคำนวณ attention และ MLP
  • เราเตอร์รายเลเยอร์จะสร้างค่าน้ำหนักสเกลาร์สำหรับแต่ละโทเคน และมีเพียงโทเคนที่อยู่ภายใน ความจุ top-k ที่กำหนดไว้ล่วงหน้าเท่านั้นที่ผ่านบล็อก ส่วนที่เหลือจะอ้อมผ่านด้วย residual connection
  • การตรึงค่า k ไว้ล่วงหน้าช่วยให้กราฟคำนวณและขนาดเทนเซอร์ยังคงเป็นแบบสแตติก แต่การกระจายการคำนวณต่อโทเคนยังเปลี่ยนได้แบบไดนามิกตามบริบท
  • ในการทดลอง MoD ให้ loss ต่ำกว่า Transformer พื้นฐานภายใต้เกณฑ์ isoFLOP เมื่อใช้เวลาฝึกเท่ากัน หรือสามารถลด FLOPs ต่อ forward pass ที่สมรรถนะเท่ากัน ทำให้ step ของการฝึกและการ sampling เร็วขึ้น
  • การตั้งค่าที่ดีที่สุดคือการใช้ บล็อกความจุ 12.5% กับบล็อกสลับกัน (every other block) และในการ sampling แบบ auto-regressive ใช้การ routing แบบอิง predictor เพื่อลดปัญหา non-causality ของ top-k

การสิ้นเปลืองการคำนวณที่ MoD มุ่งแก้

  • โมเดลภาษา Transformer ทั่วไปใช้ปริมาณการคำนวณเท่ากันกับทุกโทเคนใน forward pass
  • MoD ตั้งอยู่บนสมมติฐานว่าโทเคนและลำดับทั้งหมดไม่ได้ต้องการเวลาและความพยายามเท่ากันในการทำนาย จึงลดงบคำนวณรวมและจัดสรรการคำนวณไปยังโทเคนที่จำเป็น
  • conditional computation แบบเดิมอาจสร้างกราฟคำนวณแบบไดนามิก ซึ่งอาจไม่สอดคล้องกับ กราฟคำนวณแบบสแตติก ที่ฮาร์ดแวร์ปัจจุบันนิยม
  • เมื่อผู้ใช้กำหนดงบคำนวณรวมก่อนการฝึก โมเดลจะเรียนรู้ว่าจะใช้การคำนวณกับโทเคนและเลเยอร์ใดภายในงบนั้น
    • สามารถคาดการณ์ปริมาณการคำนวณ การใช้หน่วยความจำ และ FLOPs ต่อ forward pass ได้ล่วงหน้า
    • โทเคนที่ถูก routing จะเปลี่ยนไปตามบริบทอินพุต

โครงสร้าง routing และวิธีลดการคำนวณ

  • MoD ใช้เราเตอร์เหมือน MoE Transformer แต่แทนที่จะเลือก expert หนึ่งตัวจากหลายตัว จะส่งโทเคนไปหนึ่งในสองเส้นทาง
    • การคำนวณบล็อกมาตรฐาน ที่รวม attention และ MLP
    • residual connection ที่ส่งค่าต่อไปตามเดิม
  • residual connection มีต้นทุนการคำนวณต่ำ และเอาต์พุตของบล็อกถูกกำหนดจากค่าอินพุต
  • หากตั้ง capacity ของบล็อกให้น้อยกว่าจำนวนโทเคนทั้งหมด T จะลด FLOPs ต่อ forward pass เมื่อเทียบกับ Transformer พื้นฐาน
    • ตัวอย่างเช่น หากลด capacity เหลือ T/2 การคูณเมทริกซ์ query-key ของ attention จะเป็น (T/2)² แทน T² ทำให้ FLOPs เหลือราว 25%
  • หากตั้ง capacity เป็น 0 จะเร็วมาก แต่ไม่ได้ใช้พารามิเตอร์ส่วนใหญ่ของ Transformer ทำให้สมรรถนะแย่ลง และหากตั้ง capacity เป็น T ก็จะเท่ากับ Transformer พื้นฐาน
  • เป้าหมายคือการหาจุดกึ่งกลางระหว่างสุดโต่งทั้งสองที่ให้สมรรถนะเท่ากับหรือดีกว่า Transformer พื้นฐาน พร้อมเพิ่มความเร็วของ step ได้

การ routing แบบ top-k expert-choice

  • ในแต่ละเลเยอร์ l เราเตอร์จะสร้างค่าน้ำหนักสเกลาร์ให้กับ embedding ของแต่ละโทเคน
  • ตาม capacity C ที่ผู้ใช้กำหนด มีเพียงโทเคนที่อยู่ใน top-k ของโทเคนในลำดับเท่านั้นที่จะเข้าร่วมการคำนวณ attention และ MLP
  • โทเคนที่ไม่ถูกเลือกจะข้ามการคำนวณของบล็อกและถูกส่งต่อไปยังเลเยอร์ถัดไปผ่าน residual connection
  • วิธีนี้อิงกับ expert-choice routing
    • token-choice routing อาจเกิดปัญหา load balancing เพราะโทเคนเป็นฝ่ายเลือกเส้นทาง
    • expert-choice routing ทำให้แต่ละเส้นทางเลือกโทเคน top-k จึงเติม capacity ได้พอดี
  • ใน MoD เส้นทางมีโดยพื้นฐานสองแบบคือ “การคำนวณบล็อก” และ “การอ้อมผ่าน” จึงแบ่งโทเคนเป็นสองชุดได้ด้วยการทำ top-k เพียงครั้งเดียว
  • ค่าน้ำหนักจากเราเตอร์จะถูกคูณกับเอาต์พุตของบล็อก ทำให้อยู่บน gradient path และได้รับอิทธิพลจาก gradient descent ระหว่างการฝึก language modeling

ปัญหา non-causality ที่เกิดในการ sampling

  • การทำ top-k ของ expert-choice routing เป็นแบบ non-causal
    • การตัดสินว่าโทเคนใดเข้า top-k ต้องรู้ค่าน้ำหนักเราเตอร์ของโทเคนถัดไปด้วย
    • ในการ sampling แบบ auto-regressive ไม่สามารถใช้ข้อมูลโทเคนอนาคตได้
  • เพื่อหลีกเลี่ยงปัญหานี้ มีการทดลองสองวิธี
    • auxiliary loss แบบ binary cross-entropy เพื่อให้ sigmoid ของเอาต์พุตเราเตอร์ทำนายว่าจะถูกเลือกใน top-k หรือไม่
    • MLP predictor เสริมขนาดเล็กที่รับอินพุตเดียวกับเราเตอร์ แต่ใช้ stop-gradient
  • วิธีแรกส่งผลต่อ objective หลักของ language modeling ประมาณ 0.2~0.3%
  • วิธีที่สองซึ่งอิง predictor ไม่ส่งผลต่อ objective ของ language modeling และไม่กระทบ step speed อย่างมีนัยสำคัญ
  • งานทำนายเสริมนี้ในการทดลองไปถึง ความแม่นยำ 99% ได้อย่างรวดเร็ว

ผลการทดลองและข้อสังเกต

  • การทดลองช่วงแรกทำโดยค้นหา hyperparameter ของ MoD ภายใต้งบ 6e18 FLOPs รวม
  • การตั้งค่า MoD ที่เหมาะที่สุดคือการใช้ routing ทุกบล็อกเว้นบล็อก และตั้ง top-k เป็น 256
    • เมื่อความยาวลำดับเป็น 2048 จะมี 256 โทเคน หรือเพียง 12.5% ที่ผ่าน attention และ MLP
    • โทเคนที่เหลือ 1792 ตัว หรือ 87.5% จะอ้อมผ่านบล็อก
  • ตัวแปร MoD ขนาด 220M พารามิเตอร์ให้สมรรถนะดีกว่า baseline ที่เหมาะที่สุดตาม isoFLOP เล็กน้อย ขณะเดียวกัน step ระหว่างการฝึกเร็วขึ้นกว่า 60% และ wall-clock time รวมของการฝึกบนฮาร์ดแวร์เดียวกันแทบเท่าเดิม
  • ในการวิเคราะห์ isoFLOP ที่ 6e18, 2e19, 1e20 FLOPs โมเดล MoD ที่ FLOP-optimal ก็มีแนวโน้มมีพารามิเตอร์มากกว่า baseline
  • มีตัวแปร MoD หลายแบบที่ให้สมรรถนะเท่าเดิมหรือ loss ต่ำกว่า พร้อมใช้ FLOPs ต่อ forward pass น้อยกว่า
  • ความเร็ว step ที่เพิ่มขึ้นมาจากสองปัจจัย
    • โทเคนบางส่วนอ้อมผ่านบล็อก จึงลด FLOPs per parameter ที่ขนาดโมเดลเท่าเดิม
    • สามารถเลือกตัวแปร MoD ที่เล็กกว่าซึ่งให้สมรรถนะเท่ากับ baseline ได้
  • MoD ที่ใช้ stochastic routing มีสมรรถนะต่ำกว่า Transformer พื้นฐานและ MoD ปกติอย่างมาก แสดงให้เห็นว่า routing ที่เรียนรู้มา มีความสำคัญ
  • ในการประเมินแบบ auto-regressive ใช้ held-out data จำนวน 256,000 ลำดับ รวม 500M โทเคน และเมื่อเปลี่ยนเป็น routing แบบอิง predictor การลดลงของสมรรถนะมีเพียงเล็กน้อย
  • MoD ยังสามารถรวมกับ MoE เป็น Mixture-of-Depths-and-Experts (MoDE) ได้
    • staged MoDE จะตัดสินใจก่อน attention ว่าจะส่งโทเคนเข้าบล็อกหรือให้อ้อมผ่าน
    • integrated MoDE จะรวม routing ของ MoD โดยเพิ่ม expert แบบ “no-op” เข้าไปในกลุ่ม MLP expert เดิม
  • การวิเคราะห์ routing พบรูปแบบว่าโทเคนบางตัวผ่านหลายบล็อกบ่อยครั้ง ส่วนโทเคนอื่นจะอ้อมผ่านเมื่อทำได้
  • มีผลวิเคราะห์เบื้องต้นว่าโทเคนที่ผ่านบล็อกบ่อยสัมพันธ์กับกรณีที่ entropy ของการทำนายเอาต์พุตสูงกว่า ซึ่งเชื่อมโยงกับความเป็นไปได้ว่าเป็นโทเคนที่ทำนายได้ยากกว่า

บทความที่เกี่ยวข้อง

1 ความคิดเห็น

 
GN⁺ 2024-04-08
ความคิดเห็นบน Hacker News

  • การกำหนดเส้นทาง (routing) ที่ซับซ้อนขึ้นน่าจะพบเห็นได้มากขึ้นแน่นอน
    โดยเฉพาะอย่างยิ่ง สักวันหนึ่งน่าจะไปสู่ การกำหนดเส้นทางแบบเรียกซ้ำ (recursive routing) ที่พาให้ผ่านชุดผู้เชี่ยวชาญซ้ำอีกครั้ง ต่อไป “กระบวนการคิด (chain-of-thought)” น่าจะเกิดขึ้นแบบเรียกซ้ำภายในโมเดล

    • วัตถุสมมุติแบบนี้น่าจะเรียกว่า Recursive Neural Networks ได้
    • สิ่งที่อธิบายตรงนี้ดูคล้ายงานวิจัยตระกูล Universal Transformers อยู่บ้าง คือให้ input embedding ผ่านบล็อกทรานส์ฟอร์เมอร์เดียวกันหลายครั้ง แล้วมีโมดูลแยกมาตัดสินว่า embedding นั้น “สุก” พอหรือยังจึงดึงออกมา
      ส่วนที่ใกล้กับแนวคิด “ผู้เชี่ยวชาญ” มากกว่าคือ论文 Sparse Universal Transformers เมื่อปีที่แล้ว ซึ่งผสาน Universal Transformer กับ mixture of sparse experts โดยให้กลไก gating ตัดสินว่าจะใช้บล็อกทรานส์ฟอร์เมอร์ใด ตามลำดับใด
      นี่ไม่ใช่สาขาเชี่ยวชาญของผม แต่เท่าที่เข้าใจคือฝึกให้ดีได้ยาก และหากต้องการผลลัพธ์ใกล้เคียงกับทรานส์ฟอร์เมอร์ทั่วไป ตอน inference จะต้องใช้ปริมาณการคำนวณรวมมากกว่า ถึงอย่างนั้นก็เป็นทิศทางที่น่าสนใจ และผมมองว่าการมีเพดานจำนวนขั้นตอนการคำนวณต่อโทเคนเป็นหนึ่งในข้อเสียใหญ่ของสถาปัตยกรรมทรานส์ฟอร์เมอร์แบบคลาสสิก
    • เหตุผลที่เรื่องนี้ยังไม่เกิดขึ้น ผมคิดว่าเพราะตอนฝึกยังไม่มีวิธีตัดสินว่า ควรเรียกซ้ำกี่ครั้ง
      ถ้าเลือกจำนวนครั้งแบบสุ่ม หรือทดลองหลายระดับความลึกของการเรียกซ้ำ เอาต์พุตจะ “เบลอ” กล่าวคือจะไม่รู้ว่าเอาต์พุตของชั้นใดควรปล่อยข้อมูลที่สำคัญต่อผลลัพธ์สุดท้าย หรือควรปล่อยอินพุตที่ดีที่สุดสำหรับรอบเรียกซ้ำถัดไป
    • โดยพื้นฐานแล้ว attention ก็คือ routing และวิธี routing แบบอื่น ๆ เหล่านี้ให้ตัวเลือกที่หยาบขึ้นกับโมเดล ซึ่งอาจทำให้การเรียนรู้ง่ายขึ้น
    • กระแสไปทาง dynamic routing เพิ่มขึ้นแน่นอน แต่ผมคิดว่า MoE/MoD/MoDE ไม่ได้ทำให้ reasoning ลึกขึ้นมากนัก หากแต่ใกล้เคียงกับการช่วยบรรจุข้อเท็จจริงเพิ่มเติมไว้ใน weights โดยลดการซ้อนทับกันมากกว่า
      reasoning ที่ลึกขึ้นน่าจะมาจาก dynamic ระดับโทเคนมากกว่าระดับชั้น ตัวอย่างเช่น论文 Quiet-STaR ล่าสุดที่ให้โมเดลส่งออก rationale token ที่ถูกทิ้งไป: https://arxiv.org/abs/2403.09629

  • อาจเป็น 论文ที่สำคัญที่สุดของปี 2024
    แนวคิดที่ว่าเราอยากได้โมเดลที่ไม่ต้องใช้การคำนวณเท่ากันกับทุกโทเคนมีมานานแล้ว แต่เพิ่งเห็นกลไกที่น่าเชื่อถือสำหรับเรื่องนี้เป็นครั้งแรก

    Equipped with these new methods, we can sample autoregressively by choosing to route tokens to or around a block based on the router’s output, which does not depend on any information from future tokens. We provide empirical evidence that this is a relatively easy auxiliary task that quickly achieves 99% accuracy.
    อันนี้น่าทึ่งหน่อยไหม?

    • Sparse Universal Transformer เก่ากว่านั้น และทำ early exit แบบอิง routing อยู่แล้ว
    • ถึงขั้นสำคัญที่สุดเลยหรือ? การที่ไม่ใช่ทุกโทเคนจะต้องใช้ context window ทั้งหมด น่าจะเป็น obvious optimization อยู่แล้ว

  • เวอร์ชันอธิบายเบื้องต้นแบบง่าย ๆ:
    ลองจินตนาการว่ามีผู้ช่วยอัจฉริยะที่เข้าใจและประมวลผลคำพูดได้ ปกติผู้ช่วยคนนี้จะให้ความสนใจกับทุกคำเท่ากัน ไม่ว่าคำนั้นจะสำคัญต่อความหมายทั้งหมดแค่ไหน
    ทีนี้สมมติว่าเราพบวิธีสอนให้ผู้ช่วยใช้ “ทรัพยากรสมอง” อย่างฉลาดขึ้น แทนที่จะให้ความสนใจทุกคำเท่ากัน ก็ให้โฟกัสมากขึ้นกับคำที่สำคัญที่สุดต่อการเข้าใจความหมาย และปรับโฟกัสนั้นได้ทันทีตามบริบท
    เพื่อไม่ให้ผู้ช่วยโอเวอร์โหลด เรายังกำหนดเพดานรวมของ “ทรัพยากรสมอง” ที่ใช้ได้ในช่วงเวลาหนึ่งด้วย เหมือนให้ budget แล้วบอกว่า “ใช้ทรัพยากรกับคำได้แค่จำนวนหนึ่งในแต่ละครั้ง” ดังนั้นผู้ช่วยต้องตัดสินใจว่าคำใดสำคัญที่สุด
    แม้มีเพดานนี้ ผู้ช่วยก็ยังปรับการใช้ทรัพยากรได้อย่างยืดหยุ่น ใช้กับบางคำมากขึ้น ใช้กับคำอื่นน้อยลง และจัดลำดับความสำคัญให้เหมาะกับสถานการณ์
    เมื่อฝึกแบบนี้ ผู้ช่วยจะทำงานได้อย่างมีประสิทธิภาพพร้อมกับให้ความสนใจอย่างชาญฉลาด เข้าใจได้ดีพอ ๆ กับผู้ช่วยที่ให้ความสนใจกับทุกคำเท่ากัน แต่ใช้ทรัพยากรสมองโดยรวมน้อยกว่า จึงตอบสนองและประมวลผลข้อมูลใหม่ได้เร็วขึ้นมาก

    • เข้าใจว่านี่คือ ELI5 แต่สิ่งที่อธิบายมานี่ไม่ใช่สิ่งที่ attention ทำอยู่แล้วหรือ? เพราะมันโฟกัสเฉพาะคำที่สำคัญตามบริบทในลำดับก่อนหน้าอยู่แล้ว

  • ผมลองสรุปไว้ตรงนี้เล็กน้อยตามความเข้าใจของตัวเอง:
    https://lifeinthesingularity.com/p/googles-breakthroughs-in-...

    • บทความดีมาก ชวนให้นึกถึงสไตล์ New Scientist คือ “อธิบายไปเรื่อย ๆ แต่คงความสั้นไว้” เหมาะสำหรับจับภาพรวมของวงการนี้

  • คล้ายกับ mixture of experts มาก เพียงแต่แทนที่จะ route โทเคนไปยังผู้เชี่ยวชาญหลายราย กลับเป็นการ “จัดเข้า expert เดียวที่ข้ามแบบไดนามิกได้”

    • ถ้าผสมสองอย่างนี้เข้าด้วยกันน่าจะเจ๋งทีเดียว อาจลด ปริมาณการคำนวณของ MoE ลงได้อีกโดยยังคงประสิทธิภาพไว้

  • “This is more computationally efficient than performing a full content-based lookup across an entire memory buffer for each step in the future, and could be one step towards drastically increasing the context-length available for making a prediction.”
    นี่คือวิธีที่จะทำให้ context window 10 ล้านโทเคน เป็นไปได้หรือเปล่า? หรือหมายถึง context window ที่ยาวขึ้นในอนาคต?

  • ข้อสรุปที่ได้หลังจากพยายามทำความเข้าใจและ implement อัลกอริทึม RASP บางตัว คือ ฟังก์ชันบางอย่างต้องใช้ ชั้นทรานส์ฟอร์เมอร์ จำนวนหนึ่งจึงจะทำงานได้
    หากตามตรรกะนี้ จะเห็นว่าฟังก์ชันที่ทรานส์ฟอร์เมอร์เรียนรู้อาจกระจายอยู่ข้ามหลาย head การทำฟังก์ชันเหล่านี้ซ้ำอาจมีคุณค่ามากต่อการเข้าใจและแก้โจทย์ แต่แนวทาง inference ปัจจุบันไม่สามารถรันกลุ่ม head ที่ต่อเนื่องกันซ้ำได้论文นี้ดูเหมือนเป็นทิศทางที่มีแนวโน้มจริง ๆ
    [1] https://arxiv.org/pdf/2106.06981.pdf
    [2] https://www.youtube.com/watch?v=t5LjgczaS80

  • ข้อเสียเดียวของการที่ LLM พัฒนาเร็วเกินไป อาจเป็นตรงที่论文ออกมาเร็วกว่าความเร็วที่ใครก็ตาม อย่างน้อยก็นอก Google จะเรียนรู้และทดสอบ improvement เหล่านั้นได้
    ตอนผมเริ่มทำ deep learning นั้น ReLU กับ dropout กำลังมาแรง และสามารถแก้โค้ดหนึ่งสองบรรทัดบน 1080 ระดับผู้บริโภค แล้วทดสอบได้ภายในไม่กี่ชั่วโมงว่าดีขึ้นไหม ตอนนี้คงต้องรอหลายสัปดาห์จนกว่าที่อย่าง Mistral จะลองให้ดู

    • ยินดีต้อนรับสู่ ชนชั้นยากจนด้าน GPU
      ผมโฟกัสแนวทาง quantization และทดสอบด้วย GPU รุ่นเก่าจากเจเนอเรชันก่อน ๆ

  • hu-po รีวิว论文 AI แบบไลฟ์สตรีมเจาะลึก
    ขอแนะนำอย่างยิ่ง และในนี้พูดถึง论文 mixture-of-depths ที่กำลังคุยกันอยู่: https://www.youtube.com/watch?v=Teru_qIdB8Y