H Company ra mắt NeoMME: Bộ mã hóa đa phương thức đa ngôn ngữ mã nguồn mở

TL;DR (Tóm tắt)

Chúng tôi giới thiệu NeoMME, một dòng encoder đa phương thức đa ngôn ngữ với quy mô 260 triệu và 800 triệu tham số. Không giống như nhiều mô hình ngôn ngữ thị giác tạo sinh, NeoMME không sử dụng tháp thị giác (vision tower) tiền huấn luyện riêng biệt hay mô hình ngôn ngữ nhân quả (causal language model). Một Transformer hai chiều duy nhất sẽ xử lý cả token văn bản và các bản vá hình ảnh thô, và chúng tôi huấn luyện toàn bộ mô hình từ đầu với mục tiêu khuếch tán rời rạc có che (masked discrete-diffusion objective).

Chúng tôi đã tinh chỉnh (fine-tune) NeoMME cho tác vụ truy xuất tài liệu hình ảnh bằng phương pháp ảnh trang của ColPali. NeoMME-Retriever trả về các embedding tương tác muộn (late-interaction) và dày đặc (dense) chỉ trong một lần truyền xuôi (forward pass). Cả hai kích thước mô hình đều nằm trên đường biên Pareto của ViDoRe v3 xét về chỉ số nDCG@10 và kích thước mô hình. Với kích thước đầu vào hình ảnh 2048×2048 trên GPU NVIDIA L40S, mô hình 260M mã hóa khoảng 51 trang mỗi giây, nhanh gấp đôi thông lượng của ColModernVBERT. Việc gộp token phân cấp (hierarchical token pooling) và lượng tử hóa bất đối xứng (asymmetric quantization) giúp giảm dung lượng lưu trữ chỉ mục tương tác muộn từ khoảng 1,5 MB xuống còn 6 kB mỗi trang (nhỏ hơn 255 lần) trong khi vẫn giữ lại hơn 95% chỉ số nDCG@10 cơ sở.

NeoMME hiện đã có sẵn trên Hugging Face Transformers. Chúng tôi phát hành tất cả các checkpoint của mô hình theo giấy phép Apache 2.0.

Tại sao lại cần thêm một encoder đa phương thức nữa?

Nhiều hệ thống truy xuất tài liệu hình ảnh gần đây được chuyển đổi từ các mô hình ngôn ngữ thị giác tạo sinh đã được tiền huấn luyện. Một encoder thị giác tiền huấn luyện riêng biệt sẽ tạo ra các đặc trưng thị giác, sau đó một bộ chiếu (projector) sẽ ánh xạ chúng vào không gian đầu vào của mô hình ngôn ngữ. Một bộ giải mã nhân quả (causal decoder) sau đó sẽ xử lý các biểu diễn kết hợp giữa hình ảnh và văn bản. Các tác vụ truy xuất, phân loại và gắn nhãn token không tạo văn bản theo cơ chế tự hồi quy (autoregressive), vì vậy chúng không cần bộ giải mã nhân quả hay gánh nặng về tham số và tính toán của kiến trúc này.

ModernBERT đã mang đến những cải tiến về kiến trúc và huấn luyện hiệu quả cho các encoder hai chiều. Đối với truy xuất tài liệu hình ảnh, ModernVBERT đã áp dụng một encoder văn bản hai chiều theo phong cách ModernBERT trong khi vẫn giữ lại một tháp thị giác SigLIP2 tiền huấn luyện riêng biệt. Chúng tôi muốn tiến xa hơn nữa bằng cách thiết kế và huấn luyện một encoder đa phương thức mà không cần phải gánh chịu chi phí về tham số và tính toán của một VLM.

NeoMME (phát âm là "nee-oh-me", IPA /ˈniː.oʊ.mi/) là một encoder nền tảng đa phương thức, đa ngôn ngữ, tạo ra các biểu diễn vector cho văn bản và/hoặc hình ảnh đầu vào bằng cách sử dụng một Transformer encoder duy nhất. Nó không dựa trên bất kỳ tháp thị giác, encoder văn bản hay bộ giải mã văn bản tiền huấn luyện nào có sẵn.

