Résumé Technique : Apprentissage Emboîté Fédéré (FedNL)

Énoncé du Problème

L'apprentissage fédéré (FL) pour les modèles de langage à grande échelle (LLM) fait face à deux défis persistants et imbriqués : l'hétérogénéité des données (Non-IID) et les distributions à longue traîne. Dans des scénarios réalistes, les distributions de données des clients sont fortement biaisées (par exemple, données médicales par rapport à des données de codage), et les connaissances critiques résident souvent dans la longue traîne, ce qui les fait facilement passer au second plan par rapport aux classes principales lors de l'agrégation globale.

Les approches existantes se concentrent principalement sur la régularisation ou la personnalisation des poids statiques des modèles (par exemple, FedProx, FedSSI). Bien qu'efficaces pour atténuer l'oubli catastrophique, ces méthodes traitent le modèle global comme un conteneur de connaissances statiques. Par conséquent, lorsqu'il est déployé sur un client disposant de données non vues et hautement hétérogènes, le modèle manque de plasticité au moment du test — c'est-à-dire la capacité de s'adapter au contexte actuel sans mises à jour par gradient. Cela entraîne des performances sous-optimales sur des tâches spécifiques à un domaine où les nuances dépendent du contexte.

Méthodologie : Apprentissage Emboîté Fédéré (FedNL)

L'article propose l'Apprentissage Emboîté Fédéré (FedNL), un cadre qui reformule le FL en passant de l'agrégation de connaissances statiques à l'apprentissage collaboratif de règles d'optimisation. S'inspirant de la théorie de l'Apprentissage Emboîté (NL), FedNL traite l'inférence non pas comme un processus passif de récupération, mais comme un processus d'entraînement de « boucle interne » haute fréquence où le modèle compresse activement le contexte en un état de mémoire transitoire.

Architecture Principale

FedNL reformule le FL comme un système d'optimisation emboîté à trois niveaux :

1. Niveau 0 (Boucle Externe - Serveur) : Généralisation Collaborative. Le serveur agrège des méta-règles (paramètres $\theta$) plutôt que du contenu de mémoire. Ces règles gouvernent la construction et la mise à jour de la mémoire.
2. Niveau 1 (Boucle Intermédiaire - Client) : Apprentissage des Règles. Pendant l'entraînement local, les clients optimisent $\theta$ (par exemple, adaptateurs LoRA, mécanismes de commutation) pour garantir que la boucle interne produit des états de mémoire utiles pour les tâches en aval. Cela implique une rétropropagation à travers la trajectoire de la mémoire.
3. Niveau 2 (Boucle Interne - Client) : Adaptation au Moment du Test. Pendant l'inférence, les clients utilisent les règles globales $\theta$ pour mettre à jour en ligne un état de mémoire privé et conscient du contexte $S_t$ via une Règle Delta. Cet état est transitoire et ne quitte jamais l'appareil.

Mise en Œuvre Technique

• Attention Linéaire Basée sur Titans : FedNL exploite l'architecture Titans, qui utilise un mécanisme d'attention linéarisée. L'état de mémoire $S_t$ est mis à jour via une Règle Delta :
  $$S_t = S_{t-1} + \eta(v_t - S_{t-1}k_t)k_t^\top$$
  Cette mise à jour est mathématiquement équivalente à une seule étape de descente de gradient sur un objectif de mémoire associative, traitant efficacement l'inférence comme un processus d'optimisation en ligne.
• Découplage des Connaissances et des Règles : Le modèle se compose d'un socle figé (frozen backbone) fournissant des capacités linguistiques générales et de méta-paramètres entraînables ($\theta$). Les méta-paramètres incluent des projections de faible rang (LoRA) pour la génération de requêtes/clés/valeurs et un mécanisme de commutation de mémoire (Memory Gating, $\alpha$) pour équilibrer l'attention statique et dynamique.
• Efficacité : En n'agrégeant que le petit ensemble de méta-règles (adaptateurs) plutôt que les poids complets du modèle ou les états de mémoire, FedNL réduit considérablement la surcharge de communication. De plus, le mécanisme d'attention linéaire maintient une complexité mémoire d'inférence constante ($O(1)$) quelle que soit la longueur de la séquence, contrairement à l'attention standard qui nécessite un cache KV de $O(T)$.

Contributions Clés

