Bạn hoàn toàn có thể tự nghĩ ra kiến trúc Transformer từ ký ức

Một bài viết kỹ thuật giải thích cách các lớp trong mạng nơ-ron hoạt động như những bộ lưu trữ ký ức, từ lớp tuyến tính, MLP cho đến self-attention. Qua đó, tác giả chỉ ra rằng Transformer không phải là phát minh quá xa vời mà có thể suy luận được từ những nguyên lý cơ bản về bộ nhớ kết hợp.

Các mô hình ngôn ngữ lớn ngày nay có thể nhớ được những sự kiện hiếm gặp ở khắp mọi nơi trên thế giới. Chúng ta đã quá quen với việc hỏi mô hình bất cứ điều gì đến mức coi khối kiến thức khổng lồ của chúng là điều hiển nhiên. Nhưng bằng cách nào mô hình nhớ được nhiều thông tin đến vậy, và chúng có thể lưu trữ được bao nhiêu?

Một bài viết kỹ thuật gần đây trên alexwang.ai đã trình bày một cách tiếp cận trực quan thú vị: coi các lớp trong mạng nơ-ron — từ lớp tuyến tính, MLP cho đến self-attention — như những dạng bộ lưu trữ ký ức khác nhau. Từ góc nhìn này, ta có thể rút ra các ước lượng về dung lượng bộ nhớ của từng loại lớp.

Bộ nhớ kết hợp là gì?

Tác giả định nghĩa bộ lưu trữ ký ức là một hàm $f$ nắm giữ $M$ ký ức, mỗi ký ức là một cặp khóa – giá trị: khóa $k_i$ dùng để tra cứu dựa trên độ tương đồng, còn giá trị $v_i$ là thứ ta muốn truy xuất. Khi truy vấn bằng một vector $x$, bộ nhớ sẽ trả về $f(x)$. Một bộ nhớ tốt sẽ trả về $v_i$ khi ta truy vấn bằng $k_i$, hoặc bằng một vector tương tự $k_i$. Trong tài liệu học thuật, bộ nhớ truy xuất theo độ tương đồng như vậy được gọi là bộ nhớ kết hợp (associative memory).

Điểm mấu chốt: dù có thể lưu vô số sự kiện, nhưng nếu không truy xuất được khi cần thì cũng vô ích. Càng lưu nhiều ký ức, chúng càng gây nhiễu lẫn nhau khi truy xuất. Vì vậy, bài viết tập trung vào một đại lượng duy nhất: tỷ lệ nhiễu trên tín hiệu — đo mức độ các ký ức khác cản trở việc truy xuất. Số lượng ký ức mà bộ nhớ có thể giữ trước khi nhiễu lớn bằng tín hiệu chính là dung lượng của nó.

Từ outer product đến self-attention

Hành trình trong bài viết đi từ bộ nhớ đơn giản nhất đến phức tạp nhất:

Bộ nhớ outer product (tuyến tính): chỉ đơn giản cộng dồn các tích ngoài $W = \sum_i v_i k_i^T$. Dung lượng chỉ khoảng $d_k$ ký ức, nơi $d_k$ là số chiều của khóa. Với khóa 64 chiều, chỉ nhớ được khoảng 64 ký ức trước khi nhiễu lấn át tín hiệu.
Bộ nhớ tuyến tính tối ưu (least-squares): thay vì ghi ký ức một cách mù quáng, ta chọn trọng số $W$ sao cho việc truy xuất khớp với tất cả ký ức cùng lúc. Cách này "tẩy trắng" các khóa, loại bỏ phần dư thừa, nâng dung lượng lên gấp đôi: khoảng $2d_k$.
MLP như bộ nhớ phi tuyến: nếu ta đưa mỗi độ tương đồng qua một hàm phi tuyến $\sigma$ trước khi tính tổng có trọng số, ta có ngay một MLP hai lớp. ReLU giúp loại bỏ các khóa có độ tương đồng âm; ReLU bình phương giúp nén mạnh các khóa ít tương đồng. Với ReLU bình phương, dung lượng nhảy vọt lên khoảng $2d_k^2/3$ — với $d_k = 64$, tức khoảng 2.700 ký ức.
Self-attention như bộ nhớ phi tham số: thay vì chọn một lũy thừa cố định, ta dùng hàm mũ $e^{\beta s}$, vốn là tổng có trọng số của mọi lũy thừa. Khi chuẩn hóa, ta thu được chính công thức softmax của self-attention. Dung lượng tăng theo cấp số nhân: với $\beta = \sqrt{d_k}$, có thể lên tới hơn 1,2 triệu ký ức.

