Giới thiệu về kiến trúc Mamba LLM: Một mô hình mới trong học máy

Mô hình ngôn ngữ là một loại mô hình học máy được huấn luyện để mô phỏng phân phối xác suất trên ngôn ngữ tự nhiên. Kiến trúc của chúng chủ yếu bao gồm nhiều tầng mạng nơ-ron, như các tầng hồi quy, tầng truyền thẳng, tầng nhúng và tầng attention. Các tầng này phối hợp xử lý văn bản đầu vào và tạo ra dự đoán đầu ra.

Cuối năm 2023, các nhà nghiên cứu từ Carnegie Mellon và Đại học Princeton đã công bố một bài báo nghiên cứu giới thiệu một kiến trúc mới cho Mô hình Ngôn ngữ Lớn (LLM) có tên Mamba. Mamba là một kiến trúc mô hình không gian trạng thái mới tập trung vào bài toán mô hình hóa chuỗi. Nó được phát triển để khắc phục một số hạn chế của các mô hình transformer, đặc biệt trong xử lý các chuỗi dài, và đã cho thấy hiệu năng đầy hứa hẹn.

Hãy cùng tìm hiểu kiến trúc Mamba LLM và ý nghĩa của nó trong học máy.

Mamba là gì?

Mamba là một kiến trúc LLM mới tích hợp mô hình chuỗi Không gian Trạng thái Có Cấu trúc (S4) để quản lý các chuỗi dữ liệu dài. Kết hợp những ưu điểm của mô hình hồi quy, tích chập và thời gian liên tục, S4 có thể mô phỏng hiệu quả và hiệu suất cao các phụ thuộc dài hạn. Điều này cho phép nó xử lý dữ liệu lấy mẫu không đều, có ngữ cảnh không giới hạn và duy trì hiệu quả tính toán trong suốt quá trình huấn luyện và kiểm thử.

Mở rộng theo khuôn khổ S4, Mamba mang đến một số cải tiến đáng chú ý, đặc biệt trong xử lý các phép toán biến thiên theo thời gian. Kiến trúc của nó xoay quanh một cơ chế chọn lọc đặc biệt điều chỉnh các tham số của mô hình không gian trạng thái có cấu trúc (SSM) theo đầu vào.

Nhờ đó, Mamba có thể lọc bỏ dữ liệu kém quan trọng bằng cách chỉ tập trung vào thông tin then chốt trong chuỗi. Theo Wikipedia, "Mô hình chuyển từ khung bất biến theo thời gian sang biến thiên theo thời gian, điều này ảnh hưởng đến cả tính toán và hiệu suất của hệ thống."

Các tính năng và đổi mới chính

Việc Mamba tách khỏi các khối attention và MLP truyền thống khiến nó khác biệt. Sự đơn giản hóa này tạo ra một mô hình nhẹ hơn, nhanh hơn và mở rộng tuyến tính theo độ dài chuỗi, điều mà các thế hệ trước chưa đạt được.

Các thành phần chính của Mamba gồm:

• Selective-State-Spaces (SSM): Các mô hình hồi quy, vốn xử lý thông tin có chọn lọc tùy theo đầu vào hiện tại, là nền tảng của SSM trong Mamba. Điều này cho phép chúng lọc dữ liệu thừa và tập trung vào thông tin liên quan, có thể dẫn đến xử lý hiệu quả hơn.
• Kiến trúc đơn giản hóa: Mamba thay thế các khối attention và MLP phức tạp của Transformer bằng một khối SSM thống nhất. Mục tiêu là tăng tốc suy luận và giảm độ phức tạp tính toán.
• Song song hóa theo phần cứng: Mamba có thể hoạt động tốt hơn nữa vì nó sử dụng chế độ hồi quy với thuật toán song song được thiết kế đặc biệt cho hiệu quả phần cứng.

