Phương pháp định tuyến học được của Apple giảm độ trễ RAG bằng cách thay thế tìm kiếm trọng tâm

Các nhà nghiên cứu tại Apple Machine Learning Research đã giới thiệu một phương pháp mới nhằm đẩy nhanh quá trình tìm kiếm tích vô hướng tối đa (MIPS), động lực cốt lõi đằng sau các hệ thống gợi ý và sinh tạo tăng cường truy xuất (RAG). Được công bố trong một bài báo trình bày tại ICML 2026, nghiên cứu này thay thế chỉ mục hình học truyền thống bằng các mạng nơ-ron nhỏ, được huấn luyện ngoại tuyến, để học phân phối cụ thể của các truy vấn từ người dùng. Cách tiếp cận này giảm đáng kể độ trễ đồng thời cải thiện độ chính xác recall so với định tuyến dựa trên trọng tâm tiêu chuẩn.

Điều gì đã xảy ra

Tìm kiếm tích vô hướng tối đa tốn nhiều tài nguyên tính toán khi cơ sở dữ liệu chứa hàng triệu vector chiều cao. Các phương pháp xấp xỉ truyền thống tăng tốc quy trình này bằng cách gom cụm các khóa và sử dụng trọng tâm để định tuyến truy vấn, nhưng những giả định hình học này thường thất bại vì cấu trúc của truy vấn từ người dùng khác biệt so với các tài liệu được chọn lọc mà họ đang tìm kiếm. Nhóm của Apple nhận thấy rằng các phương pháp chung, không phụ thuộc vào truy vấn, bỏ qua phân phối ổn định và đặc trưng rõ ràng của các truy vấn trong môi trường sản xuất, khiến các lợi ích về hiệu suất chưa được khai thác.

Để giải quyết vấn đề này, các nhà nghiên cứu đã huấn luyện ngoại tuyến các mạng nơ-ron nhỏ trên một cơ sở dữ liệu khóa cố định. Những mạng này học cách ánh xạ các truy vấn đến kết quả khớp tốt nhất một cách hiệu quả hơn so với việc so sánh toàn bộ hoặc định tuyến đơn giản dựa trên trọng tâm. Nhóm đã đánh giá hai kiến trúc, SupportNet và KeyNet, sử dụng benchmark Natural Questions, ghép nối các câu hỏi kiến thức chung với khoảng 2,6 triệu đoạn văn bản duy nhất. Kết quả cho thấy các bộ định tuyến học được có thể đạt độ recall cao hơn với độ trễ thấp hơn bằng cách tận dụng hình học cụ thể của phân phối truy vấn thay vì dựa vào các xấp xỉ hình học rộng rãi.

Cách hoạt động

Insight cốt lõi dựa trên việc xem xét MIPS dưới góc độ phân tích lồi. Hàm ánh xạ một truy vấn đến tích vô hướng lớn nhất của nó với bất kỳ khóa nào trong một tập hợp được gọi là hàm hỗ trợ. Hàm này là lồi, nghĩa là nó có thể được xấp xỉ bởi các mạng nơ-ron lồi theo đầu vào (ICNNs). Hơn nữa, gradient của hàm hỗ trợ này tại bất kỳ điểm truy vấn nào sẽ trả về vector khóa khớp tốt nhất. Bằng cách huấn luyện một mạng nơ-ron để mô phỏng hàm hỗ trợ này hoặc gradient của nó, hệ thống có thể dự đoán các kết quả khớp tối ưu với chi phí tính toán thấp hơn nhiều so với việc tìm kiếm toàn bộ cơ sở dữ liệu.

Các nhà nghiên cứu đã phát triển hai mô hình cụ thể để khai thác cấu trúc toán học này. SupportNet học trực tiếp hàm hỗ trợ cho từng cụm dữ liệu, cho phép nó trả lời mức độ khớp tốt nhất của khóa bên trong một cụm với một truy vấn. Điều này cho phép định tuyến chính xác hơn so với chỉ dựa vào khoảng cách trọng tâm. Ngược lại, KeyNet học cách xuất trực tiếp vector khóa khớp tốt nhất. Nó hoạt động như một bộ viết lại truy vấn, biến đổi truy vấn gốc thành một vector gần với mục tiêu hơn trong không gian embedding trước khi chuyển nó sang một chỉ mục tìm kiếm tiêu chuẩn.

