Nvidia dùng phép toán tuyến tính đơn giản thay thế quá trình chuyển giao AI tốn kém

Nvidia vừa công bố một kỹ thuật mang tên cross-model KV cache transfer, cho phép ánh xạ trực tiếp bộ đệm KV (Key-Value) đã được điền sẵn từ mô hình nguồn sang mô hình đích. Phương pháp này dựa trên phép hồi quy tuyến tính, không cần đến mạng nơ-ron sâu tốn kém, nhằm giải quyết bài toán chi phí và độ trễ khi chuyển đổi giữa các LLM trong các quy trình agentic dài hạn.

Khi một hệ thống AI agentic chuyển một nhiệm vụ từ mô hình nhỏ sang mô hình lớn hơn — hoặc ngược lại — chi phí phải trả là rất lớn: mô hình nhận phải tính toán lại toàn bộ cuộc hội thoại từ đầu, gây tốn tài nguyên tính toán và tăng độ trễ. Đây chính là điểm nghẽn lớn cho các doanh nghiệp đang xây dựng các hệ thống đa mô hình với khối lượng ngữ cảnh lớn qua nhiều lượt tương tác.

Vì sao việc chuyển đổi mô hình giữa chừng lại đắt đỏ đến vậy?

Để hiểu rõ vấn đề, cần nắm cách LLM xử lý bộ nhớ. Khi một LLM nhận lời nhắc, nó phải thực hiện giai đoạn “prefill” — phép tính truyền xuôi ban đầu để tính toán các khóa (keys) và giá trị (values) cho tất cả token đầu vào, từ đó lấp đầy bộ đệm KV.

Sau đó, mô hình bước vào giai đoạn “decode” để sinh ra các token tiếp theo trong chuỗi. Trong giai đoạn này, mô hình đọc từ bộ đệm KV để dự đoán token mới, tránh phải đánh giá lại toàn bộ lịch sử hội thoại cho từng token.

Trong các phiên đối thoại nhiều lượt hoặc tác vụ agentic dài, ngữ cảnh ngày càng lớn dần. Vì chi phí tính toán giai đoạn prefill tỷ lệ thuận với cả kích thước mô hình lẫn độ dài đầu vào, việc xử lý các phiên dài trở nên tốn kém và dễ gây độ trễ nếu bộ đệm KV bị vô hiệu hóa.

Sự vô hiệu hóa này xảy ra khi hệ thống AI cố đổi mô hình giữa phiên làm việc, chẳng hạn như chuyển sang mô hình lớn hơn để xử lý bước suy luận phức tạp, hoặc hạ xuống mô hình nhỏ hơn để tiết kiệm chi phí. Vì các LLM khác nhau có kiến trúc khác nhau, chúng cần các định dạng cache đầu vào khác nhau.

Hệ quả là mọi lần chuyển mô hình đều buộc mô hình tiếp nhận trả lại toàn bộ chi phí prefill từ đầu để tính toán lại bộ đệm KV cho toàn bộ ngữ cảnh đã tích lũy.

Ánh xạ bộ nhớ giữa các mô hình mà không cần bắt đầu lại

Nhóm nghiên cứu tại Nvidia đã tìm hiểu khả năng truyền bộ đệm KV xuyên mô hình nhằm giúp nhà phát triển chuyển đổi bộ đệm KV của một mô hình sang đúng định dạng mà mô hình khác mong đợi, mà không phải chạy lại prefill.

Nếu giải quyết thành công, kỹ thuật này mang lại lợi ích theo cả hai hướng:

• Truyền từ mô hình nhỏ sang mô hình lớn giúp nâng cao chất lượng đầu ra. Chẳng hạn, một mô hình nhỏ giá rẻ xử lý các phần thường nhật trong quy trình agentic nhưng gặp khó với bài toán suy luận phức tạp, thì ánh xạ bộ đệm KV sang mô hình lớn hơn giúp tiếp tục quy trình liền mạch.

• Truyền từ mô hình lớn sang mô hình nhỏ giúp giảm chi phí tính toán. Một mô hình lớn có năng lực cao có thể dùng để xử lý system prompt cồng kềnh hoặc tổng hợp PDF dày đặc ở đầu phiên. Sau khi hoàn tất phần nặng nhất, bộ đệm KV của phiên được ánh xạ xuống mô hình nhỏ hơn, tiết kiệm hơn để xử lý các lượt hội thoại nhanh sau đó.

