Towards Engineering Scaling Laws with Pretraining Data Composition

Este artículo demuestra que en la física de partículas, donde los simuladores de alta fidelidad permiten la generación económica de datos sintéticos, el comportamiento de escalado de los modelos de clasificación de chorros hadrónicos puede diseñarse para priorizar la diversidad y la alineación de los datos sobre el tamaño del modelo mediante la curaduría estratégica de los conjuntos de datos de preentrenamiento.

Lo esencial

Imagina que estás intentando enseñarle a un estudiante a reconocer diferentes tipos de vehículos en una ciudad con mucho tráfico. Tienes dos formas principales de ayudarle a aprender: puedes darle un cerebro más grande (un modelo más grande) o puedes darle más problemas de práctica (más datos).

Durante mucho tiempo, los científicos que estudian la Inteligencia Artificial (IA) han creído que existe una "regla de oro" para esto. Pensaban que si tienes una cantidad fija de tiempo y dinero (presupuesto de cómputo), la mejor manera de obtener al estudiante más inteligente es dividir tus recursos aproximadamente 50/50 entre construir un cerebro más grande y darle más problemas de práctica.

Sin embargo, este nuevo artículo sugiere que, en el mundo de la física de partículas, podemos diseñar mediante ingeniería una regla mejor cambiando qué es lo primero que aprende el estudiante.

El Escenario: El Aula de Física

Los investigadores están trabajando con "jets". En la física de partículas, cuando partículas diminutas chocan entre sí, lanzan chorros de otras partículas llamadas jets. Es como un fuego artificial explotando, pero en lugar de chispas, obtienes corrientes de partículas subatómicas.

El objetivo es enseñar a una IA a mirar estas corrientes y decir: "¡Ah, esta provino de un tipo específico de explosión!".

El Experimento: Cambiando el Libro de Texto

Los investigadores probaron dos "libros de texto" diferentes (conjuntos de datos de preentrenamiento) para ver cómo cambiaban las reglas de aprendizaje:

1. El Libro de Texto Aburrido (Solo QCD): Este libro solo contenía ejemplos de explosiones "estándar". Era como una escuela de conducción que solo te enseñaba a conducir un sedán estándar.
2. El Libro de Texto Diverso (Mejorado con BSM): Este libro incluía los ejemplos estándar más explosiones complejas, raras y exóticas que no ocurren en nuestro universo normal (simulaciones de física "Más Allá del Modelo Estándar" o BSM). Era como una escuela de conducción que te enseñaba a conducir sedanes, pero también coches de carreras, camiones e incluso vehículos voladores.

El Descubrimiento: Reescribiendo las Reglas

Cuando los investigadores entrenaron a la IA usando el Libro de Texto Aburrido, la vieja regla del 50/50 se mantuvo vigente. Para obtener mejores resultados, tenías que equilibrar el hacer el cerebro más grande y darle más práctica.

Pero cuando usaron el Libro de Texto Diverso, las reglas cambiaron por completo. La IA aprendió que los problemas de práctica adicionales eran mucho más valiosos que un cerebro más grande.

• La Analogía: Imagina que la IA entrenada con el libro de texto diverso es como un estudiante que ya ha visto todo tipo de vehículos imaginables. Cuando le das una nueva prueba, no necesita un cerebro más grande para entender el nuevo coche; solo necesita ver más ejemplos de él para alcanzar la perfección. Su "cerebro" no necesita crecer tan rápido porque su "experiencia" es muy rica.

El Resultado: La Nueva Estrategia de "Prioridad a los Datos"

El artículo encontró que, al usar los datos diversos y exóticos para el entrenamiento inicial:

• La estrategia de "cerebro más grande" se volvió menos importante.
• La estrategia de "más datos" fue la ganadora.

De hecho, los investigadores descubrieron que por cada unidad de potencia de cómputo que gastas, deberías dedicar aproximadamente el 78% a obtener más datos y solo el 22% a hacer el modelo más grande. Este es un cambio enorme respecto a la antigua división 50/50.

Por qué esto importa para la Física

El artículo destaca una ventaja única de la física: Podemos crear nuestros propios datos.

