Entropy-Regularized Probabilistic Gates for Sparse Model Discovery in Scarce-Data Federated Learning

Resumen Técnico: Compuertas Probabilísticas con Regularización de Entropía para el Descubrimiento de Modelos Dispersos en Aprendizaje Federado con Datos Escasos

Planteamiento del Problema
El Aprendizaje Federado (FL) enfrenta desafíos significativos cuando opera en regímenes de alta dimensionalidad y escasez de datos, donde el número de parámetros ($d$) supera ampliamente al número total de muestras ($N$). En tales entornos ($d \gg N$), la optimización estándar suele producir configuraciones de parámetros que no logran generalizar a datos de prueba no vistos. Si bien los modelos dispersos (sparse) ofrecen beneficios en eficiencia de comunicación y computación, el descubrimiento de estas estructuras dispersas es difícil bajo condiciones de heterogeneidad de datos y participación parcial de los clientes. Los enfoques existentes, como la poda basada en magnitud o el Umbralización Dura Federada Iterativa (Fed-IHT), a menudo fallan al no explorar adecuadamente el espacio de parámetros, lo que conduce a un compromiso prematuro con soportes dispersos subóptimos. Además, las técnicas de regularización estándar (L1/L2) dependen de las magnitudes de los parámetros en lugar de imponer una densidad de parámetros específica, lo que hace que el descubrimiento de un modelo con un nivel de dispersión deseado $\rho$ sea un desafío.

Metodología
El artículo propone E-FLoPS (Aprendizaje Federado con Compuertas Probabilísticas con Regularización de Entropía), un algoritmo distribuido diseñado para aprender modelos dispersos bajo una restricción de densidad $L_0$ manteniendo la exploración de la incertidumbre.

1. Compuertas Probabilísticas y Reparametrización: El método adopta el marco de compuertas estocásticas $z \in \{0,1\}^d$, donde los parámetros efectivos se definen como $\theta = \tilde{\theta} \odot z$. Para permitir la optimización basada en gradientes, las compuertas discretas se aproximan mediante una distribución de Concreto Duro (Hard Concrete), una relajación continua de la distribución de Bernoulli. Esto permite que la pseudo-norma $L_0$ (el recuento de parámetros no nulos) se aproxime mediante el número esperado de compuertas activas, $\mathbb{E}_q[z_j]$.
2. Regularización de Entropía: Una innovación clave es la introducción de un término de regularización de entropía en el objetivo de optimización. El Lagrangiano incluye un término de divergencia de Kullback-Leibler (KL), $KL(q(z_j) || p(z_j))$, que penaliza la entropía relativa de las distribuciones de las compuertas. Este mecanismo evita que las distribuciones de las compuertas colapsen demasiado pronto en estados deterministas 0/1, fomentando así la exploración de configuraciones dispersas competidoras durante el entrenamiento. El parámetro de temperatura $T$ controla la fuerza de esta penalización.
3. Optimización Federada: El algoritmo opera bajo una orquestación de un servidor central. Los clientes realizan actualizaciones de descenso de gradiente estocástico (SGD) local sobre los parámetros brutos $\tilde{\theta}$, los parámetros de la compuerta $\phi = \log \alpha$, y el multiplicador de Lagrange $\lambda$ (que impone la densidad objetivo $\rho$).
   • El servidor agrega las actualizaciones de forma sincrónica.
   • Para garantizar la eficiencia de la comunicación, el algoritmo emplea un mecanismo de poda durante el entrenamiento. Solo se serializan y comunican los índices de los top-$m$ parámetros (donde $m = \lfloor \rho \cdot |\theta| \rfloor$) y sus estadísticas de compuerta correspondientes.
   • Al recibir los datos, el servidor reconstruye el modelo global y recupera los parámetros de la compuerta a partir de los valores agregados de las compuertas, ignorando el componente de ruido.

Contribuciones Clave

• Formulación: Los autores formulan un objetivo diferenciable, con regularización de entropía y restringido por $L_0$, diseñado específicamente para el entorno federado. Esta formulación aborda la no diferenciabilidad de la norma $L_0$ y la inestabilidad de la optimización en el régimen $d \gg N$.
• Algoritmo (E-FLoPS): El artículo introduce una variante de Promedio Federado (Federated Averaging) que integra compuertas probabilísticas con regularización de entropía. Esto permite una dispersión objetivo definida por el usuario mientras mantiene la exploración de la incertidumbre para evitar la convergencia prematura en mínimos locales pobres.
• Eficiencia de Comunicación: El método logra una compresión de comunicación significativa al transmitir únicamente el soporte disperso y las estadísticas de compuerta necesarias, de manera comparable a los métodos de umbralización dura, pero sin la sobrecarga de la poda post-entrenamiento.

Resultados Experimentales
Los autores evaluaron E-FLoPS en pruebas sintéticas y del mundo real, comparándolo con Fed-IHT y el Promedio Federado (FedAvg) denso con poda post-entrenamiento.

• Regresión Lineal Sintética: En entornos de alta dimensionalidad con variaciones en la relación muestra-parámetro ($N/|\theta|$), E-FLosps alcanzó consistentemente puntuaciones de $R^2$ de prueba más altas y una mejor precisión de recuperación de la dispersión (medida por la fracción de características verdaderas recuperadas) en comparación con Fed-IHT y FedAvg, particularmente en los regímenes de mayor escasez de datos.
• Clasificación de Imágenes MNIST: En una tarea de CNN con un 97.5% de dispersión objetivo, E-FLoPS demostró un rendimiento en tiempo de prueba y una eficiencia de comunicación similares a Fed-IHT. Sin embargo, mostró una mayor reducción en las Operaciones de Multiplicación y Acumulación (MACs) en la inferencia, indicando una eficiencia computacional superior.
• Clasificación de Leucemia: Utilizando el conjunto de datos de expresión génica de Golub (72 muestras, 3571 características, $N/|\theta| \approx 0.02$), E-FLoPS superó significativamente a Fed-IHT y FedAvg en precisión de prueba a través de 30 semillas aleatorias diferentes. Crucialmente, los genes seleccionados por E-FLoPS (por ejemplo, cystatin C y myeloperoxidase) se alinearon con investigaciones experimentales previas sobre la metástasis del cáncer, mientras que otros métodos seleccionaron características con menor relevancia biológica.

Significancia y Reivindicaciones
El artículo afirma que E-FLoPS proporciona un mecanismo robusto para el descubrimiento de modelos dispersos en FL, abordando específicamente la inestabilidad encontrada en regímenes de muestras pequeñas y alta dimensionalidad. Al utilizar la regularización de entropía, el método evita el compromiso temprano con soportes dispersos subóptimos, lo que conduce a una mejor generalización y a una recuperación más precisa de la verdadera estructura dispersa subyacente. Los autores aseguran que su enfoque logra estas ganancias estadísticas manteniendo costos de comunicación comparables a los métodos dispersos existentes como Fed-IHT, lo que lo convierte en una solución viable para entornos de aprendizaje con restricciones de recursos y preservación de la privacidad. El trabajo destaca la importancia de la exploración de la incertidumbre en la optimización federada dispersa, un factor que a menudo se pasa por alto en las técnicas de poda basadas en la magnitud.