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

Bài báo · arXiv 1701.06538

Mạng nơ-ron cực lớn: Lớp Mixture-of-Experts định tuyến thưa (Sparsely-Gated Mixture-of-Experts Layer)

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

Xuất bản 2017-01-23cs.LGcs.CLcs.NEstat.ML arXiv PDF

Đọc với bản dịch và AI

Tóm tắt trong ba dòng

🎯 Các mạng nơ-ron dày đặc (dense) gặp tình trạng bùng nổ chi phí bậc hai khi mở rộng dung lượng mô hình để tiếp nhận các tập dữ liệu khổng lồ.
🔧 Đề xuất lớp Sparsely-Gated Mixture-of-Experts (MoE) sử dụng cơ chế Noisy Top-K gating và các hàm mất mát cân bằng tải phụ trợ trên các phân vùng GPU phân tán.
📈 Mở rộng các mô hình ngôn ngữ và dịch máy lên tới 137 tỷ tham số, cải thiện perplexity mô hình ngôn ngữ tốt nhất từng công bố từ 24% đến 39% với chi phí tính toán thấp hơn.

Tóm tắt

Khả năng tiếp nhận thông tin của một mạng nơ-ron bị giới hạn bởi số lượng tham số của nó. Tính toán có điều kiện (conditional computation), trong đó các phần của mạng chỉ kích hoạt trên từng mẫu dữ liệu cụ thể, đã được đề xuất trên lý thuyết như một phương pháp tăng mạnh dung lượng mô hình mà không làm tăng chi phí tính toán tương ứng. Tuy nhiên, trên thực tế vẫn tồn tại những thách thức lớn về mặt thuật toán và hiệu năng. Trong công trình này, chúng tôi giải quyết các thách thức đó và hiện thực hóa triển vọng của tính toán có điều kiện, đạt mức cải thiện hơn 1000 lần về dung lượng mô hình trong khi chỉ giảm nhẹ hiệu quả tính toán trên các cụm GPU hiện đại. Chúng tôi giới thiệu lớp Sparsely-Gated Mixture-of-Experts (MoE), bao gồm hàng nghìn mạng con truyền thẳng (feed-forward sub-networks). Một mạng cổng (gating network) có thể huấn luyện sẽ xác định tổ hợp thưa các chuyên gia (experts) cần dùng cho mỗi mẫu dữ liệu. Chúng tôi áp dụng MoE vào các tác vụ mô hình hóa ngôn ngữ và dịch máy, nơi dung lượng mô hình đóng vai trò cốt lõi để tiếp thu lượng tri thức khổng lồ từ các kho ngữ liệu huấn luyện. Chúng tôi trình bày các kiến trúc mô hình trong đó MoE với quy mô lên tới 137 tỷ tham số được áp dụng dạng tích chập xen giữa các lớp LSTM xếp chồng. Trên các bài kiểm thử quy mô lớn về mô hình hóa ngôn ngữ và dịch máy, những mô hình này đạt kết quả vượt trội rõ rệt so với các phương pháp tiên tiến nhất (SOTA) với chi phí tính toán thấp hơn.

Tóm tắt gốc (tiếng Anh)

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.

Hướng dẫn đọc

Vấn đề và động lực

Các mô hình học sâu hưởng lợi rất lớn từ dung lượng tham số khi được huấn luyện trên các tập dữ liệu lớn. Tuy nhiên, các kiến trúc tiêu chuẩn lại kích hoạt toàn bộ mạng cho mỗi token hoặc mẫu đầu vào. Điều này dẫn đến sự gia tăng phi thực tế về chi phí tính toán khi cả kích thước mạng lẫn khối lượng dữ liệu đều mở rộng.

Mặc dù tính toán có điều kiện (chỉ kích hoạt các tập con của mô hình cho mỗi mẫu) đã được đề xuất về mặt lý thuyết, các triển khai trước đây gặp nhiều khó khăn trên các bộ tăng tốc phần cứng hiện đại do:

GPU ưu tiên các phép tính số học đồng nhất hơn là rẽ nhánh.
"Vấn đề suy giảm kích thước batch" (shrinking batch problem), khi việc định tuyến các tập con của mẫu dữ liệu khiến các mạng con riêng lẻ nhận kích thước batch rất nhỏ, làm giảm hiệu năng.
Điểm nghẽn băng thông mạng giữa các thiết bị.
Định tuyến mất cân bằng, trong đó một vài mạng con bị quá tải trong khi các mạng con khác không nhận được dữ liệu.

Các ý tưởng cốt lõi

