Các mô hình ngôn ngữ và AI tạo sinh là chủ đề nổi bật nhờ năng lực ngày càng cao. Các nhà nghiên cứu đang cải thiện hiệu quả và khả năng của chúng. Những hệ thống học sâu này được tiền huấn luyện trên lượng dữ liệu lớn, sử dụng mạng nơ-ron và cơ chế tự chú ý. Các lớp truyền thẳng, hồi quy, nhúng và chú ý xử lý văn bản đầu vào để tạo đầu ra phù hợp.

Phần lớn tham số của mô hình ngôn ngữ lớn nằm trong các lớp truyền thẳng. Các nghiên cứu cho thấy chỉ một phần nơ-ron cần tham gia tính toán đầu ra trong mỗi lần suy luận.

Bài viết giới thiệu UltraFastBERT, một biến thể dựa trên BERT có hiệu quả tương đương các mô hình BERT hàng đầu nhưng chỉ dùng 0,3% nơ-ron khi suy luận, tức 12 trong số 4.095 nơ-ron ở mỗi lớp. Chúng ta sẽ tìm hiểu kiến trúc, cách hoạt động và kết quả của mô hình.

UltraFastBERT: mô hình hóa ngôn ngữ nhanh hơn theo cấp số nhân

Mô hình ngôn ngữ thường kết hợp các lớp truyền thẳng, hồi quy, nhúng và chú ý để học mẫu trong dữ liệu và tạo đầu ra phù hợp. Trong các mô hình ngôn ngữ, lớp truyền thẳng chứa phần lớn tham số. Việc không tận dụng hợp lý số nơ-ron cần thiết cho từng đầu vào trong giai đoạn suy luận có thể làm tăng độ phức tạp, thời gian và chi phí tính toán.

UltraFastBERT thay các lớp truyền thẳng thông thường bằng mạng truyền thẳng nhanh (FFF). Nhờ vậy, mô hình chỉ sử dụng 0,3% nơ-ron nhưng vẫn cho kết quả gần với BERT có cùng quy mô và quy trình huấn luyện, đặc biệt ở các tác vụ hạ nguồn. Thiết kế này làm giảm mạnh độ phức tạp của các lớp trung gian.

Với n nơ-ron, lượt truyền xuôi của mạng truyền thẳng thông thường (FF) có độ phức tạp O(n), còn FFF là O(log₂n). FFF tổ chức nơ-ron thành cây nhị phân cân bằng và chỉ đi theo một nhánh tùy đầu vào. Khi suy luận, mô hình thực hiện phép nhân ma trận có điều kiện (CMM): mỗi hàng đầu vào lấy tích vô hướng với những cột trọng số được chọn, và kết quả trước quyết định cột tiếp theo. Toàn bộ tập nơ-ron phục vụ các đầu vào khác nhau, nhưng mỗi đầu vào chỉ cần một số ít nơ-ron. Ngược lại, phép nhân ma trận đặc (DMM) tính tích vô hướng của mọi đầu vào với mọi cột trọng số.

Tóm lại, UltraFastBERT:

• Chỉ dùng 0,3% nơ-ron khi suy luận, tương đương 12 trên 4.095 nơ-ron mỗi lớp.
• Đạt kết quả gần các mô hình BERT hàng đầu khi tinh chỉnh trên tác vụ hạ nguồn.
• Cung cấp cách triển khai CMM làm nền tảng cho FFF, với mức tăng tốc được báo cáo lên tới 78 lần so với cách triển khai DMM tối ưu được đem so sánh.

Mạng nơ-ron truyền thẳng

Mạng truyền thẳng là một dạng mạng nơ-ron nhân tạo đơn giản: thông tin chỉ đi từ đầu vào qua các nút ẩn tới đầu ra, không có vòng lặp hay chu trình. Cấu trúc dễ xây dựng hơn mạng hồi quy (RNN) và mạng tích chập (CNN). Nó gồm lớp đầu vào, các lớp ẩn và lớp đầu ra; mỗi lớp chứa nơ-ron liên kết bằng trọng số.

Nơ-ron đầu vào nhận dữ liệu rồi chuyển tới lớp kế tiếp; số nơ-ron phụ thuộc số chiều đầu vào. Các lớp ẩn thực hiện tính toán: lấy tổng có trọng số của đầu ra lớp trước, áp dụng hàm kích hoạt và chuyển kết quả tiếp. Lớp đầu ra tạo dự đoán cuối. Trong mạng truyền thẳng kết nối đầy đủ được mô tả ở phần này, mỗi nơ-ron nối với mọi nơ-ron của lớp kế tiếp. Trọng số thể hiện độ mạnh của liên kết và được cập nhật theo sai số để học các mẫu.

Quá trình hoạt động có hai giai đoạn chính: truyền xuôi và lan truyền ngược.

Giai đoạn truyền xuôi

Đầu vào được đưa vào mạng và truyền về phía trước. Các lớp ẩn tính tổng có trọng số rồi đưa qua hàm kích hoạt như ReLU, sigmoid hoặc tanh để tạo tính phi tuyến. Quá trình tiếp tục tới lớp đầu ra và mô hình đưa ra dự đoán.

Giai đoạn lan truyền ngược

Mô hình tính sai số giữa đầu ra dự đoán và đầu ra mong đợi. Sai số được truyền ngược qua mạng; thuật toán tối ưu hạ gradient điều chỉnh trọng số nhằm giảm sai số.

UltraFastBERT: kiến trúc và cách hoạt động

UltraFastBERT xây dựng trên crammedBERT và giữ các thành phần của kiến trúc này, ngoại trừ các mạng truyền thẳng ở lớp trung gian trong bộ mã hóa Transformer, vốn được thay bằng FFF. Những thay đổi chính gồm:

