Neural Field Tokenizations with Hierarchy and Spatial Locality Priors

Resumen Técnico: Tokenizaciones de Campos Neuronales con Jerarquía y Priores de Localidad Espacial

Declaración del Problema

Los campos neuronales parametrizan los datos como funciones continuas de coordenadas a valores (por ejemplo, $f_\theta: [0, 1]^2 \to \mathbb{R}^3$), ofreciendo un marco unificado para el aprendizaje de representaciones a través de diversas modalidades. Sin embargo, los enfoques agnósticos a la modalidad existentes, como Functa y sus descendientes, dependen fuertemente del meta-aprendizaje por muestra (por ejemplo, MAML o auto-decodificación) para inferir las representaciones latentes. Este paradigma presenta tres limitaciones críticas:

1. Cuello de Botella de Escalabilidad: El meta-aprendizaje requiere almacenar el grafo de computación completo para la optimización del bucle interno con el fin de calcular gradientes de segundo orden. Esto limita severamente los tamaños de lote y la eficiencia de memoria, dificultando el escalado a altas resoluciones o grandes conjuntos de datos.
2. Falta de Estructura: Los espacios latentes derivados de la optimización carecen de menudo de priores estructurales explícitos, lo que hace que el aprendizaje sea ineficiente y las representaciones resultantes sean difíciles de utilizar para tareas posteriores como la clasificación.
3. Compromiso (Trade-off): Existe una tensión fundamental entre introducir útiles priores estructurales (como localidad y jerarquía) y mantener la agnosticidad de la modalidad. La mayoría de los métodos que introducen estructura (por ejemplo, rejillas 2D, nubes de puntos) sacrifican la capacidad de manejar dominios de coordenadas arbitrarios.

Metodología: LH-NeF

Los autores proponen LH-NeF (Campos Neuronales Jerárquicos con Preservación de Localidad), un marco para aprender representaciones tokenizadas de señales continuas de propósito general sin comprometer la agnosticidad de la modalidad. El método reemplaza la optimización por muestra con una única pasada hacia adelante a través de una arquitectura de codificador-decodificador aprendida.

1. El Tokenizador LH-NeF (Codificador)
El tokenizador se basa en los Perceptores Jerárquicos (HiP) pero los modifica para respetar la localidad en el espacio de coordenadas:

• Ordenamiento con Preservación de Localidad: A diferencia del HiP estándar, que utiliza un escaneo de rasterización fijo (violando la localidad espacial en el espacio de coordenadas), LH-NeF reordena las observaciones de entrada $\{(x_i, v_i)\}$ utilizando una permutación $\pi$ que preserva la localidad. Esta se deriva de curvas de llenado de espacio (por ejemplo, Morton/Z-order para dominios euclidianos, curvas de Hilbert S2 para datos esféricos). Esto asegura que los grupos contiguos en la secuencia de tokens correspondan a regiones espaciales compactas en el dominio de coordenadas.
• Atención Agrupada Jerárquica: La secuencia reordenada se procesa a través de $L$ bloques de atención agrupada. Los tokens se fusionan progresivamente desde grupos de grano fino hacia grupos más gruesos.
• Metadatos Espaciales: El conjunto de tokens agrupados $Y^{(L)}$ resultante viene acompañado de metadatos de enrutamiento: el centroide $\mu_g$ (media de Fréchet de las coordenadas en el grupo) y la extensión $\lambda_g$ (radio espacial o caja delimitadora). Esto permite que la representación se adapte a la densidad de muestreo y la geometría de la entrada.

2. El Renderizador LH-NeF (Decodificador)
Dada la tokenización y una coordenada de consulta $x$, el renderizador reconstruye el valor del campo $f_\theta(x)$ mediante tres pasos:

• Enrutamiento Suave por k-Grupos (k-Group Soft Routing): La consulta $x$ se enruta a los $k$ grupos más cercanos basándose en la distancia entre $x$ y los centroides de los grupos $\mu_g$. Un núcleo gaussiano aprendible asigna pesos suaves $w_g(x)$ a estos grupos, asegurando la continuidad y evitando cambios abruptos entre grupos.
• Agregación por Atención Cruzada: Para cada grupo enrutado, se realiza una atención cruzada entre el embedding de la consulta y los tokens del grupo. Las salidas se agregan mediante una suma ponderada.
• Modulación Condicionada por la Geometría: El vector de características agregado se modula mediante un mecanismo de tipo FiLM condicionado en la posición relativa $\tilde{x}$ de la consulta dentro del grupo más cercano (normalizada por la extensión del grupo). Esto proporciona invarianza a la traslación y escala, permitiendo que la red aprenda patrones locales que se generalizan a través de diferentes grupos.

Contribuciones Clave

1. Identificación de un Compromiso: El artículo identifica la tensión entre la escalabilidad, los priores estructurales y la agnosticidad de la modalidad en el aprendizaje de campos neuronales, argumentando que la jerarquía y la localidad pueden introducirse sin supuestos específicos de la modalidad.
2. Marco LH-NeF: Una arquitectura novedosa que produce tokenizaciones de grupos espacialmente estructuradas. Condiciona un renderizador que decodifica campos en coordenadas arbitrarias mediante enrutamiento suave de grupos y atención cruzada, eliminando la necesidad de optimización por muestra.
3. Eficiencia y Rendimiento: El método reemplaza el bucle interno de alto consumo de memoria del meta-aprendizaje con una única pasada hacia adelante, logrando ganancias significativas en la eficiencia de entrenamiento mientras iguala o supera el rendimiento del estado del arte.

Resultados Experimentales

Los autores evaluaron LH-NeF en imágenes (CIFAR-10, CelebA-HQ, ImageNet), formas 3D (ShapeNet) y campos climáticos (ERA5).

• Reconstrucción: LH-NeF logra un rendimiento de estado del arte entre los métodos agnósticos a la modalidad. Por ejemplo, en CelebA-HQ (64²), alcanza 44.4 dB de PSNR con solo 198K parámetros, superando a Functa (38.1 dB, 2.6M params) y ENF (42.2 dB, 522K params).
• Tareas Posteriores (Downstream Tasks):
  • Generación: En CelebA-HQ, LH-NeF logra un FID de 9.7, superando a los métodos de campos neuronales generativos especializados y a las líneas base agnósticas a la modalidad.
  • Clasificación: En ShapeNet16, un clasificador ligero sobre tokens congelados alcanza un 96.8% de precisión, superando a ENF (96.6%) y Functa (90.3%).
  • Pronóstico: En ERA5, LH-NeF mejora respecto a Functa en el pronóstico de temperatura, aunque queda por detrás de ENF (que utiliza priores específicos de simetría).
• Eficiencia: Al eliminar el bucle interno de meta-aprendizaje, LH-NeF utiliza ~42× menos memoria y soporta 133× tamaños de lote más grandes (266 vs 2) en comparación con la línea base agnóstica más fuerte (ENF) en el mismo hardware.

Significado y Reivindicaciones

El artículo afirma que LH-NeF sirve como un método fundacional, escalable y agnóstico a la modalidad para el aprendizaje de representaciones de campos neuronales. Al aprovechar la jerarquía y la localidad como sesgos inductivos, el marco logra cerrar la brecha entre la eficiencia de la inferencia de pasada directa y la flexibilidad de los campos neuronales.

Los autores señalan que, si bien el método es altamente efectivo para datos con muestreo uniforme (imágenes, rejillas) y se adapta bien al muestreo irregular (nubes de puntos), la preservación de la localidad en datos altamente irregulares (por ejemplo, LiDAR) se mantiene solo en expectativa. También sugieren que trabajos futuros podrían explorar el aprendizaje de qué sesgos generalizan mejor entre tareas en lugar de diseñar manualmente los sesgos, y extender el mecanismo de condicionamiento para explotar la información de todos los niveles del codificador en lugar de solo el final.

El trabajo no pretende resolver todos los problemas de aprendizaje de representaciones, sino demostrar que la tokenización estructurada puede reemplazar el ajuste de latentes basado en optimización sin sacrificar la generalidad requerida para aplicaciones multimodales.