Scaling Laws and Symmetry, Evidence from Neural Force Fields

Este artículo demuestra que incorporar simetrías de tarea mediante arquitecturas equivariantes, particularmente aquellas con representaciones de orden superior, mejora significativamente las leyes de escalado para los potenciales interatómicos, lo que sugiere que los sesgos inductivos fundamentales deben diseñarse explícitamente en lugar de dejarse a que los modelos los descubran a medida que escalan.

Lo esencial

Imagina que estás intentando enseñar a un robot cómo predecir cómo se moverán e interactuarán los átomos en una molécula. Esto es un poco como enseñarle a un niño a entender cómo se mantiene unida una estructura compleja de Lego. Puedes darle al robot dos tipos diferentes de manuales de instrucciones:

1. El Manual "Ciego": Solo le muestras al robot millones de imágenes de estructuras de Lego y dices: "Descubre las reglas tú mismo". El robot tiene que aprender todo desde cero, incluido el hecho de que si giras toda la estructura, la física no cambia.
2. El Manual "Simetría": Le das al robot un manual que dice explícitamente: "Oye, recuerda, si giras esta estructura, sigue siendo la misma estructura. Si la volteas, las reglas permanecen iguales". Incrustas las leyes de la física (simetría) directamente en el cerebro del robot.

Durante mucho tiempo, muchos investigadores creyeron en el enfoque "Ciego". Pensaban que si simplemente le daban al robot suficientes datos y suficiente poder de cómputo (un "cerebro más grande"), eventualmente descubriría las reglas de simetría por sí mismo. Creían que enseñar explícitamente las reglas era innecesario y que un modelo simple y flexible eventualmente alcanzaría el nivel necesario.

Este artículo dice: "En realidad, no. El manual de 'Simetría' es mucho mejor, y la brecha se amplía a medida que creces".

Aquí está el desglose de sus hallazgos usando analogías simples:

1. La Carrera: Velocidad vs. Eficiencia

Los investigadores organizaron una carrera entre diferentes tipos de cerebros de robots (arquitecturas) para ver qué tan rápido podían aprender a predecir las fuerzas atómicas.

• Los Robots "Ciegos" (Sin restricciones): Son flexibles pero ineficientes. Tienen que "volver a aprender" el hecho de que una molécula rotada es la misma molécula cada vez que la ven.
• Los Robots "Simetría" (Equivariantes): Tienen las reglas de rotación y traslación integradas. No desperdician energía en volver a aprender física básica.

El Hallazgo: Cuando los robots eran pequeños, la diferencia no era enorme. Pero a medida que los investigadores hacían los robots masivos (escalando los datos y el poder de cómputo), los robots de "Simetría" no solo se mantuvieron adelante; se distanciaron dramáticamente. Los robots "Ciegos" chocaron contra un muro donde agregar más datos no les ayudaba mucho, mientras que los robots de "Simetría" seguían volviéndose más y más inteligentes.

2. El "Grado" de Simetría Importa

No todos los robots de "Simetría" son iguales. Algunos solo entienden rotaciones simples (como una moneda plana), mientras que otros entienden rotaciones complejas en 3D (como un globo giratorio).

• Simetría de Bajo Orden: Entiende reglas básicas.
• Simetría de Alto Orden: Entiende reglas muy complejas y detalladas sobre cómo interactúan las formas en el espacio 3D.

El Hallazgo: Cuanto más complejas son las reglas de simetría incrustadas en el robot, más rápido aprende. Un robot con simetría de "Alto Orden" aprendió tan rápido que la brecha entre él y el robot "Ciego" se convirtió en un abismo. Es como comparar a un estudiante que conoce el alfabeto con un estudiante que ya conoce la gramática y el vocabulario del idioma; a medida que el libro se vuelve más grueso, el segundo estudiante deja al primero muy atrás.

3. La "Lección Amarga" vs. la Realidad

Existe una idea famosa en la IA llamada la "Lección Amarga", que sugiere que deberíamos dejar de intentar codificar manualmente el conocimiento humano (como la simetría) en la IA y simplemente dejar que la IA lo aprenda de los datos brutos porque es más barato y escala mejor.

• Este artículo argumenta: En el mundo de los átomos y las moléculas, la "Lección Amarga" es incorrecta. Si intentas dejar que un modelo descubra la simetría por sí mismo, es como pedirle a un estudiante que redescubra la gravedad. Es posible, pero es increíblemente ineficiente. Para cuando el estudiante lo descubre, el estudiante al que se le enseñó la gravedad ya está volando.

4. El Equilibrio "Goldilocks"

El artículo también analizó cómo gastar dinero (poder de cómputo) de la manera más eficiente.