En campos como la medicina o el lenguaje, obtener nuevos datos es difícil, costoso o imposible (no puedes simplemente "simular" un nuevo paciente humano). Pero en la física de partículas, los científicos utilizan potentes computadoras para simular colisiones de partículas. Pueden generar cantidades infinitas de datos de alta calidad y diversos de forma gratuita (una vez que la simulación está en marcha).

La Conclusión:
Si estás construciendo una IA superinteligente para la física, no intentes solo construir el cerebro más grande posible. En su lugar, dedica tu tiempo y dinero a diseñar mediante ingeniería un currículo mejor y más diverso para que la IA aprenda primero. Una vez que la IA haya visto una amplia variedad de ejemplos "exóticos", aprenderá más rápido y mejor de la tarea específica que le asignes, y obtendrás mejores resultados si le proporcionas más datos en lugar de hacer el modelo más grande.

En resumen: Una dieta de datos de entrenamiento bien elegida y diversa es más poderosa que un cerebro más grande.

Resumen técnico

Resumen Técnico: Hacia la ingeniería de las leyes de escala con la composición de datos de preentrenamiento

Declaración del problema
Las leyes de escala neuronal describen cómo el rendimiento de un modelo mejora como una ley de potencia con respecto al cómputo, el tamaño del modelo y el tamaño del conjunto de datos. Aunque están bien establecidas para los modelos de lenguaje de gran tamaño (LLM), estas relaciones están emergiendo en la física de partículas. Una distinción clave en la física fundamental es la capacidad de generar datos sintéticos de alta fidelidad mediante simuladores a un costo relativamente bajo en comparación con el gasto computacional de entrenar modelos más grandes. Esto crea una oportunidad única para la ingeniería del propio conjunto de datos de preentrenamiento para influir en el comportamiento de la escala. La pregunta central abordada es si la composición de los datos de preentrenamiento —específicamente su diversidad y alineación con las tareas posteriores— puede ser diseñada para desplazar el régimen de escala computacionalmente óptimo a favor de modelos más grandes hacia el favor de conjuntos de datos más grandes.

Metodología
El estudio se centra en la tarea de clasificar chorros hadrónicos (hadronic jets) producidos en colisiones de partículas de alta energía. Los autores utilizan una arquitectura de transformador genérica que procesa datos de chorros como una nube de puntos, variando los tamaños de los modelos desde aproximadamente 3,000 hasta 10.5 millones de parámetros (abarcando tres órdenes de magnitud) mientras mantienen fijas la profundidad y las dimensiones de las cabezas de atención.

El diseño experimental consiste en un protocolo de entrenamiento de dos etapas:

1. Preentrenamiento: Los modelos son preentrenados en subconjuntos del conjunto de datos JetClass-II, que contiene 188 clases de chorros simulados. Los autores definen cuatro subconjuntos de preentrenamiento distintos para manipular la diversidad y la alineación:
   • QCD: Solo chorros iniciados por quarks ligeros o gluones (17 clases).
   • QCD + res2p: Chorros QCD más chorros de desintegraciones de dos cuerpos de resonancias de Más Allá del Modelo Estándar (BSM).
   • QCD + res34p: Chorjes QCD más chorjes de desintegraciones de tres o cuatro cuerpos de resonancias BSM.
   • QCD + res2p + res34p: El conjunto de datos completo que incluye todas las desintegraciones de resonancias BSM.
   • Nota: Los subconjuntos BSM introducen una mayor diversidad (más clases de procesos, cobertura de espacio de fase más amplia) y una mejor alineación con la tarea posterior (topologías de múltiples ramas o multi-prong) en comparación con los datos de solo QCD.
2. Ajuste fino (Fine-tuning): Los modelos preentrenados se ajustan finamente en el conjunto de datos original JetClass para una tarea de clasificación de chorros de 10 clases (identificación de quarks ligeros/gluones, quarks top, bosones W/Z y partículas Higgs). Esta tarea requiere identificar la multiplicidad de ramas y las escalas de masa, que están bien representadas en los datos aumentados por BSM pero pobremente representadas en los datos de solo QCD.

