Class-Grouped Normalized Momentum and Faster Hyperparameter Exploration to Tackle Class Imbalance in Federated Learning

Sintesi Tecnica: FedCGNM e FedHOO per il Disbilanciamento delle Classi nel Federated Learning

Definizione del Problema
Il disbilanciamento delle classi rimane una sfida critica nel Federated Learning (FL), in particolare quando la distribuzione globale delle etichette è a coda lunga (long-tailed). In tali scenari, le classi minoritarie soffrono di una scarsa performance predittiva a causa della sottorappresentazione nei dati di addestramento. Le tecniche centralizzate standard, come la riponderazione della perdita (loss reweighting), il campionamento avanzato o l'augmentation generativa, sono difficili da implementare nel FL. I vincoli sulla privacy impediscono lo scambio di dati e i generatori sintetici spesso non riescono a produrre campioni realistici per regimi sparsi o compiti specifici del dominio (ad esempio, il rilevamento dei difetti nei semiconduttori). Inoltre, gli approoli esistenti a livello di ottimizzazione, come la Normalizzazione Per-Classe (PCN), affrontano limitazioni negli scenari multi-classe: sono suscettibili al rumore direzionale delle classi minoritarie e introducono discrepanze di scala quando si sommano molti vettori unitari, il che destabilizza la convergenza all'aumentare del numero di classi.

Metodologia
Il documento propone un approccio a due fronti per affrontare questi problemi: un ottimizzatore lato client chiamato FedCGNM (Federated Class-Grouped Normalized Momentum) e un tuner di iperparametri ausiliario chiamato FedHOO.

1. FedCGNM (Ottimizzatore):

   • Raggruppamento delle Classi (Class Grouping): Invece di normalizzare i gradienti per ogni singola classe ($C$), FedCGNM suddivide le classi in un piccolo numero di gruppi ($H \ll C$, tipicamente maggioranza vs. minoranza). La regola di raggrupamento è derivata minimizzando la varianza intra-gruppo delle proporzioni delle classi, una strategia che si è dimostrata un proxy affidabile per una ricerca esaustiva delle soglie.
   • Momentum Normalizzato per Gruppo (Group-Wise Normalized Momentum): Per ogni gruppo, l'algoritmo mantiene un vettore di momentum. Durante l'addestramento locale, i gradienti stocastici vengono calcolati per gruppo e il momentum viene aggiornato. Fondamentalmente, il vettore di momentum di ogni gruppo viene normalizzato a lunghezza unitaria (con una costante di stabilizzazione $\delta$) prima della sommatoria.
   • Direzione di Aggiornamento: La direzione finale dell'aggiornamento è la somma di questi momentum di gruppo normalizzati. Questo design equalizza la grandezza del gradiente tra i gruppi di maggioranza e di minoranza, mitigando al contempo il rumore direzionale inerente ai gradienti delle classi rare. Operando su un numero limitato di gruppi anziché su singole classi, evita la discrepanza di scala e l'instabilità associate a PCN.
   • Allineamento del Client: Forzando l'aggiornamento di ogni gruppo ad avere una norma unitaria, le direzioni di aggiornamento aggregate tra client eterogenei diventano più allineate, affrontando una delle cause primarie del degrado delle prestazioni nel FL.

2. FedHOO (Tuner di Iperparametri):

   • Contesto: Selezionare i tassi di campionamento appropriati non è banale nel FL a causa dello spazio di ricerca combinatorio tra i client e del rischio di destabilizzare l'ottimizzazione con programmi fluttuanti.
   • Meccanismo: FedHOO è un algoritmo basato su X-armed bandit (XAB) progettato per federazioni di piccola scala (dove $K$ è piccolo). Sfrutta la natura lineare dell'aggregazione FL. In ogni round di comunicazione, il server istruisce ogni client ad addestrarsi con due candidati tassi di campionamento.
   • Efficienza: Combinando linearmente gli aggiornamenti dei client restituiti, il server può inferire la performance di validazione di tutte le $2^K$ possibili combinazioni di tassi locali senza richiedere $2^K$ round di addestramento separati. Ciò consente di identificare rapidamente i tassi di campionamento efficaci nelle prime fasi dell'addestramento mantenendo al contempo la privacy (i client restituiscono solo metriche di validazione scalari).

