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

Resumo Técnico: Redes Neurais Convolucionais Direcionáveis Bayesianas 3D

Enunciado do Problema
As Redes Neurais Convolucionais Direcionáveis (Steerable-CNNs) estabeleceram-se como um paradigma poderoso para o processamento de dados 3D, garantindo a SE(3)-equivariância ao parametrizar os núcleos de convolução como combinações lineares de funções de base direcionáveis. Esta restrição arquitetural garante que a transformação da entrada (ex: via rotação 3D) produza uma saída correspondentemente transformada, reduzindo significativamente o espaço de hipóteses e melhorando a generalização em domínios como reconhecimento de formas e predição de propriedades moleculares. No entanto, as Steerable-CNNs existentes são inteiramente determinísticas. Elas produzem estimativas de ponto único sem medidas de confiança associadas, tornando-as inadequadas para aplicações críticas de segurança (ex: imagens médicas, sistemas autônomos) onde a quantificação de incerteza (UQ) é essencial.

Um desafio primário no desenvolvimento de Bayesian Steerable-CNNs é que introduzir estocasticidade nas camadas de pesos da rede de forma ingênua normalmente quebra a propriedade de equivariância. Se os núcleos de convolução forem amostrados aleatoriamente, eles geralmente violam a restrição de "intertwiner" ou direcionabilidade, o que significa que as realizações estocásticas resultantes não seriam equivariantes, mesmo que sua média seja.

Metodologia
Os autores propõem um framework de Bayesian Steerable-CNN que resolve este conflito ao colocar distribuições posteriores sobre os coeficientes escalares das funções de base direcionáveis, em vez das próprias matrizes de núcleo.

1. Núcleos Equivariantes Estocásticos: O núcleo de convolução $K(x)$ é definido como uma combinação linear de funções de base direcionáveis determinísticas e pré-computadas $\{\kappa_b(x)\}$ (derivadas de matrizes de Wigner D para SO(3)):
   $$K(x; w) = \sum_{b=1}^{B} w_b \kappa_b(x)$$
   As funções de base $\kappa_b$ são fixas e satisfazem a restrição de direcionabilidade por construção. Os coeficientes $w = (w_1, \dots, w_B)$ são tratados como variáveis aleatórias.
2. Preservação da Equivariância: O artigo prova (Proposição 1) que se as funções de base satisfizerem a restrição de direcionabilidade, qualquer combinação linear delas — independentemente de os coeficientes serem determinísticos ou amostrados de uma distribuição de probabilidade — permanece equivariante. Isso permite núcleos estocásticos sem modificações arquiteturais para impor a equivariância.
3. Inferência Variacional: O modelo utiliza a abordagem Bayes-by-Backpropagation. Assume-se uma posterior variacional Gaussiana de campo médio $q(w|\theta)$ sobre os coeficientes, parametrizada por uma média $\mu$ e uma log-variância $\rho$. O modelo é treinado maximizando o Evidence Lower Bound (ELBO), que consiste em um termo de perda de entropia cruzada e um termo de divergência de Kullback-Leibler (KL) que regulariza a posterior contra uma priori Gaussiana.
4. Decomposição de Incerteza: Na inferência, o modelo realiza $N$ passagens estocásticas para frente (forward passes). A incerteza preditiva total é decomposta em:
   • Incerteza Aleatória: Ruído inerente aos dados (capturada pela média da entropia das predições individuais).
   • Incerteza Epistêmica: Incerteza do modelo devido a dados limitados (capturada pela variância das predições através das $N$ amostras).

Principais Contribuições

• Nova Formulação Bayesiana: A primeira integração de inferência estocástica com convoluções direcionáveis, obtendo a distribuição posterior dos coeficientes de base do núcleo mantendo a SE(3)-equivariância exata.
• Prova Teórica: Uma prova formal demonstrando que a estocasticidade no espaço dos coeficientes não viola a restrição de intertwiner, garantindo que cada núcleo amostrado permaneça equivariante.
• Incerteza Decomponível: Um framework que fornece estimativas de incerteza calibradas e decomponíveis (epistêmica vs. aleatória) para modelos equivariantes 3D.

Resultados Experimentais
O modelo foi avaliado no conjunto de dados ModelNet10 (modelos CAD 3D) utilizando uma arquitetura construída sobre a biblioteca escnn.

• Desempenho de Classificação: O modelo bayesiano alcançou uma acurácia de teste de 86,34%, sendo competitivo com seus equivalentes determinísticos.
• Calibração: O modelo demonstrou excelente calibração com um Erro de Calibração Esperado (ECE) de 0,0263, indicando que as pontuações de confiança preditas correspondem de perto à acurácia empírica.
• Robustez à Mudança de Distribuição: Sob ruído Gaussiano aditivo, o modelo bayesiano exibiu robustez superior comparado ao baseline determinístico. Em um nível alto de ruído ($\epsilon_{noise} = 0,9$), o modelo bayesiano manteve 81,83% de acurácia, superando o modelo determinístico (75,66%) em 6,17%.
• Utilidade da Incerteza:
  • Predição Seletiva: Ao filtrar amostras com alta incerteza, o modelo alcançou um ganho de acurácia de ~4% em 84% do conjunto de teste (amostras com incerteza abaixo de um limite específico).
  • Significado Semântico: Foi encontrada uma correlação negativa de Spearman estatisticamente significativa (-0,3989) entre a incerteza epistêmica e o erro de predição, confirmando que uma incerteza maior prevê de forma confiável um erro de classificação mais alto.
  • Convergência da Incerteza: Durante o treinamento, a incerteza total diminuiu de forma constante, enquanto a incerteza epistêmica permaneceu consistentemente baixa, sugerindo que o modelo tornou-se confiante em sua distribuição de pesos aprendida.

Significância e Alegações
O artigo afirma unificar a quantificação de incerteza bayesiana com o viés indutivo de CNNs equivariantes. Os autores enfatizam que este é o primeiro trabalho a integrar com sucesso a inferência estocástica e as convoluções direcionáveis em um único modelo discriminativo. A significância reside em permitir a tomada de decisão consciente de riscos em domínios 3D: o modelo não apenas respeita simetrias geométricas, mas também fornece estimativas de confiança decomponíveis e fundamentadas. Isso permite que os praticantes distingam entre a ambiguidade inerente aos dados (aleatória) e a falta de conhecimento do modelo (epistêmica), facilitando o refinamento direcionado do modelo e implementações mais seguras em aplicações sensíveis à segurança. O framework alcança esses benefícios sem sacrificar as garantias de equivariância que tornam as Steerable-CNNs eficazes para dados 3D.