Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer

Noam Shazeer; Azalia Mirhoseini; Krzysztof Maziarz; Andy Davis; Quoc Le; Geoffrey Hinton; Jeff Dean

งานวิจัย · arXiv 1701.06538

Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer

Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, Jeff Dean

เผยแพร่ 2017-01-23cs.LGcs.CLcs.NEstat.ML arXiv PDF

อ่านพร้อมคำแปลและ AI

สรุปสามบรรทัด

🎯 Dense neural networks ต้องเผชิญกับต้นทุนการคำนวณที่เพิ่มขึ้นแบบกำลังสอง เมื่อต้องขยายขนาดโมเดลเพื่อรองรับชุดข้อมูลขนาดใหญ่
🔧 นำเสนอเลเยอร์ Sparsely-Gated Mixture-of-Experts (MoE) ที่ใช้ Noisy Top-K gating พร้อมฟังก์ชันการสูญเสียเสริมเพื่อเกลี่ยภาระงานข้าม GPU shards แบบกระจายศูนย์
📈 ขยายสเกลโมเดลภาษาและโมเดลการแปลได้ถึง 1.37 แสนล้านพารามิเตอร์ และทำค่า LM perplexity ดีกว่าผลงานวิจัยที่ดีที่สุดเดิมถึง 24% ถึง 39% ด้วยต้นทุนการคำนวณที่ต่ำลง

บทคัดย่อ

ความสามารถของโครงข่ายประสาทเทียมในการซึมซับข้อมูลถูกจำกัดด้วยจำนวนพารามิเตอร์ แนวคิด conditional computation ซึ่งเปิดใช้งานบางส่วนของโครงข่ายตามอินพุตแต่ละตัวอย่าง ได้รับการเสนอในเชิงทฤษฎีเพื่อเพิ่มความจุของโมเดลอย่างมหาศาลโดยไม่ต้องเพิ่มการคำนวณในสัดส่วนเดียวกัน ทว่าในทางปฏิบัติกลับมีความท้าทายสำคัญทั้งด้านขั้นตอนวิธีและประสิทธิภาพการทำงาน ในงานวิจัยนี้ เราได้จัดการกับความท้าทายเหล่านี้และบรรลุผลสำเร็จของ conditional computation ในที่สุด โดยสามารถเพิ่มความจุของโมเดลได้มากกว่า 1,000 เท่า โดยสูญเสียประสิทธิภาพการคำนวณบนคลัสเตอร์ GPU สมัยใหม่ไปเพียงเล็กน้อยเท่านั้น เรานำเสนอเลเยอร์ Sparsely-Gated Mixture-of-Experts (MoE) ซึ่งประกอบด้วยเครือข่ายย่อยแบบ feed-forward มากถึงหลายพันเครือข่าย โดยมี gating network ที่สามารถฝึกฝนได้ทำหน้าที่เลือกกลุ่ม expert แบบเบาบาง (sparse) สำหรับแต่ละตัวอย่าง เราประยุกต์ใช้ MoE กับงาน language modeling และ machine translation ซึ่งความจุของโมเดลมีความสำคัญอย่างยิ่งต่อการซึมซับองค์ความรู้มหาศาลจากคลังข้อมูลฝึกฝน เรานำเสนอสถาปัตยกรรมโมเดลที่ใช้ MoE ขนาดพารามิเตอร์สูงสุด 1.37 แสนล้านตัว แทรกแบบคอนโวลูชันระหว่างชั้นของ LSTM ที่ซ้อนกันอยู่ ในชุดการทดสอบมาตรฐานขนาดใหญ่ของทั้ง language modeling และ machine translation โมเดลเหล่านี้ทำผลลัพธ์ได้ดีกว่าระดับ state-of-the-art อย่างมีนัยสำคัญโดยใช้ต้นทุนการคำนวณที่ต่ำกว่า

บทคัดย่อต้นฉบับ (ภาษาอังกฤษ)

