Bayesian 3D Steerable CNNs: Enabling Equivariance and Uncertainty Quantification Simultaneously

Resumen Técnico: Redes Neuronales Convolucionales 3D Direccionables Bayesianas

Planteamiento del Problema
Las Redes Neuronales Convolucionales Direccionables (Steerable-CNNs) se han consolidado como un paradigma poderoso para el procesamiento de datos 3D, garantizando la SE(3)-equivarianza al parametrizar los núcleos convolucionales como combinaciones lineales de funciones de base direccionables. Esta restricción arquitectónica asegura que la transformación de la entrada (por ejemplo, mediante una rotación 3D) produzca una salida correspondientemente transformada, reduciendo significativamente el espacio de hipótesis y mejorando la generalización en dominios como el reconocimiento de formas y la predicción de propiedades moleculares. Sin embargo, las Steerable-CNNs existentes son enteramente deterministas. Producen estimaciones de punto único sin medidas de confianza asociadas, lo que las hace inadecuadas para aplicaciones de seguridad crítica (por ejemplo, imágenes médicas, sistemas autónomos) donde la cuantificación de la incertidumbre (UQ) es esencial.

Un desafío principal en el desarrollo de Steerable-CNNs Bayesianas es que introducir estocasticidad en los pesos de la red de forma ingenua suele romper la propiedad de equivariancia. Si los núcleos de convolución se muestrean aleatoriamente, generalmente violan la restricción de "intertwiner" o de direccionabilidad, lo que significa que las realizaciones estocásticas resultantes no serían equivariantes, incluso si su media lo es.

Metodología
Los autores proponen un marco de Steerable-CNN Bayesiana que resuelve este conflicto al colocar distribuciones posteriores sobre los coeficientes escalares de las funciones de base direccionables, en lugar de sobre las matrices de los núcleos mismos.

1. Núcleos Equivariantes Estocásticos: El núcleo de convolución $K(x)$ se define como una combinación lineal de funciones de base direccionables precomputadas y deterministas $\{\kappa_b(x)\}$ (derivadas de las matrices de Wigner D para SO(3)):
   $$K(x; w) = \sum_{b=1}^{B} w_b \kappa_b(x)$$
   Las funciones de base $\kappa_b$ son fijas y satisfacen la restricción de direccionabilidad por construcción. Los coeficientes $w = (w_1, \dots, w_B)$ son tratados como variables aleatorias.
2. Preservación de la Equivariancia: El artículo demuestra (Proposición 1) que si las funciones de base satisfacen la restricción de direccionabilidad, cualquier combinación lineal de ellas —independientemente de si los coeficientes son deterministas o muestreados de una distribución de probabilidad— sigue siendo equivariante. Esto permite núcleos estocásticos sin modificaciones arquitectónicas para imponer la equivariancia.
3. Inferencia Variacional: El modelo emplea el enfoque Bayes-by-Backpropagation. Se asume una posterior variacional gaussiana de campo medio $q(w|\theta)$ sobre los coeficientes, parametrizada por una media $\mu$ y una log-varianza $\rho$. El modelo se entrena maximizando el Límber de la Evidencia (ELBO), que consiste en un término de pérdida de entropía cruzada y un término de divergencia de Kullback-Leibler (KL) que regulariza la posterior frente a una distribución previa Gaussiana.
4. Descomposición de la Incertidumbre: En la inferencia, el modelo realiza $N$ pasadas estocásticas hacia adelante. La incertidumbre predictiva total se descompone en:
   • Incertidumbre Aleatoria: Ruido inherente en los datos (capturada por la entropía promedio de las predicciones individuales).
   • Incertidumbre Epistémica: Incertidumbre del modelo debida a datos limitados (capturada por la varianza de las predicciones a través de las $N$ muestras).

Contribuciones Clave

• Nueva Formulación Bayesiana: La primera integración de la inferencia estocástica con convoluciones direccionables, obteniendo la distribución posterior de los coeficientes de la base del núcleo manteniendo la SE(3)-equivarianza exacta.
• Prueba Teórica: Una prueba formal que demuestra que la estocasticidad en el espacio de los coeficientes no viola la restricción de intertwiner, asegurando que cada núcleo muestreado permanezca equivariante.
• Incertidumbre Descomponible: Un marco que proporciona estimaciones de incertidumbre calibradas y descomponibles (epistémica vs. aleatoria) para modelos equivariantes 3D.

Resultados Experimentales
El modelo fue evaluado en el conjunto de datos ModelNet10 (modelos CAD 3D) utilizando una arquitectura construida sobre la biblioteca escnn.

• Rendimiento de Clasificación: El modelo bayesiano alcanzó una precisión de prueba del 86.34%, lo cual es competitivo con sus contrapartes deterministas.
• Calibración: El modelo demostró una excelente calibración con un Error de Calibración Esperado (ECE) de 0.0263, indicando que las puntuaciones de confianza predichas coinciden estrechamente con la precisión empírica.
• Robustez al Cambio de Distribución: Bajo ruido gaussiano aditivo, el modelo bayesiano exhibió una robustez superior comparado con la línea base determinista. Con un nivel de ruido alto ($\epsilon_{noise} = 0.9$), el modelo bayesiano retuvo un 81.83% de precisión, superando al modelo determinista (75.66%) por un 6.17%.
• Utilidad de la Incertidumbre:
  • Predicción Selectiva: Al filtrar las muestras con alta incertidumbre, el modelo logró una ganancia de precisión de ~4% en el 84% del conjunto de datos de prueba (muestras con incertidumbre por debajo de un umbral específico).
  • Significado Semántico: Se encontró una correlación de rango de Spearman negativa estadísticamente significativa (-0.3989) entre la incertidumbre epistémica y el error de predicción, confirmando que una mayor incertidumbre predice de manera fiable un mayor error de clasificación.
  • Convergencia de la Incertidumbre: Durante el entrenamiento, la incertidumbre total disminuyó constantemente, mientras que la incertidumbre epistémica se mantuvo consistentemente baja, sugiriendo que el modelo se volvió seguro en su distribución de pesos aprendida.

Significancia y Reivindicaciones
El artículo afirma unificar la cuantificación de la incertidumbre bayesiana con el sesgo inductivo de las CNNs equivariantes. Los autores enfatizan que este es el primer trabajo que integra con éxito la inferencia estocástica y las convoluciones direccionables en un único modelo discriminativo. La significancia radica en permitir la toma de decisiones consciente del riesgo en dominios 3D: el modelo no solo respeta las simetrías geométricas, sino que también proporciona estimaciones de confianza principios y descomponibles. Esto permite a los profesionales distinguir entre la ambigüedad inherente a los datos (aleatoria) y la falta de conocimiento del modelo (epistémica), facilitando el refinamiento dirigido del modelo y un despliegue más seguro en aplicaciones sensibles a la seguridad. El marco logra estos beneficios sin sacrificar las garantías de equivariancia que hacen que las Steerable-CNNs sean efectivas para datos 3D.