Lớp MoE định tuyến thưa (Sparsely-Gated MoE Layer): Các mạng con truyền thẳng ("các chuyên gia" - experts) được đặt giữa các lớp hồi quy (chẳng hạn như các lớp LSTM xếp chồng). Chỉ một nhóm top- $k$ thưa các chuyên gia được chọn để tính toán cho mỗi vị trí.
Noisy Top-K Gating: Nhiễu Gauss có thể điều chỉnh được thêm vào các logit của mạng cổng trước khi chọn top- $k$ và áp dụng softmax nhằm cung cấp gradient liên tục, có thể lan truyền ngược và khuyến khích việc khám phá cân bằng giữa các chuyên gia.
Kết hợp song song dữ liệu và song song mô hình: Các lớp tiêu chuẩn và mạng cổng chạy song song dữ liệu trên các thiết bị, trong khi các chuyên gia được phân vùng theo kiểu song song mô hình trên các GPU. Các batch trên tất cả các thiết bị được gộp lại để các chuyên gia nhận được các mini-batch tổng hợp đủ lớn.
Các ràng buộc cân bằng tải: Các hàm mất mát phụ trợ tường minh về mức độ quan trọng (importance) và tải (load) của chuyên gia giúp ngăn mạng cổng suy biến thành các giải pháp thoái hóa chỉ sử dụng lặp đi lặp lại một số ít chuyên gia.

Phương pháp (kèm các phương trình chính)

Một lớp MoE bao gồm $n$ mạng chuyên gia truyền thẳng $E_{1}, \dots, E_{n}$ và một mạng cổng $G (x)$ . Đầu ra của lớp là: $y = i = 1 \sum n G (x)_{i} E_{i} (x)$ Chi phí tính toán được triệt tiêu bất cứ khi nào $G (x)_{i} = 0$ .

Noisy Top-K Gating