Hình ảnh và văn bản sử dụng cùng một đường dẫn tính toán, vì vậy NeoMME có thể hỗ trợ việc tiền huấn luyện, tinh chỉnh, song song hóa và phục vụ (serving) trên cả hai phương thức một cách dễ dàng hơn.

Kiến trúc backbone của NeoMME

Một Transformer cho cả hình ảnh và văn bản

NeoMME có hai kích thước là 260M và 800M. Cả hai biến thể đều chia sẻ chung một kiến trúc:

Học từ hình ảnh thông qua văn bản bị che

Chúng tôi tiền huấn luyện NeoMME từ đầu như một bộ khử nhiễu văn bản khuếch tán rời rạc có che. Đối với mỗi ví dụ chỉ có văn bản, chúng tôi lấy mẫu tỷ lệ nhiễu đồng nhất trong khoảng từ 0 đến 1. Mỗi token văn bản đủ điều kiện sau đó sẽ được che độc lập theo tỷ lệ đó.

Các ví dụ đa phương thức sử dụng tỷ lệ nhiễu từ 0,3 đến 1. Các bản vá hình ảnh vẫn hiển thị trong khi NeoMME tái tạo lại văn bản bị che. Với việc che nhẹ, mô hình thường có thể khôi phục một từ bị thiếu chỉ từ văn bản xung quanh. Ví dụ, "cat" là một từ hoàn thiện hợp lý cho câu "The [MASK] sat on the mat," ngay cả khi không có hình ảnh. Tuy nhiên, việc che ở mức độ cao buộc mô hình phải học các mô tả dựa trên hình ảnh với rất ít hoặc không có tín hiệu từ các token văn bản đầu vào không bị che.

Quá trình tiền huấn luyện kết hợp văn bản đa ngôn ngữ, mã nguồn, toán học, hình ảnh tự nhiên và hình ảnh tài liệu. Mỗi mô hình xử lý khoảng 524 tỷ token đầu vào đã đóng gói, bao gồm 290 tỷ token từ các ví dụ chỉ có văn bản. Ngân sách văn bản này tương đối nhỏ so với ngân sách 2 nghìn tỷ token huấn luyện của ModernBERT. Do đó, chúng tôi đã chọn bộ tối ưu hóa NorMuon để cải thiện hiệu quả dữ liệu trong quá trình huấn luyện.

NeoMME-Retriever

Để có được đánh giá hạ nguồn (downstream evaluation) có ý nghĩa về backbone, chúng tôi tinh chỉnh NeoMME cho tác vụ truy xuất tài liệu hình ảnh bằng phương pháp ảnh trang do ColPali giới thiệu. Trong khi truy xuất dựa trên văn bản truyền thống bao gồm việc truy xuất các đoạn văn bản, NeoMME-Retriever xếp hạng các ảnh chụp màn hình trang tài liệu và bỏ qua tất cả các bước tiền xử lý OCR cần thiết để trích xuất văn bản từ PDF. Việc xử lý các trang dưới dạng hình ảnh giúp bảo toàn bố cục, biểu đồ, bảng biểu, kiểu và kích thước phông chữ, cùng các manh mối thị giác khác mà ngay cả một mô hình OCR hoàn hảo cũng không thể nắm bắt được.

Thiết kế hai đầu (dual-head) cho truy xuất dày đặc và tương tác muộn

NeoMME-Retriever tái sử dụng backbone của NeoMME nhưng bổ sung hai đầu ra được huấn luyện đồng thời để phục vụ truy xuất:

Omar Khattab, người giới thiệu khái niệm tương tác muộn (late-interaction) trong ColBERT, giải thích lý do tại sao thuật ngữ này chính xác hơn "multi-vector". Nó mô tả độ chi tiết và khả năng học của hàm tính điểm, chứ không chỉ đơn thuần là số lượng vector được lưu trữ.

Để tìm hiểu thêm về tương tác muộn, chúng tôi khuyên bạn nên đọc khóa học cấp tốc này của Amélie Chatelain.

