Hướng dẫn huấn luyện và tinh chỉnh mô hình nhúng đa vector với Sentence Transformers

Sentence Transformers là một thư viện Python dùng để sử dụng và huấn luyện các mô hình embedding và reranker cho nhiều ứng dụng khác nhau, chẳng hạn như retrieval augmented generation (RAG), tìm kiếm ngữ nghĩa, độ tương đồng văn bản ngữ nghĩa, và nhiều ứng dụng khác. Bản cập nhật v6.0 giới thiệu loại mô hình thứ tư: MultiVectorEncoder, dành cho kỹ thuật truy xuất tương tác muộn (late interaction) theo phong cách ColBERT, cùng với một phương pháp huấn luyện hoàn chỉnh cho nó. Trong bài viết này, tôi sẽ hướng dẫn bạn cách sử dụng nó để tinh chỉnh (finetune) một mô hình đa vector (multi-vector) vượt trội hơn các bộ truy xuất đa năng trên dữ liệu của riêng bạn. Phương pháp này cũng có thể huấn luyện các mô hình đa vector mới mạnh mẽ từ đầu. Mọi thứ dưới đây đều chạy được với lệnh pip install -U "sentence-transformers[train]".

Việc tinh chỉnh các mô hình đa vector bao gồm nhiều thành phần: bản thân mô hình, tập dữ liệu, hàm mất mát (loss function), các tham số huấn luyện, bộ đánh giá và lớp trainer. Tôi sẽ xem xét từng thành phần này, kèm theo các ví dụ thực tế về cách sử dụng chúng để tinh chỉnh các mô hình đa vector mạnh mẽ.

Cuối cùng, trong phần Đánh giá, tôi sẽ chứng minh rằng mô hình multi-vector-encoder/mLateOn-medical đã tinh chỉnh của tôi, được huấn luyện trong 14,5 giờ trên một chiếc RTX 3090 duy nhất cùng với bài viết này, dễ dàng vượt qua mọi mô hình truy xuất đa năng mà tôi có thể tìm thấy trong bài kiểm tra truy xuất y tế của mình: bao gồm cả mô hình dense, sparse, từ vựng (lexical) và đa vector.

Nếu bạn quan tâm đến việc tinh chỉnh các mô hình dense embedding, sparse embedding hoặc reranker, hãy cân nhắc đọc các bài viết trước của tôi về Training and Finetuning Embedding Models, Training and Finetuning Sparse Embedding Models, và Training and Finetuning Reranker Models.

Bài viết này tập trung vào việc huấn luyện các mô hình đa vector. Nếu bạn muốn tìm hiểu cách sử dụng chúng, từ việc tải và mã hóa đến lập chỉ mục trong các cơ sở dữ liệu vector, hãy xem bài viết đồng hành Multi-Vector (Late Interaction) Embedding Models with Sentence Transformers.

Mục lục

Mô hình Multi-Vector là gì?

Một mô hình dense embedding nén toàn bộ văn bản thành một vector duy nhất, và độ tương đồng là tích vô hướng giữa hai bản tóm tắt đó. Một mô hình đa vector (còn gọi là mô hình tương tác muộn hoặc phong cách ColBERT) bỏ qua quá trình nén đó. Nó giữ lại một vector nhỏ cho mỗi token và tính điểm truy vấn dựa trên tài liệu bằng toán tử MaxSim, nơi mỗi token truy vấn tìm thấy token tài liệu khớp nhất và các điểm số được cộng lại. Việc khớp ở cấp độ token bảo toàn chính xác các tín hiệu chi tiết mà một vector đơn lẻ thường phải làm mờ đi, điều này thường mang lại khả năng truy xuất mạnh mẽ hơn, với cái giá phải trả là chỉ mục (index) lớn hơn.

Bài viết đồng hành Multi-Vector Embedding Models bao gồm chi tiết về kiến trúc, mã hóa, tính điểm và lập chỉ mục, vì vậy tôi sẽ giữ phần này ngắn gọn và đi thẳng vào việc huấn luyện.

Tại sao cần tinh chỉnh?

Tinh chỉnh các mô hình đa vector giúp cải thiện đáng kể hiệu suất truy xuất trên lĩnh vực cụ thể của bạn: từ vựng, phong cách truy vấn và khái niệm về mức độ liên quan đều khác biệt giữa tìm kiếm web, khám phá pháp lý, tìm kiếm mã nguồn và đánh giá tài liệu khoa học. Vì các truy vấn và tài liệu được khớp theo từng token, các mô hình đa vector nắm bắt được các tín hiệu miền chi tiết mà các mô hình đơn vector thường làm mờ đi, và chúng phản hồi rất tốt ngay cả với lượng dữ liệu tinh chỉnh trong miền khiêm tốn.

