FedProxy: Federated Fine-Tuning of LLMs via Proxy SLMs and Heterogeneity-Aware Fusion

1. Probleemstelling

Het fine-tunen van Large Language Models (LLMs) via Federated Learning (FL) op gedecentraliseerde, privacygevoelige data stuit op een fundamenteel drievoudig dilemma (trilemma):

1. Intellectueel Eigendom (IP) Bescherming: De server wil de eigenaar van het proprietary LLM niet het volledige model blootstellen aan clients.
2. Client Privacy: Clients willen hun lokale data niet delen.
3. Prestatieverlies door Heterogeniteit: Data-distributies tussen clients zijn vaak niet-IID (niet identiek en onafhankelijk verdeeld), wat leidt tot prestatieverlies bij het samenvoegen van modellen.

Bestaande methoden zoals Offsite-Tuning (OT) lossen het IP- en privacyprobleem op door clients alleen lichte "adapters" te laten trainen terwijl het hoofdmodel (backbone) gefixeerd en verduisterd blijft. Echter, analyse in dit paper toont aan dat OT-methoden een fundamenteel prestatieknelpunt hebben. De beperkte representatiecapaciteit van alleen adapters resulteert in een significant prestatieverschil ten opzichte van centraal fine-tunen, vooral bij heterogene data.

2. Methodologie: Het FedProxy Framework

FedProxy introduceert een nieuw paradigma dat het zwakke adapter-model vervangt door een krachtige, gedeelde Proxy Small Language Model (SLM). Dit model is een gecomprimeerde versie van het originele LLM, getraind op publieke data, en fungeert als een hoogwaardig surrogaat voor collaboratief fine-tunen.

Het framework bestaat uit drie fasen:

Fase 1: Efficiënte Representatie via Server-Gestuurde Compressie

• Doel: Het creëren van een compact proxy-model ($f_\phi$) dat klein genoeg is voor client-side training, maar krachtig genoeg om de kernkennis van het LLM ($f_\theta$) te behouden.
• Methode: De server comprimeert het proprietary LLM naar een Proxy SLM met behulp van een publieke dataset (bijv. Alpaca).
• Techniek: Er wordt gebruik gemaakt van gestructureerde pruning (het verwijderen van hele transformer-blokken). De belangrijkheid van blokken wordt bepaald met de Block Influence (BI) metric, die meet hoe sterk een blok bijdraagt aan de representatiecapaciteit op de publieke data.
• Resultaat: Een uniek Proxy SLM dat als startpunt dient voor alle clients, wat de IP-blootstelling minimaliseert en de communicatiekosten verlaagt.

Fase 2: Robuuste Optimalisatie via Interferentie-Verlichtende Aggregatie

Om het probleem van heterogene data en parameterconflicten aan te pakken, introduceert FedProxy een drie-staps aggregatiestrategie:

1. Server-zijde Analyse: De server berekent per client:
   • Taakvectors: Het verschil tussen het lokale en het globale model.
   • Heterogeniteitscoëfficiënt ($h_k$): Meet hoe sterk een client afwijkt van de rest.
   • Conflictscores ($C_d$): Per parameter-dimensie wordt gemeten of updates van verschillende clients tegenstrijdige tekens hebben (positief vs. negatief).
2. Client-zijde Regularisatie (PCR - Parameter-Level Conflict-Aware Regularization):
   • Clients gebruiken de ontvangen conflictscores om hun lokale trainingsdoel aan te passen.
   • Een regularisatieterm straft afwijkingen van het globale model af op parameters waar hoge consensus is (hoge conflictscore), maar laat meer vrijheid voor personalisatie op parameters met lage consensus. Dit voorkomt dat clients "te ver" de verkeerde kant op gaan.
3. Heterogeneity-Aware TIES-Merging (H-TIES):
   • Een geavanceerde server-aggregatiestrategie die de updates van clients samenvoegt.
   • Het past de sparsiteit van updates aan op basis van heterogeniteit (hoger heterogeniteit = sparser updates).
   • Het lost parameterconflicten op door alleen updates te behouden die overeenkomen met de dominante richting (sign) van de updates, gewogen op basis van consensus. Dit voorkomt dat tegenstrijdige updates elkaar opheffen.

