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

1. Problématique : Le Trilemme du Fine-Tuning Fédéré

Le fine-tuning des grands modèles de langage (LLM) dans un cadre fédéré (Federated Learning - FL) se heurte à un trilemme difficile à résoudre simultanément :

1. Protection de la Propriété Intellectuelle (IP) : Les propriétaires de LLM (serveurs) ne veulent pas révéler leurs modèles propriétaires complets aux clients.
2. Confidentialité des Données Clients : Les données des clients doivent rester locales et privées.
3. Performance sur Données Hétérogènes : Les méthodes existantes souffrent d'une perte de performance significative due à l'hétérogénéité des données et aux limitations de capacité de représentation.

Limites des approches actuelles (Offsite-Tuning - OT) :
Les méthodes comme Offsite-Tuning (OT) tentent de résoudre ce problème en envoyant aux clients un émulateur léger et des adaptateurs (adapters) entraînables, tout en gardant le modèle de base figé. Cependant, l'analyse des auteurs révèle que ces adaptateurs légers ont une capacité de représentation intrinsèquement faible, créant un goulot d'étranglement de performance qui les empêche d'approcher les résultats du fine-tuning centralisé.

2. Méthodologie : Le Framework FedProxy

Pour combler l'écart de performance tout en respectant les contraintes de confidentialité et de ressources, les auteurs proposent FedProxy, un cadre d'adaptation fédéré en trois étapes. L'idée centrale est de remplacer les faibles adaptateurs par un Proxy Small Language Model (SLM) puissant, compressé à partir du LLM propriétaire, servant de véhicule de haute fidélité pour l'apprentissage collaboratif.

Étape 1 : Représentation Efficace via Compression (Server-Guided Compression)

• Objectif : Créer un proxy SLM ($f_\phi$) compact mais puissant à partir du LLM propriétaire ($f_\theta$).
• Mécanisme : Le serveur utilise un ensemble de données publiques (ex: Alpaca) pour compresser le LLM via un élagage structuré (suppression de blocs entiers de transformeurs).
• Critère d'élagage : Une métrique appelée Block Influence (BI) est utilisée pour quantifier l'importance de chaque bloc. Les blocs ayant le plus fort impact sur les représentations d'état caché sont conservés.
• Résultat : Un SLM unique et unifié, beaucoup plus petit que le LLM original ($|\phi| \ll |\theta|$), qui conserve la structure et les connaissances de base du modèle, prêt à être envoyé aux clients.

Étape 2 : Optimisation Fédérée Robuste (Interference-Mitigating Aggregation)

Cette étape gère l'hétérogénéité des données clients et les conflits de paramètres lors de l'agrégation. Elle se divise en trois sous-processus :

1. Analyse Serveur : Le serveur calcule les vecteurs de tâche de chaque client et évalue leur similarité pour déterminer :

   • Un coefficient d'hétérogénéité ($h_k$) : mesure l'écart d'un client par rapport à la population.
   • Un poids d'agrégation ($w_k$) : favorise les clients en accord avec le réseau global.
   • Un score de conflit paramétrique ($C_d$) : identifie les dimensions de paramètres où les mises à jour des clients ont des signes opposés.

2. Régularisation Côté Client (PCR - Parameter-Level Conflict-Aware Regularization) :

   • Les clients intègrent les scores de conflit reçus dans leur fonction de perte locale.
   • Cela pénalise la divergence des paramètres où il y a un fort consensus (conflit élevé) tout en permettant une personnalisation sur les paramètres conflictuels.

3. Fusion Serveur (H-TIES - Heterogeneity-aware TIES-merging) :

   • Adaptation de la méthode TIES-Merging pour le contexte fédéré.
   • Sparsification adaptative : Les mises à jour des clients très hétérogènes sont rendues plus clairsemées.
   • Résolution de conflits : Seules les mises à jour conformes au signe dominant (positif ou négatif) sont conservées et moyennées, éliminant ainsi les interférences destructrices.

Étape 3 : Fusion de Connaissances sans Entraînement (Training-Free "Plug-in")

• Objectif : Transférer les connaissances apprises par le proxy SLM fédéré vers le LLM original.
• Mécanisme : Une fusion directe dans l'espace des paramètres. Comme le proxy SLM est un sous-réseau du LLM original, les poids mis à jour du proxy remplacent directement les couches correspondantes du LLM.
• Avantage : Aucune ré-entraînement coûteux ni latence d'inférence supplémentaire n'est nécessaire.

3. Contributions Clés

1. Framework FedProxy : Une solution holistique résolvant le trilemme IP/Confidentialité/Performance via un proxy SLM unifié.
2. Stratégies H-TIES et PCR : Des mécanismes novateurs pour gérer l'hétérogénéité des données et les interférences de paramètres, guidant à la fois l'entraînement local et la fusion serveur.
3. Performance Supérieure : Démonstration expérimentale que FedProxy surpasse significativement les méthodes basées sur OT (Offsite-Tuning) et atteint des performances proches du fine-tuning centralisé.

4. Résultats Expérimentaux

Les expériences ont été menées sur LLaMA2-7B et Mistral-7B avec des taux de compression de 30% et 50%.

• Tâches Homogènes (IID) : FedProxy (50% compression) atteint 90,1% de la performance du fine-tuning centralisé sur les tâches de QA (Question Answering) et 90,7% sur GLUE, surpassant largement FedOT (Offsite-Tuning fédéré) de +11,7 à +22,1 points de pourcentage.
• Tâches Hétérogènes (Non-IID) : Dans des scénarios plus réalistes avec des données très variées, FedProxy maintient une supériorité marquée, atteignant 91,3% de la performance centrale sur QA.
• Analyse Coût-Performance : Bien que FedProxy nécessite plus de calcul côté client et de communication itérative (mise à jour de LoRA sur un SLM plus grand) que FedOT, le gain de performance justifie cet investissement. Le coût initial de communication (téléchargement du proxy) reste comparable.
• Études d'ablation :
  • La suppression de la régularisation PCR ou de la fusion H-TIES entraîne une chute significative des performances.
  • L'architecture du Proxy SLM est intrinsèquement supérieure à l'architecture adaptateur/émulateur de l'OT, même si l'émulateur est entraîné.
  • Le proxy SLM initial est trop faible pour être utilisé seul (protection IP efficace), mais devient puissant après fusion.

5. Signification et Impact

FedProxy établit une nouvelle référence pour l'adaptation fédérée des LLM.

• Pratique : Il permet aux entreprises de fine-tuner des modèles propriétaires sur des données sensibles distribuées sans compromettre la propriété intellectuelle ni la qualité du modèle.
• Théorique : Il démontre que le remplacement des adaptateurs légers par un proxy SLM compressé est une voie plus efficace pour surmonter les limitations de capacité de représentation dans le FL.
• Limitations : La méthode dépend de la qualité du jeu de données public utilisé pour la compression initiale et la complexité de l'agrégation croît quadratiquement avec le nombre de clients ($O(K^2)$), ce qui pourrait nécessiter des optimisations pour des déploiements massifs.

En résumé, FedProxy offre une solution robuste et performante pour débloquer le potentiel du fine-tuning fédéré des LLMs, en équilibrant efficacement sécurité, confidentialité et qualité des résultats.