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

Resumo Técnico: Portões Probabilísticos com Regularização de Entropia para Descoberta de Modelos Esparsos em Aprendizado Federado com Dados Escassos

Declaração do Problema
O Aprendizado Federado (FL) enfrenta desafios significativos ao operar em regimes de alta dimensionalidade e escassez de dados, onde o número de parâmetros ($d$) excede amplamente o número total de amostras ($N$). Nesses cenários ($d \gg N$), a otimização padrão frequentemente produz configurações de parâmetros que falham em generalizar para dados de teste não observados. Embora modelos esparsos ofereçam benefícios em eficiência de comunicação e computação para o FL, a descoberta dessas estruturas esparsas é difícil sob condições de heterogeneidade de dados e participação parcial de clientes. Abordagens existentes, como a poda baseada em magnitude ou o Federated Iterative Hard Thresholding (Fed-IHT), frequentemente falham em explorar adequadamente o espaço de parâmetros, levando ao compromisso precoce com suportes esparsos subótimos. Além disso, técnicas de regularização padrão (L1/L2) dependem de magnitudes de parâmetros em vez de impor uma densidade de parâmetros específica, tornando desafiadora a descoberta de um modelo com um nível de esparsidade desejado $\rho$.

Metodologia
O artigo propõe o E-FLoPS (Entropy-Regularized Federated Learning with Probabilistic Gates), um algoritmo distribuído projetado para aprender modelos esparsos sob uma restrição de densidade $L_0$ enquanto mantém a exploração de incerteza.

1. Portões Probabilísticos e Reparametrização: O método adota o framework de portões estocásticos $z \in \{0,1\}^d$, onde os parâmetros efetivos são definidos como $\theta = \tilde{\theta} \odot z$. Para permitir a otimização baseada em gradiente, os portões discretos são aproximados usando uma distribuição Hard Concrete, uma relaxação contínua da distribuição de Bernoulli. Isso permite que a pseudo-norma $L_0$ (a contagem de parâmetros não nulos) seja aproximada pelo número esperado de portões ativos, $\mathbb{E}_q[z_j]$.
2. Regularização de Entropia: Uma inovação fundamental é a introdução de um termo de regularização de entropia no objetivo de otimização. O Lagrangiano inclui um termo de divergência de Kullback-Leibler (KL), $KL(q(z_j) || p(z_j))$, que penaliza a entropia relativa das distribuições dos portões. Este mecanismo evita que as distribuições dos portões colapsem precocemente em estados determinísticos 0/1, incentivando assim a exploração de configurações esparsas concorrentes durante o treinamento. O parâmetro de temperatura $T$ controla a força dessa penalidade.
3. Otimização Federada: O algoritmo opera sob uma orquestração de servidor central. Os clientes realizam atualizações de gradiente descendente estocástico (SGD) local nos parâmetros brutos $\tilde{\theta}$, nos parâmetros dos portões $\phi = \log \alpha$ e no multiplicador de Lagrange $\lambda$ (que impõe a densidade alvo $\rho$).
   • O servidor agrega as atualizações de forma síncrona.
   • Para garantir a eficiência de comunicação, o algoritmo emprega um mecanismo de poda durante o treinamento. Apenas os índices top-$m$ (onde $m = \lfloor \rho \cdot |\theta| \rfloor$) dos parâmetros e suas respectivas estatísticas de portão são serializados e comunicados.
   • Ao receber os dados, o servidor reconstrói o modelo global e recupera os parâmetros dos portões a partir dos valores agregados dos portões, ignorando o componente de ruído.

Principais Contribuições

• Formulação: Os autores formulam um objetivo diferenciável, com regularização de entropia e restrito por $L_0$, especificamente para o ambiente federado. Esta formulação aborda a não diferenciabilidade da norma $L_0$ e a instabilidade da otimização no regime $d \gg N$.
• Algoritmo (E-FLoPS): O artigo introduz uma variante de Federated Averaging que integra portões probabilísticos com regularização de entropia. Isso permite uma esparsidade alvo definida pelo usuário enquanto mantém a exploração de incerteza para evitar a convergência prematura para mínimos locais ruins.
• Eficiência de Comunicação: O método alcança compressão significativa de comunicação ao transmitir apenas o suporte esparso e as estatísticas necessárias dos portões, comparável aos métodos de limiarização iterativa (iterative hard-thresholding), mas sem o overhead de poda pós-treinamento.

Resultos Experimentais
Os autores avaliaram o E-FLoPS em benchmarks sintéticos e do mundo real, comparando-o com o Fed-IHT e o Federated Averaging denso (FedAvg) com poda pós-treinamento.

• Regressão Linear Sintética: Em cenários de alta dimensionalidade com variando razões de amostra-para-parâmetro ($N/|\theta|$), o E-FLoPS consistentemente alcançou pontuações de $R^2$ de teste mais altas e melhor precisão de recuperação de esparsidade (medida pela fração de características verdadeiras recuperadas) em comparação ao Fed-IHT e FedAvg, particularmente nos regimes de maior escassez de dados.
• Classificação de Imagens MNIST: Em uma tarefa de CNN com 97,5% de esparsidade alvo, o E-FLoPS demonstrou desempenho de tempo de teste e eficiência de comunicação semelhantes ao Fed-IHT. No entanto, mostrou uma redução maior nas Operações de Acúmulo e Multiplicação (MACs) na inferência, indicando eficiência computacional superior.
• Classificação de Leucemia: Utilizando o conjunto de dados de expressão gênica Golub (72 amostras, 3571 características, $N/|\theta| \approx 0,02$), o E-FLoPS superou significativamente o Fed-IHT e o FedAvg em precisão de teste em 30 sementes aleatórias diferentes. Crucialmente, os genes selecionados pelo E-FLoPS (ex: cystatin C e myeloperoxidase) alinharam-se com pesquisas experimentais prévias sobre metástase de câncer, enquanto outros métodos selecionaram características menos relevantes biologicamente.

Significância e Alegações
O artigo afirma que o E-FLoPS fornece um mecanismo robusto para a descoberta de modelos esparsos em FL, abordando especificamente a instabilidade encontrada em regimes de pequenas amostras e alta dimensionalidade. Ao utilizar a regularização de entropia, o método evita o compromisso precoce com suportes esparsos subótimos, levando a uma melhor generalização e uma recuperação mais precisa da verdadeira estrutura esparsa subjacente. Os autores asseguram que sua abordagem alcança esses ganhos estatísticos mantendo custos de comunicação comparáveis aos métodos esparsos existentes como o Fed-IHT, tornando-a uma solução viável para ambientes de aprendizado preservação de privacidade e com recursos limitados. O trabalho destaca a importância da exploração de incerteza na otimização esparsa federada, um fator frequentemente negligenciado pelas técnicas de poda baseadas em magnitude.