Giải thích chi tiết cơ chế hoạt động của Mixture of Experts (MoE)

Mixture of Experts (MoE) là gì

Bạn muốn cải thiện hiệu suất của model? => Tăng size nó lên. Đó là lý do vì sao bạn thấy LLMs ngày càng lớn. Nhưng size thì không thể nào cứ tăng mãi được, vì model càng lớn thì train càng mệt, càng tốn tiền. Với một ngân sách cố định (và nghèo nàn như số dư trong tài khoản mình) thì bạn có hai lựa chọn để cải thiện hiệu suất model:

• Tăng size model lên nhưng train ít steps lại
• Giảm size model xuống nhưng train nhiều steps lên

Chỉ có trẻ con thì mới lựa chọn, người lớn lấy cả hai. Mixture of Experts là cơ chế giúp bạn vừa tăng size model, vừa giữ được chi phí training không quá lớn. MoE có thể đạt được hiệu suất tương đương với dense model nhưng inference time thì nhanh hơn rất nhiều

Vậy MoE là gì? Nó có 2 thành phần chính:

• Sparse MoE layers: Thay vì sử dụng các lớp feed-forward dense (FFN) rất lớn tốn kém chi phí tính toán, MoE sẽ chia nó thành các "cục" nhỏ hơn, mỗi "cục" được gọi là một "expert". Mỗi expert sẽ thực hiện một nhiệm vụ khác nhau. Dựa vào data đầu vào của bạn mà chỉ những experts liên quan mới được dùng để xử lý input. Ví dụ như input là những đoạn văn về màu sắc, thì những experts chuyên xử lý màu sắc sẽ được chọn, thay vì những experts chuyên xử lý về âm nhạc hoặc điện ảnh Ấn Độ. Mỗi cục "expert" này có thể là một FFN nhỏ hơn, hoặc một network phức tạp hơn, hoặc thậm chí cũng có thể là một MoE layer khác - tức là MoE lồng trong MoE luôn!!

• Router/Gate Network: Này có tác dụng giống như một cảnh sát đang đứng phân luồng giao thông vậy. Router sẽ quyết định xem token nào sẽ được gửi đến expert nào để xử lý. Kiểu sếp nào chạy quá tốc độ thì lùa vô làn 1, sếp nào bốc 3 thì bị lùa vô làn 2. Mỗi làn đó sẽ có một người (expert) chuyên xử lý những ca đó. Vậy làm sao để Router biết được token nào thì đi vô expert nào? Nó sẽ được train cùng với các phần khác của models trong quá trình training.

MoE có hai phần chính: Sparse MoE layer chứa các experts (là FFN 1, FFN 2, ... trong hình) và Router có nhiệm vụ phân chia token vào expert tương ứng

MoE được ứng dụng vào Transformer như thế nào?

Chắc hẳn chúng ta đã không còn xa lạ với cái tên Transformer - mô hình đứng đằng sau mọi thành tựu mà LLM đạt được hiện tại. Transformer gồm có hai phần: encoder và decoder như hình bên dưới.

Trong cả encoder và decoder, các bạn sẽ thấy nó đều có một thành phần gọi là "Feed Forward", đó chính là FFN layer mà chúng ta nói đến ở trên. Đây chính là chỗ mà ta có thể thay thế bằng MoE layer.

Tóm lại: Ta thay mỗi lớp FFN của Transformer bằng một lớp MoE (là gồm có Sparse MoE layer và Router) 

Điểm mạnh của MoE:

• Training time và Inference time nhanh hơn so với mô hình cùng size sử dụng dense FFN

Điểm yếu của MoE:

• Thực nghiệm cho thấy MoE có vẻ hơi khó khăn trong việc tổng quát hóa (generalize) quá trình huấn luyện và dễ bị overfit
• Inference: Các bạn thấy là mỗi lần inference chỉ có một số experts là được kích hoạt thôi, nên inference time nhanh hơn. Nhưng tham số của tất cả các experts phải được load vào RAM => yêu cầu về RAM để chạy đc MoE sẽ cao hơn. 

Sparsity trong MoE

Một trong những khái niệm quan trọng nhất của MoE đó là sparsity (tính thưa). Sparsity sử dụng ý tưởng conditional computation (tính toán có điều kiện), nghĩa là khi đưa một input vào, thì chỉ có một phần nhỏ tham số của model là được sử dụng để xử lý input, trong khi các mô hình dense bình thường thì phải dùng hết tham số để xử lý input.

Trong MoE, sparsity được thực hiện thông qua router, nó quyết định experts nào của model sẽ được kích hoạt cho một input cụ thể. Vậy làm sao để train cho router quyết định được experts nào là phù hợp với input nào?

Cơ chế gating của Router

Có rất nhiều cách để cài đặt cơ chế gating của router, cách cơ bản nhất là sử dụng hàm softmax. Cụ thể như sau:

Giả sử input là $x$, ta có các experts $E_1, E_2, ..., E_n$ với $n$ là tổng số experts. Output của MoE sẽ được tính bằng công thức sau:

$$y = \sum^n_{i=1}\alpha_i E_i(x)$$

$\alpha_i$ chính là trọng số của expert $E_i$ - tượng trưng cho mức độ liên quan của expert $E_i$ với input hiện tại. Nếu $\alpha_i = 0$, nghĩa là expert $E_i$ chả có liên quan gì tới input hiện tại hết, nên khi các bạn thực hiện phép nhân $\alpha_i E_i(x)$ thì kết quả sẽ bằng 0, tương tự như việc expert $E_i$ đã được tắt.

