Обучение и дообучаване на многовекторни модели за вграждане със Sentence Transformers

Актуализация в GitHub

Sentence Transformers е библиотека на Python за използване и обучение на модели за вграждане и преранжиране за широк спектър от приложения, като генериране с допълнено извличане, семантично търсене, семантична текстова близост и други. Актуализацията ѝ до версия 6.0 въвежда четвърти тип модел: MultiVectorEncoder, предназначен за извличане с късно взаимодействие в стил ColBERT, заедно с цялостен подход за обучението му. В тази публикация ще ви покажа как да го използвате за дообучаване на мултивекторен модел, който превъзхожда универсалните извличащи модели върху вашите данни. Този метод може също да обучава силни нови мултивекторни модели от нулата. Всичко по-долу работи с pip install -U "sentence-transformers[train]".

Дообучаването на мултивекторни модели включва няколко компонента: самия модел, набори от данни, функции на загубата, аргументи за обучение, оценители и клас на обучаващия модул. Ще разгледам всеки от тези компоненти, придружен от практически примери за това как могат да се използват за дообучаване на силни мултивекторни модели.

Накрая, в раздела „Оценяване“, ще ви покажа, че моят дообучен модел multi-vector-encoder/mLateOn-medical, обучен за 14,5 часа на една RTX 3090 едновременно с написването на тази публикация, лесно превъзхожда всеки универсален модел за извличане, който успях да намеря при оценката си за медицинско извличане: плътен, разреден, лексикален и мултивекторен.

Ако вместо това се интересувате от дообучаване на плътни модели за вграждане, разредени модели за вграждане или преранжиращи модели, прочетете предишните ми публикации Обучение и дообучаване на модели за вграждане, Обучение и дообучаване на разредени модели за вграждане и Обучение и дообучаване на преранжиращи модели.

Тази публикация е посветена на обучението на мултивекторни модели. Ако искате да научите как да ги използвате — от зареждането и кодирането до индексирането във векторни бази данни — вижте придружаващата публикация Мултивекторни модели за вграждане (късно взаимодействие) със Sentence Transformers.

Съдържание

• Какво представляват мултивекторните модели?
• Защо да ги дообучаваме?
• Компоненти на обучението
• Модел
  • Дообучаване на съществуващ мултивекторен модел
  • Създаване на модел от базов трансформатор
  • Коя начална точка да изберете?
• Набор от данни
  • Данни в Hugging Face Hub
  • Локални данни
  • Формат на набора от данни
• Функция на загубата
• Аргументи за обучение
• Оценител
• Обучаващ модул
  • Обратни извиквания
  • Обучение с множество набори от данни
• Оценяване
  • Оптимизиране на индекса
• Благодарности
• Допълнителни ресурси
  • Примери за обучение
  • Документация

Какво представляват мултивекторните модели?

Плътният модел за вграждане компресира цял текст в един-единствен вектор, а сходството е едно скаларно произведение между две такива обобщения. Мултивекторният модел (наричан още модел с късно взаимодействие или модел в стил ColBERT) пропуска това компресиране. Той запазва един малък вектор за всеки токен и оценява заявка спрямо документ с оператора MaxSim, при който всеки токен на заявката намира най-добре съвпадащия токен в документа, а резултатите се сумират. Съпоставянето на ниво токен запазва точно фините сигнали, които един-единствен вектор трябва да усредни, което обикновено означава по-силно извличане за сметка на по-голям индекс.

Придружаващата публикация Мултивекторни модели за вграждане разглежда подробно архитектурата, кодирането, оценяването и индексирането, затова ще запазя този раздел кратък и ще премина към обучението.

Защо да ги дообучаваме?

Дообучаването на мултивекторни модели значително подобрява ефективността им при извличане във вашия конкретен домейн: речникът, стилът на заявките и представата за релевантност се различават при уеб търсене, правно откриване на документи, търсене в код и преглед на научна литература. Тъй като заявките и документите се съпоставят токен по токен, мултивекторните модели улавят фини домейн-специфични сигнали, които едновекторните модели са склонни да усредняват, и реагират много добре дори на умерени количества домейн-специфични данни за дообучаване.

