Sintesi Tecnica: PreLort – Prefix-Nested LoRA per il Fine-Tuning Federato sotto Eterogeneità di Rank

1. Definizione del Problema

L'Apprendimento Federato (FL) consente l'addestramento collaborativo di Large Language Models (LLM) senza condividere dati grezzi, ma affronta sfide significative quando si adatta i modelli di base tramite metodi di parameter-efficient fine-tuning (PEFT) come Low-Rank Adaptation (LoRA). Un ostacolo primario è l'eterogeneità di rank: in scenari cross-device realistici, i client possiedono vincoli computazionali e di memoria variabili, che li portano a inizializzare gli adattatori LoRA con diversi rank ($r$).

Le soluzioni esistenti per il federated LoRA eterogeneo si concentrano principalmente sulla fase di aggregazione, impiegando spesso lo zero-padding per allineare gli adattatori di diversi rank prima della media. Tuttavia, gli autori identificano un limite fondamentale in questi approaggi: il disallineamento delle rappresentazioni. Anche quando l'aggregazione è matematicamente definita, la combinazione di modelli addestrati indipendentemente con diversi rank porta a inconsistenze. Nello specifico, i client a basso rank non beneficiano degli aggiornamenti dei client ad alto rank perché l'informazione non è distribuita gerarchicamente attraverso le dimensioni di rank. Lo zero-padding, inoltre, diluisce i contributi dei client ad alto rank introducendo segmenti con valori nulli, risultando in un uso subottimale delle rappresentazioni a basso rank condivise.

2. Metodologia: PreLort

Gli autori propongono PreLort (Prefix-Nested LoRA), un framework che affronta l'eterogeneità di rank strutturando congiuntamente la dinamica di addestramento e le regole di aggregazione. La filosofia centrale è organizzare le dimensioni dell'adattatore in una gerarchia di prefissi, garantendo che le dimensioni a basso rank codifichino informazioni rilevanti per il task condivise tra tutti i client, mentre le dimensioni ad alto rank catturino la capacità incrementale.

2.1 Struttura Low-Rank Annidata (Nested)

PreLort impone un'organizzazione a struttura di prefisso sulle dimensioni dell'adattatore. Per un client con rank $r_k$, solo le prime $r_k$ dimensioni sono attive. Fondamentalmente, durante l'addestramento locale, l'adattatore non è ottimizzato esclusivamente al suo rank completo. Invece, viene ottimizzato sotto configurazioni di troncamento multipli (sotto-rank di prefisso).

• Meccanismo: Un client con rank $r_k$ è addestrato per performare bene non solo a $r_k$, ma anche a tutti i livelli di rank inferiori presenti nella federazione ($r \in \{r \mid r \le r_k\}$).
• Obiettivo: Questo costringe il modello a concentrare i segnali rilevanti per il task nelle dimensioni principali (a basso rank), assicurando che queste dimensioni rimangano informative e utili anche in caso di troncamento.

2.2 Strategia di Addestramento Locale

L'obiettivo locale per un client $k$ è definito come la media delle perdite di cross-entropy calcolate su tutti i validi sotto-rank di prefisso $R_k$:
$$\mathcal{L}_k = \frac{1}{|R_k|} \sum_{r \in R_k} \mathcal{L}_{CE}(f(\cdot; \theta, r), \mathcal{D}_k)$$
In pratica, ogni step di addestramento comporta un passaggio forward/backward completo e passaggi aggiuntivi per tutti i sotto-rank. Questa strategia assicura che i segmenti di prefisso siano esplicitamente ottimizzati alla stessa granularità utilizzata durante l'aggregazione, incoraggiando una gerarchia coerente in cui le dimensioni inferiori catturano le informazioni più critiche.

2.3 Aggregazione per Segmenti

Per evitare la diluizione causata dallo zero-padding, PreLort introduce una regola di aggregazione per segmenti che opera nello spazio del gradiente (aggiornamento).

• Processo: Invece di mediare tutti i client per ogni dimensione, il server aggrega i gradienti per ogni segmento di rank $[r_{l-1}, r_l)$ solo sul sottoinsieme di client il cui rank $r_k$ è sufficiente a contribuire a quel segmento (ovvero $r_k \ge r_l$).
• Formula: Per il $l$-esimo segmento, il gradiente aggregato è:
  $$\nabla A_{global}^{r_{l-1}:r_l} = \frac{\sum_{k: r_k \ge r_l} n_k \nabla A_k^{r_{l-1}:r_l}}{\sum_{k: r_k \ge r_l} n_k}$$