Analisi Teorica
Il documento fornisce un'analisi della convergenza per FedCGNM che tiene esplicitamente conto dei tassi di campionamento variabili nel tempo. Sotto le assunzioni standard di regolarità (smoothness), gradienti limitati e imparzialità (unbiasedness), gli autori dimostrano che FedCGNM raggiunge un tasso di punto stazionario standard di $O(T^{-1/2})$. L'analisi evidenzia che, affinché la convergenza sia garantita, la variazione cumulativa dei tassi di campionamento tra i round ($V_T$) deve soddisfare condizioni specifiche, implicando che un campionamento aggressivo ma instabile può ostacre la convergenza.

Contributi Chiave

1. Ottimizzatore FedCGNM: Il primo ottimizzatore lato client che raggruppa le classi e applica un momentum a norma unitaria per gruppo. Bilancia l'influenza di maggioranza/minoranza, riduce il rumore e mitiga la discrepanza di scala.
2. Raggruppamento Consapevole della Varianza (Variance-Aware Grouping): Una strategia basata su principi e leggera che inquadra la partizione delle classi come la minimizzazione della varianza intra-gruppo, producendo split ottimali comparabili alla ricerca esaustiva.
3. Analisi della Convergenza: Un quadro teorico innovativo che combina FL, normalizzazione di gruppo e tassi di campionamento dinamici, dimostrando l'impatto degli schemi di campionamento sulla convergenza.
4. FedHOO: Il primo algoritmo di ottimizzazione per determinare i tassi di campionamento locali combinatori nel FL utilizzando un framework X-armed bandit che sfrutta il parallelismo federato per un'esplorazione efficiente.

Risultati Sperimentali
I metodi sono stati valutati su quattro benchmark long-tailed pubblici (CIFAR-10-LT, CIFAR-100-LT, Adult Income, UNSW-NB15) e su un dataset proprietario di difetti di chip semiconduttori (CDD).

• Performance: FedCGNM ha superato costantemente i baseline (inclusi FedAvg, FedProx, CLIMB, FedGraB e varianti di PCN) in tutti gli scenari. Nei regimi con pochi client ($K=5$), la combinazione di FedCGNM con FedHOO ha prodotto i risultati migliori, con miglioramenti fino al 29% rispetto a FedAvg con campionamento.
• Applicazione Industriale: Sul dataset CDD (tasso di difetto dell'1,7%), il metodo proposto ha ottenuto un miglioramento del 16% rispetto al miglior baseline, dimostrando l'efficacia in contesti industriali reali e sensibili alla privacy.
• Robustezza: L'approccio ha mostrato resilienza in regimi con molti client con partecipazione parziale e sotto diversi gradi di eterogeneità dei dati (Non-IID).
• Ablation: Gli esperimenti hanno confermato che due gruppi ($H=2$) offrono il miglior compromesso bias-varianza e che fattori di momentum moderati ($\beta \approx 0.5$) producono prestazioni ottimali.

Significatività e Rivendicazioni
Il documento sostiene che FedCGNM e FedHOO formino un framework pratico per mitigare il disbilanciamento globale delle classi in contesti federati che preservano la privacy, senza fare affidamento sullo scambio di dati o sulla generazione sintetica. La significatività risiede nel risolvere il disbilanciamento delle classi a livello di ottimizzazione, il che è più generalmente applicabile rispetto ai trucchi a livello di dati limitati dalla privacy o dalle caratteristiche del dataset. Gli autori sottolineano che il loro approccio è particolarmente efficace nelle federazioni di piccola scala (comuni nelle implementazioni industriali) dove il tuning congiunto degli iperparametri è fattibile e critico, mentre la strategia di raggruppamento garantisce stabilità e allineamento negli scenari multi-classe in cui le precedenti tecniche di normalizzazione falliscono. Il lavoro fa progredire il campo fornendo sia un ottimizzatore robusto che un metodo efficiente e teoricamente fondato per il tuning delle strategie di campionamento in presenza di disbilanciamento delle classi.