jFoF: GPU Cluster Finding with Gradient Propagation

Resumen Técnico de "jFoF: Búsqueda de Cúmulos de GPU con Propagación de Gradientes"

Planteamiento del Problema
Las simulaciones cosmológicas dependen en gran medida del algoritmo de Amigos de los Amigos (FoF, por sus siglas en inglés) para identificar halos de materia oscura, los componentes fundamentales de la estructura cósmica. Aunque es efectivo, las implementaciones tradicionales de FoF enfrentan dos limitaciones primarias en los flujos de trabajo astrofísicos modernos:

1. Cuellos de botella de rendimiento: Las implementaciones estándar suelen depender de estructuras de datos espaciales basadas en CPU (por ejemplo, árboles $k$-d) o marcos de trabajo de memoria distribuida MPI. Estos enfoques incurren en una sobrecarga significativa debido a las transferencias de datos entre el host y el dispositivo, y tienen dificultades para explotar el paralelismo masivo de las arquitecturas GPU modernas, particularmente cuando se integran en flujos de trabajo de simulación nativos de GPU (por ejemplo, PMWD, GRAM-X).
2. No diferenciabilidad: La búsqueda de halos implica operaciones discretas —umbralización, selección de vecinos y conectividad de grafos— que son inherentemente no diferenciables. Esto impide la propagación de gradientes a través del paso de identificación de halos, obstaculizando el uso de métodos de optimización y de inferencia basados en gradientes (como el Monte Carlo de Hamilton) para calibrar modelos cosmológicos o reconstruir condiciones iniciales directamente a partir de catálogos de halos observados.

Metodología
Los autores presentan jFoF, un buscador de halos de Amigos de los Amigos totalmente nativo de GPU implementado en JAX. El código está diseñado para operar enteramente en el dispositivo, eliminando las costosas transferencias de datos CPU-GPU y permitiendo la integración en flujos de trabajo astrofísicos diferenciables.

• Estrategias de Búsqueda de Vecinos: jFoF implementa dos algoritmos complementarios de búsqueda de vecinos optimizados para la ejecución en GPU:
  1. Árbol $k$-d: Un método de partición espacial jerárquico adaptado para GPUs mediante el uso de arreglos planos y formas fijas para evitar la recursión y la indirección de punteros. Utiliza la biblioteca jaxkd, ofreciendo versiones tanto de JAX puro como compiladas en CUDA.
  2. Rejilla de Celdas Vinculadas (Linked-Cell Grid): Un método que discretiza el volumen de simulación en una rejilla uniforme de celdas cúbicas. Este enfoque restringe las búsquedas de vecinos a un estarcido (stencil) local de celdas adyacentes, reduciendo la complejidad de $O(N^2)$ a $O(N)$ y manejando naturalmente las condiciones de contorno periódicas mediante el envoltorio de imagen mínima (minimum-image wrapping).
• Construcción de Grupos: Ambos métodos emplean un esquema de propagación de etiquetas iterativo. Las partículas adoptan la etiqueta mínima entre ellas y sus vecinos vinculados, convergiendo a componentes conectados que definen los halos.
• Propagación de Gradientes: Para abordar la no diferenciabilidad del agrupamiento discreto, jFoF introduce tres modos para la estimación de gradientes:
  1. Asignación Congelada (Frozen Assignment): La pertenencia al halo se trata como fija durante la retropropagación. Los gradientes fluyen solo a través de variables continuas (por ejemplo, posiciones de partículas), permitiendo la diferenciación de las posiciones de los halos y de las propiedades de los halos definidas como funciones analíticas de las coordenadas de las partículas (por ejemplo, proxies "decorados" como la varianza de la partícula para la masa).
  2. Optimización Topológica (REINFORCE): Un enfoque estocástico donde el criterio de vinculación se relaja en un proceso de Bernoulli probabilístico. Utilizando el estimador REINFORCE (de función de puntuación), los gradientes se propagan a través de la topología discreta del grafo del halo. Esto permite la optimización de la conectividad, la pertenencia y el número de halos en expectativa, uniendo campos continuos con catálogos discretos.

Resultados Clave

• Pruebas de Rendimiento (Benchmarks): Los autores compararon jFoF con cFoF, una implementación basada en CPU paralela que utiliza MPI. En simulaciones que van desde $64^3$ hasta $512^3$ partículas, las implementaciones de jFoF nativas de GPU (tanto el árbol $k$-d como la rejilla de celdas vinculadas) lograron aceleraciones de hasta un orden de magnitud en horas-núcleo en comparación con la base de CPU, particularmente en conteos de partículas grandes.
• Exactitud: Los catálogos de halos generados por jFoF muestran una concordancia casi completa con la implementación estándar de cFoF. Los coeficientes de correlación cruzada para las posiciones de los halos superan el $0.995$ en todas las escalas. Se observaron pequeñas discrepancias solo en el extremo de muy baja masa ($\leq 40$ partículas), atribuidas a diferencias de precisión de punto flotante.
• Validación de Gradientes:
  • En las pruebas de asignación congelada, los autores demostraron que los gradientes de las posiciones de los halos y de las propiedades de masa de los proxies (varianza de la partícula) pueden computarse eficientemente. Estos gradientes guiaron con éxito a un optimizador de descenso de gradiente para recuperar un parámetro cosmológico ($\sigma_8$) a partir de una simulación de fase inicial fija, convergiendo al valor real en pocas evaluaciones de la función.
  • En las pruebas de optimización topológica, el enfoque basado en REINFORCE reconstruyó con éxito una distribución de campo de halos objetivo a partir de condiciones iniciales aleatorias. El optimizador aprendió a fragmentar halos fuera de la región objetivo y a fusionar partículas dentro de ella, recuperando la morfología de la estructura a gran escala con altos coeficientes de correlación.

Significado y Reivindicaciones
El artículo sostiene que jFoF establece una base para la identificación de estructuras diferenciable. Al permitir la optimización de extremo a extremo, basada en gradientes, de los modelos de formación de estructuras dentro de flujos de trabajo acelerados por GPU, jFoF permite:

• Inferencia Directa: La capacidad de realizar inferencia sobre parámetros cosmológicos o modelos de física de sub-malla mediante la diferenciación directa a través del proceso de búsqueda de halos, en lugar de depender de modelos de sesgo aproximados o emuladores.
• Flujos de Trabajo Unificados: La integración de la simulación, el análisis y la inferencia en un único grafo computacional diferenciable, facilitando las restricciones conjuntas sobre parámetros físicos.
• Escalabilidad: La demostración de que la búsqueda de halos por FoF puede reimplementarse para arquitecturas de aceleradores modernos sin sacrificar la exactitud, allanando el camino para el análisis de simulaciones que superan los $10^9$ partículas.

Los autores mantienen un tono modesto respecto al enfoque de optimización topológica, señalando que, si bien proporciona una estimación de gradiente no sesgada, presenta una alta varianza y requiere un ajuste cuidadoso de los parámetros de relajación y las estrategias de muestreo. Posicionan a jFoF no como una solución final para todos los problemas dinámicos, sino como una herramienta de prueba de concepto que cierra la brecha entre los catálogos observacionales discretos y los campos físicos continuos.