Resumen Técnico: HERMES: Un sustrato de etiquetado de multi-granularidad para mezclas de datos de pre-entrenamiento

1. Planteamiento del problema

Las estrategias actuales de mezcla de datos para el pre-entrenamiento de Modelos de Lenguaje de Gran Escala (LLM) enfrentan un cuello de botella en el sistema de etiquetado más que en los propios algoritmos de mezcla. Los enfoques existentes suelen particionar los corpus en grupos utilizando taxonomías fijas (procedencia, temas, formatos) o clústeres de embeddings planos (por ejemplo, KMeans). Estos métodos sufren de tres limitaciones primarias:

1. Granularidad/Eje único: Se comprometen con un único eje semántico a una resolución fija. Cambiar la granularidad (por ejemplo, de temas amplios a clústeres finos) requiere reconstruir todo el conjunto de etiquetas.
2. Costo de re-computación: Los métodos de agrupamiento plano (como KMeans) requieren re-computación para cambiar el número de clústeres ($K$), lo que hace que los barridos sistemáticos de granularidad sean computacionalmente costosos y desarticulados.
3. Inflexibilidad: La elección de los grupos dicta lo que un mezclador puede expresar, limitando la capacidad de probar interacciones entre la resolución de los datos y las estrategias de muestreo.

Los autores argumentan que el cuello de botella es la infraestructura de etiquetado, no la maquinaria de mezcla. Existe la necesidad de un sustrato reutilizable, derivado de los datos, que permita a los profesionales navegar por una jerarquía de granularidades sin necesidad de re-agrupamiento.

2. Metodología: HERMES

HERMES (Hierarchical Embedding-based RVQ for Mixture Selection) es un sustrato de etiquetado jerárquico de un solo paso (one-shot) diseñado para anotar documentos con códigos de grueso a fino. Consta de dos etapas principales:

A. Transformación Semántica Aprendida (LST)

Para preparar los embeddings para la cuantización, HERMES emplea una Transformación Semántica Aprendida.

• Entrada: Embeddings de documentos congelados ($e_i \in \mathbb{R}^{1024}$) de un codificador pre-entrenado.
• Mecanismo: Una única capa lineal con normalización ($h_i = \text{normalize}(W e_i + b)$) rota el espacio de embeddings.
• Objetivo de entrenamiento: La transformación se entrena conjuntamente con los libros de códigos de cuantización para optimizar tres pérdidas:
  1. Preservación de la estructura por pares: Mantiene la geometría de coseno entre los embeddings originales y los transformados.
  2. Reconstrucción consciente de la cuantización: Minimiza el error de reconstrucción del cuantizador vectorial.
  3. Ortogonalidad: Evita el colapso de la representación imponiendo la casi-ortogonalidad de la matriz de transformación $W$.

B. RVQ de 3 etapas (Residual Vector Quantization)

Los embeddings transformados ($h_i$) son codificados por un RVQ en cascada de 3 etapas.

• Arquitectura: Tres cuantificadores vectoriales secuenciales, cada uno con un tamaño de libro de códigos $K=256$.
• Proceso:
  1. Etapa 1 (Gruesa): Codifica $h_i$ para encontrar el centroide más cercano $c_1$, produciendo un residuo $r_2 = h_i - q_1$.
  2. Etapa 2 (Intermedia): Codifica $r_2$ para encontrar $c_2$, produciendo el residuo $r_3$.
  3. Etapa 3 (Fina): Codifica $r_3$ para encontrar $c_3$.
• Salida: Cada documento recibe un código jerárquico $(c_1, c_2, c_3)$.
• Control de Granularidad: La granularidad se define por la longitud del prefijo del código:
  • L1: Código $(c_1)$ $\rightarrow$ 256 cubetas (buckets).
  • L12: Código $(c_1, c_2)$ $\rightarrow$ $\sim$65k cubetas.
  • L123: Código $(c_1, c_2, c_3)$ $\rightarrow$ $\sim$130k cubetas observadas (de un total nominal de $1.67 \times 10^7$).

C. Taxonomía de Muestreo

El artículo define un muestreador como una composición de dos etapas que operan sobre el sustrato HERMES:

1. Etapa 1 (Peso Exterior): Asigna pesos a los ancestros de L1. El artículo utiliza Uniforme (entropía máxima) o DoReMi (group-DRO sobre un modelo proxy).
2. Etapa 2 (Muestreador de Sub-cubetas): Decide la elegibilidad y selección dentro de los ancestros de L1.
   • Distribución de Masa: Igual (entropía máxima) vs. Proporcional al tamaño.
   • Elegibilidad por Documento: Todos los documentos vs. Top-30% por puntuación de calidad (FineWeb-Edu) dentro de la sub-cubeta.

3. Contribuciones Clave