Để kết hợp tính thưa với cơ chế cổng khả vi, nhiễu được đưa vào các phép chiếu tuyến tính: $H (x)_{i} = (x \cdot W_{g})_{i} + StandardNormal () \cdot Softplus ((x \cdot W_{noise})_{i})$ $KeepTopK (v, k)_{i} = {v_{i} - \infty n \overset{ˊ}{\overset{e}{ˆ}} u v_{i} n \overset{ˋ}{\overset{a}{˘}} m trong top k ph \overset{ˋ}{\overset{a}{ˆ}} n tử lớn nh \overset{ˊ}{\overset{a}{ˆ}} t của v ng ư ợc lại$ $G (x) = Softmax (KeepTopK (H (x), k))$

Các hàm mất mát cân bằng tải (Load-Balancing Losses)

Để áp đặt tổng trọng số cổng đồng đều, hàm mất mát tầm quan trọng (importance loss) được định nghĩa trên một batch $X$ bằng hệ số biến thiên ( $C V$ ): $Importance (X) = x \in X \sum G (x)$ $L_{importance} (X) = w_{importance} \cdot C V (Importance (X))^{2}$ Để đảm bảo số lượng mẫu đồng đều giữa các chuyên gia, ước lượng tải trơn sử dụng hàm phân phối tích lũy chuẩn (normal 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}$

Kết quả (kèm các số liệu quan trọng nhất)

1 Billion Word Benchmark: Mô hình MoE-4096-h có cùng mức chi phí tính toán đạt perplexity kiểm thử là 34.1 (giảm 24% so với các mô hình LSTM baseline tương đương về tính toán). Mô hình MoE với ngân sách tính toán cao đạt perplexity 28.0 sau 47 giờ trên 32 GPU K40, vượt qua kết quả SOTA trước đó là 34.7 (vốn mất 59 giờ trên 32 GPU K40).
Kho ngữ liệu 100 Billion Word: Mở rộng quy mô lên 65,536 chuyên gia (68 tỷ tham số) đã giảm 39% perplexity kiểm thử so với baseline dạng dày đặc (28.9 so với 47.0), trong khi vẫn duy trì 0.72 TFLOPS/GPU.
Dịch máy: GNMT tích hợp MoE đạt 40.56 BLEU trên WMT'14 En $\to$ Fr (+1.34 BLEU so với GNMT không dùng RL) và 26.03 BLEU trên WMT'14 En $\to$ De (+1.12 BLEU) trong khi yêu cầu số phép tính trên mỗi bước thời gian ít hơn một nửa.
Dịch đa ngữ: Một mô hình MoE đa ngữ duy nhất đã vượt trội hơn baseline GNMT đa ngữ trên 11/12 cặp ngôn ngữ (lên tới +5.84 BLEU ở cặp Hàn $\to$ Anh) và giảm perplexity trên tập dev 19%.

Hạn chế và các câu hỏi mở

Tính thưa cực hạn có thể dẫn đến hiệu suất giảm dần hoặc suy giảm độ chính xác; perplexity bị giảm sút khi tăng từ 65,536 lên 131,072 chuyên gia trên tập dữ liệu 100 tỷ từ.
Chi phí truyền thông phân tán vẫn phụ thuộc vào tỷ lệ giữa số chiều ẩn và đầu vào/đầu ra, khiến các dung lượng chuyên gia nhỏ kém hiệu quả về băng thông.
Việc áp dụng MoE đệ quy bên trong các bước hồi quy của RNN phá vỡ tối ưu hóa giải cuộn batch (batch-unrolling) qua các vị trí trong chuỗi, đòi hỏi các cơ chế tính toán lại kích hoạt (activation recomputation).

Thuật ngữ chính

Sparsely-Gated Mixture-of-Experts Layer (MoE): Một lớp mạng nơ-ron chứa nhiều mạng con chuyên gia song song, trong đó một mạng cổng có thể huấn luyện sẽ chọn một tập hợp con thưa các chuyên gia cho mỗi mẫu đầu vào.
Conditional Computation: Một mô hình kiến trúc trong đó các phần của mạng nơ-ron được kích hoạt hoặc vô hiệu hóa động theo từng mẫu dữ liệu nhằm tăng dung lượng tham số mà không làm tăng tuyến tính khối lượng tính toán.
Noisy Top-K Gating: Một cơ chế chọn cổng bổ sung nhiễu Gauss có thể điều chỉnh vào các logit trước khi giữ lại k giá trị lớn nhất và áp dụng softmax, cho phép lan truyền ngược và ngăn chặn sự sụp đổ định tuyến.
Shrinking Batch Problem: Sự suy giảm kích thước batch hiệu dụng mà từng mạng con nhận được khi dữ liệu đầu vào bị phân vùng thưa thớt qua nhiều chuyên gia, làm giảm hiệu quả tính toán số học của phần cứng.
Importance Loss: Một hàm mất mát phụ trợ phạt bình phương hệ số biến thiên của tổng giá trị cổng trên mỗi chuyên gia trong một batch nhằm khuyến khích phân bổ đồng đều giữa các chuyên gia.
Load Loss: Một mục tiêu phụ trợ giảm thiểu phương sai trong số lượng token được phân bổ theo xác suất ước tính giữa các chuyên gia để cân bằng tải phần cứng.
Hierarchical Mixture of Experts: Cấu trúc MoE hai cấp, trong đó mạng cổng chính định tuyến dữ liệu đầu vào đến các cụm MoE thứ cấp nhằm giảm hệ số phân nhánh và chi phí truyền thông phân tán.

Câu hỏi thường gặp

Mạng cổng cho phép lan truyền ngược qua phép định tuyến top-k rời rạc như thế nào?

Vì phép chọn top-k được áp dụng trước softmax, các giá trị cổng khác không của các chuyên gia được chọn vẫn giữ đạo hàm liên tục theo trọng số đầu vào, cho phép gradient truyền ngược qua mạng cổng.

Vấn đề suy giảm kích thước batch được giải quyết như thế nào trên các thiết bị phân tán?

Mô hình kết hợp song song dữ liệu và song song mô hình: các lớp tiêu chuẩn xử lý các mini-batch độc lập trên các thiết bị, nhưng các token từ mọi thiết bị được gom lại với nhau trước khi định tuyến đến các chuyên gia cố định trên các GPU khác nhau.

Tại sao việc thêm nhiễu vào mạng cổng là cần thiết?

Nhiễu Gauss được đưa vào giúp thúc đẩy việc khám phá các chuyên gia ở giai đoạn đầu huấn luyện và cung cấp một bộ ước lượng trơn cho hàm mất mát cân bằng tải, ngăn chặn việc mô hình tự củng cố và sụp đổ vào một vài chuyên gia duy nhất.

Điều gì xảy ra khi mô hình được mở rộng lên 131,072 chuyên gia?

Perplexity bị giảm nhẹ từ 28.9 (ở 65,536 chuyên gia) xuống 29.2 (ở 131,072 chuyên gia), và hiệu suất tính toán giảm xuống 0.30 TFLOPS/GPU, nguyên nhân nhiều khả năng do tính thưa quá lớn và kích thước batch cố định không được mở rộng tương ứng.

Phần giải thích trên trang này do AI viết dựa trên bài báo và có thể có sai sót. Hãy đọc bài báo gốc để biết chi tiết.

Đọc với bản dịch và AI