Discovering Crystal Structure Prediction Algorithms with an AI Co-Scientist

Resumen Técnico: Descubrimiento de Algoritmos de Predicción de Estructura Cristalina con un Co-Científico de IA

Planteamiento del Problema
La Predicción de la Estructura Cristalina (CSP, por sus siglas en inglés) es un desafío central en el descubrimiento de materiales computacionales, que requiere la generación de estructuras cristalinas estables a partir de composiciones químicas. El problema está limitado por la periodicidad (los cristales se repiten en el espacio), las coordenadas fraccionarias que viven en un toro, la simetría del grupo espacial y la existencia de múltiples polimorfos para una sola composición. Si bien los avances recientes en modelos generativos (por ejemplo, modelos de difusión, flow matching) han mejorado la CSP, estos suelen depender de formulaciones continuas que requieren un muestreo estocástico o iterativo, careciendo de un espacio de decisión discreto finito para una decodificación codiciosa (greedy decoding) confiable. Además, los agentes de investigación de IA existentes suelen operar dentro de especificaciones de problemas fijas, optimizando código o hiperparámetros para una familia de modelos dada, en lugar de descubrir nuevos principios de modelado a través de diversos dominios.

Metodología: El Sistema HACO
Los autores presentan HACO (sistema de co-descubrimiento Humano-IA), un marco diseñado para descubrir algoritmos científicos mediante la búsqueda trans-dominio y el guiado humano disperso. HACO opera bajo dos principios fundamentales:

1. Transferencia Trans-dominio: En lugar de perturbar modelos específicos de un dominio, el agente busca a través de marcos de modelado generativo de diversos campos (visión, lenguaje, etc.) para identificar principios transferibles.
2. Interacción Humana Dispersa: Los humanos proporcionan objetivos o mecanismos de alto nivel (por ejemplo, "mejorar la cobertura de polimorfos" o "incorporar simetría"), mientras que el agente se encarga de la implementación, el entrenamiento, la evaluación y el refinamiento.

HACO emplea una búsqueda estructurada en árbol orquestada por un LLM (Claude Opus 4.7). La búsqueda procede en tres etapas:

• Etapa de Descubrimiento: Busca la metodología generativa (presupuesto de 2 horas por candidato).
• Etapa de Escalamiento: Optimiza la metodología seleccionada con presupuestos mayores (presupuesto de 12 horas).
• Optimización de Muestreo: Ajusta el algoritmo de muestreo.

Los operadores clave incluyen idea (proponer marcos trans-dominio), draft (instanciar código), improve (realizar cambios coherentes) y debug. El operador idea identificó específicamente MaskGIT, un modelo generativo enmascarado de visión computacional, como un candidato prometedor para la CSP.

El Algoritmo Resultante: MaskGXT
A través del proceso HACO, el sistema desarrolló MaskGXT (Masked Generative Crystal Transformer). A diferencia de los modelos de difusión continuos, MaskGXT formula la CSP como una tarea de generación de tokens discretos.

• Tokenización: Un cristal se representa como una secuencia de $1 + 6 + 5N$ tokens:
  • 1 token de Grupo Espacial.
  • 6 tokens de Red (longitudes y ángulos cuantizados).
  • $5N$ tokens de sitio (3 tokens de coordenadas cuantizadas, 1 token de posición de Wyckoff y 1 token de tipo de átomo).
  • Los tipos de átomos se proporcionan como información lateral (condicionamiento).