• La Vieja Forma: Quizás deberías comprar un cerebro más grande (más parámetros) o conseguir más libros de texto (más datos).
• El Nuevo Hallazgo: Resulta que necesitas comprar ambos al mismo tiempo. Si duplicas tus datos, también deberías duplicar el tamaño de tu modelo. Este "escalado en tándem" funciona mejor para todos los tipos de robots, pero los robots de "Simetría" son simplemente mucho más eficientes al usar ese poder combinado.

5. ¿Qué pasa con "hacer trampa" con las funciones de pérdida?

Algunos investigadores intentaron engañar a los robots "Ciegos" agregando una puntuación de penalización si cometían un error sobre la simetría (por ejemplo: "Si dices que una molécula rotada es diferente, recibes una mala calificación").

• El Hallazgo: Esto no funcionó bien. Es como decirle a un estudiante: "No olvides las reglas", pero en realidad no enseñarle las reglas. El robot aún tenía que luchar para aprender el patrón. Era mucho mejor simplemente construir la regla en el cerebro del robot desde el principio.

La Conclusión

Si quieres construir una IA superinteligente para entender moléculas, no solo lances más datos a un modelo simple y flexible y esperes que descubra las leyes de la física. Construye las leyes de la física directamente en el diseño del modelo.

A medida que escalas a tamaños masivos, los modelos que respetan las simetrías fundamentales del universo (rotación, traslación) no solo serán ligeramente mejores; serán exponencialmente más poderosos que aquellos que intentan aprender estas reglas desde cero. El enfoque de "Simetría" cambia la naturaleza misma de la curva de aprendizaje, haciendo la tarea más fácil y los resultados mejores.

Resumen técnico

Enunciado del Problema
El artículo aborda el comportamiento de escalado de los Potenciales Interatómicos de Redes Neuronales (NNIPs), que son modelos de aprendizaje profundo diseñados para predecir propiedades mecánico-cuánticas (específicamente energía potencial y fuerzas atómicas) de sistemas atómicos. Aunque la literatura reciente en lenguaje natural y visión sugiere que las leyes de escalado (relaciones de ley de potencias entre el rendimiento y los datos/parámetros/computación) son en gran medida independientes de la arquitectura —lo que implica que los modelos pueden aprender por sí mismos los sesgos inductivos necesarios, como la simetría, a medida que escalan—, esta visión es cuestionada en dominios geométricos. Los autores investigan si la equivarianza arquitectónica explícita (imponiendo simetrías rotacionales y de permutación) proporciona una ventaja distintiva en las leyes de escalado para los NNIPs, o si modelos más simples, no equivariantes, pueden lograr un rendimiento comparable dado suficiente cómputo.

Metodología
Los autores realizan un estudio empírico exhaustivo sobre el conjunto de datos de moléculas neutras OpenMol (aprox. 34 millones de muestras de entrenamiento, ~9.2 × 10⁸ tokens). Comparan cuatro familias arquitectónicas distintas que representan diversos grados de restricciones de simetría:

1. MPNN sin restricciones: Una red neuronal de paso de mensajes (Message Passing Neural Network) estándar que procesa características geométricas (posiciones relativas) sin restricciones de simetría.
2. Escalares invariantes (GemNet-OC): Utiliza características invariantes (distancias, ángulos, dihedros) pero aproxima funciones equivariantes mediante paso de mensajes basado en aristas; se clasifica como una arquitectura de 4 cuerpos, orden tensorial $\ell=0$.
3. Vectores cartesianos (EGNN): Una GNN $E(n)$-equivariante que utiliza canales vectoriales (orden tensorial $\ell=1$).
4. Tensores esféricos de alto orden (eSEN): Una red equivariante que utiliza representaciones irreducibles de orden superior del grupo de rotación ($\ell \ge 2$), empleando alineación de marcos para espaciar los productos tensoriales.

El estudio emplea un régimen de entrenamiento de una sola época para alinearse con la literatura teórica de leyes de escalado, utilizando optimizadores AdamW sin programador de tasas de aprendizaje para mitigar artefactos de la programación de la tasa de aprendizaje. Las leyes de escalado se ajustan frente a tres métricas:

• Cómputo: Tanto FLOPs teóricos ($C$) como tiempo de entrenamiento en tiempo real (horas de GPU, $H$).
• Datos: Número de tokens de entrenamiento ($D$).
• Parámetros: Tamaño del modelo ($N$).

Los autores también investigan los efectos de la regularización por pérdida de simetría (penalizando desviaciones de la equivarianza en modelos no equivariantes), el entrenamiento de múltiples épocas con aumento de datos y el promediado de grupos en tiempo de prueba.

Contribuciones Clave

