Allineare prima dell’aggregazione: LoRA federata coerente con la base in presenza di ranghi eterogenei

Abstract degli autori

Il fine-tuning federato dei modelli linguistici di grandi dimensioni (LLM) in presenza di capacità eterogenee dei client è sempre più importante, e l’adattamento a basso rango (LoRA) con ranghi eterogenei lo rende praticabile. Gli aggregatori server-side esistenti compongono gli aggiornamenti locali delle matrici $ΔW_i$ a partire dai fattori LoRA specifici di ciascun client $(A_i,B_i)$, prima di aggregarli. Sebbene questi aggiornamenti delle matrici appartengano allo stesso spazio dei pesi $\mathbb{R}^{m\times n}$ e possano essere mediati direttamente, i fattori LoRA che li generano possono essere espressi in basi locali specifiche dei client. Nel tipico ciclo di composizione, aggregazione, fattorizzazione e distribuzione, questa inconsistenza tra le basi può rendere instabile, dopo la rifattorizzazione lato server, il prefisso di riferimento di rango-$r_i$ distribuito ai client con rango inferiore. In questo lavoro proponiamo \textbf{FedAbA}, un framework di aggregazione LoRA federata coerente con la base, per ranghi eterogenei. FedAbA estrae innanzitutto la base locale di ciascun client, la allinea al prefisso corrispondente della base di riferimento globale e costruisce, prima dell’aggregazione, una ricostruzione coerente con la base di ciascun aggiornamento locale della matrice. Aggrega quindi queste ricostruzioni con pesi che combinano la dimensione dei dati del client e il punteggio di coerenza, seguiti da una rifattorizzazione esatta tramite SVD nello spazio dei coefficienti e dal fissaggio della gauge prima della distribuzione dei prefissi di rango. La nostra analisi teorica chiarisce il ruolo della coerenza tra le basi e supporta il meccanismo di aggregazione proposto, basato sulla coerenza. Esperimenti approfonditi mostrano che FedAbA supera le baseline rappresentative di LoRA federata, con i miglioramenti maggiori in presenza della più forte eterogeneità dei ranghi.