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

Technical Summary: MOSAIC

Problem Statement
The paper addresses the challenge of adapting large-scale, general-domain text embedding models to specialized domains. While general-domain models are trained on vast corpora using two-stage approaches (pre-training followed by supervised fine-tuning), they often lack domain-specific knowledge, limiting their performance in specialized areas. Existing adaptation methods face a dilemma:

1. Vocabulary Expansion without Joint Training: Simply adding domain-specific tokens to a pre-trained model and continuing Masked Language Modeling (MLM) training degrades the model's contrastive properties, as the encoder loses its sentence-level discrimination capabilities under token prediction objectives.
2. Contrastive-Only Adaptation: Continuing training with only a contrastive objective provides insufficient signals to update new domain tokens effectively, as the embedding matrix receives diluted signals due to pooling functions.
3. Joint Optimization Challenges: Attempting to jointly optimize MLM and contrastive objectives often leads to conflicting training signals. The MLM loss typically dominates due to its larger vocabulary size and higher gradient magnitudes, overwhelming the contrastive objective and resulting in suboptimal performance.

Methodology: MOSAIC
The authors propose MOSAIC (Masked Objective with Selective Adaptation for In-domain Contrastive Learning), a multi-stage framework designed to balance token-level and sentence-level supervision. The method relies on a simple prerequisite: vocabulary expansion.

The framework consists of three stages:

1. Stage 1: Vocabulary Augmentation. The model's tokenizer is expanded with domain-specific tokens identified from a domain corpus. These new tokens are added to the embedding matrix and initialized as the average of their base model subword embeddings. Crucially, the original encoder weights are preserved to maintain sentence-level discrimination capabilities.
2. Stage 2: Joint Optimization with Selective Masking. The model undergoes joint training combining MLM and contrastive objectives. To mitigate the dominance of the MLM loss, the authors restrict the MLM objective to only the newly introduced domain-specific tokens (domain-restricted vocabulary). This reduces the denominator in the InfoNCE-style loss, balancing the gradient magnitudes between the token-level (MLM) and sentence-level (contrastive) objectives. The joint loss is formulated as $L = \alpha \cdot L_{domain}^{MLM} + L_{CL}$, where $\alpha$ balances the contributions.
3. Stage 3: Contrastive-Only Training. After the domain tokens are learned, the model undergoes a final stage of training using only the contrastive objective. This serves as a corrective step to reinforce sentence-level discrimination, which may have been diluted during the joint training phase.

Key Contributions

1. Novel Adaptation Approach: The paper introduces MOSAIC, a method that jointly optimizes MLM and contrastive objectives within a mutual information maximization framework specifically for domain adaptation of existing strong text embedding models.
2. Selective Adaptation Mechanism: The authors propose restricting MLM supervision to domain-specific tokens only, a strategy that aligns gradient magnitudes and prevents the MLM objective from overwhelming the contrastive signal.
3. Empirical Validation Across Resources: The method is validated on both high-resource (Biomedical) and low-resource (Islamic) domains, demonstrating robustness even with limited in-domain data.
4. Comprehensive Analysis: The study includes extensive ablation studies analyzing the impact of the joint objective, the weighting hyperparameter ($\alpha$), masking strategies, and training stage ordering.

Results

• High-Resource Domain (Biomedical): On the MTEB biomedical subset, MOSAIC achieved improvements up to 13.4% in NDCG@10 over strong general-domain baselines. The full three-stage pipeline (MOSAIC-Bio-Stage3) outperformed naive vocabulary expansion and contrastive-only adaptation, achieving the highest average score across retrieval, clustering, and semantic textual similarity tasks.
• Low-Resource Domain (Islamic): In the Islamic domain, characterized by scarce English data, MOSAIC-ID-Stage3 achieved an NDCG@10 of 36.8, surpassing general-domain baselines and a contrastive-only variant. The model showed a significant reduction in queries with zero retrieval scores compared to baselines.
• Ablation Findings:
  • The weighting factor $\alpha$ is critical; $\alpha=0.3$ yielded optimal balance, while higher values caused MLM dominance and performance drops.
  • Restricting MLM to domain tokens was more effective than applying it to the full vocabulary, even with very low weights.
  • The order of stages matters; performing joint training before pure contrastive training yielded better results than the reverse.
  • The method is computationally efficient, requiring approximately 3x the compute of standard contrastive fine-tuning, significantly less than full domain-adaptive pre-training (DAPT).

Significance and Claims
The paper claims that MOSAIC provides a principled and effective mechanism for adapting text embeddings to specialized domains without requiring complex architectural changes or large-scale retraining. By leveraging the complementary strengths of token-level and sentence-level supervision through selective masking, the method enables models to learn domain-relevant representations while preserving the robust semantic discrimination of the original general-domain model. The authors emphasize the practical value of this approach for low-resource domains, where data scarcity often precludes traditional adaptation methods, thereby facilitating equitable access to state-of-the-art language technologies. The work is presented as a model-driven approach that operates after language model pre-training, distinguishing it from prior data-driven or adapter-based techniques.