Một thành phần quan trọng khác là Tính Bất Biến Theo Thời Gian Tuyến Tính (LTI); LTI là một trong những đặc trưng cốt lõi của các mô hình S4. Đặc điểm này gợi ý rằng các tham số của mô hình giữ nguyên trong mọi bước thời gian, duy trì tính nhất quán trong động lực của mô hình. Việc xây dựng mô hình chuỗi trở nên dễ dàng và hiệu quả hơn với LTI, vốn là nền tảng của hồi quy và tích chập.

Chi tiết kiến trúc Mamba LLM

Kiến trúc của Mamba nhấn mạnh hơn nữa tầm quan trọng của những tiến bộ trong học máy. Nó điều chỉnh cách các mô hình xử lý chuỗi bằng cách đưa vào một tầng mô hình không gian trạng thái có chọn lọc (SSM). Điều này cho phép Mamba thực hiện hai điều cực kỳ quan trọng:

1. Tập trung vào thông tin liên quan – Mamba có thể ưu tiên dữ liệu giàu sức dự báo cho tác vụ bằng cách gán trọng số khác nhau cho từng đầu vào.

2. Thích ứng động với đầu vào – Nhờ khả năng thích ứng với đầu vào, Mamba có thể dễ dàng xử lý nhiều nhiệm vụ mô hình hóa chuỗi khác nhau.

Kết quả là, Mamba có thể xử lý các chuỗi với hiệu suất chưa từng có, trở thành lựa chọn lý tưởng cho các tác vụ liên quan đến chuỗi dữ liệu dài.

Triết lý thiết kế của Mamba dựa trên sự thấu hiểu năng lực phần cứng hiện đại. Nó được thiết kế để tận dụng tối đa sức mạnh tính toán của GPU, đảm bảo:

• Tối ưu hóa sử dụng bộ nhớ: Giảm thời gian truyền dữ liệu và tăng tốc xử lý bằng cách thiết kế phần mở rộng trạng thái của Mamba phù hợp với bộ nhớ băng thông cao (HBM) của GPU.

• Tối đa hóa xử lý song song: Bằng cách điều phối các phép tính phù hợp với bản chất song song của tính toán GPU, Mamba đạt mức hiệu năng thiết lập tiêu chuẩn mới cho các mô hình chuỗi.

Mamba so với Transformers

Sự ra đời của Transformers, như GPT-4, đã đưa lĩnh vực xử lý ngôn ngữ tự nhiên (NLP) tiến lên và thiết lập các chuẩn mực cho nhiều tác vụ ngôn ngữ. Các chuỗi dài từ lâu là điểm yếu của transformer vì chúng làm giảm đáng kể hiệu quả của mô hình.

Đây chính là điểm Mamba vượt trội. Cụ thể, Mamba có thể xử lý các chuỗi dài nhanh hơn transformer và làm điều đó một cách đơn giản hơn nhờ kiến trúc độc đáo.

Kiến trúc Transformer

Transformer rất giỏi làm việc với các chuỗi dữ liệu, như văn bản cho mô hình ngôn ngữ. Chúng xử lý toàn bộ chuỗi đồng thời, khác với các mô hình trước đây xử lý tuần tự. Đặc tính bẩm sinh này cho phép chúng nắm bắt các mối quan hệ phức tạp trong dữ liệu.

Chúng sử dụng cơ chế attention cho phép mô hình tập trung vào các đoạn khác nhau của chuỗi khi tạo dự đoán. Ba bộ trọng số được dùng để tính attention: values, keys và queries được lấy từ dữ liệu đầu vào.

Mỗi phần tử trong chuỗi được gán trọng số tương quan với mọi phần tử khác để chỉ ra mức độ — hay "attention" — mà nó nên có để dự báo phần tử tiếp theo trong dãy.

Transformer gồm hai khối chính: decoder, tạo đầu ra, và encoder, xử lý dữ liệu đầu vào.