Освен това повечето публикувани модели за извличане са конфигурирани за кратки откъси. Класическите контролни точки на ColBERT съкращават документите до 180 или 300 токена, а много популярни плътни модели — до 256 или 512, тъй като данните им за обучение в стил MS MARCO рядко надхвърлят тези стойности. Ако документите ви са дълги, тези модели безшумно изхвърлят по-голямата част от всеки документ, преди да го оценят. При медицинската ми оценка, с откъси средно от 941 токена, измерих, че това съкращаване струва до 0,24 NDCG@10 — значително повече от всяка разлика между архитектурите на моделите. Когато обучавате собствен модел, задавате дължината на документа, от която вашите данни се нуждаят.

LightOn се сблъска със същата динамика при извличането на код, където универсалният LateOn не беше достатъчен и те обучиха LateOn-Code. Вашият домейн — независимо дали е медицински, правен, финансов или представлява вътрешните документи на компанията ви — няма да получи официален модел. Тази публикация показва как да го създадете сами за няколко часа, на един потребителски графичен процесор.

Компоненти на обучението

Обучението на модели MultiVectorEncoder включва следните компоненти:

1. Модел: Моделът за дообучаване или архитектурата, която ще изградите от нулата.
2. Набор от данни: Данните, използвани за обучение и оценяване.
3. Функция на загубата: Функция, която измерва ефективността на модела и насочва процеса на оптимизация.
4. Аргументи за обучение (по избор): Параметри, които влияят върху ефективността на обучението, проследяването и отстраняването на грешки.
5. Оценител (по избор): Клас за оценяване на модела преди, по време на или след обучението.
6. Обучаващ модул: Обединява всички компоненти на обучението.

Нека разгледаме по-подробно всеки компонент.

Модел

Обучението на мултивекторни модели ви дава реален избор на начална точка и това е по-важно, отколкото може да очаквате.

Дообучаване на съществуващ мултивекторен модел

Ако искате допълнително да дообучите съществуващ мултивекторен модел, изобщо не е нужно да се тревожите за архитектурата:

from sentence_transformers import MultiVectorEncoder


model = MultiVectorEncoder(
    "lightonai/mLateOn-unsupervised",
    model_kwargs={"torch_dtype": "float32"},
    processor_kwargs={"model_max_length": 8192},  
)

Контролната точка носи собствената си рецепта: маркерите за заявка и документ, проекционната глава и списъка за пропускане при оценяване. При дообучаване обикновено искате да запазите всичко това и да промените само онова, което вашите данни изискват. Първото нещо за проверка е конфигурацията на дължината, тъй като много публикувани контролни точки ограничават документите до 180–512 токена (вижте „Защо да ги дообучаваме?“), а моите медицински откъси достигат 1400 токена. Семейството mLateOn вече използва пълния контекст от 8192 токена на базовия модел, но ако началната ви контролна точка има ограничения, премахнете ги:

model[0].query_length = None
model[0].document_length = None

Когато ограниченията за конкретната задача не са зададени, съкращаването се връща към model_max_length на токенизатора, поради което конфигурирам тази граница при зареждането по-горе.

Направих още една промяна: добавих списък за пропускане на пунктуацията, който изключва пунктуационните токени от оценяването и съхранението от страната на документа. При аблация в 4 варианта (без, пунктуация, стоп-думи, и двете) това даде умерено по-добро качество и безплатно намали индекса на документите с 9,6% върху тези данни:

import string


model[2].skiplist_words = list(string.punctuation)
model[2].resolve_with_tokenizer(model.tokenizer)  

Създаване на модел от базов трансформатор

Можете също да насочите MultiVectorEncoder към всеки базов трансформатор, като към него автоматично се добавя нова, произволно инициализирана проекция на ниво токен:

from sentence_transformers import MultiVectorEncoder

model = MultiVectorEncoder("answerdotai/ModernBERT-base", model_kwargs={"torch_dtype": "float32"})