Fase 3: Trainingsvrije "Plug-in" Kennisfusie

• Doel: De geleerde kennis uit het geaggregeerde Proxy SLM terugvoeren naar het originele LLM zonder dure hertraining.
• Methode: Omdat het Proxy SLM een direct sub-netwerk is van het LLM, worden de bijgewerkte lagen van het Proxy SLM direct overgeschreven op de corresponderende lagen in het volledige LLM.
• Voordeel: Dit is een "plug-in" mechanisme dat geen extra inferentie-latentie of retraining vereist.

3. Belangrijkste Bijdragen

1. FedProxy Framework: Een holistische oplossing voor het IP-privacy-prestatie dilemma, die een Proxy SLM gebruikt als surrogaat in plaats van beperkte adapters.
2. H-TIES en PCR: Nieuwe algoritmen voor heterogene data. H-TIES lost parameterinterferentie op tijdens het samenvoegen, terwijl PCR clients helpt om lokaal te trainen zonder de globale cohesie te schaden.
3. Prestatiebewijs: Uitgebreide experimenten tonen aan dat FedProxy aanzienlijk beter presteert dan bestaande OT-methoden (zoals FedOT en FedBiOT) en de prestaties van centraal fine-tunen (CentSFT) benadert.

4. Experimentele Resultaten

De auteurs testten FedProxy op LLaMA2-7B en Mistral-7B met compressieverhoudingen van 30% en 50%.

• Homogene Taken (IID data): FedProxy (50% compressie) bereikte 90,1% van de prestatie van centraal fine-tunen op QA-taken en 90,7% op GLUE-taken. Dit is een enorme verbetering ten opzichte van FedOT (bijv. +22,1% op GLUE).
• Heterogene Taken (Niet-IID data): In dit uitdagende scenario, waar OT-methoden vaak falen door parameterconflicten, behaalde FedProxy 91,3% van de CentSFT-prestatie op QA en 89,6% op GLUE.
• Kosten-Prestatie Afweging: Hoewel FedProxy meer client-computatie en iteratieve communicatie vereist dan FedOT (door het trainen van een groter Proxy SLM in plaats van alleen adapters), is de initiële communicatiekosten vergelijkbaar. De investering in rekenkracht wordt ruimschoots gecompenseerd door de prestatiewinst.
• Ablatie Studies:
  • Zonder PCR of H-TIES daalt de prestatie aanzienlijk, wat bewijst dat beide componenten essentieel zijn.
  • FedProxy presteert beter dan een variant waarbij ook de emulator wordt getraind (FedOT-ET), wat aantoont dat de architectuur van een geïntegreerde Proxy SLM superieur is aan losse adapters.
  • De IP-bescherming is effectief: het geïsoleerde Proxy SLM presteert slecht als standalone model, maar na fusie verbetert het LLM aanzienlijk.

5. Betekenis en Conclusie

FedProxy stelt een nieuwe standaard voor voor veilig en hoogpresterend federated fine-tuning van LLMs. Het overbrugt de kloof tussen de theoretische voordelen van federated learning en de praktische beperkingen van IP-bescherming en data-heterogeniteit.

Door de beperkingen van adapters te omzeilen en te kiezen voor een krachtig, gecomprimeerd surrogaatmodel met geavanceerde conflict-oplossing, maakt FedProxy het mogelijk om LLMs te verbeteren op gedecentraliseerde data zonder in te leveren op privacy of modelkwaliteit. De methode is schaalbaar en biedt een praktische route voor industrieel gebruik van federated learning met foundation models.

Beperkingen: De auteurs erkennen dat de prestaties afhangen van de compressieverhouding (hoge compressie leidt tot kwaliteitsverlies) en dat de schaalbaarheid van de aggregatie (O(K²)) bij zeer grote aantallen clients verbeterd moet worden, bijvoorbeeld via hiërarchische aggregatie.