MIPIC: Matryoshka Representation Learning via Self-Distilled Intra-Relational and Progressive Information Chaining

1. Declaración del Problema

El Aprendizaje de Representación Matryoshka (MRL) tiene como objetivo crear incrustaciones anidadas donde un único vector de alta dimensión contiene prefijos significativos de menor dimensión. Esto permite una inferencia adaptativa, donde un modelo puede truncar las incrustaciones para ajustarse a diferentes presupuestos computacionales sin necesidad de reentrenamiento.

Sin embargo, los métodos existentes de MRL enfrentan dos desafíos críticos:

1. Incoherencia Estructural: Los enfoques actuales a menudo dependen de la alineación a nivel de oración o de una supervisión independiente en diferentes niveles de truncamiento. No logran organizar explícitamente cómo se dispone la información semántica internamente a través de las dimensiones de la incrustación y las capas de la red. En consecuencia, los prefijos de baja dimensión a menudo pierden relaciones estructurales críticas encontradas en el espacio de dimensión completa.
2. Cuello de Botella de Compresión Semántica: Comprimir conocimientos ricos y específicos de la tarea desde capas profundas hacia prefijos estrechos es difícil. Los métodos estándar a menudo aplican supervisión solo en la capa final, obligando al modelo a realizar un "salto semántico" que puede provocar inestabilidad y pérdida de información en las representaciones tempranas y de baja capacidad.

2. Metodología: Marco MIPIC

Los autores proponen MIPIC, un marco de entrenamiento unificado que sincroniza el aprendizaje de representaciones anidadas con el flujo de información interno del modelo. Emplea dos mecanismos sinérgicos: Alineación de Relaciones Intra Auto-Distilada (SIA) y Encadenamiento Progresivo de Información (PIC).

A. Alineación de Relaciones Intra Auto-Distilada (SIA)

SIA aborda la consistencia estructural cruzada-dimensiones. En lugar de simplemente alinear vectores de oraciones, reorganiza la jerarquía interna de características para priorizar las características estructurales centrales dentro de subespacios de baja dimensión.

• Coincidencia de Distribución de Atención: La representación de dimensión completa actúa como un "maestro". Alinea el orden de importancia relativa de los tokens (derivado de las puntuaciones de atención) entre la dimensión completa y los prefijos truncados. Esto asegura que los tokens más informativos permanezcan enfatizados incluso después de la compresión.
• Alineación de Estados Ocultos Top-k mediante CKA: Para evitar la saturación de información en cuellos de botella pequeños, SIA utiliza una estrategia de Selección Dura Top-k. Selecciona los $k$ tokens más importantes (basados en puntuaciones de atención) y alinea la estructura geométrica de estos subconjuntos entre el estudiante (truncado) y el maestro (completo) utilizando Alineación de Núcleo Centrado (CKA).
  • Se elige CKA por su invarianza a transformaciones ortogonales y escalado isotrópico, lo que lo hace robusto para comparar representaciones de diferentes dimensiones.
  • Esto crea una "doble estructura Matryoshka": los núcleos de baja dimensión capturan un esqueleto semántico grueso, mientras que las dimensiones superiores añaden detalles más finos.

B. Encadenamiento Progresivo de Información (PIC)

PIC aborda la consolidación semántica en profundidad. Guía el flujo de información específica de la tarea desde capas más profundas (más expresivas) hacia capas anteriores.

• Alineación Andamiada: A diferencia de los métodos que supervisan solo la salida final, PIC construye una tubería donde cada capa actúa como un puente semántico. Propaga señales relevantes para la tarea desde las capas superiores hacia las inferiores.
• Objetivo InfoNCE: El método maximiza la información mutua entre puntos de control adyacentes (capas) en la red. Utiliza un proyector ligero para mapear representaciones aguas arriba de menor dimensión al espacio objetivo aguas abajo, minimizando una pérdida InfoNCE.
• Supervisión Selectiva: La supervisión se restringe al subespacio compacto reorganizado por SIA. Esto permite que las dimensiones restantes continúen modelando características lingüísticas de bajo nivel (sintaxis, semántica general) sin sobre-regularización, asegurando una base flexible para el refinamiento específico de la tarea.

C. Objetivo General de Entrenamiento

La función de pérdida total combina la pérdida de tarea estándar de MRL con las pérdidas de alineación propuestas:
$$L_{MIPIC} = \alpha L_{MRL} + (1 - \alpha)[L_{SIA} + L_{PIC}]$$
Donde $\alpha$ equilibra el rendimiento de la tarea y la alineación estructural/semántica.

3. Contribuciones Clave

1. Marco Unificado: MIPIC es el primer marco que coordina explícitamente tanto la alineación cruzada-dimensiones (vía SIA) como la consolidación de información en profundidad (vía PIC) para resolver la incoherencia estructural y semántica en MRL.
2. Mecanismos Nuevos:
   • SIA: Introduce auto-distilación Top-k CKA para imponer consistencia topológica, forzando estructuras semánticas críticas hacia prefijos de baja dimensión.
   • PIC: Introduce una estrategia de aprendizaje andamiado que permite a las capas tempranas adquirir información relevante para la tarea progresivamente, actuando como un "puente semántico".
3. Validación Empírica Exhaustiva: El método fue probado en diversas arquitecturas base (TinyBERT, BERT-base, BGEM3, Qwen3-0.6B) y puntos de referencia (STS, NLI, Clasificación).

4. Resultados Experimentales

• Rendimiento: MIPIC logra un rendimiento competitivo o superior en comparación con las líneas base más avanzadas (MRL, Incrustaciones de Oración Espresso, SimCSE) en todas las capacidades.
• Ventaja de Baja Dimensión: Las ganancias más significativas se observan en regímenes de extrema baja dimensión (16 y 32 dimensiones). Por ejemplo, en el conjunto de datos Banking77 con TinyBERT, MIPIC logró una precisión del 48.93% a 16 dimensiones, superando significativamente a MRL (40.64%) y ESE (43.30%).
• Generalización: El método demuestra robustez tanto en configuraciones dentro del dominio como fuera del dominio (OOD), manteniendo una alta densidad semántica incluso bajo un truncamiento severo.
• Estudios de Ablación:
  • Eliminar SIA causa caídas bruscas de rendimiento en tamaños pequeños, confirmando que el orden geométrico es clave para la densidad de información.
  • Eliminar PIC conduce a pérdida semántica en etapas tempranas, destacando la necesidad de una guía progresiva.
  • La escalada progresiva de dimensiones (usando diferentes dimensiones en diferentes capas) supera a mantener todas las capas en tamaño completo.

5. Significado y Compensaciones

• Eficiencia de Inferencia: MIPIC introduce cero latencia adicional de inferencia ni parámetros. Todos los proyectores auxiliares y módulos de alineación se descartan después del entrenamiento. El modelo se comporta exactamente como una arquitectura base estándar durante la inferencia, pero con representaciones truncadas superiores.
• Costo de Entrenamiento: El método incurre en una sobrecarga de entrenamiento mayor (aproximadamente un 45% menos de rendimiento que ESE) debido al cálculo de pérdidas auxiliares CKA e InfoNCE a través de múltiples capas. Sin embargo, los autores argumentan que esto es una inversión única justificable dados las ganancias masivas en la calidad de la representación y la eficiencia de compresión, especialmente para modelos desplegados en entornos con recursos limitados.
• Impacto: MIPIC proporciona un enfoque principiado para organizar la información semántica, haciendo que las incrustaciones Matryoshka sean viables para escenarios de compresión extrema donde los métodos anteriores no lograron mantener la expresividad semántica.