Los exponentes de escala se extraen ajustando leyes de potencia al tamaño del modelo computacionalmente óptimo ($N^*$) y al tamaño del conjunto de datos ($D^*$) como una función del cómputo total ($C$). El estudio compara estos exponentes entre el entrenamiento "desde cero" (scratch) y las diversas configuraciones de preentrenamiento.

Resultados Clave
El estudio demuestra que la composición de los datos de preentrenamiento altera significamente los exponentes de escala computacionalmente óptimos:

• Entrenamiento desde cero (Scratch): El entrenamiento desde cero produce exponentes de $a \approx 0.52$ (tamaño del modelo) y $b \approx 0.48$ (tamaño del conjunto de datos), lo que indica una asignación de recursos computacionales aproximadamente equilibrada entre el tamaño del modelo y los datos, consistente con los hallazgos en los LLM.
• Preentrenamiento de solo QCD: El preentrenamiento únicamente en chorjes QCD resulta en un cambio marginal ($a \approx 0.53, b \approx 0.47$), lo que sugiere que el preentrenamiento por sí solo, sin una alineación o diversidad específica, no cambia fundamentalmente el régimen de escala.
• Preentrenamiento aumentado con BSM: La inclusión de desintegraciones de resonancia BSM en el corpus de preentrenamiento provoca un cambio drástico. Con el conjunto de datos completo aumentado por BSM, los exponentes cambian a $a \approx 0.22$ y $b \approx 0.78$.
  • Esto indica un régimen donde la estrategia computacionalmente óptima favorece fuertemente el aumento del tamaño del conjunto de datos sobre el aumento del tamaño del modelo.
  • El cambio representa una reducción de un factor de 2.3 en el exponente de escala para el tamaño del modelo en comparación con la línea base de entrenamiento desde cero.
  • Las curvas de pérdida de ajuste fino confirman que el preentrenamiento aumentado por BSM reduce consistentemente la pérdida en todos los tamaños de modelo, con beneficios que aumentan para los modelos más grandes.

Contribuciones Clave

1. Ingeniería de las Leyes de Escala: El artículo proporciona el primer estudio sistemático que muestra que la composición de los datos de preentrenamiento puede ser diseñada para desplazar los exponentes de escala en la física fundamental. Demuestra que la diversidad y la alineación con la tarea posterior en el corpus de preentrenamiento pueden mover el régimen de escala óptimo hacia estrategias que favorecen los datos.
2. Cambio Cuantitativo: El trabajo cuantifica el cambio de un régimen de escala equilibrado ($a \approx b \approx 0.5$) a un régimen que favorece fuertemente los datos ($a \approx 0.22, b \approx 0.78$) mediante la incorporación de la física BSM en el preentrenamiento.
3. Implicaciones para los Modelos Fundacionales: Los resultados sugieren que los modelos fundacionales preentrenados en datos sintéticos diversos y alineados pueden alcanzar un rendimiento óptimo con recuentos de parámetros más pequeños, permitiendo que los presupuestos de cómputo ahorrados se redirijan a la generación de datos de entrenamiento adicionales.

Significancia y Reivindicaciones
Los autores afirman que este trabajo identifica un nuevo espacio de diseño para el aprendizaje automático científico: los inputs físicos para el entrenamiento de modelos fundacionales. A diferencia de los dominios del lenguaje natural o de las imágenes, donde la curación de datos está limitada por la disponibilidad, la física fundamental puede aprovechar simuladores de alta fidelidad y bajo costo para construir corpus de preentrenamiento que den forma activamente a las leyes de escala.

El artículo concluye modestamente que, si bien el preentrenamiento en corpus bien compuestos permite un régimen donde el cómputo posterior se emplea mejor en más datos que en modelos más grandes, se requiere más trabajo para verificar si estos resultados se generalizan a diferentes tareas de ajuste fino, escalas de modelos más grandes y diferentes tamaños de conjuntos de datos. El estudio no pretende haber resuelto todos los desafíos de escala, sino que destaca la ingeniería de la composición del preentrenamiento como una palanca poco explorada para maximizar el potencial de descubrimiento de los modelos fundacionales científicos.