Ngoài ra, hầu hết các mô hình truy xuất được phát hành đều được cấu hình cho các đoạn văn ngắn. Các checkpoint ColBERT cổ điển cắt bớt tài liệu ở mức 180 hoặc 300 token, và nhiều mô hình dense phổ biến ở mức 256 hoặc 512, vì dữ liệu huấn luyện theo phong cách MS MARCO của chúng hiếm khi vượt quá mức đó. Nếu tài liệu của bạn dài, các mô hình này sẽ âm thầm loại bỏ phần lớn tài liệu trước khi tính điểm. Trong bài kiểm tra y tế của tôi với các đoạn văn trung bình 941 token, tôi đo lường được rằng việc cắt bớt này làm giảm tới 0,24 NDCG@10, đáng kể hơn bất kỳ sự khác biệt nào giữa các kiến trúc mô hình. Khi bạn tự huấn luyện mô hình của mình, bạn có thể cấu hình độ dài tài liệu phù hợp với nhu cầu dữ liệu của bạn.

LightOn đã gặp phải tình trạng tương tự với việc truy xuất mã nguồn, nơi mô hình LateOn tổng quát không đủ đáp ứng và họ đã huấn luyện LateOn-Code. Lĩnh vực của bạn, dù là y tế, pháp lý, tài chính hay tài liệu nội bộ của công ty, đều không có sẵn một mô hình chính thức. Bài viết này chỉ cho bạn cách tự xây dựng nó, chỉ trong vài giờ, trên một GPU tiêu dùng duy nhất.

Các thành phần huấn luyện

Việc huấn luyện các mô hình MultiVectorEncoder bao gồm các thành phần sau:

Hãy cùng xem xét kỹ hơn từng thành phần.

Mô hình

Huấn luyện đa vector cho bạn một lựa chọn thực sự về điểm bắt đầu, và nó quan trọng hơn bạn nghĩ.

Tinh chỉnh một mô hình đa vector hiện có

Nếu bạn muốn tinh chỉnh thêm một mô hình đa vector hiện có, bạn không cần phải lo lắng về kiến trúc:

Checkpoint mang theo công thức riêng của nó: các token đánh dấu truy vấn và tài liệu, lớp projection head, và danh sách bỏ qua (skiplist) tính điểm. Để tinh chỉnh, bạn thường muốn giữ lại tất cả những thứ đó và chỉ thay đổi những gì dữ liệu của bạn yêu cầu. Điều đầu tiên cần kiểm tra là cấu hình độ dài, vì nhiều checkpoint được phát hành giới hạn tài liệu từ 180 đến 512 token (xem phần Tại sao cần tinh chỉnh?), trong khi các đoạn văn y tế của tôi lên tới 1.400 token. Họ mLateOn đã hỗ trợ ngữ cảnh đầy đủ 8192 token của backbone, nhưng nếu checkpoint bắt đầu của bạn có giới hạn, hãy gỡ bỏ chúng:

Với các giới hạn theo tác vụ đã được gỡ bỏ, việc cắt bớt sẽ quay về giá trị model_max_length của tokenizer, đó là lý do tại sao tôi cấu hình giới hạn đó tại thời điểm tải ở trên.

Tôi đã thực hiện thêm một thay đổi, thêm một danh sách bỏ qua dấu câu (punctuation skiplist) để loại trừ các token dấu câu khỏi việc tính điểm và lưu trữ phía tài liệu. Trong một thử nghiệm cắt bỏ 4 chiều (không, dấu câu, từ dừng, cả hai), nó đã thắng nhẹ về chất lượng và giúp thu nhỏ chỉ mục tài liệu thêm 9,6% trên dữ liệu này mà không mất phí:

Xây dựng từ một transformer cơ sở

Bạn cũng có thể trỏ MultiVectorEncoder vào bất kỳ transformer cơ sở nào, và một lớp projection cấp token mới, được khởi tạo ngẫu nhiên sẽ được thêm vào cho bạn:

Đó là pipeline ColBERT cổ điển: một Transformer tạo ra các embedding token theo ngữ cảnh, một lớp Dense cấp token chiếu từng token xuống 128 chiều, một MultiVectorMask quyết định token nào được tính trong quá trình tính điểm, và một lớp Normalize cấp token. Quá trình chiếu bắt đầu ngẫu nhiên, vì vậy cần phải huấn luyện trước khi mô hình này hữu ích. Thú vị là, điều này cũng hoạt động với các backbone dense embedding mạnh mẽ. Một lớp chiếu mới trên Alibaba-NLP/gte-modernbert-base đã đạt kết quả trong khoảng 0,03 so với các điểm bắt đầu từ checkpoint hiện có trong các thử nghiệm của tôi, chỉ với lớp chiếu và 25 nghìn cặp huấn luyện.

