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

Riepilogo Tecnico: MOSAIC

Problematica
Il documento affronta la sfida di adattare modelli di embedding testuali di grandi dimensioni e di dominio generale a domini specializzati. I modelli di dominio generale, pur essendo addestrati su vasti corpora attraverso approcci a due stadi (pre-addestramento seguito da fine-tuning supervisionato), spesso mancano di conoscenze specifiche del dominio, limitandone le prestazioni in aree specializzate. I metodi di adattamento esistenti affrontano un dilemma:

1. Espansione del Vocabolario senza Addestramento Congiunto: L'aggiunta di semplici token specifici del dominio a un modello pre-addestrato e il proseguimento dell'addestramento tramite Masked Language Modeling (MLM) degrada le proprietà contrastive del modello, poiché l'encoder perde le sue capacità di discriminazione a livello di frase sotto gli obiettivi di predizione dei token.
2. Adattamento Solo Contrastivo: Proseguire l'addestramento utilizzando solo un obiettivo contrastivo fornisce segnali insufficienti per aggiornare efficacemente i nuovi token del dominio, poiché la matrice di embedding riceve segnali diluiti a causa delle funzioni di pooling.
3. Sfide dell'Ottimizzazione Congiunta: Tentare di ottimizzare congiuntamente gli obiettivi MLM e contrastivo porta spesso a segnali di addestramento conflittuali. La perdita MLM tende tipicamente a dominare a causa della sua dimensione del vocabolario più ampia e delle magnitudo dei gradienti più elevate, sovrastando l'obiettivo contrastivo e risultando in prestazioni subottimali.

Metodologia: MOSAIC
Gli autori propongono MOSAIC (Masked Objective with Selective Adaptation for In-domain Contrastive Learning), un framework multi-stadio progettato per bilanciare la supervisione a livello di token e a livello di frase. Il metodo si basa su un semplice prerequisito: l'espansione del vocabolario.

Il framework consiste in tre stadi:

1. Stadio 1: Augmentation del Vocabolario. Il tokenizer del modello viene espanso con token specifici del dominio identificati da un corpus di dominio. Questi nuovi token vengono aggiunti alla matrice di embedding e inizializzati come la media degli embedding subword del modello base. Fondamentalmente, i pesi originali dell'encoder vengono preservati per mantenere le capacità di discriminazione a livello di frase.
2. Stadio 2: Ottimizzazione Congiunta con Mascheramento Selettivo. Il modello subisce un addestramento congiunto che combina gli obiettivi MLM e contrastivi. Per mitigare il dominio della perdita MLM, gli autori limitano l'obiettivo MLM solo ai nuovi token introdotti dal dominio specifico (vocabolario ristretto al dominio). Ciò riduce il denominatore nella perdita di tipo InfoNCE, bilanciando le magnitudo dei gradienti tra gli obiettivi a livello di token (MLM) e a livello di frase (contrastivo). La perdita congiunta è formulata come $L = \alpha \cdot L_{domain}^{MLM} + L_{CL}$, dove $\alpha$ bilancia i contributi.
3. Stadio 3: Addestramento Solo Contrastivo. Dopo che i token del dominio sono stati appresi, il modello subisce uno stadio finale di addestramento utilizzando solo l'obiettivo contrastivo. Questo funge da passo correttivo per rafforzare la discriminazione a livello di frase, che potrebbe essere stata diluita durante la fase di addestramento congiunto.

Contributi Chiave

1. Nuovo Approccio di Adattamento: Il documento introduce MOSAIC, un metodo che ottimizza congiuntamente gli obiettivi MLM e contrastivi all'interno di un framework di massimizzazione dell'informazione mutua, specificamente per l'adattamento di dominio di esistenti e forti modelli di embedding testuale.
2. Meccanismo di Adattamento Selettivo: Gli autori propongono di limitare la supervisione MLM solo ai token specifici del dominio, una strategia che allinea le magnitudo dei gradienti e impedisce all'obiettivo MLM di sovrastare il segnale contrastivo.
3. Validazione Empirica su Diverse Risorse: Il metodo è validato sia su domini ad alte risorse (Biomedicale) che a basse risorse (Islamico), dimostrando robustezza anche con dati limitati del dominio.
4. Analisi Esaustiva: Lo studio include estesi studi di ablazione che analizzano l'impatto dell'obiettivo congiunto, del parametro di ponderazione ($\alpha$), delle strategie di mascheramento e dell'ordine degli stadi di addestramento.

Risultati

• Dominio ad Alte Risorse (Biomedicale): Sul subset biomedico di MTEB, MOSAIC ha ottenuto miglioramenti fino al 13,4% in NDCG@10 rispetto a forti baseline di dominio generale. L'intero pipeline a tre stadi (MOSAIC-Bio-Stage3) ha superato l'espansione del vocabolario ingenua e l'adattamento solo contrastivo, ottenendo il punteggio medio più alto nei compiti di retrieval, clustering e similitudine testuale semantica.
• Dominio a Basse Risorse (Islamico): Nel dominio islamico, caratterizzato da scarsi dati in inglese, MOSAIC-ID-Stage3 ha raggiunto un NDCG@10 di 36,8, superando le baseline di dominio generale e una variante solo contrastiva. Il modello ha mostrato una riduzione significativa delle query con punteggi di retrieval pari a zero rispetto alle baseline.
• Risultati dell'Ablazione:
  • Il fattore di ponderazione $\alpha$ è critico; $\alpha=0,3$ ha prodotto il bilanciamento ottimale, mentre valori più elevati hanno causato la dominanza del MLM e cali di prestazione.
  • Limitare l'MLM ai soli token del dominio è stato più efficace rispetto all'applicarlo all'intero vocabolario, anche con pesi molto bassi.
  • L'ordine degli stadi è importante; eseguire l'addestramento congiunto prima dell'addestramento puramente contrastivo ha prodotto risultati migliori rispetto al contrario.
  • Il metodo è computazionalmente efficiente, richiedendo circa 3 volte la computazione del normale fine-tuning contrastivo, significativamente meno del DAPT (Domain-Adaptive Pre-training) completo.

Significatività e Rivendicazioni
Il documento afferma che MOSAIC fornisce un meccanismo fondato ed efficace per adattare gli embedding testuali a domini specializzati senza richiedere cambiamenti architettonici complessi o ri-addestramenti su larga scala. Sfruttando i punti di forza complementari della supervisione a livello di token e a livello di frase attraverso il mascheramento selettivo, il metodo consente ai modelli di apprendere rappresentazioni rilevanti per il dominio preservando al contempo la robusta discriminazione semantica del modello originale di dominio generale. Gli autori sottolineano il valore pratico di questo approccio per i domini a basse risorse, dove la scarsità di dati spesso preclude i metodi di adattamento tradizionali, facilitando così l'accesso equo alle tecnologie linguistiche allo stato dell'arte. Il lavoro è presentato come un approccio guidato dal modello che opera dopo il pre-addestramento del modello linguistico, distinguendosi dalle precedenti tecniche basate sui dati o basate su adapter.