Encoder gồm nhiều lớp – mỗi lớp có hai lớp con: một mạng truyền thẳng đầy đủ kết nối theo vị trí và một cơ chế tự-attention đa đầu. Để hỗ trợ huấn luyện mạng sâu, các kết nối tắt (residual) và chuẩn hóa được dùng ở mỗi lớp con.

Tương tự encoder, decoder gồm hai lớp với hai lớp con, nhưng còn bổ sung lớp con thứ ba xử lý attention đa đầu trên đầu ra của encoder. Tính tự hồi quy của decoder được duy trì do tính tuần tự, giới hạn dự đoán tại một vị trí chỉ xét các vị trí trước đó.

Vì vậy, Transformer cố gắng giải quyết vấn đề chuỗi dài bằng cách sử dụng các cơ chế attention phức tạp hơn, còn Mamba chọn một hướng tiếp cận khác.

Kiến trúc Mamba

Mamba tận dụng các không gian trạng thái có chọn lọc. Cách tiếp cận này giải quyết sự kém hiệu quả tính toán của Transformer với các chuỗi dài.

Kiến trúc của Mamba cho phép suy luận nhanh hơn và mở rộng tuyến tính theo độ dài chuỗi, tạo nên một mô hình mới cho mô hình hóa chuỗi có thể hiệu quả hơn khi chuỗi ngày càng dài.

Vì ở trên chúng ta đã đi sâu vào kiến trúc của Mamba, nên phần này sẽ không nhắc lại.

Dưới đây là một biểu đồ từ Wikipedia giúp hình dung rõ hơn sự so sánh giữa Mamba và Transformer:

Cũng cần lưu ý rằng, dù SSM có nhiều ưu thế so với Transformer, thì Transformer có thể xử lý các chuỗi dài hơn nhiều so với khả năng lưu trữ trong bộ nhớ của SSM, cần ít dữ liệu hơn để học các tác vụ tương tự, và vượt trội trong các nhiệm vụ cần truy xuất hoặc sao chép từ ngữ cảnh đầu vào, ngay cả với ít tham số hơn. 

Bắt đầu với Mamba

Nếu bạn muốn thử nghiệm Mamba hoặc tận dụng nó trong một dự án, bạn cần có:

• Linux
• GPU NVIDIA
• PyTorch 1.12+
• CUDA 11.6+

Để cài đặt các gói cần thiết từ kho Mamba, dùng một vài lệnh pip đơn giản:

• [Tùy chọn] pip install causal-conv1d>=1.2.0: triển khai hiệu quả một lớp Conv1d nhân quả đơn giản được dùng trong khối Mamba.
• pip install mamba-ssm: gói Mamba lõi.

Bạn cũng có thể build từ mã nguồn với pip install . từ kho này.

Nếu phiên bản PyTorch gây ra vấn đề tương thích, có thể dùng pip với tùy chọn --no-build-isolation để hỗ trợ. Các mô hình này được huấn luyện trên những bộ dữ liệu lớn như Pile và SlimPajama, và được xây dựng để đáp ứng nhiều yêu cầu tính toán và tiêu chí hiệu năng.

Mô hình Mamba có một số lớp giao diện, nhưng mô-đun chính là khối kiến trúc Mamba bao bọc SSM chọn lọc.

Khi đã cài đặt xong, có thể sử dụng như sau:

# Source: Mamba Repository
import torch
from mamba_ssm import Mamba

batch, length, dim = 2, 64, 16
x = torch.randn(batch, length, dim).to("cuda")
model = Mamba(
    # This module uses roughly 3 * expand * d_model^2 parameters
    d_model=dim, # Model dimension d_model
    d_state=16,  # SSM state expansion factor
    d_conv=4,    # Local convolution width
    expand=2,    # Block expansion factor
).to("cuda")
y = model(x)
assert y.shape == x.shapeĐược hỗ trợ bởi

Ứng dụng của Mamba

