Technische Samenvatting: PreLort – Prefix-Nested LoRA voor Federated Fine-Tuning onder Rang-heterogeniteit

1. Probleemstelling

Federated Learning (FL) maakt collaboratieve training van Large Language Models (LLMs) mogelijk zonder het delen van ruwe data, maar het staat voor aanzienlijke uitdagingen bij het aanpassen van fundamentele modellen via parameter-efficiënte fine-tuning (PEFT) methoden zoals Low-Rank Adaptation (LoRA). Een primair obstakel is rang-heterogeniteit: in realistische cross-device omgevingen beschikken cliënten over variërende computationele en geheugenbeperkingen, wat leidt tot het initialiseren van LoRA-adapters met verschillende rangen ($r$).

Bestaande oplossingen voor heterogene federated LoRA richten zich voornamelijk op de aggregatiefase, waarbij vaak zero-padding wordt toegepast om adapters van verschillende rangen uit te lijnen voor middeling. De auteurs identificeren echter een fundamentele beperking in deze benaderingen: representatie-misalignement. Zelfs wanneer aggregatie wiskundig gedefinieerd is, leidt het combineren van onafhankelijk getrainde modellen met verschillende rangen tot inconsistenties. Specifiek profiteren cliënten met een lagere rang niet van de updates van cliënten met een hogere rang, omdat de informatie niet hiërarchisch over de rang-dimensies is verdeeld. Zero-padding verzwakt bovendien de bijdragen van high-rank cliënten door het introduceren van segmenten met een waarde van nul, wat resulteert in een suboptimale benutting van gedeelde low-rank representaties.

2. Methodologie: PreLort

De auteurs stellen PreLort (Prefix-Nested LoRA) voor, een framework dat rang-heterogeniteit aanpakt door de trainingsdynamiek en de aggregatieregels gezamenlijk te structureren. De kernfilosofie is het organiseren van adapter-dimensies in een prefix-hiërarchie, waardoor wordt gewaarborgd dat dimensies met een lagere rang taakrelevante informatie coderen die gedeeld wordt tussen alle cliënten, terwijl dimensies met een hogere rang incrementele capaciteit vastleggen.

2.1 Geneste Low-Rank Structuur

PreLort legt een prefix-gestructureerde organisatie op aan de adapter-dimensies. Voor een cliënt met rang $r_k$ zijn alleen de eerste $r_k$ dimensies actief. Cruciaal is dat tijdens de lokale training de adapter niet uitsluitend wordt geoptimaliseerd op zijn volledige rang. In plaats daarvan wordt deze geoptimaliseerd onder meerdere getrunceerde configuraties (prefix sub-rangen).

• Mechanisme: Een cliënt met rang $r_k$ wordt getraind om goed te presteren niet alleen bij $r_k$, maar ook bij alle lagere rang-niveaus die aanwezig zijn in de federatie ($r \in \{r \mid r \le r_k\}$).
• Doelstelling: Dit dwingt het model om taakrelevante signalen te concentreren in de leidende (low-rank) dimensies, waardoor wordt gegarandeerd dat deze dimensies informatief en bruikbaar blijven, zelfs wanneer ze worden getrunceerd.

2.2 Lokale Trainingsstrategie

De lokale doelfunctie voor een cliënt $k$ wordt gedefinieerd als het gemiddelde van de cross-entropy losses berekend over alle geldige prefix sub-rangen $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 de praktijk omvat elke trainingsstap één full-rank forward/backward pass en aanvullende passes voor alle sub-rangen. Deze strategie zorgt ervoor dat de prefix-segmenten expliciet worden geoptimaliseerd op dezelfde granulariteit als gebruikt tijdens de aggregatie, wat een consistente hiërarchie stimuleert waarbij lagere dimensies de meest kritieke informatie vastleggen.

2.3 Segment-wijze Aggregatie

Om de verdunning veroorzaakt door zero-padding te voorkomen, introduceert PreLort een segment-wijze aggregatieregel die opereert in de gradiënt-ruimte (update space).

• Proces: In plaats van alle cliënten te middelen voor elke dimensie, aggregeert de server gradiënten voor elk rang-segment $[r_{l-1}, r_l)$ uitsluitend over de subset van cliënten waarvan de rang $r_k$ voldoende is om bij te dragen aan dat segment (d.w.z. $r_k \ge r_l$).
• Formule: Voor het $l$-de segment is de geaggregeerde gradiënt:
  $$\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}$$
