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

Sintesi Tecnica: Gate Probabilistici Regolarizzati dall'Entropia per la Scoperta di Modelli Sparsi nel Federated Learning con Scarsità di Dati

Problematica
Il Federated Learning (FL) affronta sfide significative quando opera in regimi ad alta dimensionalità e scarsità di dati, dove il numero di parametri ($d$) supera di gran lunga il numero totale di campioni ($N$). In tali contesti ($d \gg N$), l'ottimizzazione standard spesso produce configurazioni di parametri che non riescono a generalizzare su dati di test non visti. Sebbene i modelli sparsi offrano vantaggi in termini di efficienza di comunicazione e computazionale per l'FL, la scoperta di queste strutture sparse è difficile in condizioni di eterogeneità dei dati e partecipazione parziale dei client. Gli approcci esistenti, come la potatura basata sull'ampiezza (magnitude-based pruning) o il Federated Iterative Hard Thresholding (Fed-IHT), spesso falliscono nell'esplorare adeguatamente lo spazio dei parametri, portando a un impegno precoce verso supporti sparsi subottimali. Inoltre, le tecniche di regolarizzazione standard (L1/L2) si affidano all'ampiezza dei parametri piuttosto che imporre una specifica densità di parametri, rendendo difficile la scoperta di un modello con un livello di sparsità desiderato $\rho$.

Metodologia
Il documento propone E-FLoPS (Entropy-Regularized Federated Learning with Probabilistic Gates), un algoritmo distribuito progettato per apprendere modelli sparsi sotto un vincolo di densità $L_0$ mantenendo l'esplorazione dell'incertezza.

1. Gate Probabilistici e Riperparametrizzazione: Il metodo adotta il framework di gate stocastici $z \in \{0,1\}^d$, dove i parametri effettivi sono definiti come $\theta = \tilde{\theta} \odot z$. Per consentire l'ottimizzazione basata sul gradiente, i gate discreti sono approssimati utilizzando una distribuzione Hard Concrete, una rilassazione continua della distribuzione di Bernoulli. Ciò consente di approssimare la pseudo-norma $L_0$ (il conteggio dei parametri non nulli) come il numero atteso di gate attivi, $\mathbb{E}_q[z_j]$.
2. Regolarizzazione dell'Entropia: Un'innovazione chiave è l'introduzione di un termine di regolarizzazione dell'entropia nell'obiettivo di ottimizzazione. Il Lagrangiano include un termine di divergenza di Kullback-Leibler (KL), $KL(q(z_j) || p(z_j))$, che penalizza l'entropia relativa delle distribuzioni dei gate. Questo meccanismo impedisce alle distribuzioni dei gate di collassare troppo presto in stati deterministici 0/1, incoraggiando così l'esplorazione di configurazioni sparse concorrenti durante l'addestramento. Il parametro di temperatura $T$ controlla la forza di tale penalità.
3. Ottimizzazione Federata: L'algoritmo opera sotto un'orchestrazione di un server centrale. I client eseguono aggiornamenti di discesa del gradiente stocastica (SGD) locale sui parametri grezzi $\tilde{\theta}$, sui parametri dei gate $\phi = \log \alpha$ e sul moltiplicatore di Lagrange $\lambda$ (che impone la densità target $\rho$).
   • Il server aggrega gli aggiornamenti in modo sincrono.
   • Per garantire l'efficienza della comunicazione, l'algoritmo impiega un meccanismo di potatura durante l'addestramento. Solo gli indici top-$m$ (dove $m = \lfloor \rho \cdot |\theta| \rfloor$) dei parametri e le relative statistiche dei gate vengono serializzati e comunicati.
   • Al ricevimento, il server ricostruisce il modello globale e recupera i parametri dei gate dai valori dei gate aggregati, ignorando la componente di rumore.

Contributi Chiave

• Formulazione: Gli autori formulano un obiettivo differenziabile, regolarizzato dall'entropia e vincolato da $L_0$, specificamente per l'ambiente federato. Questa formulazione affronta la non-differenziabilità della norma $L_0$ e l'instabilità dell'ottimizzazione nel regime $d \gg N$.
• Algoritmo (E-FLoPS): Il documento introduce una variante di Federated Averaging che integra gate probabilistici con la regolarizzazione dell'entropia. Ciò consente una sparsità target definita dall'utente pur mantenendo l'esplorazione dell'incertezza per evitare la convergenza prematura verso minimi locali poveri.
• Efficienza di Comunicazione: Il metodo ottiene una significativa compressione della comunicazione trasmettendo solo il supporto sparso e le statistiche dei gate necessarie, paragonabile ai metodi di hard-thresholding iterativo ma senza l'overhead di potatura post-addestramento.

Risultati Sperimentali
Gli autori hanno valutato E-FLoPS su benchmark sintetici e reali, confrontandolo con Fed-IHT e FedAvg con potatura post-addestramento.

• Regressione Lineare Sintetica: In scenari ad alta dimensionalità con variabili rapporti campione-parametro ($N/|\theta|$), E-FLoPS ha costantemente raggiunto punteggi $R^2$ di test più elevati e una migliore accuratezza nel recupero della sparsità (misurata come frazione di feature reali recuperate) rispetto a Fed-IHT e FedAvg, particolarmente nei regimi con maggiore scarsità di dati.
• Classificazione di Immagini MNIST: In un compito di classificazione CNN con una sparsità target del 97,5%, E-FLoPS ha dimostrato prestazioni al tempo di test ed efficienza di comunicazione simili a Fed-IHT. Tuttavia, ha mostrato una maggiore riduzione delle operazioni di Moltiplica e Accumula (MAC) all'inferenza, indicando una superiore efficienza computazionale.
• Classificazione della Leucemia: Utilizzando il dataset di espressione genica Golub (72 campioni, 3571 feature, $N/|\theta| \approx 0,02$), E-FLoPS ha superato significativamente Fed-IHT e FedAvg in termini di accuratezza di test su 30 diversi seed casuali. Fondamentalmente, i geni selezionati da E-FLoPS (ad esempio, cystatin C e myeloperoxidase) si allineavano con la ricerca sperimentale pregressa sulla metastasi del cancro, mentre altri metodi selezionavano feature meno biologicamente rilevanti.

Significatività e Rivendicazioni
Il documento afferma che E-FLoPS fornisce un meccanismo robusto per la scoperta di modelli sparsi nel FL, affrontando specificamente l'instabilità riscontrata nei regimi con pochi campioni e alta dimensionalità. Utilizzando la regolarizzazione dell'entropia, il metodo evita l'impegno precoce verso supporti sparsi subottimali, portando a una migliore generalizzazione e a un recupero più accurato della vera struttura sparsa sottostante. Gli autori sostengono che il loro approccio ottenga questi guadagni statistici mantenendo costi di comunicazione comparabili ai metodi sparsi esistenti come Fed-IHT, rendendolo una soluzione valida per ambienti di apprendimento che preservano la privacy e sono limitati nelle risorse. Il lavoro evidenzia l'importanza dell'esplorazione dell'incertezza nell'ottimizzazione federata sparsa, un fattore spesso trascurato dalle tecniche di potatura basate sull'ampiezza.