[Bản dịch] MoD (Mixture-of-Depths): Cách tiếp cận tối ưu hóa tính toán cho mô hình ngôn ngữ dựa trên Transformer, và MoDE (MoD+MoE)

(discuss.pytorch.kr)

7 điểm bởi ninebow 2024-04-13 | Chưa có bình luận nào. | Chia sẻ qua WhatsApp

PyTorchKR

Gần đây, phương pháp MoE (Mixture-of-Experts) — chọn một hoặc nhiều khối “chuyên gia” để xử lý token đầu vào — đang được sử dụng rất nhiều, như trong Jamba, Qwen1.5-MoE, hay DBRX. Cách làm MoE này cho thấy hiệu năng tốt hơn nhờ phân bổ tài nguyên tính toán hiệu quả bằng cách điều tiết động một phần của các LLM khổng lồ.
Nếu MoE điều chỉnh động “độ rộng (width) của mô hình” bằng cách chọn một số chuyên gia trong số nhiều chuyên gia, thì MoD (Mixture-of-Depths) được giới thiệu lần này lại chọn cách điều chỉnh động “độ sâu (depth) của mô hình”. Nói cách khác, với một số token nhất định, mô hình không tính toán qua toàn bộ layer mà sẽ động bỏ qua một số layer. Hãy cùng xem chi tiết.

MoD (Mixture-of-Depths): Cách tiếp cận tối ưu hóa tính toán cho mô hình ngôn ngữ dựa trên Transformer (Mixture-of-Depths: Dynamically allocating compute in transformer-based language models)

Giới thiệu

Theo cách truyền thống, mô hình Transformer phân bổ cùng một lượng tài nguyên tính toán cho mọi phần của chuỗi đầu vào. Tuy nhiên, bài báo về MoD (Mixture-of-Depth) Transformer được giới thiệu lần này thử nghiệm việc phân bổ tài nguyên tính toán một cách động cho các vị trí chuỗi cụ thể trong Transformer, đồng thời thay đổi cách phân bổ đó qua các layer khác nhau. Để sử dụng chiến lược Mixture-of-Depths (MoD), các tác giả dùng cơ chế định tuyến top-k để quyết định token nào sẽ được xử lý, từ đó kiểm soát tổng lượng tính toán. Kỹ thuật này yêu cầu ít FLOPs hơn rất nhiều nhưng vẫn đạt hiệu năng tương đương hoặc tốt hơn mô hình hiện có, đồng thời mang lại mức tăng tốc tới 50% trong giai đoạn lấy mẫu sau huấn luyện.

Cấu trúc Transformer được giới thiệu trong bài báo Attention is All You Need

Các mô hình Transformer và LLM (Large Language Model) gần đây đang không ngừng tăng về quy mô và độ phức tạp. Những mô hình này sử dụng hàng tỷ tham số để tiến gần tới khả năng hiểu và sinh ngôn ngữ ở mức con người, nhờ đó thể hiện hiệu năng ấn tượng. Tuy nhiên, đi cùng với sự phát triển đó là chi phí tính toán và mức tiêu thụ năng lượng đáng kể; khi quy mô mô hình tăng lên, thời gian huấn luyện và suy luận cũng tăng mạnh.

Điều này xuất phát từ việc Transformer thực hiện cùng một lượng tính toán cho mọi token trong chuỗi đầu vào. Nhưng ý tưởng ở đây là: không phải mọi token đều cung cấp lượng thông tin như nhau; có token đóng vai trò cốt lõi trong việc hiểu ngữ cảnh, trong khi có token tương đối ít quan trọng hơn.

MoD Transformer cho phép mô hình tập trung tài nguyên vào nơi cần thiết bằng cách phân bổ tính toán động cho từng token. Điều này giúp giảm chi phí tính toán tổng thể của mô hình, đồng thời nâng cao hiệu quả xử lý thông tin bằng cách dành nhiều sự chú ý hơn cho các token quan trọng. Đặc biệt với các mô hình lớn như LLM, cách tiếp cận của MoD có thể trở thành một phương pháp quan trọng để tiết kiệm thời gian và tài nguyên trong cả huấn luyện lẫn suy luận.

Giới thiệu kỹ thuật MoD (Mixture-of-Depths)

Kỹ thuật MoD tích hợp cơ chế định tuyến top-k để động chọn token nào sẽ được xử lý ở từng layer của Transformer. Nhờ đó, chỉ một phần token được chọn mới đi qua toàn bộ quá trình tính toán, còn các token khác sẽ bỏ qua layer (hoặc nhiều layer).