The capacity of a neural network to absorb information is limited by its number of parameters. Conditional computation, where parts of the network are active on a per-example basis, has been proposed in theory as a way of dramatically increasing model capacity without a proportional increase in computation. In practice, however, there are significant algorithmic and performance challenges. In this work, we address these challenges and finally realize the promise of conditional computation, achieving greater than 1000x improvements in model capacity with only minor losses in computational efficiency on modern GPU clusters. We introduce a Sparsely-Gated Mixture-of-Experts layer (MoE), consisting of up to thousands of feed-forward sub-networks. A trainable gating network determines a sparse combination of these experts to use for each example. We apply the MoE to the tasks of language modeling and machine translation, where model capacity is critical for absorbing the vast quantities of knowledge available in the training corpora. We present model architectures in which a MoE with up to 137 billion parameters is applied convolutionally between stacked LSTM layers. On large language modeling and machine translation benchmarks, these models achieve significantly better results than state-of-the-art at lower computational cost.

คู่มือการอ่าน

ปัญหาและแรงจูงใจ

โมเดลการเรียนรู้เชิงลึก (Deep learning) ได้รับประโยชน์อย่างมหาศาลจากความจุของพารามิเตอร์เมื่อฝึกบนชุดข้อมูลขนาดใหญ่ อย่างไรก็ตาม สถาปัตยกรรมแบบมาตรฐานจะเปิดใช้งานโครงข่ายประสาททั้งหมดสำหรับทุกๆ โทเคนหรือตัวอย่างอินพุต ซึ่งส่งผลให้ต้นทุนการคำนวณพุ่งสูงเกินความเป็นจริงเมื่อทั้งขนาดโมเดลและปริมาณข้อมูลขยายตัวขึ้น

แม้ว่า conditional computation (การเลือกเปิดใช้งานเฉพาะส่วนย่อยของโมเดลตามแต่ละตัวอย่าง) จะได้รับการนำเสนอในเชิงทฤษฎี แต่การนำมาใช้งานจริงในยุคแรกต้องประสบปัญหาบนตัวเร่งความเร็วฮาร์ดแวร์สมัยใหม่เนื่องจาก:

GPU ถูกออกแบบมาให้ประมวลผลทางคณิตศาสตร์แบบสม่ำเสมอ (uniform) มากกว่าการแตกแขนงเงื่อนไข (branching)
"ปัญหา shrinking batch" ซึ่งการส่งตัวอย่างอินพุตไปยังเครือข่ายย่อยเฉพาะกลุ่ม ทำให้แต่ละเครือข่ายย่อยได้รับ batch ขนาดเล็กมากจนทำงานได้ไม่มีประสิทธิภาพ
คอขวดของแบนด์วิดท์เครือข่ายระหว่างอุปกรณ์ (inter-device network bandwidth)
การจัดสรรเส้นทางที่ไม่สมดุล (imbalanced routing) ซึ่งมีเครือข่ายย่อยเพียงไม่กี่ตัวที่ถูกใช้งานอย่างหนักในขณะที่ตัวอื่นๆ ไม่ได้ถูกใช้งาน

แนวคิดหลัก

Sparsely-Gated MoE Layer: เครือข่ายย่อยแบบ feed-forward ("experts") จะถูกวางไว้ระหว่าง recurrent layers (เช่น stacked LSTMs) โดยในแต่ละตำแหน่งจะมีการประเมินและคำนวณเฉพาะกลุ่ม experts ที่ได้รับการคัดเลือกแบบ sparse top- $k$ เท่านั้น
Noisy Top-K Gating: การเพิ่มสัญญาณรบกวนแบบ Gaussian (Gaussian noise) ที่สามารถปรับจูนได้เข้าไปยัง gating logits ก่อนเข้าสู่กระบวนการคัดเลือก top- $k$ และ softmax เพื่อสร้างการไล่ระดับ (gradients) ที่ต่อเนื่องและหาอนุพันธ์ย้อนกลับได้ รวมถึงกระตุ้นให้เกิดการสำรวจและกระจายการใช้งาน experts อย่างสมดุล
การผสมผสาน Data และ Model Parallelism: เลเยอร์มาตรฐานและ gating จะทำงานแบบขนานตามข้อมูล (data-parallel) ข้ามอุปกรณ์ต่างๆ ในขณะที่ experts จะถูกแบ่งขนานตามโมเดล (model-parallel) ข้าม GPU จากนั้นจึงนำแบตช์จากทุกอุปกรณ์มารวมกันเพื่อให้ experts แต่ละตัวได้รับ mini-batch รวมที่มีขนาดใหญ่เพียงพอ
Balancing Constraints: ใช้ฟังก์ชันการสูญเสียเสริม (auxiliary losses) กำกับความสำคัญ (importance) และภาระงาน (load) ของ expert อย่างชัดเจน เพื่อป้องกันไม่ให้ gating network ลู่เข้าสู่คำตอบแบบ degenerate ที่พึ่งพา expert เพียงไม่กี่ตัวมากเกินไป