Това е класическият конвейер на ColBERT: Transformer, който създава контекстуализирани вграждания на токените; Dense, който проектира всеки от тях до 128 измерения; MultiVectorMask, който определя кои токени участват в оценяването; и Normalize на ниво токен. Проекцията започва случайно, така че преди този модел да стане полезен, е необходимо обучение. Интересното е, че това работи добре и със силни базови модели за плътни вграждания. Нова проекция върху Alibaba-NLP/gte-modernbert-base достигна разлика от само 0,03 спрямо началните точки със съществуващи контролни точки в моите експерименти, използвайки единствено проекцията и 25 хиляди обучаващи двойки.

Класическите трикове за токенизация на ColBERT ([MASK] разширяване на заявката, префиксни токени [Q] / [D], ограничение на дължината на документа и списък за пропускане на пунктуацията) са изключени по подразбиране и могат да се конфигурират. Вижте Създаване на персонализирани модели за пълния набор. Между другото, тествах разширяването на заявката с [MASK] в четири конфигурации за дообучаването на моя домейн и нито една не даде измерима разлика, така че не се чувствайте длъжни да използвате класическата рецепта.

Коя начална точка да изберете?

Измерих това директно при подготовката на тази публикация, като взех шест начални точки и обучих всяка с идентична рецепта върху 25 хиляди двойки медицински въпрос–откъс от MIRIAD, след което оцених върху 1000 въпроса, задържани за проверка, спрямо корпус от 50 хиляди откъса:

Резултатът ме изненада и се повтори при две семейства модели. Контролните точки с -unsupervised се адаптират към нов домейн много по-добре от завършените си събратя, като ги изпреварват въпреки по-ниския старт. Тези контролни точки са след мащабно контрастивно предварително обучение, но преди надзиравано дообучаване за универсално извличане, така че носят цялата структура на късното взаимодействие, без универсалната настройка, която домейн обучението после трябва да премахне. За разлика от тях завършените контролни точки почти не се промениха или дори регресираха при всички скорости на обучение, които изпробвах.

Затова, ако семейството модели, което харесвате, публикува контролна точка преди надзираваното обучение, започнете от нея. Ако няма такава, нова проекция върху силен бекбоун, предварително обучен за извличане, е близка втора възможност. Продължаването от напълно завършена контролна точка е най-слабият вариант за домейн адаптация, въпреки че изглежда най-естественият.

Набор от данни

MultiVectorEncoderTrainer използва обекти datasets.Dataset или datasets.DatasetDict за обучение и оценяване. Можете да зареждате данни от Hugging Face Datasets Hub или да използвате локални данни в предпочитан от вас формат (например CSV, JSON, Parquet, Arrow или SQL).

Забележка: Много публични набори от данни, които работят директно със Sentence Transformers, са маркирани с sentence-transformers в Hugging Face Hub, така че лесно можете да ги откриете на https://huggingface.co/datasets?other=sentence-transformers. Разгледайте ги, за да намерите готови набори от данни, които може да са полезни за вашите задачи, домейни или езици.

Данни в Hugging Face Hub

Можете да използвате функцията load_dataset, за да зареждате данни от набори в Hub:

from datasets import load_dataset

train_dataset = load_dataset("tomaarsen/miriad-4.4M-split", split="train")

print(train_dataset)
"""
Dataset({
    features: ['question', 'passage_text'],
    num_rows: 4467542
})
"""

Това е наборът от данни, върху който ще обучавам в тази публикация: 4,4 милиона медицински въпроса от MIRIAD, всеки свързан с изходния откъс, съдържащ отговора му (със средно 941 токена). Простите двойки (заявка, релевантен откъс) като тези са най-лесните за събиране обучаващи данни за извличане в собствен домейн и, както ще видите, са всичко необходимо.

Локални данни

Можете също да използвате load_dataset за зареждане на локални данни в разпространени файлови формати:

from datasets import load_dataset

dataset = load_dataset("csv", data_files="my_file.csv")

dataset = load_dataset("json", data_files="my_file.json")

А ако локалните ви данни изискват предварителна обработка, можете да използвате datasets.Dataset.from_dict, за да инициализирате набора си от данни с речник от списъци:

from datasets import Dataset

queries = []
documents = []



dataset = Dataset.from_dict({
    "query": queries,
    "document": documents,
})

Формат на набора от данни