Hình 1 / Trái: Tổng quan về Transformer MoD (Mixture-of-Depths): tương tự cấu trúc MoE (Mixture-of-Experts), đường đi tính toán được chọn thông qua router. Tuy nhiên, khác với MoE, ở đây router quyết định có thực hiện khối chuẩn (Self-Attention và MLP) hay không. Nếu token đầu vào ($X_i$) không kích hoạt toàn bộ phép tính theo quyết định của router, token đó sẽ không đi qua tính toán của layer tương ứng.
Hình 1 / Phải: Kết quả quyết định định tuyến trên một chuỗi ngắn rút gọn còn 64 token. (Trục X: chuỗi, Trục Y: layer) Màu tím là thực hiện toàn bộ tính toán của layer đó, màu cam là bỏ qua tính toán của layer đó.

Như phần mô tả trong hình trên, MoD triển khai một router dự đoán trong kiến trúc Transformer để quyết định với mỗi token xem có tiếp tục tính toán tại layer đó hay không. Qua đó, tài nguyên tính toán có thể được tập trung vào các token tương đối “quan trọng”.

Huấn luyện MoD Transformer

Quá trình huấn luyện MoD Transformer về cơ bản tương tự Transformer tiêu chuẩn. Chuỗi đầu vào được đưa vào mô hình, đi qua từng layer và thực hiện các phép tính cần thiết. Tuy nhiên, MoD Transformer bổ sung thêm một bước, trong đó router quyết định mỗi token sẽ đi theo đường tính toán nào.

Tính trọng số router: Router tính trọng số cho từng token ở mỗi layer của mô hình. Trọng số này biểu thị lượng tính toán mà token đó nên nhận, dựa trên mức độ quan trọng của token.
Chọn các token top: Dựa trên các trọng số đã tính, router chọn ra nhóm token đứng đầu cho mỗi layer. Những token này sẽ được đưa vào đường thực hiện đầy đủ tính toán.
Phân bổ đường tính toán: Các token được chọn sẽ đi vào đường thực hiện tính toán chuẩn (Self-Attention và MLP). Các token còn lại được chuyển tiếp thông qua kết nối residual và đi sang layer tiếp theo mà không cần tính toán bổ sung.
Tính loss và lan truyền ngược: Đầu ra của mô hình được đánh giá bằng hàm loss cuối cùng, sau đó gradient được lan truyền ngược qua mô hình để cập nhật trọng số. Trong quá trình này, router cũng được học để theo thời gian có thể đưa ra quyết định định tuyến hiệu quả hơn.

Cách định tuyến của MoD Transformer (Routing Schemes)

Triển khai định tuyến là phần cốt lõi của MoD Transformer. Các tác giả so sánh 3 cách định tuyến và xác nhận rằng định tuyến token theo kiểu Expert-Choice MoD cho hiệu năng tốt nhất. Hãy xem so sánh với các cách còn lại:

Định tuyến chọn token (Token-choice)

Trong định tuyến chọn token, mỗi token tự chọn đường mà nó sẽ đi qua. Mô hình tạo ra một phân phối xác suất cho nhiều đường tính toán khác nhau của từng token (ví dụ: các layer hoặc khối xử lý khác nhau), rồi gửi token theo phân phối đó.

Định tuyến chọn token cho phép mỗi token chọn đường tính toán phù hợp nhất, nên rất linh hoạt. Ngoài ra, mô hình có thể cân nhắc chính xác hơn mức độ quan trọng hoặc ngữ cảnh của từng token, từ đó tối ưu ở cấp độ token riêng lẻ.

Định tuyến chọn chuyên gia (Expert-choice)

Trong định tuyến chọn chuyên gia, mô hình trực tiếp chọn các token mà mỗi đường (hoặc “chuyên gia”) sẽ xử lý. Việc này được thực hiện bằng cơ chế top- để chọn ra nhóm token đứng đầu dựa trên trọng số router của từng token.

Cách định tuyến chọn chuyên gia giúp các token có thể được phân bổ đồng đều hơn giữa các đường, nhờ đó tải tính toán được cân bằng tốt hơn và tài nguyên tính toán của mô hình được tận dụng hiệu quả hơn. Ngoài ra, vì các token top được xác định rõ ràng, khả năng dự đoán và độ ổn định của mô hình cũng được cải thiện.

Expert-choice MoD

Expert-choice MoD là việc áp dụng khái niệm định tuyến chọn chuyên gia vào cấu trúc của MoD Transformer. Trong cách này, mô hình chọn ra nhóm token đứng đầu cho mỗi layer hoặc mỗi bước xử lý, và chỉ các token được chọn mới đi theo đường có tính toán. Các token còn lại đi theo đường chi phí thấp hơn thông qua kết nối residual.