Điều đáng chú ý: self-attention không dùng tham số huấn luyện cho khóa và giá trị, mà tính chúng từ chính các token trong ngữ cảnh. Mỗi token lưu một ký ức, nên số ký ức chính là độ dài ngữ cảnh.

Bảng tóm tắt dung lượng các loại bộ nhớ

Loại bộ nhớ	Dung lượng	Ở $d_k = 64$
Outer product	$d_k$	64
Tuyến tính tối ưu	$2d_k$	128
ReLU MLP	$2d_k$	128
ReLU bình phương (GLU MLP)	$2d_k^2/3$	~2.700
SwiGLU MLP	~$2d_k^2/3$	~2.900
ReLU lập phương	$2d_k^3/15$	~35.000
Self-attention, $\beta = \sqrt{d_k}$	$e^{2\sqrt{d_k} - 2}$	~1,2 triệu
Self-attention, $\beta \propto \log M$	không đạt giới hạn	hơn 2 triệu

Ví dụ thực tế: mẹo mở rộng ngữ cảnh dài

Vì $\beta$ quyết định mức độ triệt tiêu các khóa ít tương đồng, ta có thể tăng $\beta$ khi số cặp khóa – giá trị tăng lên. Bài viết điểm qua một số kỹ thuật đang được dùng trong các mô hình thực tế:

Su (2021) / Scalable-Softmax: nhân logit với $\log M / \log L$, giữ entropy của attention không đổi khi $M$ tăng.
LogN-Scaling của Qwen: áp dụng hệ số tương tự nhưng chỉ khi ngữ cảnh dài hơn độ dài tham chiếu $L$.
YaRN (dùng trong Qwen2.5-1M): nhân logit với $(0.1 \log(M/L) + 1)^2$, tăng nhanh hơn theo $(\log M)^2$.

Kết quả: khi $\beta$ tỉ lệ với $\log M$, số ký ức có thể lưu gần như không giới hạn, với các khóa phân biệt được.

Góc nhìn dành cho bạn đọc Việt Nam

Với cộng đồng AI và lập trình viên Việt Nam đang ngày càng tiếp cận sâu với các mô hình ngôn ngữ lớn, cách nhìn "mỗi lớp là một bộ nhớ" mang lại nhiều giá trị thực tiễn:

Giúp hiểu rõ hơn tại sao các kiến trúc như SwiGLU, GLU hay attention lại được thiết kế như hiện nay, thay vì chỉ học thuộc công thức.
Cung cấp trực giác để đánh giá chi phí và giới hạn khi mở rộng ngữ cảnh cho các mô hình tiếng Việt — vốn thường cần ngữ cảnh dài do đặc thù ngôn ngữ.
Mở ra hướng nghiên cứu về bộ nhớ ngắn hạn cập nhật trong lúc suy luận (fast weights, Mamba, DeltaNet, Titans…) và bộ nhớ ngoài trọng số (kNN-LM, RETRO, Memorizing Transformers).

Tác giả kết luận bằng một suy nghĩ khá thú vị: ông luôn bị cuốn hút bởi ký ức, có lẽ vì cảm nhận rõ giới hạn của trí nhớ chính mình. Biết đâu, bằng cách hiểu cách máy móc ghi nhớ, ta cũng học được điều gì đó về bản thân.