Vậy làm sao để router tính được các trọng số $\alpha_i$ cho từng expert? Ta sẽ dùng hàm softmax

$$\alpha_i = \frac{e^{\mathbf{w}_i^\top \mathbf{x}}}{\sum_{j=1}^{N} e^{\mathbf{w}_j^\top \mathbf{x}}}$$

Nếu bạn chưa biết hàm softmax là gì thì sẽ khá khó hiểu khúc này, các bạn cần phải tìm hiểu thêm. Đại khái thì hàm softmax sẽ giúp tạo ra một phân phối xác suất trên các outputs, ở đây chính là trọng số của từng expert. Với input hiện tại, giả sử ta nên dùng $80\%$ của expert $E_1$, $2\%$ của expert $E_2$, ... thì sử dụng hàm softmax bên trên, nó sẽ tìm ra cho chúng ta $\alpha_1=0.8$, $\alpha_2=0.02$, ... Hàm softmax đảm bảo được hai điều:

• $\alpha_i>0$ 
• $\sum_{i=1}^n \alpha_i = 1$
  

Trong công thức trên, $\mathbf{w}_i$ là weight vector của expert $E_i$, còn $x$ là input. Các weight vectors này sẽ được train để hàm softmax tính được $\alpha_i$ tốt nhất cho từng input $x$

-----------------------------------------------------------------------------------------------------

Ngoài cách gating cơ bản bằng softmax ra thì còn các cách khác như Noisy Top-K Gating:

Bước 1: Thêm nhiễu (noise)

$$H_i(x) = (\mathbf{w}_i^\top \mathbf{x}) + \mathcal{N}(0, 1) \cdot softplus(\mathbf{w_{noise}}_i^\top \mathbf{x})$$

• $x$ là input
• $\mathbf{w}_i^\top \mathbf{x}$ là điểm liên quan của expert $E_i$ đối với input $x$
• $\mathcal{N}(0, 1)$ là phân phối chuẩn (Standard Normal Distribution) - là phân phối có trung bình bằng 0 và độ lệch chuẩn = 1
• $softplus$ là một hàm kích hoạt nhằm đảm bảo giá trị nhiễu luôn dương. Softplus là một hàm xấp xỉ của hàm ReLU
• $\mathbf{w_{noise}}_i^\top \mathbf{x}$ là giá trị của nhiễu mà ta muốn thêm vào cho expert $E_i$

Trong bước này, chúng ta tính toán điểm số cho mỗi expert, nhưng thêm nhiễu vào để tạo ra một chút biến đổi trong quá trình chọn expert. Nhiễu giúp cho model "explore" được nhiều hơn và sẽ đỡ bị overfit hơn.

Bước 2: Chọn K experts có điểm số cao nhất

$$\text{TopK}(H, k)_i = \begin{cases} H_i & \text{nếu } H_i \text{ thuộc top } k \text{ giá trị lớn nhất }, \\ -\infty & \text{ngược lại.} \end{cases}$$

• $H_i$ là điểm số đã thêm nhiễu của expert $E_i$ được tính ở bước 1
• $k$ là số lượng experts mà mình muốn chọn

Bước này đảm bảo rằng chỉ có $k$ experts tốt nhất được chọn để tính tiếp. Các experts còn lại bị loại bỏ bằng cách gán với giá trị $-\infty$.

Bước 3: Sử dụng Softmax lên $k$ experts được chọn ở bước 2

$$\alpha_i = \frac{e^{H_i(x)}}{\sum_{j=1}^k e^{H_j(x)}}$$

Sau khi chọn ra $k$ experts tốt nhất đối với input hiện tại, ta phải tiếp tục tìm ra trọng số phù hợp cho từng expert. Có nghĩa là trong $k$ thằng tốt nhất đó, không phải thằng nào cũng được bật như nhau. Mà mỗi thằng sẽ có một trọng số xác định đóng góp của nó vào kết quả cuối cùng.

Bằng cách chọn số $k$ đủ nhỏ (VD: 1 hoặc 2), ta có thể train và chạy inference nhanh hơn so với việc có nhiều experts được kích hoạt hơn (VD: 5 hoặc 6). Vậy sao không chọn 1 expert (thằng tốt nhất) thôi cho nó siêu nhanh luôn? Bởi vì sử dụng ít nhất là 2 experts là cần thiết để router học được cách gating đến các experts khác nhau.

Kết bài

Trong bài viết sau, chúng ta sẽ cùng tìm hiểu về Switch Transformer, một mô hình mở rộng của MoE Transformer. Switch Transformer chỉ sử dụng duy nhất một expert trong MoE layer (k = 1) chứ không sử dụng top-k experts, điều này đi ngược với nhận định ít nhất k phải bằng 2 mà chúng ta nói đến bên trên.

Tài liệu tham khảo:

• https://huggingface.co/blog/moe
• https://en.wikipedia.org/wiki/Mixture_of_experts
• https://www.ibm.com/topics/mixture-of-experts
• https://abvijaykumar.medium.com/mixture-of-experts-2d0b681f506d
• https://medium.com/@developer.yasir.pk/understanding-mixture-of-experts-in-nlp-a-technical-deep-dive-a02b5a73e025