Cách làm này giúp giảm đáng kể tổng mức sử dụng FLOP vì chỉ các token được chọn mới thực hiện tính toán. Đồng thời, với các token quan trọng, mô hình vẫn đảm bảo tính toán trên toàn bộ layer, còn các phép tính không cần thiết thì được giảm thiểu để duy trì hoặc cải thiện hiệu năng. Cuối cùng, nhờ định tuyến động qua router, mô hình có ưu điểm là có thể phân bổ tài nguyên tính toán linh hoạt theo đặc điểm của chuỗi đầu vào.

Cách triển khai định tuyến

Việc triển khai định tuyến có thể thực hiện qua 3 bước chính sau:

Xác định embedding token và trọng số router: Với mỗi layer của mô hình, ta có tập embedding token của chuỗi đã cho là $X^i$. Với mỗi embedding token $x^l_i$, router tính một trọng số vô hướng $r^l_i$ thông qua phép biến đổi tuyến tính; trọng số này biểu thị khả năng token đó sẽ tham gia vào tính toán.
Xác định các trọng số top: Từ tập trọng số router của token $R^l$, tìm phân vị xác suất của trọng số thứ $\beta$, ký hiệu là $P_\beta(R^l)$. Ở đây $\beta$ quyết định tỷ lệ token mà mô hình sẽ xử lý ở mỗi layer. Vì vậy, các token có trọng số router lớn hơn $P_\beta(R^l)$ sẽ tham gia vào tính toán.
Chọn đường tính toán: Mỗi token được so sánh trọng số router $r^l_i$ với $P_\beta(R^l)$. Nếu lớn hơn hoặc bằng, token đó sẽ được xử lý tại layer tương ứng; nếu nhỏ hơn, nó sẽ được chuyển thẳng sang layer tiếp theo. Nhờ đó, chỉ các token được chọn mới được tính toán, còn các token khác có thể đi theo đường không phát sinh chi phí tính toán.

Khi triển khai định tuyến như vậy, cần đặc biệt lưu ý các điểm sau:

Quyết định động dựa trên trọng số: Quá trình định tuyến được quyết định động dựa trên trọng số của từng token, phản ánh tầm quan trọng của token đó. Token quan trọng nhận nhiều tính toán hơn để tối ưu hiệu năng mô hình, trong khi token ít quan trọng hơn sẽ tiết kiệm chi phí bằng cách đi theo đường residual.
Tối ưu hiệu quả tính toán và hiệu năng: Thông qua cơ chế định tuyến này, MoD Transformer có thể sử dụng tài nguyên tính toán hiệu quả mà vẫn duy trì hoặc cải thiện hiệu năng mô hình. Đặc biệt, bằng cách chỉ xử lý các token quan trọng ở những layer có chi phí tính toán cao, tổng mức sử dụng FLOP có thể được giảm xuống.
Đơn giản và hiệu quả trong triển khai: Việc định tuyến được triển khai thông qua biến đổi tuyến tính và tính phân vị, nên tương đối đơn giản và hiệu quả. Điều này cho phép phân bổ tài nguyên tính toán động mà không tạo thêm gánh nặng tính toán đáng kể trong quá trình huấn luyện và suy luận.

Ví dụ cách MoD Transformer hoạt động

Dựa trên những gì đã xem, ta hãy giải thích cách MoD Transformer hoạt động bằng một ví dụ.

Ví dụ, giả sử chuỗi đầu vào có 100 token và router đặt giá trị là 10. Khi đó, ở mỗi layer, router sẽ chọn 10 token quan trọng nhất, và chỉ những token này mới được xử lý theo đường tính toán chuẩn.

90 token còn lại sẽ bỏ qua tính toán và được chuyển thẳng sang layer tiếp theo thông qua kết nối residual. Nhờ quy trình này, MoD Transformer có thể giảm tổng mức sử dụng FLOP mà vẫn đảm bảo thực hiện các phép tính cần thiết cho các token quan trọng.

So sánh hiệu năng của MoD Transformer

Hiệu năng huấn luyện

Tối ưu siêu tham số: Thông qua thí nghiệm huấn luyện một biến thể MoD Transformer trong đó các khối định tuyến dung lượng 12,5% được xen kẽ với các khối attention đầy đủ, các tác giả đã tìm ra cấu hình siêu tham số tối ưu. Cấu hình này cho hiệu năng tốt hơn cả đường cơ sở tối ưu trong phân tích isoFLOP.

Hiệu quả tính toán: Một số biến thể MoD nhất định đạt cùng mức hiệu năng với mô hình hiện có nhưng dùng ít FLOP hơn, đồng thời cho tốc độ xử lý nhanh hơn 66%.

Phân tích định tuyến

Tính thưa của quyết định định tuyến: Khi dùng định tuyến dung lượng 12,5%, phần lớn token đi vòng qua khối, và chỉ một số ít token quan trọng mới đi xuyên qua khối. Điều này cũng có thể quan sát được trong phân bố trọng số router; mạng thậm chí còn có xu hướng ưu tiên định tuyến một số token cụ thể theo độ sâu.

