MOSAIC: Masked Objective with Selective Adaptation for In-domain Contrastive Learning

Technische Samenvatting: MOSAIC

Probleemstelling
Het artikel behandelt de uitdaging van het aanpassen van grootschalige, algemene tekst-embeddingmodellen aan gespecialiseerde domeinen. Algemene modellen worden getraind op enorme corpora met twee fasen (pre-training gevolgd door supervised fine-tuning), maar ze missen vaak domeinspecifieke kennis, wat hun prestaties in gespecialiseerde gebieden beperkt. Bestaande adaptatiemethoden kampen met een dilemma:

1. Vocabulaire-expansie zonder gezamenlijke training: Het simpelweg toevoegen van domeinspecifieke tokens aan een vooraf getraind model en het voortzetten van Masked Language Modeling (MLM) training, verslechtert de contrastieve eigenschappen van het model, omdat de encoder zijn vermogen tot zin-niveau discriminatie verliest onder token-voorspellingsdoelstellingen.
2. Alleen contrastieve adaptatie: Het voortzetten van training met alleen een contrastieve doelstelling levert onvoldoende signalen om nieuwe domeinspecifieke tokens effectief bij te werken, aangezien de embedding-matrix verdunde signalen ontvangt door pooling-functies.
3. Uitdagingen bij gezamenlijke optimalisatie: Het pogen tot het gezamenlijk optimaliseren van MLM- en contrastieve doelstellingen leidt vaak tot conflicterende trainingssignalen. De MLM-loss domineert doorgaans vanwege de grotere vocabulaire-omvang en hogere gradiënt-magnitudes, waardoor de contrastieve doelstelling wordt overstemd, wat resulteert in suboptimale prestaties.

Methodologie: MOSAIC
De auteurs stellen MOSAIC voor (Masked Objective with Selective Adaptation for In-domain Contrastive Learning), een meerfasig framework dat is ontworpen om de balans te vinden tussen token-niveau en zin-niveau supervisie. De methode rust op een eenvoudige vereiste: vocabulaire-expansie.

Het framework bestaat uit drie fasen:

1. Fase 1: Vocabulaire-augmentatie. De tokenizer van het model wordt uitgebreid met domeinspecifieke tokens die zijn geïdentificeerd uit een domeincorpus. Deze nieuwe tokens worden toegevoegd aan de embedding-matrix en geïnitialiseerd als het gemiddelde van hun basismodel subword-embeddings. Cruciaal is dat de oorspronkelijke encoder-gewichten behouden blijven om de zin-niveau discriminatiecapaciteiten te handhaven.
2. Fase 2: Gezamenlijke optimalisatie met selectieve masking. Het model ondergaat gezamenlijke training waarbij MLM en contrastieve doelstellingen worden gecombineerd. Om de dominantie van de MLM-loss te beperken, beperken de auteurs de MLM-doelstelling tot alleen de nieuw geïntroduceerde domeinspecifieke tokens (domein-beperkte vocabulaire). Dit vermindert de noemer in de InfoNCE-stijl loss, waardoor de gradiënt-magnitudes tussen de token-niveau (MLM) en de zin-niveau (contrastieve) doelstellingen worden gebalanceerd. De gezamenlijke loss wordt geformuleerd als $L = \alpha \cdot L_{domain}^{MLM} + L_{CL}$, waarbij $\alpha$ de bijdragen balanceert.
3. Fase 3: Alleen contrastieve training. Nadat de domeintokens zijn geleerd, ondergaat het model een laatste trainingsfase met alleen de contrastieve doelstelling. Dit dient als een corrigerende stap om de zin-niveau discriminatie te versterken, die tijdens de gezamenlijke trainingsfase mogelijk is verdund.

Belangrijkste Bijdragen

1. Nieuw aanpak voor adaptatie: Het artikel introduceert MOSAIC, een methode die MLM- en contrastieve doelstellingen gezamenlijk optimaliseert binnen een raamwerk van wederzijdse informatie-maximalisatie, specifiek voor de domeinadaptatie van bestaande sterke tekst-embeddingmodellen.
2. Mechanisme voor selectieve adaptatie: De auteurs stellen voor om MLM-supervisie te beperken tot enkel domeinspecifieke tokens, een strategie die gradiënt-magnitudes afstemt en voorkomt dat de MLM-doelstelling het contrastieve signaal overstemt.
3. Empirische validatie over diverse bronnen: De methode wordt gevalideerd op zowel rijk-resource (Biomedisch) als arm-resource (Islamitisch) domeinen, waarbij de robuustheid zelfs met beperkte in-domein data wordt aangetoond.
4. Uitgebreide analyse: De studie bevat uitgebreide ablatie-studies die de impact van de gezamenlijke doelstelling, de weging-hyperparameter ($\alpha$), masking-strategieën en de volgorde van de trainingsfasen analyseren.

Resultaten

• Rijk-resource domein (Biomedisch): Op de MTEB biomedische subset behaalde MOSAIC verbeteringen tot wel 13,4% in NDCG@10 ten opzichte van sterke algemene domein-baselines. De volledige drie-fasen pipeline (MOSAIC-Bio-Stage3) presteerde beter dan naïeve vocabulaire-expansie en alleen contrastieve adaptatie, en behaalde de hoogste gemiddelde score voor retrieval, clustering en semantische tekstuele gelijkenis-taken.
• Arm-resource domein (Islamitisch): In het Islamitische domein, gekenmerkt door schaarse Engelse data, behaalde MOSAIC-ID-Stage3 een NDCG@10 van 36,8, waarmee het algemene domein-baselines en een contrastieve-alleen variant overtrof. Het model vertoonde een significante reductie in queries met een retrieval-score van nul vergeleken met de baselines.
• Ablatie-bevindingen:
  • De wegingfactor $\alpha$ is cruciaal; $\alpha=0,3$ leverde de optimale balans op, terwijl hogere waarden leidden tot MLM-dominantie en prestatiedalingen.
  • Het beperken van MLM tot domeintokens was effectiever dan het toepassen ervan op de volledige vocabulaire, zelfs met zeer lage gewichten.
  • De volgorde van de fasen is van belang; het uitvoeren van gezamenlijke training vóór pure contrastieve training leverde betere resultaten op dan de omgekeerde volgorde.
  • De methode is computationeel efficiënt en vereist ongeveer 3x de rekenkracht van standaard contrastieve fine-tuning, aanzienlijk minder dan volledige domain-adaptive pre-training (DAPT).

Betekenis en Claims
Het artikel claimt dat MOSAIC een principieel en effectief mechanisme biedt voor het aanpassen van tekst-embeddings aan gespecialiseerde domeinen zonder complexe architecturale wijzigingen of grootschalige hertraining te vereisen. Door gebruik te maken van de complementaire sterktes van token-niveau en zin-niveau supervisie via selectieve masking, stelt de methode modellen in staat om domein-relevante representaties te leren terwijl de robuuste semantische discriminatie van het oorspronkelijke algemene domeinmodel behouden blijft. De auteurs benadrukken de praktische waarde van deze aanpak voor arm-resource domeinen, waar datascarciteit traditionele adaptatiemethoden vaak onmogelijk maakt, waardoor zij zo de toegang tot state-of-the-art taaltechnologieën vergemakkelijken. Het werk wordt gepresenteerd als een model-gestuurde aanpak die opereert na de taalmodel pre-training, wat het onderscheidt van eerdere data-gestuurde of adapter-gebaseerde technieken.