• Dùng GELU cho cả nút lá và nút không phải lá, bổ sung trọng số đầu ra cho các nút và loại bỏ độ lệch đầu ra. Kích thước lá được cố định là 1.
• Cho phép nhiều cây FFF chạy song song. Đầu ra trung gian được tính bằng tổng đầu ra của từng cây.

Quy trình huấn luyện theo crammedBERT: tắt dropout khi tiền huấn luyện và dùng lịch tốc độ học tam giác một chu kỳ. Sau đó mô hình được tinh chỉnh trong 5 epoch trên nhiều tác vụ, chủ yếu thuộc bộ GLUE.

Suy luận

Suy luận là khía cạnh quan trọng của FFF và có tiềm năng tăng tốc lớn. Ví dụ, lớp truyền thẳng trong mỗi lớp Transformer của GPT-3 có hơn 49.100 nơ-ron. Nếu có thể huấn luyện một FFF độ sâu tối đa 15 để thay thế, mạng đó sẽ có hơn 65.000 nơ-ron nhưng chỉ dùng 16 nơ-ron mỗi lần suy luận, khoảng 0,03% tổng số nơ-ron. Đây là ví dụ minh họa tiềm năng, không phải khẳng định GPT-3 đã được thay thế theo cách này.

Thuật toán và khả năng tương thích

UltraFastBERT dùng thuật toán mã giả đệ quy cho suy luận truyền thẳng nhanh, như hình dưới.

Trong đó B là kích thước lô, H là độ rộng đầu vào và M chỉ các cột. Một vấn đề cần xem xét là CMM có tương thích với quy trình DMM và các khung học sâu hiện có hay không. CMM không tạo ra sự không tương thích trong cách triển khai này, nhưng làm tăng độ phức tạp của việc dùng bộ nhớ đệm.

DMM đơn luồng thực hiện các lệnh nhân và tích lũy (MAC). CMM có lợi cho CPU vì cần ít lệnh MAC hơn để tính đầu ra của lớp cho mỗi phần tử. Dù có điều kiện rẽ nhánh, “phân nhánh nơ-ron” được thực hiện qua độ lệch bộ nhớ của các con trỏ liên quan. Do đó không cần dùng đầy đủ cơ chế dự đoán nhánh lệnh để thực thi điều kiện CMM; chỉ tải riêng những cột trọng số cần thiết. Vì vẫn tính tích vô hướng hàng-cột, xử lý song song véc-tơ SIMD cũng có thể tăng tốc suy luận trên phần cứng phù hợp.

UltraFastBERT: hiệu năng và kết quả

Phần này đánh giá cả tinh chỉnh và suy luận để so sánh UltraFastBERT với các mô hình ngôn ngữ hàng đầu.

Kết quả tinh chỉnh

Hình dưới trình bày kết quả trên GLUE-dev. N là số nơ-ron sẵn có để huấn luyện và Avg là điểm trung bình các tác vụ.

Sau khoảng 24 giờ huấn luyện trên GPU A6000, UltraFastBERT giữ gần 96% hiệu năng dự đoán của BERT gốc trên các tác vụ GLUE. Khi cây FFF sâu hơn, hiệu năng giảm, nhưng phần lớn mức giảm tập trung ở CoLA. Nếu loại CoLA khỏi phép tính, hiệu năng còn khoảng 98,6% so với mô hình tham chiếu.

Kết quả suy luận

Các cách triển khai FF và FFF được so sánh ở ba cấp:

• Cấp 1: dùng BLAS cấp 1, gồm tích vô hướng-véc-tơ và tích vô hướng giữa các véc-tơ.
• Cấp 2: dùng BLAS cấp 2, gồm các phép tính theo lô và tích ma trận-véc-tơ theo cách triển khai được mô tả.
• Cấp 3: dùng phép nhân ma trận-ma trận BLAS cấp 3 không chia lô. Đây là cách nhanh nhất được xét cho FF, nhưng không có cách tương đương cho FFF trong thư viện vì không hỗ trợ tính thưa ở mức véc-tơ của CMM.

Các cách triển khai GPU sử dụng kernel CUDA tùy chỉnh hoặc kernel PyTorch.

Bảng so sánh mức tăng tốc tương đối của FFF so với FF khi hai bên sử dụng cùng các phép toán đại số tuyến tính cơ bản.

Cần lưu ý đây là phép so sánh dùng các phép toán cơ bản giống nhau nhằm đảm bảo công bằng, không phải mức tăng tốc của toàn bộ mô hình so với mọi cách triển khai tối ưu. Ở cấp 1 và cấp 2, FFF suy luận nhanh hơn lần lượt 48 và 78 lần so với cách triển khai FF nhanh nhất trong nhóm so sánh tương ứng.

Kết luận

UltraFastBERT là biến thể BERT thay lớp truyền thẳng bằng mạng truyền thẳng nhanh để chỉ sử dụng một phần nhỏ nơ-ron cho mỗi đầu vào. Nó dùng 0,3% nơ-ron nhưng vẫn đạt kết quả gần với các mô hình BERT có quy mô và quy trình huấn luyện tương tự, đặc biệt trên tác vụ hạ nguồn.

Thiết kế này làm các lớp trung gian nhanh hơn đáng kể. Kết quả cho thấy mô hình ngôn ngữ lớn có thể duy trì hiệu năng cao dù mỗi lần suy luận chỉ huy động một phần tham số. Mức tăng tốc tới 78 lần được báo cáo thuộc các phép so sánh triển khai lớp nêu trên.