Hiệu năng suy luận

Suy luận tự hồi quy: Khi chuyển từ sơ đồ định tuyến top- phi nhân quả (non-causal) trong huấn luyện sang cách tiếp cận dựa trên dự đoán nhân quả (causal) trong suy luận, gần như không có suy giảm hiệu năng. Có vẻ như nguyên nhân là bài toán dự đoán tương ứng này khá dễ học, và mô hình đã đạt độ chính xác trên 97% ngay từ giai đoạn đầu huấn luyện.

Mixture-of-Depths-and-Experts (MoDE) = MoD + MoE

MoDE (Mixture-of-Depths-and-Experts) là sự tích hợp giữa kỹ thuật MoD (Mixture-of-Depths) và kỹ thuật MoE (Mixture of Experts). Mục tiêu của cách làm này là kết hợp ưu điểm của hai hướng tiếp cận để đồng thời tối ưu hiệu quả tính toán và hiệu năng của mô hình Transformer. Nếu MoD chọn để mỗi token được tính toán sâu hơn hoặc nông hơn, thì MoE xử lý token thông qua “chuyên gia” phù hợp nhất trong số nhiều chuyên gia. MoDE tích hợp hai cách tiếp cận này để token vừa có thể bỏ qua một số khối hoặc chọn thực hiện tính toán sâu hơn một cách có chọn lọc, vừa có thể chọn ra phương thức xử lý tối ưu nhất từ nhiều chuyên gia khác nhau.

Cách triển khai MoDE

MoDE có hai cách triển khai chính: MoDE theo giai đoạn (Staged) và MoDE tích hợp (Integrated).

Staged MoDE (MoDE theo giai đoạn): MoDE theo giai đoạn trước hết áp dụng kỹ thuật MoD để quyết định đường tính toán của từng token, sau đó áp dụng kỹ thuật MoE để định tuyến các token trên đường đã chọn vào các khối chuyên gia tương ứng. Nhờ đó, mô hình trước tiên điều chỉnh độ sâu tính toán, rồi sau đó chọn chuyên gia tối ưu ở từng bước.
Integrated MoDE (MoDE tích hợp): MoDE tích hợp triển khai MoD và MoE trong cùng một bước định tuyến. Tức là khi định tuyến token, mô hình đưa ra một quyết định duy nhất để đồng thời chọn độ sâu của đường tính toán và chuyên gia sẽ xử lý token đó. Cách này có thể đơn giản và hiệu quả hơn trong triển khai, đồng thời có ưu điểm là quyết định được cả đường tính toán lẫn chuyên gia phù hợp nhất cho từng token trong một lần.

Ưu điểm của MoDE

Hiệu quả tính toán: MoDE quản lý hiệu quả chi phí tính toán tổng thể của mô hình bằng cách quyết định động lượng tính toán cần thiết và chuyên gia xử lý cho từng token.
Tối ưu hiệu năng: Bằng cách gán mỗi token cho đường tính toán và chuyên gia phù hợp nhất, MoDE tối đa hóa hiệu năng mô hình. Điều này đặc biệt hữu ích khi giải quyết các bài toán phức tạp hoặc xử lý nhiều loại dữ liệu khác nhau.
Tính linh hoạt và khả năng mở rộng: MoDE mang lại tính linh hoạt để áp dụng cho nhiều cấu trúc mô hình và tác vụ khác nhau, đồng thời cung cấp một khung mở rộng có thể cải thiện đồng thời cả hiệu quả tính toán lẫn hiệu năng mô hình.

MoDE gợi mở hướng phát triển tương lai cho các mô hình Transformer và đóng góp quan trọng cho nghiên cứu về cách tối đa hóa hiệu năng mô hình trong khi vẫn giảm thiểu chi phí tính toán.

Đọc thêm

https://arxiv.org/abs/2404.02258

Bài viết này được tổng hợp dựa trên nội dung do mô hình GPT tóm lược, nên có thể có những điểm được diễn đạt khác với nội dung hoặc ý đồ của nguyên bản. Nếu bạn quan tâm, hãy tham khảo thêm cả bài gốc. Nếu trong lúc đọc bạn phát hiện nội dung gượng gạo hoặc sai sót, mong bạn để lại bình luận để góp ý.

⚠️Quảng cáo⚠️: Bạn thấy bài viết do Cộng đồng người dùng PyTorch Hàn Quốc tổng hợp này hữu ích chứ? Nếu đăng ký thành viên, bạn sẽ nhận được các bài viết nổi bật qua email! (Mặc định là Weekly, nhưng cũng có thể đổi sang Daily.)