• Discretización y Suavizado: Las coordenadas fraccionarias continuas se mapean a bins discretos en un círculo (módulo 1). Para preservar la estructura métrica, el modelo utiliza un suavizado de etiquetas ordinal y circular para las coordenadas y un suavizado ordinal estándar para los parámetros de red, en lugar de tratar los bins como categorías no ordenadas.
• Manejo de Simetría: El modelo aprende a asignar tokens consistentes de grupo espacial y Wyckoff a través de diferentes orígenes de la celda unidad y ordenamientos atómicos mediante aumentos en línea (desplazamientos de origen normalizadores euclidianos y permutaciones intra/inter-órbita).
• Arquitectura: Un tronco de Transformer bidireccional con atención de normalización QK y redes feedforward SwiGLU. Incluye una cabeza de regresión ligera para predecir desplazamientos de sub-bin (sub-bin offsets), refinando la precisión geométrica más allá de la resolución del bin discreto.
• Muestreo:
  • Muestreo Estratificado de Grupos Espaciales: El modelo predice una distribución posterior sobre los grupos espaciales. Para la cobertura de polimorfos, la generación se ramifica en flujos separados para los grupos espaciales con mayor probabilidad posterior.
  • Decodificación Codiciosa Clasificada por Confianza: En lugar de un muestreo estocástico, el modelo desenmascara determinísticamente los tokens con mayor confianza (argmax), permitiendo la generación eficiente de estructuras de alta probabilidad.

Contribuciones Clave

1. Marco HACO: Demuestra que un co-científico de IA, guiado por objetivos humanos dispersos, puede identificar un principio de modelado transferible (MaskGIT de la visión) y adaptarlo a un dominio científico (CSP) para producir un algoritmo de vanguardia.
2. MaskGXT: Introduce un transformador generativo enmascarado discreto para cristales que unifica la red, las coordenadas, los grupos espaciales y las posiciones de Wyckoff en una única secuencia de tokens, permitiendo una decodificación codiciosa confiable y un manejo explícito de la simetría.
3. Ganancias de Rendimiento: Muestra que un solo modelo MaskGXT logra tasas de coincidencia (match rates) de vanguardia en bancos de pruebas estándar y mejora significativamente la cobertura de polimorfos.

Resultados Experimentales
El modelo fue evaluado en los bancos de pruebas MP-20 y MPTS-52, así como en el split de polimorfos MP-20 (que prueba la cobertura de múltiples estructuras por composición).

• CSP Estándar (Uno-a-Uno): MaskGXT logró la tasa de coincidencia (MR) más alta en MP-20 (73.79% sin filtrar) y MPTS-52 (36.75% sin filtrar), superando bases sólidas como Crystalite, MCFlow y FlowMM. También logró el RMSE de pares emparejados más bajo.
• Cobertura de Polimorfos (METRe): En el split de polimorfos MP-20, MaskGXT alcanzó un 79.06% de precisión de Match-Everyone-to-Reference (METRe), comparado con el 70.87% de la base más fuerte (Crystalite). Esta mejora se atribuye al muestreo estratificado de grupos espaciales, que asigna el presupuesto de generación a distintos grupos espaciales de alta probabilidad.
• Estudios de Ablación:
  • Eliminar los tokens de simetría causó una caída significativa en el METRe en el split de polimorfos, confirmando la necesidad de un modelado de simetría explícito para la cobertura.
  • Eliminar la regresión de desplazamiento de sub-bin tuvo un efecto mínimo en la tasa de coincidencia pero aumentó sustancialmente el RMSE, probando su función en la precisión geométrica.
  • La decodificación codiciosa mejoró la precisión de la estructura única, mientras que el muestreo estratificado fue crítico para la cobertura de polimorfos.

Significado y Reivindicaciones
El artículo afirma que en dominios caracterizados por una "validación barata, rápida y bien alineada" (como la CSP con emparejamiento de estructuras automatizado), los co-científicos de IA interactivos guiados por transferencia pueden contribuir al descubrimiento de algoritmos científicos. Específicamente, el sistema logró con éxito:

• Identificar un principio de modelado de un campo no relacionado (modelado generativo enmascarado en visión) y adaptarlo a las restricciones de la cristalografía.
• Combinar esta transferencia con la guía humana dirigida (por ejemplo, solicitar la generación consciente de la simetría) para refinar el algoritmo.
• Producir un algoritmo competitivo que avanza en los estándares de la industria sin requerir que el humano especifique detalles de implementación de bajo nivel.

Los autores señalan que el éxito de este bucle depende de la estructura específica del problema de la CSP (conjuntos de datos fijos, ciclos de entrenamiento rápidos y métricas de validación fiables). Sugieren que extender este enfoque a dominios con experimentos físicos o indicadores indirectos (proxy objectives) débiles requeriría una evaluación más estructurada y un guiado humano más fuerte.