• Vantaggio: Questo previene la diluizione del segnale dai client a basso rank con zero-padding e assicura che le dimensioni ad alto rank siano aggiornate solo dai client che possiedono effettivamente la capacità di apprenderle, mentre le dimensioni inferiori beneficiano dell'intera popolazione.

3. Contributi Chiave

Il documento delinea tre contributi primari:

1. Identificazione del Disallineamento delle Rappresentazioni: Gli autori individuano che la mancanza di un allineamento strutturato tra le dimensioni di rank è un limite chiave nei metodi attuali di federated LoRA eterogeneo, portando a rappresentazioni inconsistenti.
2. Formulazione Annidata e Strategia Unificata: Propongono una formulazione low-rank annidata che impone una comune struttura di prefisso. Questa è accoppiata a una strategia di addestramento e aggregazione unificata che consente un'aggregazione coerente tra rank eterogenei senza zero-padding.
3. Validazione Empirica: Il metodo è dimostrato essere costantemente in grado di migliorare la qualità e la stabilità del modello rispetto ai precedenti metodi di federated LoRA eterogeneo su benchmark di instruction-tuning.

4. Risultati Sperimentali

Gli autori hanno valutato PreLort su due modelli base (TinyLlama-1.1B e Qwen2.5-0.5B) su tre dataset: Alpaca, Databricks-dolly-15 e 20 Newsgroups.

• Metriche di Performance: PreLort ha superato costantemente i baseline (ZeroPad, HetLoRA, FLoRA) in Accuratezza (MMLU, 20 Newsgroups) e punteggi ROUGE-L.
  • Su TinyLlama-1.1B, PreLort ha raggiunto il 33,86% di accuratezza su MMLU (Dolly) e il 30,76% su Alpaca, superando i migliori baseline.
  • Su Qwen2.5-0.5B, ha raggiunto il 39,85% di accuratezza su MMLU (Dolly).
• Perplessità: PreLort ha ottenuto la perplessità più bassa su Dolly e NewsGroup per entrambi i modelli base, dimostrando una superiore verosimiglianza a livello di token.
• Robustezza del Rank: L'analisi ha mostrato che PreLort mantiene prestazioni stabili attraverso diversi livelli di troncamento del prefisso. Al contrario, i baseline come "Nested Aggregation Only" hanno mostrato instabilità, in particolare ad alti rank, confermando che l'aggregazione da sola è insufficiente senza l'obiettivo di addestramento annidato.
• Importanza per Dimensione: La visualizzazione dell'importanza delle dimensioni ha rivelato che PreLort concentra il segnale rilevante per il task nelle dimensioni principali (una distribuzione strutturata), mentre l'addestramento standard produce una distribuzione uniforme. Questa prioritizzazione gerarchica facilita aggiornamenti stabili tra rank eterogenei.

5. Significato e Rivendicazioni

Il paper sostiene che PreLort rappresenti un cambio di paradigma nel gestire l'eterogeneità di rank: invece di trattarla come un problema di aggregazione post-hoc, la risolve strutturando il processo di addestramento stesso.

• Argomento Centrale: Gli autori sostengono che affrontare l'eterogeneità di rank richieda non solo un'aggregazione migliorata, ma anche un allineamento strutturato delle rappresentazioni tra i rank.
• Efficienza: Il metodo consente ai client a basso rank di beneficiare delle informazioni più ricche fornite dai client ad alto rank perché le dimensioni del prefisso sono apprese e aggregate in modo coerente.
• Limitazioni: Gli autori riconoscono che la strategia di addestramento annidato introduce circa un 40% di overhead computazionale aggiuntivo a causa dei molteplici passaggi forward/backward per batch. Notano inoltre che l'efficacia dipende da un numero sufficiente di step di ottimizzazione locale; con pochissimi epoch locali, il modello potrebbe non apprendere appieno le rappresentazioni di prefisso a basso rank ben allineate.

In conclusione, PreLort dimostra che l'eterogeneità di rank nel federated LoRA è gestita al meglio strutturando congiuntamente il modo in cui gli adattatori vengono addestrati (tramite l'ottimizzazione del prefisso annidato) e il modo in cui i loro aggiornamenti vengono aggregati (tramite regole per segmenti), ottenendo rappresentazioni di prefisso allineate e robuste che superano i metodi di aggregazione eterogenei esistenti.