Controllable Diversity in Normalization-Based Implicit Ensembles via Softmax-Temperature Modulation

Resumen Técnico: Diversidad Controlable en Ensambles Implícitos Basados en Normalización mediante Modulación de Temperatura Softmax

Planteamiento del Problema
Los ensambles profundos (deep ensembles) siguen siendo el estándar de oro para la cuantificación de la incertidumbre en el aprendizaje profundo, ofreciendo una calibración fiable y detección de datos fuera de distribución (OOD) mediante el entrenamiento de múltiples modelos independientes. Sin embargo, sus costes computacionales y de memoria escalan linealamente con el número de miembros del ensamble, lo que resulta prohibitivo para modelos a gran escala. Los ensambles implícitos abordan esto compartiendo un único backbone entre los miembros, introduciendo una modulación o partición por miembro para simular el comportamiento de un ensamble en una sola ejecución de entrenamiento.

Una limitación crítica de los métodos actuales de ensamble implícito es la falta de control sobre la diversidad de los miembros. La diversidad es un determinante primario de la calidad del ensamble, pero los enfoques actuales la fijan en la inicialización (por ejemplo, métodos basados en adaptadores como BatchEnsemble) o la codifican de forma rígida en la arquitectura (por ejemplo, métodos de partición como Packed Ensembles). En consecuencia, no existe un mecanismo para moldear o navegar el compromiso (trade-off) entre precisión y calibración durante el entrenamiento. Además, los métodos existentes suelen requerir rediseños arquitectónicos específicos, lo que dificulta su aplicación a modelos preentrenados.

Metodología: $\sigma$N-Ens
Los autores presentan $\sigma$N-Ens, un ensamble implícito basado en normalización que trata a cada miembro como una tarea distinta dentro de un marco de aprendizaje multitarea (MTL).

1. Arquitectura ($\sigma$Norm): El método reemplaza las capas de normalización estándar (BatchNorm o LayerNorm) en un backbone compartido por capas específicas de cada tarea. Para cada miembro $t$, la normalización es modulada por un escalador acotado por una sigmoide:
   $$\sigma\text{Norm}_t(x) = \sigma(\gamma_t) \odot \text{norm}_t(x)$$
   donde $\sigma(\gamma_t) = \frac{1}{1 + e^{-\gamma_t}}$. Estos escaladores actúan como importancias por característica. Debido a que solo se replican las capas de normalización, los pesos del backbone (convolucionales o de transformadores) permanecen totalmente compartidos, lo que permite al método envolver cualquier arquitectura estándar y adaptar modelos preentrenados mediante un breve ajuste fino (fine-tuning).

2. Diversidad Controlable mediante Temperatura Softmax: Para moldear activamente la diversidad durante el entrenamiento, los autores introducen un regularizador de temperatura softmax aplicado a los vectores de importancia ($\sigma(\gamma)$) entre miembros para cada característica.
   $$\mathcal{L}_{div} = \lambda \sum_{\ell \in \mathcal{R}} \sum_{i} \sum_{m=1}^{M} \log \left[ \text{softmax}_{\tau}(\sigma(\gamma_{\cdot, \ell, i})) \right]_m$$
   Aquí, $\tau$ actúa como un hiperparámetro de temperatura que establece el nivel de equilibrio de la compartición.

   • $\tau$ Bajo: Empuja la distribución hacia una función one-hot, forzando a que las características pertenezcan a un solo miembro (diversidad máxima, similar a la partición rígida).
   • $\tau$ Alto: Empuja la distribución hacia una distribución uniforme, fomentando la compartición total de las características.
   • $\tau$ Intermedio: Crea una asignación de tipo "$M$ de $k$ suave", donde las características se comparten entre un subconjunto de $k$ miembros.
     A diferencia de la regularización estándar ($\lambda$), que determina qué tan fuerte se empuja a la red hacia un objetivo, $\tau$ selecciona en qué equilibrio se establece el sistema. Esto permite a los profesionales trazar una frontera de precisión-calibración mediante el barrido de $\tau$.