• Voordeel: Dit voorkomt signaalverdunning door zero-padded low-rank cliënten en zorgt ervoor dat high-rank dimensies alleen worden bijgewerkt door cliënten die daadwerkelijk de capaciteit hebben om ze te leren, terwijl low-rank dimensies profiteren van de volledige populatie.

3. Belangrijkste Bijdragen

Het artikel schetst drie primaire bijdragen:

1. Identificatie van Representatie-misalignment: De auteurs wijzen erop dat het gebrek aan gestructureerde uitlijning over de rang-dimensies een sleutelbeperking is in huidige heterogene federated LoRA-methoden, wat leidt tot inconsistente representaties.
2. Geneste Formulering en Verenigde Strategie: Zij stellen een geneste low-rank formulering voor die een gemeenschappelijke prefix-structuur oplegt. Dit wordt gekoppeld aan een verenigde trainings- en aggregatiestrategie die consistente aggregatie over heterogene rangen mogelijk maakt zonder zero-padding.
3. Empirische Validatie: De methode wordt consequent aangetoond de modelkwaliteit en stabiliteit te verbeteren ten opzichte van eerdere heterogene federated LoRA-methoden op instructie-tuning benchmarks.

4. Experimentele Resultaten

De auteurs hebben PreLort geëvalueerd op twee basismodellen (TinyLlama-1.1B en Qwen2.5-0.5B) over drie datasets: Alpaca, Databricks-dolly-15, en 20 Newsgroups.

• Prestatie-metrieken: PreLort presteerde consequent beter dan de baselines (ZeroPad, HetLoRA, FLoRA) in Nauwkeurigheid (MMLU, 20 Newsgroups) en ROUGE-L scores.
  • Op TinyLlama-1.1B behaalde PreLort 33,86% nauwkeurigheid op MMLU (Dolly) en 30,76% op Alpaca, waarmee het de op één na beste baselines overtrof.
  • Op Qwen2.5-0.5B behaalde het 39,85% nauwkeurigheid op MMLU (Dolly).
• Perplexity: PreLort behaalde de laagste perplexity op Dolly en NewsGroup voor beide basismodellen, wat duidt op een superieure token-level likelihood.
• Rang-robuustheid: Analyse toonde aan dat PreLort een stabiele prestatie behoudt over verschillende prefix-truncatie niveaus. In tegenstelling hiertoe vertoonden baselines zoals "Nested Aggregation Only" instabiliteit, met name bij hogere rangen, wat bevestigt dat aggregatie alleen onvoldoende is zonder de geneste trainingsdoelstelling.
• Belang per Dimensie: Visualisatie van de dimensie-belangrijkheid onthulde dat PreLort het taakrelevante signaal concentreert in de leidende dimensies (een gestructureerde distributie), terwijl standaard training een uniforme distributie oplevert. Deze hiërarchische prioritering faciliteert stabiele updates over heterogene rangen.

5. Betekenis en Claims

Het artikel claimt dat PreLort een verschuiving vertegenwoordigt in het aanpakken van rang-heterogeniteit: in plaats van het te behandelen als een post-hoc aggregatieprobleem, wordt het opgelost door het trainingsproces zelf te structureren.

• Kernargument: De auteurs betogen dat het aanpakken van rang-heterogeniteit niet alleen verbeterde aggregatie vereist, maar ook een gestructureerde uitlijning van representaties over rangen heen.
• Efficiëntie: De methode stelt low-rank cliënten in staat om te profiteren van rijkere informatie bijgedragen door high-rank cliënten, omdat de prefix-dimensies consistent worden geleerd en geaggregeerd.
• Bescheidenheid over Beperkingen: De auteurs erkennen dat de geneste trainingsstrategie ongeveer 40% extra computationele overhead introduceert door meerdere forward/backward passes per batch. Ze merken ook op dat de effectiviteit afhangt van voldoende lokale optimalisatiestappen; bij zeer beperkte lokale epochs leert het model mogelijk de goed uitgelijnde low-rank prefix-representaties niet volledig.

Concluderend toont PreLort aan dat rang-heterogeniteit in federated LoRA het best kan worden afgehandeld door zowel de manier waarop adapters worden getraind (via geneste prefix-optimalisatie) als de manier waarop hun updates worden geaggregeerd (via segment-wijze regels) gezamenlijk te structureren, wat resulteert in uitgelijnde, robuuste prefix-representaties die bestaande heterogene aggregatiemethoden overtreffen.