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

Resumen Técnico: MOSAIC

Planteamiento del Problema
El artículo aborda el desafío de adaptar modelos de incrustación (embedding) de texto de dominio general a gran escala hacia dominios especializados. Los modelos de dominio general suelen entrenarse en corpus vastos mediante enfoques de dos etapas (preentrenamiento seguido de ajuste fino supervisado), pero a menudo carecen de conocimiento específico del dominio, lo que limita su rendimiento en áreas especializadas. Los métodos de adaptación existentes enfrentan un dilema:

1. Expansión del Vocabulario sin Entrenamiento Conjunto: Simplemente añadir tokens específicos del dominio a un modelo preentrenado y continuar el entrenamiento de Modelado de Lenguaje con Máscara (MLM) degrada las propiedades contrastivas del modelo, ya que el codificador pierde sus capacidades de discriminación a nivel de oración bajo objetivos de predicción de tokens.
2. Adaptación Únicamente Contrastiva: Continuar el entrenamiento utilizando solo un objetivo contrastivo proporciona señales insuficientes para actualizar eficazmente los nuevos tokens del dominio, ya que la matriz de incrustación recibe señales diluidas debido a las funciones de agrupación (pooling).
3. Desafíos de la Optimización Conjunta: Intentar optimizar conjuntamente los objetivos MLM y contrastivo suele derivar en señales de entrenamiento conflictivas. El error (loss) de MLM suele dominar debido a su mayor tamaño de vocabulario y mayores magnitudes de gradiente, abrumando al objetivo contrastivo y resultando en un rendimiento subóptimo.

Metodología: MOSAIC
Los autores proponen MOSAIC (Masked Objective with Selective Adaptation for In-domain Contrastive Learning), un marco de trabajo de múltiples etapas diseñado para equilibrar la supervisión a nivel de token y a nivel de oración. El método se basa en un prerrequisito simple: la expansión del vocabulario.

El marco consta de tres etapas:

1. Etapa 1: Aumentación del Vocabulario. El tokenizador del modelo se expande con tokens específicos del dominio identificados a partir de un corpus del dominio. Estos nuevos tokens se añaden a la matriz de incrustación e inicializan como el promedio de las incrustaciones de subpalabras del modelo base. Crucialmente, se preservan los pesos originales del codificador para mantener las capacidades de discriminación a nivel de oración.
2. Etapa 2: Optimización Conjunta con Máscara Selectiva. El modelo se somete a un entrenamiento conjunto que combina los objetivos MLM y contrastivos. Para mitigar el dominio del error MLM, los autores restringen el objetivo MLM únicamente a los nuevos tokens específicos del dominio (vocabulario restringido al dominio). Esto reduce el denominador en la pérdida de tipo InfoNCE, equilibrando las magnitudes de gradiente entre el objetivo a nivel de token (MLM) y el objetivo a nivel de oración (contrastivo). La pérdida conjunta se formula como $L = \alpha \cdot L_{domain}^{MLM} + L_{CL}$, donde $\alpha$ equilibra las contribuciones.
3. Etapa 3: Entrenamiento Únicamente Contrastivo. Una vez aprendidos los tokens del dominio, el modelo se somete a una etapa final de entrenamiento utilizando únicamente el objetivo contrastivo. Esto sirve como un paso correctivo para reforzar la discriminación a nivel de oración, que pudo haber sido diluida durante la fase de entrenamiento conjunto.

Contribuciones Clave

1. Enfoque de Adaptación Novedoso: El artículo introduce MOSAIC, un método que optimiza conjuntamente los objetivos MLM y contrastivos dentro de un marco de maximización de información mutua, específicamente para la adaptación de dominio de modelos de incrustación de texto potentes existentes.
2. Mecanismo de Adaptación Selectiva: Los autores proponen restringir la supervisión MLM solo a los tokens específicos del dominio, una estrategia que alinea las magnitudes de gradiente y evita que el objetivo MLM abrume la señal contrastiva.
3. Validación Empírica en Diversos Recursos: El método se valida tanto en dominios de altos recursos (Biomédico) como de bajos recursos (Islámico), demostrando robustez incluso con datos limitados del dominio.
4. Análisis Exhaustivo: El estudio incluye extensos estudios de ablación que analizan el impacto del objetivo conjunto, el hiperparámetro de ponderación ($\alpha$), las estrategias de enmascaramiento y el orden de las etapas de entrenamiento.

Resultados

• Dominio de Altos Recursos (Biomédico): En el subconjunto biomédico de MTEB, MOSAIC logró mejoras de hasta un 13.4% en NDCG@10 sobre bases de comparación de dominio general robustas. El flujo completo de tres etapas (MOSAIC-Bio-Stage3) superó la expansión de vocabulario ingenua y la adaptación únicamente contrastiva, logiendo la puntuación promedio más alta en tareas de recuperación, agrupamiento (clustering) y similitud semántica de texto.
• Dominio de Bajos Recursos (Islámico): En el dominio islámico, caracterizado por la escasez de datos en inglés, MOSAIC-ID-Stage3 alcanzó un NDCG@10 de 36.8, superando a las bases de comparación de dominio general y a una variante únicamente contrastiva. El modelo mostró una reducción significativa en las consultas con puntuaciones de recuperación de cero en comparación con las bases de comparación.
• Hallazgos de Ablación:
  • El factor de ponderación $\alpha$ es crítico; $\alpha=0.3$ produjo el equilibrio óptimo, mientras que valores más altos causaron el dominio de MLM y caídas en el rendimiento.
  • Restringir el MLM a los tokens del dominio fue más efectivo que aplicarlo al vocabulario completo, incluso con pesos muy bajos.
  • El orden de las etapas importa; realizar el entrenamiento conjunto antes del entrenamiento puramente contrastivo produjo mejores resultados que el proceso inverso.
  • El método es computacionalmente eficiente, requiriendo aproximadamente 3 veces la computación de un ajuste fino contrastivo estándar, significativamente menos que el preentrenamiento de adaptación de dominio completo (DAPT).

Significancia y Reivindicaciones
El artículo afirma que MOSAIC proporciona un mecanismo basado en principios y efectivo para adaptar las incrustaciones de texto a dominios especializados sin requerir cambios arquitectónicos complejos o un reentrenamiento a gran escala. Al aprovechar las fortalezas complementarias de la supervisión a nivel de token y a nivel de oración mediante el enmascaramiento selectivo, el método permite que los modelos aprendan representaciones relevantes del dominio mientras preservan la robusta discriminación semántica del modelo original de dominio general. Los autores enfatizan el valor práctico de este enfoque para dominios de bajos recursos, donde la escasez de datos a menudo impide los métodos de adaptación tradicionales, facilitando así el acceso equitativo a tecnologías de lenguaje de vanguardia. El trabajo se presenta como un enfoque impulsado por el modelo que opera después del preentrenamiento del modelo de lenguaje, distinguiéndose de las técnicas previas basadas en datos o basadas en adaptadores.