3. Incertidumbre de Modulación: Los autores reformulan la incertidumbre epistémica expresada por este ensamble como incertidumbre de modulación. Dado que el backbone es compartido y determinista, la incertidumbre surge únicamente del reponderado de las características compartidas por los parámetros de modulación. Esto es estructuralmente distinto de la incertidumbre derivada de extractores de características diferentes en los ensambles profundos completos.

Contribuciones Clave

• Diversidad Controlable: El artículo presenta el primer mecanismo en la familia de ensambles implícitos para moldear la diversidad de los miembros durante el entrenamiento mediante un único hiperparámetro ($\tau$), en lugar de fijarla en la inicialización o construcción. Esto permite la navegación de la frontera de Pareto de precisión-calibración.
• Incertidumbre de Modulación: Los autores identifican y definen la "incertidumbre de modulación" como el componente epistémico específico expresado por los ensambles de backbone compartido. Establecen que esta incertidumbre mantiene la calibración bajo corrupción de entrada, pero exhibe una detección de OOD más débil en comparación con los ensambles completos, ya que el backbone compartido no puede divergir en entradas que no codifica.
• Agnosticismo de Arquitectura y Eficiencia: Al operar únicamente en las capas de normalización, $\sigma$N-Ens se aplica tanto a backbones convolucionales como de transformadores sin rediseños invasivos. Puede convertir modelos preentrenados existentes en ensambles calibrados mediante un breve ajuste fino, evitando la necesidad de entrenar los miembros desde cero.

Resultados
El método fue evaluado en CIFAR-10/100, ImageNet y SST-2 utilizando ResNets, WideResNets y Vision Transformers (ViT).

• Rendimiento vs. Coste: $\sigma$N-Ens iguala o supera a los ensambles profundos en términos de precisión y Log-Verosimilitud Negativa (NLL) utilizando una fracción del número de parámetros (por ejemplo, en CIFAR-100 con WideResNet, logra una precisión de vanguardia con significativamente menos parámetros).
• Escalabilidad: A diferencia de los métodos de partición (como Packed Ensembles) que se degradan y colapsan a medida que aumenta el número de miembros ($M$) debido a las restricciones de capacidad, $\sigma$N-Ens escala de forma fluida. A medida que $M$ crece, el mecanismo de compartición suave agrupa la capacidad de manera efectiva, permitiendo que el rendimiento mejore o se sature en lugar de colapsar.
• Calibración bajo Desplazamiento: El método mantiene una calibración robusta bajo desplazamiento de distribución (corrupción de entrada), con el Error de Calibración Esperado (ECE) manteniéndose casi plano a medida que aumenta la severidad de la corrupción.
• Detección de OOD: Consistente con el análisis teórico, la detección de OOD es más débil que en los ensambles profundos completos, particularmente cuando se utilizan backbones preentrenados donde la representación de características es fija. Sin embargo, el rendimiento de detección mejora cuando se entrena desde cero, ya que los miembros moldean la representación compartida bajo la presión de la diversidad.
• Sensibilidad a Hiperparámetros: El rendimiento final es mayormente insensible a la fuerza de regularización $\lambda$ una vez convergido, lo que convierte a $\tau$ en el único hiperparámetro crítico para que los profesionales ajusten.

Significado
El artículo argumenta que la diversidad controlable transforma los ensambles implícitos de una elección arquitectónica estática en un sistema dinámico y ajustable. Al desacoplar la diversidad de la inicialización y la arquitectura, $\sigma$N-Ens ofrece una alternativa práctica para entrenar ensambles profundos completos, especialmente a medida que cre el tamaño de los modelos. La capacidad de convertir un único modelo preentrenado en un ensamble calibrado mediante un breve ajuste fino aborda una brecha significativa en el despliegue de aprendizaje profundo consciente de la incertidumbre en dominios con recursos limitados o sensibles al riesgo. El trabajo clarifica los compromisos inherentes a los ensambles de backbone compartido, vinculando explícitamente la limitación estructural (extractores de características compartidos) con la debilidad observada en la detección de OOD, mientras destaca la fortaleza en la robustez de la calibración.