Từ GPT-2 đến Kimi K3: Cách các mô hình ngôn ngữ học cách quản lý bộ nhớ

Sáu kiến trúc, một ý tưởng: từ bộ nhớ đệm KV (KV cache) đến sự kết hợp lai giữa bộ nhớ đệ quy, cơ chế chú ý toàn cục và các chuyên gia thưa thớt (sparse experts) của Kimi K3, mỗi kiến trúc là một cách thông minh hơn để quản lý bộ nhớ hạn chế. Tóm tắt: Năm 2019, GPT-2 (phiên bản nhỏ) có 124 triệu tham số. Đến năm 2026, Kimi K3 có 2,8 nghìn tỷ tham số, gấp khoảng 22.580 lần. Tuy nhiên, bước nhảy vọt này không chỉ nằm ở kích thước. Một phần lớn câu chuyện ít được nhắc đến hơn là cách các mô hình học cách ghi nhớ, cập nhật và truy xuất thông tin hiệu quả hơn. Bài viết này sẽ theo dõi một mạch ý tưởng đó qua sáu kiến trúc, bắt đầu với mô hình giúp ý tưởng này dễ hình dung nhất, GPT-2. Bài viết này được truyền cảm hứng từ bài phân tích chuyên sâu xuất sắc của Ali Taha "22.580: GPT-2 đến Kimi K3, giải thích" trên blog của Baseten (cấu trúc sáu giai đoạn và phép so sánh 22.580 lần tuân theo bài viết này một cách chặt chẽ), và từ bài báo khoa học về Kimi K3, "Kimi K3: Trí tuệ Tiên phong Mở". 1. GPT-2: Kiến trúc Transformer tiêu chuẩn GPT-2 là một mô hình Transformer chỉ có bộ giải mã (decoder-only Transformer), kiến trúc từ bài báo "Attention Is All You Need" năm 2017. "Chỉ có bộ giải mã" có nghĩa là nó thực hiện chính xác một nhiệm vụ: dự đoán mã thông báo (token) tiếp theo từ các mã thông báo đứng trước nó. Không có giai đoạn "bộ mã hóa" (encoder) riêng biệt để đọc câu nguồn như cách một mô hình dịch thuật làm; GPT-2 đơn giản là tiếp tục một chuỗi. Khi mô hình xử lý bất kỳ vị trí nào, nó chỉ được phép nhìn vào vị trí đó và các vị trí đứng trước nó, không bao giờ nhìn vào tương lai. GPT-2 hoạt động theo cơ chế tự hồi quy (autoregressively): nó dự đoán một mã thông báo, thêm mã thông báo đó vào chuỗi, sau đó sử dụng toàn bộ chuỗi đã cập nhật để dự đoán mã thông báo tiếp theo, và lặp lại. Mỗi dự đoán đều dựa trên tất cả những gì đã xảy ra trước đó: lời nhắc (prompt) ban đầu cộng với mỗi mã thông báo đã được tạo ra cho đến nay. Hãy xem xét lời nhắc "The cat sat on the". Để chọn từ tiếp theo, mô hình xử lý năm mã thông báo đó và, từ biểu diễn tại mã thông báo cuối cùng (the), dự đoán một từ tiếp nối có khả năng, ví dụ "mat". Mã thông báo đó được thêm vào, mô hình chạy lại (giờ đây có thể chú ý đến cả sáu mã thông báo), và cứ thế tiếp diễn. Có hai điều đáng chú ý. Thứ nhất, GPT-2 không nén quá khứ vào một bản tóm tắt duy nhất đang chạy như các mạng thần kinh hồi quy (RNN) cũ hơn; nó giữ một biểu diễn riêng biệt, có thể truy cập trực tiếp cho mỗi mã thông báo trong cửa sổ. Thứ hai, nó xử lý lời nhắc trong một lượt song song ("prefill"), sau đó tạo ra từng mã thông báo một ("decoding"). Cơ chế cho phép một mã thông báo "đọc" các mã thông báo khác được gọi là cơ chế chú ý (attention). Tại mỗi vị trí, mô hình đánh giá mức độ liên quan của mỗi mã thông báo trước đó (và chính vị trí hiện tại), sau đó kết hợp thông tin của chúng theo tỷ lệ với các điểm số đó, để các mã thông báo liên quan đóng góp nhiều hơn và các mã thông báo không liên quan đóng góp ít hơn. Chính biểu diễn tại vị trí hiện tại mang cái nhìn tổng hợp có trọng số này về phía trước để dự đoán mã thông báo tiếp theo. Cái nhìn tổng hợp có trọng số này là thao tác quan trọng nhất trong một mô hình Transformer, và đây là thao tác mà các kiến trúc sau này nỗ lực nhất để làm cho nó rẻ hơn. Hình 1. Dự đoán từ sau "The cat sat on the." Cơ chế chú ý chạy từ **mã thông báo đầu vào cuối cùng** (the), mã thông báo này có thể nhìn vào chính nó và tất cả các mã thông báo trước đó; biểu diễn mà nó tạo ra là thứ dự đoán mã thông báo tiếp theo, "mat". Lưu ý rằng "mat" không có biểu diễn riêng cho đến khi nó được tạo ra và đưa trở lại vào. Hiệu quả tiềm ẩn Sự chú ý tại một vị trí nhất định cần một gói thông tin nhỏ từ chính nó và mọi token trước đó. (Chính xác gói thông tin đó chứa gì sẽ được đề cập sau.) Nếu mô hình xây dựng lại gói thông tin đó cho mỗi token đã qua từ đầu ở mỗi bước, nó sẽ lặp lại gần như cùng một công việc nhiều lần, và công việc dư thừa sẽ tích tụ nhanh hơn khi văn bản dài ra. Hình 2. Mỗi ô màu là key/value của một token đang được tính toán trong một bước nhất định. Chỉ có một ô thực sự mới mỗi bước (màu xanh lá cây); phần còn lại là các phép tính toán lại các phép chiếu key/value đã thực hiện (màu đỏ). Việc xây dựng lại chúng từ đầu tốn khoảng 1 + 2 + 3 + 4 = 10 phép tính chiếu chỉ để phát ra 4 token, và tăng theo cấp số nhân n². (Điều này tính công việc chiếu key/value, không tính việc so sánh sự chú ý, sẽ được đề cập dưới đây.) Giải pháp: bộ nhớ đệm KV (KV cache) Giải pháp bắt đầu bằng việc đặt tên cho "gói thông tin" đó. Bên trong mỗi lớp chú ý (attention layer), biểu diễn ẩn (hidden representation) của mỗi token được truyền qua ba phép biến đổi nhỏ đã học để tạo ra ba vector: Q, truy vấn (query): "tôi đang tìm kiếm điều gì?" K, khóa (key): "tôi nên được người khác tìm thấy như thế nào?" V, giá trị (value): "tôi truyền thông tin gì nếu tôi được chọn?" Sự chú ý so sánh Q của token hiện tại với K của mọi token trước đó (và của chính nó) để quyết định mức độ lấy thông tin từ mỗi token, sau đó trộn các V của chúng một cách phù hợp. K và V được gọi là các biểu diễn trung gian (intermediate representations) vì chúng không phải là đầu vào (các từ) và không phải là đầu ra cuối cùng (dự đoán). Chúng là các vector được tính toán bên trong lớp chú ý, trên đường từ đầu vào đến đầu ra. Và đây là điểm mấu chốt của toàn bộ giải pháp: K và V của một token đã qua không bao giờ thay đổi khi một token mới đến. Câu "The cat sat" tạo ra cùng K và V cho từ "cat" dù từ tiếp theo là "on", "quietly" hay "down". Tại sao? Sự chú ý có tính nhân quả (causal): biểu diễn nội bộ của một token chỉ được xây dựng từ chính nó và các token trước đó (mặt nạ chặn mọi việc nhìn trước), và K và V được tính toán trực tiếp từ biểu diễn đó. Vì vậy, một token đến sau không có cách nào để quay lại và thay đổi chúng. Vì các vector đó ổn định, không cần phải tính toán lại chúng. K và V của mỗi token được tính toán một lần và lưu trữ; nơi lưu trữ đó là bộ nhớ đệm KV (KV cache). Mỗi bước mới sau đó tính toán một Q, K và V mới cho riêng token hiện tại, thêm K và V của nó vào bộ nhớ đệm, và tái sử dụng mọi thứ đã lưu, điều này loại bỏ tam giác tính toán lại màu đỏ trong Hình 2. Mặc dù bộ nhớ đệm KV giúp mô hình không phải xây dựng lại các khóa và giá trị đã qua (sự lãng phí trong Hình 2), nhưng nó không giúp mô hình không phải đọc chúng. Để tạo ra mỗi token mới, mô hình vẫn phải quét toàn bộ bộ nhớ đệm: so sánh token mới với mọi khóa đã lưu cho đến nay và trộn vào mọi giá trị. Và bộ nhớ đệm đó chỉ tăng lên. Việc ghi token thứ 10 có nghĩa là quét khoảng 10 mục; token thứ 1.000 quét khoảng 1.000 mục; token thứ một triệu.