PSSA: Mô hình ngôn ngữ phi Transformer viết hoàn toàn bằng Rust

Trong bối cảnh các mô hình ngôn ngữ lớn ngày càng phụ thuộc vào kiến trúc Transformer và các framework học máy cồng kềnh, một dự án mã nguồn mở mang tên PSSA (Plastic State-Space Architecture) đang thu hút sự chú ý của cộng đồng lập trình viên khi chứng minh rằng vẫn còn những hướng đi hoàn toàn khác.

PSSA là gì?

PSSA là một mô hình ngôn ngữ nhỏ không phải Transformer. Nó đọc văn bản từng token một thông qua một lớp state-space hồi tiếp, duy trì một ngân hàng bộ nhớ episodic để tra cứu thông tin, và có khả năng tự viết lại một phần trọng số của chính mình trong lúc chạy.

Điểm đáng chú ý: toàn bộ dự án được viết bằng Rust từ con số không, không có PyTorch, không TensorFlow, không bất kỳ framework học máy nào bên dưới. Đại số tuyến tính được viết tay, kèm một nhánh CUDA cho huấn luyện và một tham chiếu CPU dạng scalar để kiểm tra từng gradient.

Sơ đồ kiến trúc PSSA

Khác biệt cốt lõi so với Transformer

Cách tiếp cận của PSSA khác căn bản ở điểm sau:

• Transformer chấm điểm mọi cặp token trong ngữ cảnh, nên chi phí mỗi bước tăng theo bình phương độ dài chuỗi, và toàn bộ ngữ cảnh phải được đọc lại ở mỗi bước.
• PSSA mang theo một trạng thái kích thước cố định xuyên suốt chuỗi trong một lượt từ trái sang phải, và tra cứu trong ngân hàng bộ nhớ thay vì đọc lại ngữ cảnh. Nhờ đó, chi phí tăng tuyến tính theo độ dài.

Kiến trúc này gồm bốn thành phần chính:

• Lõi state-space hồi tiếp: các ma trận trạng thái liên tục được học để mang thông tin về phía trước trong một trạng thái kích thước cố định.
• Ngân hàng bộ nhớ episodic: 512 khe với truy xuất kiểu hyperbolic (Poincaré), tìm kiếm top-4 có giới hạn, được ghi và đọc trong lúc chạy.
• Trọng số dẻo (plastic weights): cập nhật nhanh củng cố những gì hiệu quả, tính mới lạ thúc đẩy tăng trưởng, và một cổng refractory giới hạn tốc độ ghi đè để dữ liệu mâu thuẫn lặp lại ít gây hư hại.
• Hợp nhất dạng đóng: một bước hồi quy ridge gấp các cập nhật dẻo nhanh trở lại ma trận chuyển tiếp cơ sở — giống như giấc ngủ củng cố kiến thức trong ngày.

Kết quả thực nghiệm

Hai mô hình được huấn luyện trên cùng corpus, cùng tokenizer, cùng lịch trình optimizer, cùng seed và cùng số tham số: một là PSSA, một là Transformer tiêu chuẩn. Trên 12,7 triệu token WikiText-103 đã làm sạch:

• PSSA kết thúc ở cross-entropy huấn luyện 3,98
• Transformer dừng ở 4,43

Khoảng cách 0,45 nats tương ứng perplexity 53,7 so với 83,7. Đáng chú ý hơn, Transformer đã tiêu tốn toàn bộ ngân sách 12,7 triệu token mà vẫn chưa đạt được mức loss mà PSSA đã vượt qua chỉ quanh mốc 2 triệu token. Hai đường cong không bao giờ cắt nhau.

Đường cong loss huấn luyện của PSSA và Transformer

Khả năng tổng quát hóa trên dữ liệu chưa thấy

Loss huấn luyện chỉ cho biết mô hình khớp với luồng dữ liệu được đưa vào. Vì vậy, cả hai checkpoint được chấm điểm trên một lát cắt 198.939 token lấy từ phần corpus mà cả hai lần chạy chưa từng chạm tới:

Chỉ số	PSSA	Transformer
Cross-entropy	3,997	4,429
Perplexity	54,4	83,8
Độ chính xác token kế tiếp	24,1%	18,0%

Khoảng cách trên dữ liệu held-out là 0,43 nats, gần như trùng khớp với khoảng cách huấn luyện. Điều này cho thấy PSSA không ghi nhớ giỏi hơn, mà tổng quát hóa tốt hơn.

Đường cong loss trên tập held-out

Tốc độ sinh văn bản

Sinh 200 token trên cùng CPU, cùng prompt, cùng sampler:

	PSSA	Transformer
200 token	226 ms	2.735 ms
Tương quan	nhanh hơn 12 lần	mốc chuẩn

Lý do rất rõ ràng: mô hình hồi tiếp mang theo trạng thái kích thước cố định, nên chi phí mỗi token mới không tăng theo độ dài phần đã sinh trước đó. Transformer thì phải đọc lại toàn bộ ngữ cảnh ở mỗi bước.

Những điều cần thẳng thắn nhìn nhận

Tác giả dự án khá minh bạch về giới hạn của kết quả:

"Đây là các mô hình 1,5 triệu tham số trên 12,7 triệu token. Đó là một nguyên mẫu nghiên cứu, không phải đối thủ của bất kỳ thứ gì bạn từng nghe tới."

Chất lượng văn bản ở quy mô này còn kém với cả hai mô hình. PSSA sinh ra "a barget of the Prian Academy", còn Transformer tạo "a material circulation of the United States". So sánh ở đây là về hiệu quả học tập, không phải độ lưu loát.

Ngoài ra, so sánh tốc độ ở trên là CPU-với-CPU nên công bằng, nhưng các con số thông lượng huấn luyện sâu hơn trong tài liệu không khớp phần cứng và không nên đọc như một kết quả kiến trúc. Hai thí nghiệm vẫn chưa được đo: khả năng giữ lại kỹ năng cũ sau khi đổi corpus, và liệu việc loại bỏ ngân hàng bộ nhớ có làm thay đổi loss hay không.

Dùng thử và đóng góp

Cài đặt khá đơn giản với Rust toolchain:

git clone https://github.com/Sparticle62ops/pssa.git
cd pssa
cargo build --release
./target/release/oxide_ai_pssa

Chạy không kèm tham số sẽ hiện màn hình chính liệt kê mọi lệnh cùng checkpoint và corpus tìm thấy trong thư mục làm việc. Dự án cung cấp CLI đầy đủ cho huấn luyện, sinh văn bản, chat tương tác, đánh giá và benchmark, cùng khả năng huấn luyện theo chuỗi (chained training) để vượt qua giới hạn thời gian của phiên notebook.

Dự án đang cần hỗ trợ ở ba mảng: hiệu năng kernel, baseline hồi tiếp hiện đại để so sánh, và đánh giá ngoài next-token loss. Tác giả cũng nêu rõ: toàn bộ kết quả được huấn luyện trên notebook miễn phí với một GPU entry-level, và thứ có giá trị nhất mà cộng đồng có thể trao cho dự án là compute thực sự.

Góc nhìn cho cộng đồng công nghệ Việt Nam

Với các lập trình viên và nhà nghiên cứu AI tại Việt Nam — nơi nguồn lực GPU còn hạn chế — PSSA là một ví dụ đáng tham khảo về việc tối ưu hiệu quả trên phần cứng khiêm tốn. Việc một mô hình phi Transformer có thể học hiệu quả hơn trên CPU không chỉ là chuyện học thuật: nó mở ra khả năng thử nghiệm AI ngay trên máy cá nhân, không cần thuê hạ tầng đám mây đắt đỏ. Dự án cũng là minh chứng cho giá trị của việc hiểu sâu toán học thay vì chỉ dựa vào framework có sẵn.

Mã nguồn PSSA được công khai trên GitHub theo giấy phép của dự án, kèm tài liệu CLI, quy trình làm sạch corpus WikiText, và bộ kiểm thử tích hợp đầy đủ.