Các thủ thuật token hóa ColBERT cổ điển (mở rộng truy vấn [MASK], token tiền tố [Q] / [D], giới hạn độ dài tài liệu, danh sách bỏ qua dấu câu) đều bị tắt theo mặc định và có thể cấu hình được. Xem phần Tạo mô hình tùy chỉnh để biết bộ đầy đủ. Để bạn tham khảo, tôi đã thử nghiệm mở rộng truy vấn [MASK] trong bốn cấu hình cho việc tinh chỉnh miền của mình và không có cấu hình nào tạo ra sự khác biệt đáng kể, vì vậy đừng cảm thấy bắt buộc phải sử dụng công thức cổ điển.

Bạn nên chọn điểm bắt đầu nào?

Tôi đã đo lường trực tiếp điều này trong khi chuẩn bị bài viết này, lấy sáu điểm bắt đầu và huấn luyện từng cái với cùng một công thức trên 25 nghìn cặp câu hỏi-đoạn văn y tế từ MIRIAD, sau đó đánh giá trên 1.000 câu hỏi tách biệt so với kho tài liệu 50.000 đoạn văn:

Kết quả làm tôi ngạc nhiên, và nó lặp lại trên hai dòng mô hình. Các checkpoint -unsupervised thích nghi với miền mới tốt hơn nhiều so với các phiên bản đã hoàn thiện, vượt qua chúng mặc dù bắt đầu thấp hơn. Các checkpoint này nằm sau quá trình tiền huấn luyện tương phản quy mô lớn nhưng trước khi tinh chỉnh có giám sát trên truy xuất tổng quát, vì vậy chúng mang tất cả cấu trúc tương tác muộn mà không có bất kỳ sự tinh chỉnh đa năng nào mà việc huấn luyện miền sau đó phải xóa bỏ. Ngược lại, các checkpoint đã hoàn thiện hầu như không thay đổi hoặc thậm chí thoái lui ở mọi tốc độ học (learning rate) mà tôi đã thử.

Vì vậy, nếu dòng mô hình bạn thích có xuất bản checkpoint tiền giám sát, hãy bắt đầu từ đó. Nếu không, một lớp chiếu mới trên một backbone tiền huấn luyện truy xuất mạnh mẽ là lựa chọn thay thế gần nhất. Tiếp tục từ một checkpoint đã hoàn thiện hoàn toàn là lựa chọn yếu nhất cho việc thích nghi miền, mặc dù đó là lựa chọn có vẻ tự nhiên nhất.

Tập dữ liệu

MultiVectorEncoderTrainer sử dụng các instance datasets.Dataset hoặc datasets.DatasetDict để huấn luyện và đánh giá. Bạn có thể tải dữ liệu từ Hugging Face Datasets Hub hoặc sử dụng dữ liệu cục bộ ở bất kỳ định dạng nào bạn thích (ví dụ: CSV, JSON, Parquet, Arrow hoặc SQL).

Lưu ý: Rất nhiều tập dữ liệu công khai hoạt động ngay lập tức với Sentence Transformers đã được gắn thẻ sentence-transformers trên Hugging Face Hub, vì vậy bạn có thể dễ dàng tìm thấy chúng tại https://huggingface.co/datasets?other=sentence-transformers. Hãy cân nhắc duyệt qua các tập dữ liệu này để tìm các tập dữ liệu sẵn sàng sử dụng có thể hữu ích cho các tác vụ, miền hoặc ngôn ngữ của bạn.

Dữ liệu trên Hugging Face Hub

Bạn có thể sử dụng hàm load_dataset để tải dữ liệu từ các tập dữ liệu trên Hub:

Đây là tập dữ liệu tôi sẽ huấn luyện trong bài viết này: 4,4 triệu câu hỏi y tế từ MIRIAD, mỗi câu hỏi được ghép nối với đoạn văn nguồn chứa câu trả lời của nó (trung bình 941 token). Các cặp (truy vấn, đoạn văn liên quan) đơn giản như thế này là dữ liệu huấn luyện truy xuất dễ thu thập nhất cho miền của riêng bạn, và như bạn sẽ thấy, chúng là tất cả những gì bạn cần.