Technische Samenvatting: Federated Nested Learning (FedNL)

Probleemstelling

Federated Learning (FL) voor Large Language Models (LLM's) staat voor twee aanhoudende, met elkaar verweven uitdagingen: data-heterogeniteit (Non-IID) en verdeling met lange staart. In realistische scenario's zijn de data-distributies van clients sterk scheef (bijvoorbeeld medische versus coderingsdata), en bevindt kritieke kennis zich vaak in de lange staart, die tijdens globale aggregatie gemakkelijk wordt overschaduwd door hoofdclasses.

Bestaande benaderingen richten zich voornamelijk op het regulariseren of personaliseren van statische modelgewichten (bijvoorbeeld FedProx, FedSSI). Hoewel deze methoden effectief zijn in het mitigeren van catastrofaal vergeten, behandelen ze het globale model als een container voor statische kennis. Bijgevolg mist het model, wanneer het wordt ingezet bij een client met onbekende, sterk heterogene data, plasticiteit op het moment van testen – het vermogen om zich aan te passen aan de huidige context zonder gradiëntupdates. Dit leidt tot suboptimale prestaties op domeinspecifieke taken waarbij nuances afhankelijk zijn van de context.

Methodologie: Federated Nested Learning (FedNL)

Het artikel stelt Federated Nested Learning (FedNL) voor, een raamwerk dat FL herformuleert van het aggregeren van statische kennis naar het gezamenlijk leren van optimalisatieregels. Geïnspireerd door de theorie van Nested Learning (NL), behandelt FedNL inferentie niet als een passief ophaalproces, maar als een hoogfrequente "inner-loop" trainingsprocess waarbij het model context actief comprimeert tot een transient geheugentoestand.

Kernarchitectuur

FedNL reformuleert FL als een drie-niveau genest optimalisatiesysteem:

1. Niveau 0 (Buitenste lus - Server): Gezamenlijke Generalisatie. De server aggregeert meta-regels (parameters $\theta$) in plaats van geheugeninhoud. Deze regels bepalen hoe geheugen wordt opgebouwd en bijgewerkt.
2. Niveau 1 (Intermediaire lus - Client): Regel leren. Tijdens lokale training optimaliseren clients $\theta$ (bijvoorbeeld LoRA-adapters, gating-mechanismen) om ervoor te zorgen dat de inner-loop bruikbare geheugentoestanden produceert voor downstream-taken. Dit omvat backpropagatie door de geheugentrajekt.
3. Niveau 2 (Binnenste lus - Client): Aanpassing op het moment van testen. Tijdens inferentie gebruiken clients de globale regels $\theta$ om een privé, contextbewuste geheugentoestand $S_t$ online bij te werken via een Delta Rule. Deze toestand is transient en verlaat het apparaat nooit.

Technische Implementatie

• Titans-gebaseerde lineaire aandacht: FedNL maakt gebruik van de Titans-architectuur, die een gelijneerde aandachtsmechanisme gebruikt. De geheugentoestand $S_t$ wordt bijgewerkt via een Delta Rule:
  $$S_t = S_{t-1} + \eta(v_t - S_{t-1}k_t)k_t^\top$$
  Deze update is wiskundig equivalent aan een enkele stap van gradiëntafstijging op een associatief geheugen-doel, waarbij inferentie effectief wordt behandeld als een online optimalisatieproces.
• Ontkoppeling van kennis en regels: Het model bestaat uit een bevroren backbone (die algemene linguïstische capaciteiten biedt) en trainbare meta-parameters ($\theta$). De meta-parameters omvatten Low-Rank Projecties (LoRA) voor query/key/value-generatie en een Memory Gating-mechanisme ($\alpha$) om statische en dynamische aandacht in evenwicht te brengen.
• Efficiëntie: Door alleen de kleine set meta-regels (adapters) te aggregeren in plaats van volledige modelgewichten of geheugentoestanden, vermindert FedNL de communicatie-overhead drastisch. Bovendien behoudt het lineaire aandachtsmechanisme constante inferentiememory-complexiteit ($O(1)$), ongeacht de sequentielengte, in tegenstelling tot standaard aandacht die een $O(T)$ KV-cache vereist.

Belangrijkste Bijdragen

1. Theoretische Herformulering: Het artikel introduceert het Nested Learning-perspectief in FL, waarbij het probleem wordt geformaliseerd als het gezamenlijk trainen van optimalisatieregels in plaats van statische representaties. Dit biedt een theoretische basis voor het aanpakken van Non-IID-problemen via aanpassing op het moment van testen.
2. Het FedNL-raamwerk: Een praktisch algoritme dat Titans-gebaseerde lineaire aandacht integreert in de FL-pijplijn. Door de Delta Rule te behandelen als een online gradiëntafstijgingsstap, maakt het Zero-Shot Test-Time Adaptation mogelijk op onbekende domeinen zonder zware computereisen.
3. Empirische Superioriteit: Experimenten tonen aan dat FedNL concurrerende prestaties behaalt in redeneren met korte context en superieure winsten in ophalen met lange context en streamingdiagnostiek vergeleken met sterke federale baselines.

Experimentele Resultaten

De auteurs evalueerden FedNL in twee federale settings: een Non-IID MMLU-verdeling met vijf clients en een Non-IID Needle In A Haystack (NIAH)-benchmark voor lange context met zeven clients.

• Non-IID MMLU (Redeneren met korte context):

  • Op de Qwen2.5-1.5B-backbone behaalde FedNL de hoogste gemiddelde nauwkeurigheid (58,88%), net iets beter dan de beste baseline, FedSSI (58,70%).
  • Er werden significante winsten waargenomen in STEM (+2,0 pp) en Wiskunde/Informatica (+3,8 pp) domeinen.
  • Op de Llama-3.2-1B-backbone bereikte FedNL 42,64% versus 42,10% voor FedSSI.
  • Cruciaal minimaliseerde FedNL de "aggregatie-daling" (prestatieverlies bij de overstap van lokale naar federale modellen), waarbij deze dicht bij nul bleef over domeinen, terwijl statische baselines significante degradatie vertoonden.

• Ophalen met lange context (NIAH):

  • FedNL behaalde de hoogste gemiddelde nauwkeurigheid (29,7%) over zeven ophaalsjablonen, beter dan FedALA (28,6%).
  • Winsten waren het meest uitgesproken in MK-NIAH, MV-NIAH en UUID-code taken.
  • Streaming Cross-Entropy (CE) Diagnostiek: In een 16K-token streaming-probe nam de genormaliseerde CE van FedAvg toe met 3,1% naarmate de context zich ontvouwde (accumulatie van onzekerheid), terwijl FedNL afnam met 2,1%, wat aangeeft dat de recurrente geheugentoestand succesvol nuttige context over lange streams absorbeerde.
  • Geheugenefficiëntie: FedNL behield constante inferentiememory, waardoor Out-Of-Memory-fouten bij 16K tokens werden vermeden die statische-aandacht baselines troffen.

• Communicatie-efficiëntie:

  • FedNL verlaagde de communicatie per ronde met ongeveer 350x vergeleken met FedAvg (0,26 MB versus 22,5 MB voor de trainbare parameters in de geteste setup).

Betekenis en Claims

Het artikel stelt dat FedNL een nieuwe richting voor Federated Learning vestigt, waarbij modellen evolueren van een repository van statische kennis naar een paradigma van continu, contextbewust leren.

• Paradigmaverschuiving: Door algemene linguïstische capaciteiten (bevroren backbone) te ontkoppelen van regels voor geheugenopbouw (trainbare adapters), stelt FedNL diverse clients in staat om "leercapaciteiten" te delen in plaats van conflicterende "kennisstaten".
• Robuustheid tegen Heterogeniteit: De globale meta-parameters hoeven niet de conflicten tussen verschillende domeinen te coderen (bijvoorbeeld Code versus Medisch). In plaats daarvan coderen ze een universele regel: "Als de context X betreft, update het geheugen om Y op te slaan." Dit verlicht het Non-IID-conflict in de parameterruimte.
• Zero-Shot Adaptatie: Het raamwerk stelt modellen in staat om zich tijdens de eerste paar tokens van inferentie aan te passen aan specifieke lokale distributies, puur door uitgevoerde geleerde geheugenupdate-regels, waarmee het "cold-start"-probleem wordt opgelost dat vaak voorkomt bij Test-Time Training (TTT)-methodes.

De auteurs merken op dat hoewel de resultaten consistent zijn op de schaal van 1B–1,5B, het uitbreiden van FedNL naar grotere foundation-modellen een noodzakelijke volgende stap blijft om de generaliteit van geheugenregel-aggregatie te valideren.