Chi tiết chính

• Nghiên cứu được trình bày tại ICML 2026 bởi Apple Machine Learning Research.
• Các mô hình SupportNet đạt độ recall định tuyến 95% nhanh gấp 1,9 lần so với định tuyến trọng tâm (1,74 ms so với 3,36 ms).
• KeyNet giảm thời gian đạt đến độ recall 99% lên tới 8 lần so với tìm kiếm truy vấn thô (0,23 ms so với 1,8 ms).
• Các mô hình được huấn luyện trên tập dữ liệu Natural Questions chứa khoảng 2,6 triệu đoạn văn bản duy nhất.
• SupportNets được kiểm thử với kích thước dao động từ 1% đến 20% kích thước cơ sở dữ liệu, với độ sâu 16 lớp.
• Các truy vấn được biến đổi bởi KeyNet đạt độ recall 99% trong khi chỉ quét 0,009% cơ sở dữ liệu, so với 0,10% đối với các truy vấn thô.

Tại sao điều này quan trọng

Đối với các kỹ sư xây dựng hệ thống RAG, khoảng cách giữa các câu hỏi ngắn của người dùng và các embedding tài liệu dài tạo ra một nút thắt cổ chai dai dẳng. Các chỉ mục tệp đảo ngược tiêu chuẩn dựa vào trọng tâm không tính đến sự phân tán của các khóa trong một cụm, dẫn đến việc bỏ sót các lân cận và lãng phí tài nguyên tính toán không cần thiết. Bằng cách học phân phối truy vấn cụ thể, các nhà phát triển có thể tái phân bổ tài nguyên tính toán hiệu quả hơn, đảm bảo rằng các tài liệu liên quan nhất được truy xuất mà không cần tìm kiếm toàn bộ.

Phương pháp này cũng cung cấp một con đường thực tế để tích hợp với hạ tầng hiện có. KeyNet có thể hoạt động như một bước tiền xử lý cho các chỉ mục tìm kiếm xấp xỉ tiêu chuẩn như FAISS, không yêu cầu thay đổi hệ thống lưu trữ cơ bản. Điều này có nghĩa là các nhóm có thể đạt được giảm độ trễ đáng kể và cải thiện độ chính xác bằng cách thêm một lớp mạng nơ-ron nhẹ phía trước pipeline truy xuất hiện tại, thay vì xây dựng lại toàn bộ kiến trúc tìm kiếm vector của họ.

Bạn có thể làm gì

• Đánh giá xem phân phối truy vấn trong môi trường sản xuất của bạn có ổn định và đặc trưng rõ ràng hay không, vì phương pháp này dựa vào huấn luyện ngoại tuyến trên dữ liệu đại diện.
• Cân nhắc triển khai SupportNet để thay thế định tuyến dựa trên trọng tâm trong chỉ mục tệp đảo ngược của bạn nếu lỗi định tuyến ảnh hưởng đến độ recall.
• Sử dụng KeyNet như một lớp viết lại truy vấn trước chỉ mục tìm kiếm vector hiện có để giảm số lượng khóa được quét trong quá trình truy xuất.
• Huấn luyện các mạng nơ-ron nhỏ ngoại tuyến bằng cách hồi quy trên các truy vấn mẫu và các kết quả khớp tốt nhất thực tế (ground-truth) để phân bổ chi phí tính toán.
• Benchmark sự đánh đổi giữa kích thước mô hình và độ trễ, lưu ý rằng ngay cả các mô hình nhỏ (1% kích thước cơ sở dữ liệu) cũng mang lại tốc độ tăng đáng kể so với các phương pháp cơ sở.
• Kiểm tra loss nhất quán (consistency loss) được suy ra từ định lý Euler khi huấn luyện các mô hình KeyNet để giữ cho các dự đoán có nền tảng hình học vững chắc.