FlexAct: Why Learn when you can Pick?

Resumen Técnico: Flex-Act

Planteamiento del Problema
A pesar del papel crítico que desempeñan las funciones de activación en la propagación de señales, la capacidad de entrenamiento y la generalización dentro de las redes neuronales profundas, estas suelen tratarse como hiperparámetros fijos (por ejemplo, ReLU, GELU, Swish) aplicados uniformemente en todas las capas y tareas. Este paradigma estático asume que una única no linealidad es óptima independientemente de la profundidad de la capa o del régimen de datos, lo que potencialmente conduce a un sesgo inductivo subóptimo, fragmentación de gradientes (shattered gradients) o problemas de saturación. Aunque existen funciones de activación parametrizadas (como PReLU, SPLASH) y aproximadores de funciones aprendibles, estas suelen introducir una alta complejidad de optimización, una reducida interpretabilidad o requieren comprometerse con una familia funcional específica. Los autores argumentan que ninguna no linealidad es universalmente óptima y que las redes profundas se beneficiarían de la capacidad de seleccionar dinámicamente la función de activación más apropiada para cada capa sin necesidad de un ajuste manual de hiperparámetros.

Metodología
El artículo presenta Flex-Act, un marco que permite la selección discreta, por capa, de funciones de activación a partir de un conjunto predefinido durante el entrenamiento. El mecanismo central se basa en el truco de Gumbel-Softmax para crear una relajación diferenciable de las elecciones discretas.

1. Enrutamiento Discreto mediante Gumbel-Softmax:
   Para cada capa $i$, el modelo computa una combinación convexa de funciones de activación candidatas $\{\sigma^{(1)}, \dots, \sigma^{(p)}\}$ aplicadas a la transformación afín $h_i$. Las probabilidades de selección $p_i$ se derivan de los logits entrenables $\pi_i$ utilizando la distribución Gumbel-Softmax:
   $$X_i = \sum_{j=1}^{p} p_i^{(j)} \cdot \sigma^{(j)}(h_i)$$
   donde $p_i \sim \text{GumbelSoftmax}(\log \pi_i, \tau)$. A medida que la temperatura $\tau$ se enfría (annealed), el modelo transiciona de explorar activaciones mezcladas a comprometerse con una función discreta específica por capa. Esto permite que el proceso de selección sea optimizable de extremo a extremo mediante el descenso de gradiente.

2. Corrección de Sesgo mediante Normalización de Gradientes:
   Los autores identifican un fallo crítico en el enrutamiento ingenuo de Gumbel-Softmax: un sesgo de escala hacia funciones de activación no acotadas (como ReLU). Debido a que la retropropagación escala los gradientes según la magnitud de la activación, las funciones con salidas más grandes dominan el proceso de selección independientemente de su idoneidad para los datos.
   Para contrarrestar esto, Flex-Act introduce un regularizador basado en la norma del gradiente. Calcula las normas de los gradientes de cada función de activación candidata con respecto a la entrada y las convierte en pseudo-probabilidades. La función de pérdida total combina la pérdida de la tarea principal ($L_{task}$) con un término de divergencia de Kullback-Leibler ($L_{KL}$) que alinea las probabilidades de enrutamiento aprendidas con estas pseudo-etiquetas basadas en gradientes:
   $$L = L_{task} + \alpha L_{KL}$$
   Esta regularización evita que el modelo se decante por funciones no acotadas simplemente debido a diferencias de magnitud, asegurando que la selección esté impulsada por la idoneidad funcional.

Contribuciones Clave

• Enrutamiento de Activación Discreta: Un mecanismo novedoso que utiliza la reparametrización Gumbel-Softmax para permitir la optimización de extremo a extremo de las elecciones de funciones de activación discretas por capa.
• Derivaciones Analíticas de Gradientes: Derivación formal de las ecuaciones de retropropagación para redes que emplean enrutamiento de activación discreto, resolviendo la no diferenciabilidad del proceso de selección.
• Regularización de la Norma del Gradiente: Una técnica específica para corregir el sesgo inherente hacia las activaciones no acotadas en el enrutamiento discreto, mejorando significativamente la precisión de la selección.
• Interpretabilidad y Modularidad: Al tratar la selección de la activación como una variable discreta y aprendible, el marco proporciona información sobre qué no linealidades se utilizan en diferentes profundidades, ofreciendo una alternativa modular a las familias parametrizadas.

Resultados Experimentales
El artículo evalúa Flex-Act tanto en tareas de regresión sintéticas como en entornos de referencia del mundo real:

• Regresión Sintética: En tareas donde la función de activación real es conocida (ReLU, Sigmoid, Tanh, LeakyReLU, Identidad), Flex-Act converge consistentemente a la función correcta.
  • Sin regularización ($\alpha=0$), el modelo falla al seleccionar funciones acotadas (Sigmoid/Tanh), decantándose por ReLU debido al sesgo de magnitud del gradiente.
  • Con regularización ($\alpha=0.3$), Flex-Act logra un Error Cuadrático Medio (MSE) cercano a cero en todos los escenarios de verdad fundamental, igualando o superando a las líneas base de activación fija sin requerir múltiples ejecuciones de entrenamiento o barridos de activación.
• Benchmarks del Mundo Real:
  • Clasificación de Imágenes (CIFAR-10/100): Integrado en ResNet-18 y ResNet-34, Flex-Act (aplicado a la capa penúltima y a todas las capas) logra mejoras estadísticamente significativas sobre las líneas base de ReLU fijo y otros métodos parametrizados (PReLU, Swish, AReLU). Por ejemplo, ResNet-34 con Flex-Act (todas las capas) alcanza un 95.75% de precisión en CIFAR-10 y 79.42% en CIFAR-100, superando a la línea base de ReLU.
  • Clasificación de Texto: El marco también fue probado en los benchmarks de GLUE utilizando modelos BERT-base, demostrando compatibilidad con arquitecturas de transformadores.

Significancia y Reivindicaciones
El artículo afirma que Flex-Act ofrece una solución "plug-and-play" para el aprendizaje profundo adaptativo. Argumenta que, al permitir que las redes recuperen dinámicamente la no linealidad apropiada para una tarea o capa específica, los profesionales pueden evitar la carga computacional de los barridos manuales de hiperparámetros o el entrenamiento de múltiples modelos con diferentes activaciones fijas. El marco cierra la brecha entre la flexibilidad teórica y la utilidad práctica, proporcionando un enfoque robusto, interpretable y modular para el diseño de no linealidades que escala desde tareas sintéticas superficiales hasta arquitecturas profundas del mundo real. Los autores enfatizan que su método logra estas ganancias sin aumentar significativamente el número de parámetros, apoyándose en su lugar en la selección inteligente de formas funcionales existentes y bien caracterizadas.