ระเบียบวิธี (พร้อมสมการสำคัญ)

เลเยอร์ MoE ประกอบด้วยเครือข่าย expert แบบ feed-forward จำนวน $n$ เครือข่าย ได้แก่ $E_{1}, \dots, E_{n}$ และ gating network $G (x)$ โดยผลลัพธ์ของเลเยอร์คือ: $y = i = 1 \sum n G (x)_{i} E_{i} (x)$ การคำนวณจะไม่เกิดขึ้นเมื่อใดก็ตามที่ $G (x)_{i} = 0$

Noisy Top-K Gating

เพื่อผสานคุณสมบัติความเบาบาง (sparsity) เข้ากับ gating ที่สามารถหาอนุพันธ์ได้ จึงมีการแทรกสัญญาณรบกวนเข้าไปในการแปลงเชิงเส้น: $H (x)_{i} = (x \cdot W_{g})_{i} + StandardNormal () \cdot Softplus ((x \cdot W_{noise})_{i})$ $KeepTopK (v, k)_{i} = {v_{i} - \infty if v_{i} is in the top k elements of v otherwise$ $G (x) = Softmax (KeepTopK (H (x), k))$

Load-Balancing Losses

เพื่อบังคับให้ค่าน้ำหนัก gating รวมมีค่าเท่าๆ กัน จึงมีการนิยาม importance loss ตลอดทั้งแบตช์ $X$ โดยใช้สัมประสิทธิ์การแปรผัน ( $C V$ ): $Importance (X) = x \in X \sum G (x)$ $L_{importance} (X) = w_{importance} \cdot C V (Importance (X))^{2}$ และเพื่อรับประกันว่า expert แต่ละตัวจะได้รับจำนวนตัวอย่างที่เท่ากัน ตัวประมาณค่า smooth load จะใช้ฟังก์ชัน CDF ปกติ $Φ$ : $P (x, i) = Φ (\frac{( x \cdot W _{g} ) _{i} - kth_excluding ( H ( x ) , k , i )}{Softplus (( x \cdot W _{noise} ) _{i} )})$ $Load (X)_{i} = x \in X \sum P (x, i)$ $L_{load} (X) = w_{load} \cdot C V (Load (X))^{2}$

ผลการทดลอง (พร้อมตัวเลขที่สำคัญที่สุด)

1 Billion Word Benchmark: โมเดล MoE-4096-h ที่จำกัดงบประมาณการคำนวณให้เท่ากัน สามารถทำค่า test perplexity ได้ที่ 34.1 (ลดลง 24% เมื่อเทียบกับ baseline LSTM ที่ใช้การคำนวณเท่ากัน) ส่วนโมเดล MoE ที่ใช้งบประมาณสูงทำ test perplexity ได้ถึง 28.0 ภายใน 47 ชั่วโมงบน GPU K40 จำนวน 32 ตัว ซึ่งเหนือกว่าผลงาน state-of-the-art ก่อนหน้าซึ่งทำได้ 34.7 (และใช้เวลา 59 ชั่วโมงบน GPU K40 จำนวน 32 ตัว)
100 Billion Word Corpus: เมื่อขยายขนาดขึ้นเป็น 65,536 experts (6.8 หมื่นล้านพารามิเตอร์) ค่า test perplexity ลดลงถึง 39% เมื่อเทียบกับ dense baseline (28.9 เทียบกับ 47.0) โดยยังคงประสิทธิภาพการประมวลผลไว้ที่ 0.72 TFLOPS/GPU
Machine Translation: แบบจำลอง GNMT ที่เสริมด้วย MoE ทำคะแนนได้ 40.56 BLEU บนชุดทดสอบ WMT'14 En $\to$ Fr (+1.34 BLEU เหนือ GNMT ที่ไม่ใช้ RL) และ 26.03 BLEU บน WMT'14 En $\to$ De (+1.12 BLEU) ขณะที่ใช้จำนวนการคำนวณ (ops) ต่อ timestep ไม่ถึงครึ่งหนึ่ง
Multilingual Translation: โมเดล MoE แบบหลายภาษา (multilingual) เพียงโมเดลเดียว สามารถเอาชนะ baseline GNMT หลายภาษาได้ถึง 11 จาก 12 คู่ภาษา (เพิ่มขึ้นสูงสุดถึง +5.84 BLEU ในคู่ภาษาเกาหลี $\to$ อังกฤษ) และลด dev perplexity ลงได้ 19%

ข้อจำกัดและคำถามปลายเปิด

สภาวะความเบาบางขั้นขั้ว (Extreme sparsity) อาจนำไปสู่ภาวะผลตอบแทนลดน้อยถอยลง (diminishing returns) หรือประสิทธิภาพที่ถดถอยลง โดยค่า perplexity มีประสิทธิภาพแย่ลงเมื่อขยายจาก 65,536 experts เป็น 131,072 experts บนชุดข้อมูล 100B-word
ภาระโอเวอร์เฮดในการสื่อสารแบบกระจายศูนย์ยังคงขึ้นอยู่กับอัตราส่วนระหว่าง hidden ต่อ input/output ส่งผลให้ expert ที่มีขนาดเล็กมีประสิทธิภาพการใช้แบนด์วิดท์ต่ำกว่า
การนำ MoE ไปใช้แบบวนซ้ำ (recurrently) ภายในขั้นตอน RNN recurrence จะทำลายการปรับแต่ง batch-unrolling ข้ามตำแหน่งของลำดับข้อมูล ทำให้จำเป็นต้องใช้โครงร่างการคำนวณ activation ซ้ำใหม่ (activation recomputation)

คำศัพท์สำคัญ

Sparsely-Gated Mixture-of-Experts Layer (MoE): เลเยอร์ของโครงข่ายประสาทที่ประกอบด้วยเครือข่ายย่อย expert แบบขนานหลายตัว โดยมี gating network ที่สามารถฝึกฝนได้ทำหน้าที่เลือกกลุ่ม expert เพียงส่วนย่อยแบบเบาบางสำหรับอินพุตแต่ละตัวอย่าง
Conditional Computation: กระบวนทัศน์เชิงสถาปัตยกรรมที่ส่วนต่างๆ ของโครงข่ายประสาทจะถูกเปิดหรือปิดการทำงานแบบพลวัตสำหรับแต่ละตัวอย่าง เพื่อเพิ่มความจุของพารามิเตอร์โดยไม่ทำให้ต้นทุนการคำนวณเพิ่มขึ้นในสัดส่วนเดียวกัน
Noisy Top-K Gating: กลไก gating ที่เพิ่ม Gaussian noise แบบปรับค่าได้ลงใน logits ก่อนที่จะเก็บไว้เฉพาะค่า top-k และใช้ฟังก์ชัน softmax ช่วยให้ส่งผ่านการแพร่ค่าย้อนกลับ (backpropagation) ได้และป้องกันไม่ให้การกระจายเส้นทางลู่ไปที่ expert เพียงบางตัว
Shrinking Batch Problem: การลดลงของขนาดแบตช์จริงที่เครือข่ายย่อยแต่ละตัวได้รับ เมื่อตัวอย่างอินพุตถูกแบ่งส่งไปยัง experts จำนวนมากอย่างกระจัดกระจาย ซึ่งส่งผลให้ประสิทธิภาพการคำนวณทางคณิตศาสตร์ของฮาร์ดแวร์ลดลง
Importance Loss: ฟังก์ชันการสูญเสียเสริมที่ลงโทษค่ากำลังสองของสัมประสิทธิ์การแปรผันของผลรวมค่าน้ำหนัก gate ต่อ expert ตลอดทั้งแบตช์ เพื่อกระตุ้นให้เกิดการจัดสรรทรัพยากรไปยัง expert แต่ละตัวอย่างเท่าเทียมกัน
Load Loss: วัตถุประสงค์เสริมที่ลดความแปรปรวนของจำนวนโทเคนโดยประมาณที่ถูกส่งไปยัง experts ต่างๆ ตามน้ำหนักความน่าจะเป็น เพื่อสร้างสมดุลให้กับภาระการทำงานของฮาร์ดแวร์
Hierarchical Mixture of Experts: โครงสร้าง MoE แบบสองระดับ โดยมี gating network ระดับหลักคอยกระจายอินพุตไปยังกลุ่ม MoE ย่อยระดับรอง เพื่อลดค่า branching factor และลดโอเวอร์เฮดในการสื่อสารแบบกระจายศูนย์

คำถามที่พบบ่อย

Gating network สามารถทำการแพร่ค่าย้อนกลับ (backpropagation) ผ่านการคัดเลือก top-k ซึ่งเป็นแบบไม่ต่อเนื่องได้อย่างไร?

เนื่องจากการคัดเลือก top-k ถูกนำมาใช้ก่อนการคำนวณ softmax ค่า gate ที่ไม่เป็นศูนย์สำหรับ experts ที่ได้รับการคัดเลือกจึงยังคงมีอนุพันธ์แบบต่อเนื่องเทียบกับค่าน้ำหนักอินพุต ทำให้เกรเดียนต์สามารถไหลย้อนกลับผ่าน gating network ได้

ปัญหา shrinking batch ได้รับการแก้ไขอย่างไรในการทำงานบนอุปกรณ์แบบกระจายศูนย์?

โมเดลจะผสาน data parallelism และ model parallelism เข้าด้วยกัน โดยเลเยอร์มาตรฐานจะประมวลผล mini-batches อิสระบนแต่ละอุปกรณ์ แต่โทเคนจากทุกอุปกรณ์จะถูกนำมารวมเข้าด้วยกันก่อนที่จะถูกส่งต่อไปยัง experts ประจำตำแหน่งซึ่งกระจายอยู่บน GPU ต่างๆ

เหตุใดการใส่สัญญาณรบกวน (noise) จึงจำเป็นใน gating network?

Gaussian noise ที่ถูกแทรกเข้าไปช่วยส่งเสริมให้เกิดการสำรวจ experts ต่างๆ ในช่วงต้นของการฝึก และยังเป็นตัวประมาณค่าที่ราบเรียบสำหรับฟังก์ชัน load-balancing loss ซึ่งช่วยป้องกันไม่ให้เกิดปรากฏการณ์ที่โมเดลลู่เข้าสู่การเลือกใช้เฉพาะ expert เดิมซ้ำๆ

จะเกิดอะไรขึ้นเมื่อโมเดลถูกขยายสเกลไปจนถึง 131,072 experts?

ค่า perplexity มีประสิทธิภาพแย่ลงเล็กน้อยจาก 28.9 (เมื่อใช้ 65,536 experts) กลายเป็น 29.2 (เมื่อใช้ 131,072 experts) และประสิทธิภาพการคำนวณลดลงเหลือ 0.30 TFLOPS/GPU ซึ่งน่าจะมีสาเหตุมาจากความเบาบางขั้นขั้ว (extreme sparsity) และขนาดของแบตช์ที่ถูกกำหนดคงที่ไว้โดยไม่ได้ขยายตาม

คำอธิบายในหน้านี้เขียนโดย AI จากงานวิจัยและอาจมีข้อผิดพลาด โปรดอ่านต้นฉบับเพื่อดูรายละเอียด

อ่านพร้อมคำแปลและ AI