Sự xuất hiện của Mamba LLM là một thay đổi tiềm năng đáng kể trong không gian kiến trúc LLM. Nhanh hơn, hiệu quả hơn và có khả năng mở rộng, Mamba xử lý mượt mà các chuỗi dài với tiêu chuẩn hiệu năng cao, giải thích vì sao nó sẽ đóng vai trò quan trọng trong việc định hình tương lai của các hệ thống AI tinh vi.

Cụ thể, làn sóng đổi mới AI tiếp theo có thể được thúc đẩy bởi hiệu quả và hiệu năng của Mamba, mở đường cho việc tạo ra các mô hình và ứng dụng ngày càng phức tạp. Ảnh hưởng tiềm năng của nó là rất lớn, bao gồm ứng dụng xử lý âm thanh và giọng nói, phân tích văn bản dài, sáng tạo nội dung, dịch ngôn ngữ theo thời gian thực và nhiều lĩnh vực khác.

Các ngành có thể được cách mạng hóa gồm:

• Y tế: Mamba có thể đẩy nhanh quá trình phát triển thuốc cá thể hóa bằng cách phân tích nhanh dữ liệu di truyền.

• Tài chính: Mamba có thể được triển khai để phân tích xu hướng thị trường dài hạn, từ đó dự báo cổ phiếu chính xác hơn.

• Dịch vụ khách hàng: Mamba có thể vận hành chatbot theo dõi các cuộc hội thoại dài, nhờ đó cải thiện giao tiếp với khách hàng.

Chặng đường phía trước cho Mamba

Mamba nổi bật như một ngọn hải đăng đổi mới, mở ra cơ hội mới để giải quyết các bài toán mô hình hóa chuỗi phức tạp. Sự ra mắt của nó là một bước tiến tới các hệ thống thông minh hơn, có khả năng mở rộng và hiệu quả hơn, có thể hiểu và xử lý chuỗi với độ sâu chưa từng có.

Dù Mamba đã chứng minh là một cột mốc kỹ thuật quan trọng, thành công của nó không chỉ phụ thuộc vào năng lực công nghệ – nghiên cứu hợp tác, tài nguyên chia sẻ và đóng góp mã nguồn mở đều đóng vai trò thiết yếu:

• Đóng góp mã nguồn mở: Khuyến khích các nhà nghiên cứu và nhà phát triển đóng góp vào codebase của Mamba có thể tạo ra các mô hình bền vững và linh hoạt hơn.

• Tài nguyên chia sẻ: Cộng đồng có thể tăng tốc tiến bộ bằng cách chia sẻ kiến thức, tài nguyên và các mô hình tiền huấn luyện.

• Nghiên cứu hợp tác: Sự hợp tác giữa các viện học thuật và doanh nghiệp có thể mở rộng năng lực của Mamba.

Kết luận

Mamba không chỉ mang lại một cải tiến nhỏ cho các mô hình chuỗi hiện tại; nó tái định nghĩa những gì có thể đạt được. Với sự ra đời của nó, lịch sử trí tuệ nhân tạo sẽ chứng kiến một chương mới, nơi những kém hiệu quả tính toán và hạn chế độ dài chuỗi dần trở nên lỗi thời.

Vài năm qua, chúng ta đã chứng kiến sự tiến hóa của AI từ RNN đến Transformer và nay là Mamba, mỗi bước tiến gần hơn tới việc hiện thực hóa AI có khả năng tư duy sâu và xử lý thông tin tương đương con người. Mamba thể hiện tinh thần cách mạng thúc đẩy lĩnh vực AI tiến lên với cách tiếp cận không gian trạng thái có chọn lọc và khả năng mở rộng theo thời gian tuyến tính.

Mamba đánh dấu khởi đầu cho một bước phát triển đầy hứa hẹn trong trí tuệ nhân tạo. Đây là một mô hình cho tương lai và sẽ tạo ra tác động đáng kể đến AI với tiềm năng gần như không giới hạn.

Để tiếp tục học, hãy xem:

• Cách xây dựng ứng dụng LLM với Langchain
• Cách huấn luyện một LLM với PyTorch