Những nỗ lực trước đây nhằm giải quyết bài toán truyền bộ đệm KV gặp phải một số hạn chế chính, như đòi hỏi phải huấn luyện dựa trên gradient (gradient-based training) đắt đỏ hoặc các ràng buộc kiến trúc rất nghiêm ngặt.

Trong nghiên cứu ban đầu này, các tác giả giới hạn phạm vi trong các cặp mô hình cùng họ (within-family), ví dụ chuyển giữa các phiên bản có kích thước khác nhau trong gia đình Qwen, Llama hoặc Ministral. Các mô hình này dùng chung tokenizer, dữ liệu huấn luyện và phong cách kiến trúc cốt lõi, nhưng khác nhau về kích thước và độ sâu.

Điểm mấu chốt trong nghiên cứu của Nvidia là cấu trúc bộ đệm KV xuyên mô hình mang tính tuyến tính đáng kể, nghĩa là có thể thực hiện ánh xạ bằng các thủ thuật đại số đơn giản mà không cần huấn luyện mạng nơ-ron lớn. Thử nghiệm trên cặp Qwen3 từ 14 tỷ tham số lên 32 tỷ tham số cho thấy một phép hồi quy tuyến tính từ một lớp nguồn sang một lớp đích có thể phục hồi 56% phương sai (variance) của khóa và 32% phương sai của giá trị trong mô hình đích. Khi kết hợp nhiều lớp nguồn, các con số này tăng lên lần lượt 79% và 65%.

Kiến trúc bộ ánh xạ: ba thành phần chính

Để biến mối quan hệ tuyến tính này thành một hệ thống thực dụng, nhóm nghiên cứu đã thiết kế một bộ ánh xạ ridge regression (hồi quy ridge) khép kín theo từng đầu (per-head) với ba thành phần chính:

• Hồi quy ridge theo từng đầu: Thay vì dùng deep learning phức tạp, họ khớp một hồi quy tuyến tính đơn giản với tập dữ liệu hiệu chuẩn nhỏ chỉ vài trăm chuỗi văn bản. Kỹ thuật này giải bài toán đường hồi quy tối ưu cho từng attention head một cách độc lập.

• Chọn lớp nguồn chéo: Vì mô hình nguồn và mô hình đích có số lớp khác nhau, bộ ánh xạ đánh giá và chọn các lớp nguồn có khả năng dự đoán tốt nhất để đưa vào từng lớp đích cụ thể. Như vậy, hệ thống chỉ lấy những phần bộ nhớ hữu ích nhất từ mô hình cũ để xây dựng bộ nhớ mới.

• Ánh xạ không gian nội dung: Trước khi dịch dữ liệu, bộ ánh xạ gỡ bỏ mã hóa RoPE (Rotary Position Embedding) — cơ chế tiêu chuẩn áp dụng phép xoay toán học phụ thuộc vị trí lên dữ liệu để mô hình hiểu được thứ tự token. Việc gỡ bỏ RoPE giúp bộ ánh xạ tổng quát hóa cho các chuỗi có độ dài vượt ra ngoài dữ liệu huấn luyện.

Thử nghiệm thực chiến bộ ánh xạ tuyến tính

Để kiểm chứng hiệu quả, các nhà nghiên cứu đánh giá quy trình truyền trên sáu gia đình mô hình “matched-KV” — tức mô hình nguồn và mô hình đích có cùng số đầu KV và cùng kích thước chiều mỗi đầu, đặc trưng cho các phiên bản khác kích thước trong cùng gia đình.

Các gia đình mô hình gồm Qwen3, Llama 3.1 và Ministral 3, với các thử nghiệm truyền trên kích thước từ 3 tỷ đến 70 tỷ tham số — bao gồm bước nhảy tham số lớn gấp 8,8 lần từ Llama 3.1 8B lên 70B.

Để phủ rộng nhiều loại tác vụ, họ đánh giá trên năm chuẩn đo độ chính xác cốt lõi (ARC-Challenge, HellaSwag, WinoGrande, MMLU và GSM8K), chuẩn perplexity về mô hình ngôn ngữ trên WikiText-2 và tác vụ hội thoại nhiều lượt CoQA. Tập dữ liệu hiệu chuẩn chỉ vỏn vẹn 500 chuỗi văn bản, mỗi chuỗi 1.024 token.