1. Exponentes de Escalado Dependientes de la Arquitectura: El artículo demuestra que los exponentes de escalado no son constantes entre arquitecturas. A medida que aumenta el "grado" de equivarianza (de sin restricciones a de bajo orden a de alto orden), los exponentes de ley de potencias para los datos ($\beta$) y los parámetros ($\alpha$) aumentan significativamente.
2. Escalado Superior de Modelos Equivariantes: Las arquitecturas equivariantes, particularmente aquellas con representaciones tensoriales de alto orden (eSEN), exhiben curvas de escalado más pronunciadas. Esto implica que la brecha de rendimiento entre modelos equivariantes y no equivariantes se amplía a medida que escalan el cómputo y los datos, contradiciendo la noción de que los modelos pueden simplemente "aprender" la simetría más tarde.
3. Asignación Óptima de Cómputo: El estudio encuentra que, para un entrenamiento óptimo en cómputo, el tamaño del modelo ($N$) y el tamaño del conjunto de datos ($D$) deben escalar al unísono ($N \propto D$) en todas las arquitecturas, reflejando hallazgos en el modelado de lenguaje (escalado Chinchilla). Sin embargo, la constante de proporcionalidad y la reducción de pérdida resultante difieren según el sesgo de simetría de la arquitectura.
4. Ineficacia de la Pérdida de Simetría: Imponer simetría a través de un término de pérdida (regularización) en modelos sin restricciones no produce los mismos beneficios de escalado que integrar la equivarianza en la arquitectura. Si bien mejora ligeramente la eficiencia de los datos, no logra igualar los exponentes de escalado de los modelos nativamente equivariantes.
5. Insights sobre Múltiples Épocas y Aumento: En configuraciones de múltiples épocas con pocos datos, el aumento de datos es necesario para los modelos sin restricciones para evitar el sobreajuste y recuperar el escalado de ley de potencias. Sin embargo, incluso con aumento, los modelos sin restricciones no igualan los exponentes de escalado de los modelos equivariantes.

Resultados

• Leyes de Escalado: La pérdida de validación sigue una ley de potencias $L \propto C^{-\gamma}$. El exponente $\gamma$ aumenta con la complejidad arquitectónica:
  • MPNN sin restricciones: $\gamma \approx 0.14$
  • EGNN: $\gamma \approx 0.17$
  • GemNet-OC: $\gamma \approx 0.25$
  • eSEN (Alto orden): $\gamma \approx 0.40$
• Escalado de Datos y Parámetros:
  • Los exponentes de escalado de datos ($\beta$) oscilan entre 0.31 (sin restricciones) y 0.75 (eSEN).
  • Los exponentes de escalado de parámetros ($\alpha$) oscilan entre 0.28 (sin restricciones) y 0.82 (eSEN).
• Pérdida de Simetría: Añadir un término de pérdida de simetría a un modelo sin restricciones aumenta ligeramente el exponente de datos ($\beta$) pero disminuye el exponente de parámetros ($\alpha$), resultando en ninguna ganancia neta en la pendiente del frente óptimo de cómputo en comparación con la línea base sin restricciones.
• Profundidad: Para modelos equivariantes, la profundidad óptima de la red aumenta con el orden de la representación de rotación, mientras que los modelos sin restricciones sufren de sobre-suavizado a profundidades mayores.

Significado y Afirmaciones
El artículo argumenta que, contrariamente a la hipótesis de la "lección amarga" (que sugiere que los modelos deben aprender sesgos inductivos de los datos), la simetría arquitectónica explícita es crítica para el escalado en tareas geométricas. Los autores afirman que la simetría no es meramente una técnica de reducción de datos, sino que altera fundamentalmente la dificultad inherente de la tarea y sus leyes de escalado.

El significado principal radica en el hallazgo de que las representaciones equivariantes de alto orden se traducen en mejores exponentes de escalado. Esto sugiere que, para NNIPs a gran escala, invertir en arquitecturas complejas y conscientes de la simetría (como eSEN) es más efectivo que escalar modelos más simples y no equivariantes. Los autores concluyen que los sesgos inductivos fundamentales como la simetría deben codificarse en la arquitectura en lugar de dejarse para que el modelo los descubra, ya que cambian la propia trayectoria de escalado.

El artículo mantiene una postura moderada respecto a su alcance, señalando limitaciones como el enfoque en el entrenamiento de una sola época, el conjunto de datos específico utilizado (moléculas neutras) y la exclusión de estrategias de preentrenamiento de eliminación de ruido utilizadas en otros trabajos recientes. Hace un llamado a futuros trabajos teóricos para explicar por qué la simetría cambia los exponentes de escalado y sugiere extender estos estudios a tipos moleculares más diversos y regímenes de múltiples épocas.