Важно е форматът на набора ви от данни да съответства на функцията на загубата (или да изберете функция на загубата, която съответства на формата на данните). Проверката дали даден формат работи с функция на загубата включва две стъпки:

1. Ако функцията ви на загубата изисква етикет според таблицата Преглед на функциите на загубата, наборът ви от данни трябва да има колона с име „label“ или „score“. Тази колона автоматично се използва като етикет.
2. Всички колони, които не се наричат „label“ или „score“, се считат за входове според таблицата Преглед на функциите на загубата. Броят на останалите колони трябва да съответства на броя на валидните входове за избраната функция на загубата. Имената на тези колони са без значение — важен е само редът.

Освен това има две специфични за мултивекторните модели конвенции:

• Позиционно разпределяне на заявка и документи: първата колона се вгражда като заявка, а всички следващи — като документи, независимо от имената на колоните. Това поведение по подразбиране може да бъде заменено за всяка колона чрез стандартния аргумент за обучение router_mapping.
• Формат за дистилация на знания: по една колона за всеки кандидат-документ, т.е. (query, document_1, ..., document_N, scores), където scores е списък от N оценки на учителя за всеки ред. За набори от данни за KD, които съхраняват идентификатори на заявките и документите заедно с отделни текстови набори (например lightonai/ms-marco-en-bge), можете да използвате resolve_ids, за да преобразувате идентификаторите в текстове в движение.

Функция на загубата

Функциите на загубата измерват колко добре се представя моделът за дадена партида данни, като позволяват на оптимизатора да актуализира теглата му, така че да генерира по-благоприятни (т.е. по-ниски) стойности на загубата. Подходящата функция на загубата за вашата задача зависи от наличните данни и от целта ви. Пълен списък с възможностите ще намерите в Преглед на функциите на загубата.

При обичайния случай на двойки въпрос–отговор или въпрос–откъс основният метод е обучение с отрицателни примери в рамките на партидата чрез MultiVectorMultipleNegativesRankingLoss, при което всеки друг документ в партидата служи като отрицателен пример за всяка заявка. По-големите партиди означават повече отрицателни примери и по-силно обучение, така че на практика ще искате варианта му с GradCache — CachedMultiVectorMultipleNegativesRankingLoss, който отделя ефективния размер на партидата от това, което се побира във вашия графичен процесор:

from sentence_transformers import MultiVectorEncoder
from sentence_transformers.multi_vector_encoder.losses import CachedMultiVectorMultipleNegativesRankingLoss

model = MultiVectorEncoder("lightonai/mLateOn-unsupervised", model_kwargs={"torch_dtype": "float32"})

loss = CachedMultiVectorMultipleNegativesRankingLoss(
    model=model,
    mini_batch_size=16,  
)

Параметърът mini_batch_size ограничава паметта, като кодира документите на части с този размер, докато ефективният контрастивен размер на партидата (128 в моя случай, а при аблациите ми по-големите партиди не донесоха допълнителна полза) остава свободен избор. GradCache гарантира идентични резултати независимо от размера на частта, така че можете да го намалите за по-малки графични процесори с цената единствено на време. Когато дължините на документите се различават много, помислете за сродния параметър mini_batch_num_tokens, който запълва всяка част до общ бюджет токени вместо до определен брой документи, така че необичайно дълга група документи никога да не предизвика пик в потреблението на памет (моят mini_batch_size=16 при около 940 токена на документ съответства на mini_batch_num_tokens=15_000).

Един специфичен за мултивекторните модели капан е, че контрастивните функции на загубата използват по подразбиране scale=1.0, за разлика от еквивалента за плътни вграждания, чиято стойност по подразбиране е scale=20.0. Стойността 20.0 съществува, защото косинусното сходство е една стойност в [-1, 1] — твърде тесен диапазон за остър softmax. За разлика от него резултатът MaxSim сумира едно най-добро сходство за всеки токен на заявката, така че вече обхваща приблизително [0, query_length]: заявка от 32 токена може да получи резултат до 32. Затова не пренасяйте scale=20.0 от скрипт за обучение на плътни модели, тъй като това ще насити softmax и ще унищожи градиентите.