Kết quả đáng chú ý:

• Với bốn trong sáu cặp được thử nghiệm, bộ hồi quy ridge tuyến tính khép kín giữ được từ 73% đến 98% độ chính xác so với prefill tiêu chuẩn của mô hình đích — bao gồm cả bước nhảy khổng lồ Llama 3.1 8B → 70B, giữ được 72,8% độ chính xác.

• Bộ ánh xạ chạy nhanh hơn từ 2,7 đến 25 lần so với re-prefill. Ví dụ, khi dịch bộ đệm KV 32.768 token từ Qwen3 14B sang 32B, quá trình truyền chỉ mất 278 mili giây, so với gần 7 giây cho re-prefill tiêu chuẩn.

• Hệ thống duy trì độ ổn định cao trên các tác vụ nhiều bước. Khi thử nghiệm với hội thoại nhiều lượt, mức trôi (drift) giữa baseline của mô hình đích và cache truyền qua vẫn cực nhỏ sau 10 lượt, chứng tỏ nó sẽ không gây lỗi dây chuyền trong các phiên agentic dài.

Tuy nhiên, cách tiếp cận tuyến tính đơn giản gặp giới hạn ở một số cặp mô hình cụ thể. Với hai cấu hình Ministral, bộ ánh xạ tuyến tính suy giảm mạnh vì phép khớp tuyến tính không thể ngoại suy ra ngoài dữ liệu hiệu chuẩn. Để khắc phục, nhóm nghiên cứu thay bằng một mạng perceptron đa tầng (MLP) phi tuyến với hai lớp ẩn 1.024 đơn vị, huấn luyện trên cùng dữ liệu. Điều này tăng thêm chi phí độ phức tạp và huấn luyện, nhưng đã phục hồi độ chính xác lên trên 90%.

Bài toán lớn hơn một bài báo có thể giải quyết

Việc giới thiệu cross-model KV cache transfer nằm trong một nỗ lực rộng khắp ngành nhằm giải quyết điểm nghẽn KV cache — một trong những rào cản chính khi mở rộng quy mô AI doanh nghiệp. Khi các nhà phát triển đẩy LLM xử lý tài liệu hoặc mã nguồn khổng lồ và thực hiện các tác vụ suy luận kéo dài, quản lý tầng bộ nhớ này đang trở nên quan trọng ngang bằng chính các mô hình.

Trong năm qua, giới nghiên cứu đã tấn công bài toán tính toán và bộ nhớ này từ nhiều góc độ:

• Nvidia giới thiệu kỹ thuật dynamic memory sparsification (DMS), trực tiếp loại bỏ các token ít quan trọng khỏi bộ đệm KV để cắt giảm chi phí suy luận tới 8 lần.

• MIT phát triển kỹ thuật Attention Matching, nén bộ đệm KV tới 50 lần mà không suy giảm chất lượng.

• Nvidia cũng giới thiệu KV Cache Transform Coding (KVTC), mượn ý tưởng nén phương tiện truyền thông để giảm bộ nhớ 20 lần mà không thay đổi trọng số mô hình.

• Ngoài nén, các nhà nghiên cứu còn tấn công chi phí tính toán khi truy xuất bộ nhớ. Các trình tối ưu như IndexCache loại bỏ các phép tính lớp thừa, giúp rút ngắn đáng kể thời gian đến token đầu tiên (time-to-first-token) trong ứng dụng ngữ cảnh dài. Các mô hình như DeepSeek và dòng GLM cũng đang tối ưu KV cache thông qua đổi mới kiến trúc.

Khi các hệ thống AI đảm nhận những tác vụ dài hơi và có kiến trúc phức tạp hơn, hạ tầng bộ nhớ bên dưới đang trở nên quan trọng ngang bằng chính các mô hình. Cross-model KV cache transfer trao cho nhà phát triển thêm một công cụ để kiểm soát chi phí suy luận khi mở rộng quy mô các hệ thống agentic đa mô hình — một bước tiến đáng chú ý cho các doanh nghiệp đang muốn tối ưu hóa vận hành AI của mình.