1. Sustrato HERMES: Un sistema de RVQ jerárquico de un solo paso que anota $\sim$50M de documentos en hasta $\sim$130k celdas. Expone la granularidad como la longitud de un prefijo en una única pila de libros de códigos entrenada, eliminando la necesidad de re-agrupar cuando se cambia la resolución.
2. Interacción Granularidad-Muestreador: Los autores demuestran que la efectividad de una regla de muestreo (por ejemplo, selección basada en calidad) depende del nivel de granularidad.
3. Diagnóstico de Competencia de Candidatos: El artículo identifica un límite estructural donde aumentar la granularidad reduce el tamaño del grupo de candidatos a un punto donde el ranking dentro de la cubeta se vuelve inestable.

4. Resultados Experimentales

Los experimentos se realizaron en un modelo de estilo LLaMA de 1B de parámetros entrenado durante 25B de tokens, evaluado en un macro-promedio de 16 tareas (Avg.).

A. La Meseta de Agrupamiento (L1=256)

• Hallazgo: En la granularidad más gruesa (L1=256), HERMES rinde a la par con los métodos de agrupamiento estándar (KMeans, MiniBatchKMeans, Plain RVQ).
• Métricas: En métricas intrínsecas (compacidad, equilibrio de masa, entropía), HERMES se sitúa en una meseta con otros agrupadores de 256 vías.
• Downstream: Al combinarse con un muestreador Uniforme + L1 de máxima entropía, HERMES y KMeans difieren por solo 0.0002 en la precisión promedio (Avg.).
• Implicación: La ganancia de rendimiento de HERMES no se debe a ser un algoritmo de agrupamiento superior, sino al sustrato mismo.

B. Hallazgo A: Regla de Etapa-2 en Granularidad Intermedia (L12)

• Configuración: Pesos exteriores DoReMi fijos en L1, libro de códigos HERMES fijo, granularidad en L12 ($\sim$65k cubetas).
• Comparación: Cambiar el muestreador de la Etapa-2 de cobertura de máxima entropía (masa igual en sub-cubetas, todos los documentos) a top-30% de calidad (masa proporcional al tamaño, top-30% por puntuación FineWeb-Edu).
• Resultado: Este cambio aumentó el promedio de las 16 tareas en +0.0253 (z-score +2.09).
• Interpretación: En L12, las sub-cubetas son lo suficientemente grandes (pool mediano $\sim$2,271 documentos) como para que el ranking de calidad sea un proxy estable de la calidad global.

C. Hallazgo B: Colapso en Granularidad Fina (L123)

• Configuración: Misma regla aplicada a L123 ($\sim$130k cubetas).
• Observación: El pool de candidatos mediano se contrae 5.3x de L12 (2,271 docs) a L123 (429 docs).
• Resultado: La ventaja de la regla de calidad top-30% colapsa. La brecha de rendimiento entre la calidad top-30% y la cobertura aleatoria local se reduce a +0.0002 (0.3988 vs. 0.3986), lo cual es estadísticamente indistinguible bajo evaluación de semilla única.
• Implicación: Esto respalda la hipótesis de la competencia de candidatos: cuando las sub-cubetas se vuelven demasiado pequeñas, el ranking dentro de la cubeta ya no es un proxy fiable de la calidad global, anulando el beneficio de la selección basada en calidad.

D. Mejor Configuración

La configuración de mayor rendimiento en el barrido fue pesos exteriores DoReMi + calidad top-30% de L12 (lector corregido), logrando un promedio de 0.4222.

5. Significado y Reivindicaciones

El artículo reivindica modestamente que HERMES redefine el diseño de la mezcla de datos. Su principal significancia radica en desplazar el enfoque de "elegir entre conjuntos de etiquetas fijos" hacia "navegar por una jerarquía de granularidad reutilizable y derivada de los datos".

• Sustrato sobre Agrupador: La contribución es la infraestructura que permite que la granularidad sea una variable libre (un dial), no el rendimiento del agrupamiento en sí.
• Determinación Conjunta: La granularidad y el muestreador de la Etapa-2 son ejes determinados conjuntamente. Una regla que funciona en una granularidad (L12) puede fallar en otra (L123) debido a restricciones estructurales como el tamaño del pool de candidatos.
• Limitaciones: Los autores reconocen que las conclusiones están limitadas a un régimen de 1B/25B en un corpus interno. La profundidad de granularidad óptima está limitada por la "competencia de candidatos", lo que podría cambiar con modelos y corpus más grandes. También señalan que la elección de los pesos exteriores (DoReMi vs. Uniforme) en la regla óptima de la Etapa-2 (L12 calidad top-30%) no fue totalmente validada cruzadamente (no se realizó la ejecución de Uniforme + L12 calidad top-30%).

En resumen, HERMES proporciona un mecanismo para tratar la granularidad de los datos como un parámetro continuo y ajustable, revelando que la eficacia de las estrategias de selección de datos está intrínsecamente ligada a la resolución del sistema de etiquetas subyacente.