За дистилация от по-силен учител — така се обучават най-силните универсални модели с късно взаимодействие — вижте MultiVectorDistillKLDivLoss и раздела „Дистилация на знания“ в документацията Преглед на обучението.

Аргументи за обучение

Можете да персонализирате процеса на обучение чрез класа MultiVectorEncoderTrainingArguments. Този клас ви позволява да настройвате параметри, които могат да повлияят на скоростта на обучение и да ви помогнат да разберете какво се случва по време на него.

За повече информация относно най-полезните аргументи за обучение вижте Multi-Vector Encoder > Преглед на обучението > Аргументи за обучение. Струва си да го прочетете, за да извлечете максимума от обучението си.

Ето пример, използващ стойностите от действителното ми обучение:

from sentence_transformers import MultiVectorEncoderTrainingArguments
from sentence_transformers.base.sampler import BatchSamplers

args = MultiVectorEncoderTrainingArguments(
    
    output_dir="models/mLateOn-medical",
    
    num_train_epochs=1,
    per_device_train_batch_size=128,  
    per_device_eval_batch_size=16,
    learning_rate=1e-4,
    warmup_steps=0.05,
    prompts={"question": "[Q] ", "passage_text": "[D] "},  
    fp16=False,  
    bf16=True,  
    batch_sampler=BatchSamplers.NO_DUPLICATES,  
    
    eval_strategy="steps",
    eval_steps=0.1,
    save_strategy="steps",
    save_steps=0.05,
    logging_steps=0.01,
    run_name="mLateOn-medical",  
)

Няколко от тях заслужават коментар:

• prompts: обучението не прилага автоматично подсказките, съхранени в модела, затова изрично ги съпоставете с колоните за обучение. Тук това са маркерът [Q] на контролната точка за колоната с въпроси и [D] за колоната с откъси, което запазва съответствието между обучението и извеждането.
• max_length (умишлено незададен): този аргумент ограничава токенизацията само по време на обучение, когато искате по-евтино обучение от пълната дължина при обслужване на модела. Измерих цената на този компромис върху тези данни. Обучението при 512 токена загуби около 0,015 NDCG@10 срещу приблизително двойно по-висока скорост, а дефицитът не намаля с повече данни, защото моделът просто никога не вижда изрязаната част. Оставете го незададен, така че обучението да съответства на извеждането, освен ако не се нуждаете от ускорението повече, отколкото от качеството.
• learning_rate=1e-4: след изпробване на стойности от 5e-6 до 2e-4 имах най-добри резултати с тази по-висока от обичайната скорост на обучение.

Оценител

За да проследявате ефективността на модела по време на обучението, можете да подадете eval_dataset на обучаващия модул за загуба при оценяване, но конкретните метрики за извличане са много по-информативни. Sentence Transformers включва следните вградени оценители за мултивекторни модели:

При домейн дообучаване най-важен е MultiVectorInformationRetrievalEvaluator, изграден от ваши собствени задържани данни. Един съвет при създаването му е корпусът да е достатъчно труден, за да могат моделите да бъдат разграничени. В моя случай въпросите от MIRIAD са генерирани от собствените им изходни откъси, което прави извличането необичайно лесно. Само спрямо 10-те хиляди златни откъса почти всеки модел получи над 0,97 NDCG@10. Ако оценката ви се насища по този начин, добавете разсейващи откъси (аз използвам дедупликирани откъси от обучаващото разделение), докато резултатите се разпределят:

Дообученият модел оглавява таблицата, изпреварвайки най-силния модел без обучение от която и да е архитектура с +0,062 NDCG@10. С други думи, най-силният модел без обучение връща правилния откъс като първи резултат при 75,8% от заявките, докато дообученият прави това при 84,9%, намалявайки грешките на първа позиция с повече от една трета.

Моделът от въведението е дообучен върху един милион двойки медицински въпрос–откъс, със същата конфигурация, описана по-горе. Той беше оценен спрямо над 50 конфигурации на модели за извличане, като използва 1000 медицински въпроса и корпус от 200 хиляди уникални откъса.

Благодарности

Благодаря на общността на Hugging Face и на авторите и поддържащите Sentence Transformers за инструментите, които направиха това обучение възможно.

Допълнителни ресурси

• Примери за обучение
• Документация