Một lần truyền xuôi của NeoMME-Retriever trả về cả hai biểu diễn, mang lại sự linh hoạt bất kể trường hợp sử dụng và cơ sở hạ tầng của bạn là gì. Chúng tôi khuyên bạn nên sử dụng các embedding tương tác muộn nói chung vì chúng mạnh mẽ hơn và có thể dễ dàng sử dụng với các thư viện mã nguồn mở như NextPlaid. Tuy nhiên, nếu bạn có một tập dữ liệu rất lớn, bạn có thể chạy một lần truyền xuôi với NeoMME-Retriever để lấy embedding dày đặc, truy xuất một số lượng nhỏ tài liệu thông qua chỉ mục ANN, sau đó sử dụng tương tác muộn để xếp hạng lại các ứng viên đã truy xuất.

Truy xuất cạnh tranh ở kích thước mô hình nhỏ gọn

Chúng tôi báo cáo chỉ số nDCG@10 trên ViDoRe v3. NeoMME-Retriever-260M đạt 0,523, điểm số cao nhất trong số các mô hình được đánh giá có dưới 800 triệu tham số. Nó chỉ kém ColQwen2.5 0,002 điểm nDCG@10 trong khi sử dụng ít hơn khoảng 14 lần tham số. NeoMME-Retriever-800M đạt 0,556, chỉ kém 0,009 điểm nDCG@10 so với Vultron Retriever Flash (0,8B) có kích thước tương đương. Cả hai mô hình NeoMME-Retriever đều nằm trên đường biên Pareto về kích thước mô hình.

ViDoRe v1 và v2 sử dụng nDCG@5. Trên cả hai tiêu chuẩn này, NeoMME-Retriever-260M đều vượt trội hơn ColModernVBERT và ColSmol-500M lớn gấp đôi. NeoMME-Retriever-800M vượt trội hơn ColPali v1.3 trong khi sử dụng ít hơn 3,6 lần tham số.

† Điểm số từ MTEB. ‡ Kết quả từ các đánh giá của riêng chúng tôi.

Hiện thực hóa truy xuất độ phân giải cao cho tương tác muộn

Dung lượng lưu trữ tương tác muộn tỷ lệ thuận với số lượng vector trong embedding đầu ra. Hình ảnh có độ phân giải cao hơn chứa nhiều bản vá hơn, do đó chúng tạo ra các embedding lớn hơn. Ví dụ, một trang vuông 2048×2048 tạo ra các embedding chứa 4.200 vector với NeoMME-Retriever, tương đương khoảng 2,1 MB ở định dạng float32. Trên toàn bộ tiêu chuẩn ViDoRe v3, mức trung bình đo được là khoảng 1,5 MB mỗi tài liệu.

Để giảm dung lượng lưu trữ của chỉ mục tương tác muộn, chúng tôi kết hợp hai phương pháp nén bổ sung cho nhau:

Chúng tôi đã thử nghiệm thiết lập này trên ViDoRe v3. Với hệ số gộp (pooling factor) là 10 cùng truy vấn và tài liệu ở định dạng int8, dung lượng lưu trữ giảm từ khoảng 1,5 MB xuống còn 39 kB mỗi trang, giảm 39 lần, trong khi vẫn giữ lại hơn 99% chỉ số nDCG@10 cơ sở. Một cấu hình mạnh mẽ hơn sử dụng hệ số gộp 8, truy vấn int8 và tài liệu nhị phân. Phiên bản đó sử dụng 6 kB mỗi trang (nhỏ hơn 255 lần) và giữ lại hơn 95% chất lượng truy xuất ban đầu.

Người dùng có thể chọn cài đặt nén từ đường biên đó dựa trên ngân sách lưu trữ và chất lượng truy xuất yêu cầu.

Suy luận nhanh để lập chỉ mục tập dữ liệu đa phương thức với chi phí thấp hơn

Trước khi có thể tìm kiếm trong một tập dữ liệu, mô hình truy xuất phải chuyển đổi tài liệu của bạn thành các embedding, sau đó sẽ được lưu trữ trong một kho lưu trữ vector như Qdrant, Weaviate hoặc Milvus. Việc mã hóa nhanh hơn giúp quá trình xây dựng và thêm tài liệu mới vào chỉ mục trở nên nhanh chóng hơn, từ đó giảm thời gian hoạt động của GPU và chi phí tính toán cần thiết.