Resumen Técnico: SPOKES – Optimización para la Selección de Datos de Preentrenamiento Diversos

1. Declaración del Problema

La diversidad de los datos es un factor crítico en la construcción de corpus de preentrenamiento, ya que reduce la redundancia y mejora el rendimiento en tareas posteriores bajo presupuestos de datos fijos. Sin embargo, optimizar la diversidad es inherentemente difícil porque es una propiedad a nivel de conjunto que depende de las interacciones entre puntos de datos, más que de ejemplos individuales. La optimización directa de la diversidad es computacionalmente intratable debido a la naturaleza combinatoria de la selección de subconjuntos, particularmente a la escala del preentrenamiento.

Los enfoques existentes suelen depender de proxies o aproximaciones, como el agrupamiento (clustering) de grano grueso (por temas, habilidades o clústeres no supervisados) o la deduplicación semántica. Aunque estos métodos aseguran la representación de los clústeres, dependen de particiones de baja resolución y a menudo no logran capturar la información en la "larga cola" (long-tail) o en el conocimiento subrepresentado. Además, los métodos actuales suelen tratar las métricas de diversidad (como la puntuación G-Vendi) como herramientas de evaluación post-hoc en lugar de objetivos de optimización. El desafío central abordado por este trabajo es cómo extraer de manera fiable subconjuntos diversos de corpus de preentrenamiento a gran escala para mejorar los resultados posteriores sin depender de heurísticas que no escalan o que no capturan la diversidad fina.

2. Metodología: SPOKES

Los autores presentan SPOKES (Spokes: Optimizing for Diverse Pretraining Data Selection), un marco escalable que optimiza directamente la diversidad utilizando la puntuación G-Vendi dentro de un marco de diversificación probabilística.

2.1 Objetivo Central

El método formula la selección de datos como un problema de optimización global que maximiza conjuntamente la calidad y la diversidad. Dado un conjunto de datos de $n$ ejemplos, el objetivo es seleccionar un subconjunto $S$ de tamaño $k$ que maximice:
$$\max_{S \subseteq [n], |S|=k} \alpha \ln \left( \frac{1}{k} \sum_{i \in S} q_i \right) + (1 - \alpha) \ln \text{Vendi}(K_S)$$
Donde:

• $q_i$ representa la puntuación de calidad del ejemplo $i$ (por ejemplo, de un clasificador FineWeb-Edu).
• $K_S$ es el núcleo de similitud (similarity kernel) restringido al subconjunto seleccionado.
• $\alpha \in [0, 1]$ es un parámetro ajustable que controla el equilibrio entre calidad y diversidad.

2.2 Diversidad Basada en Gradientes (G-Vendi)

A diferencia de los métodos tradicionales que utilizan embeddings estáticos, SPOKES opera en el espacio de gradientes. Calcula los gradientes $\nabla_\theta \ell(x_i)$ con respecto a los parámetros del modelo $\theta$ para cada punto de datos utilizando un modelo proxy. La puntuación G-Vendi mide la diversidad examinando la entropía del espectro de la matriz de similitud construida a partir de estos gradientes. Este enfoque fomenta la ortogonalidad entre los vectores de representación, reduciendo la redundancia en las direcciones de actualización y promoviendo pasos de optimización más independientes e informativos.

2.3 Optimización mediante Descenso de Gradiente Exponenciado

Para resolver el problema de selección de subconjuntos combinatorios, SPOKES emplea una estrategia de optimización relajada:

1. Relajación Continua: En lugar de una selección discreta, se asigna un vector de pesos continuo no negativo $w$ a cada punto de datos, restringido al simplex ($w \in \Delta_n$).
2. Construcción del Núcleo: Se define un núcleo dependiente de los pesos $K(w)$ donde las entradas se escalan por $\sqrt{w_i w_j}$.
3. Optimización: El objetivo se optimiza mediante descenso de gradiente exponencial. En cada iteración, los pesos se actualizan mediante $w_i^{(t+1)} \propto w_i^{(t)} \exp(-\eta g_i^{(t)})$ y se normalizan para evitar el colapso sobre unos pocos ejemplos de alto peso.
4. Recuperación Discreta: Tras la convergencia, el subconjunto final se recupera seleccionando las $k$ entradas superiores del vector de pesos aprendido.

2.4 Mejoras de Escalabilidad

Para hacer que este enfoque sea viable a la escala del preentrenamiento, los autores implementan varias estrategias de eficiencia:

• Gradientes Truncados: En lugar de calcular los gradientes de todo el modelo, el método los aproxima utilizando únicamente las últimas $n$ capas del transformer (típicamente 2 capas). Esto reduce significativamente el coste computacional preservando la estructura del núcleo de similitud.
• Proyección de Johnson-Lindenstrauss (JL): Para manejar vectores de gradiente de alta dimensión (por ejemplo, 200 millones de parámetros), los autores aplican proyecciones aleatorias (matrices de Rademacher) para reducir la dimensionalidad. Los resultados empíricos muestran que incluso las proyecciones de baja dimensión (por ejemplo, $k=1024$) mantienen una alta estabilidad del subconjunto (95.5% de superposición con las referencias de alta dimensión).
• Paralelización: El conjunto de datos puede particionarse en subconjuntos disjuntos optimizados de forma independiente, lo que genera una aceleración lineal.

3. Contribuciones Clave

1. Marco SPOKES: Introducción de una técnica escalable y fundamentada para la optimización de la diversidad que opera eficientemente a la escala de los conjuntos de datos de preentrenamiento modernos, yendo más allá del agrupamiento heurístico.
2. Optimización Directa de G-Vendi: Demostración de que la puntuación G-Vendi puede optimizarse directamente mediante descenso de gradiente exponencial para extraer subconjuntos altamente diversos, en lugar de servir meramente como una métrica post-hoc.
3. Optimización Conjunta de Calidad-Diversidad: Prueba de que la optimización conjunta de calidad y diversidad produce resultados superiores en comparación con la optimización de un solo objetivo o el uso de baselines estándar.
4. Mejoras de Eficiencia: Validación de que el cálculo de gradientes truncados y la reducción de dimensionalidad mediante la proyección JL permiten escalar el método a corpus masivos sin costes computacionales prohibitivos.

4. Resultados Experimentales

El método fue evaluado en dos corpus de preentrenamiento a gran escala: DCLM (subconjunto Dolmino, filtrado agresivamente) y FineWeb (filtrado menos agresivamente). Los modelos fueron preentrenados desde cero utilizando una arquitectura LLaMA-1B con 100 mil millones de tokens.

4.1 Métricas de Diversidad

• Incremento de la Puntuación G-Vendi: En un subconjunto de 500k muestras, SPOKES logró un incremento de +489 puntos en la puntuación G-Vendi en comparación con el muestreo aleatorio. En contraste, la deduplicación semántica (SemDeDup) solo logró un incremento de +7 puntos.
• Diversidad de Lote (Batch): Durante el entrenamiento, los lotes seleccionados por SPOKES mantuvieron puntuaciones de diversidad promedio consistentemente más altas (aumentando de 315 a 520 en DCLM) en comparación con el muestreo aleatorio.

4.2 Rendimiento en Tareas Posteriores (Downstream)

La evaluación se realizó en 10 benchmarks de inglés de la suite OLMES.

• Solo Diversidad: SPOKES (solo diversidad) mejoró el rendimiento promedio en tareas posteriores en +0.4 (DCLM) y +0.5 (FineWeb) puntos sobre el muestreo aleatorio.
• Optimización Conjunta: Los resultados más sólidos se obtuvieron optimizando conjuntamente la calidad y la diversidad. SPOKES logró ganancias de +1.5 (DCLM) y +1.4 (FineWeb) puntos sobre el muestreo aleatorio, superando a todos los baselines, incluyendo la deduplicación semántica y el filtrado de solo calidad.
• Fortalezas Complementarias: Los resultados indican que la calidad y la diversidad capturan propiedades distintas pero complementarias. Mientras que el filtrado de calidad destaca en benchmarks intensivos en conocimiento (ej. MMLU), la optimización de la diversidad mejora el razonamiento de sentido común (ej. HellaSwag). La optimización conjunta cierra esta brecha, rindiendo de forma competitiva en todas las tareas.

5. Significado y Reivindicaciones

El artículo sostiene que la optimización directa de la diversidad es tanto tratable a escala como prácticamente beneficiosa. Al alejarse de las aproximaciones de grano grueso y las heurísticas, SPOKES demuestra que:

• La diversidad es un eje crítico: Incluso sin supervisión explícita de calidad, optimizar la diversidad por sí sola produce ganancias consistentes, lo que sugiere que una representación de gradiente diversa de la distribución de datos mejora la generalización.
• Las métricas de calidad son insuficientes: Las métricas de calidad estándar no capturan totalmente la utilidad de los datos. La optimización conjunta de la diversidad puede proporcionar señales de entrenamiento complementarias que conducen a un mejor rendimiento posterior, incluso cuando el subconjunto seleccionado tiene una puntuación de calidad promedio menor que un baseline de solo calidad.
• Escalabilidad: El método propuesto extrae con éxito subconjuntos densos y altamente diversos de corpus masivos de preentrenamiento, ofreciendo una alternativa viable a los pipelines de selección existentes que dependen de la deduplicación semántica o el muestreo aleatorio.

Los autores posicionan a SPOKES como una herramienta para construir subconjuntos de alta calidad y diversos, enfatizando que los beneficios de la diversidad se traducen de forma fiable a través de diferentes conjuntos de datos con distintos regímenes de filtrado y perfiles de ruido.