1. Reformulation Théorique : L'article introduit la perspective de l'Apprentissage Emboîté dans le FL, formalisant le problème comme l'entraînement collaboratif de règles d'optimisation plutôt que de représentations statiques. Cela fournit une base théorique pour résoudre les problèmes Non-IID via l'adaptation au moment du test.
2. Le Cadre FedNL : Un algorithme pratique intégrant l'attention linéaire basée sur Titans dans le pipeline FL. En traitant la Règle Delta comme une étape de descente de gradient en ligne, il permet une Adaptation au Moment du Test Zero-Shot sur des domaines non vus sans effort computationnel lourd.
3. Supériorité Empirique : Les expériences démontrent que FedNL atteint des performances compétitives en raisonnement à court contexte et des gains supérieurs en récupération à long contexte et en diagnostics en flux par rapport à de solides références fédérées.

Résultats Expérimentaux

Les auteurs ont évalué FedNL sur deux configurations fédérées : une division Non-IID MMLU à cinq clients et un benchmark de contexte long Non-IID Needle In A Haystack (NIAH) à sept clients.

• MMLU Non-IID (Raisonnement à Court Contexte) :

  • Sur le socle Qwen2.5-1.5B, FedNL a obtenu la précision moyenne la plus élevée (58,88 %), surpassant légèrement la meilleure référence, FedSSI (58,70 %).
  • Des gains significatifs ont été observés dans les domaines STEM (+2,0 points de pourcentage) et Math/Informatique (+3,8 points de pourcentage).
  • Sur le socle Llama-3.2-1B, FedNL a atteint 42,64 % contre 42,10 % pour FedSSI.
  • Crucialement, FedNL a minimisé la « chute d'agrégation » (perte de performance lors du passage des modèles locaux aux modèles fédérés), la maintenant proche de zéro à travers les domaines, tandis que les références statiques montraient une dégradation significative.

• Récupération à Long Contexte (NIAH) :

  • FedNL a obtenu la précision moyenne la plus élevée (29,7 %) sur sept modèles de récupération, surpassant FedALA (28,6 %).
  • Les gains étaient les plus marqués dans les tâches MK-NIAH, MV-NIAH et UUID code.
  • Diagnostics de l'Entropie Croisée (CE) en Flux : Dans une sonde en flux de 16K tokens, l'entropie croisée normalisée de FedAvg a augmenté de 3,1 % à mesure que le contexte se déroulait (accumulation d'incertitude), alors que celle de FedNL a diminué de 2,1 %, indiquant que l'état de mémoire récurrent a absorbé avec succès le contexte utile sur de longs flux.
  • Efficacité Mémoire : FedNL a maintenu une mémoire d'inférence constante, évitant les erreurs de mémoire insuffisante (Out-Of-Memory) à 16K tokens qui affectaient les références à attention statique.

• Efficacité de Communication :

  • FedNL a réduit la communication par tour d'environ 350 fois par rapport à FedAvg (0,26 Mo contre 22,5 Mo pour les paramètres entraînables dans la configuration testée).

Importance et Revendications

L'article revendique que FedNL établit une nouvelle direction pour l'apprentissage fédéré où les modèles évoluent d'un répertoire de connaissances statiques vers un paradigme d'apprentissage continu et conscient du contexte.

• Changement de Paradigme : En découplant les capacités linguistiques générales (socle figé) des règles de construction de mémoire (adaptateurs entraînables), FedNL permet à des clients divers de partager des « capacités d'apprentissage » plutôt que des « états de connaissances » conflictuels.
• Robustesse à l'Hétérogénéité : Les méta-paramètres globaux n'ont pas besoin de coder le conflit entre des domaines disparates (par exemple, Code vs Médical). Au lieu de cela, ils codent une règle universelle : « Si le contexte implique X, mettez à jour la mémoire pour stocker Y. » Cela atténue le conflit Non-IID dans l'espace des paramètres.
• Adaptation Zero-Shot : Le cadre permet aux modèles de s'adapter à des distributions locales spécifiques au cours des premiers tokens d'inférence, simplement en exécutant des règles apprises de mise à jour de la mémoire, résolvant ainsi le problème du « démarrage à froid » souvent rencontré par les méthodes d'entraînement au moment du test (TTT).

Les auteurs notent que si les résultats sont cohérents à l'échelle 1B–1,5B, l'extension de FedNL à des modèles de base plus grands reste une prochaine étape nécessaire pour valider la généralité de l'agrégation de règles de mémoire.