Résumé technique : PreLort – LoRA à préfixe imbriqué pour le fine-tuning fédéré sous hétérogénéité de rang

1. Énoncé du problème

L'apprentissage fédéré (FL) permet l'entraînement collaboratif de modèles de langage de grande taille (LLM) sans partage de données brutes, mais il fait face à des défis importants lors de l'adaptation de modèles de fondation via des méthodes de fine-tuning efficace en paramètres (PEFT) comme la Low-Rank Adaptation (LoRA). Un obstacle majeur est l'hétérogénéité de rang : dans des contextes réels de type cross-device, les clients possèdent des contraintes de calcul et de mémoire variables, ce qui les amène à initialiser des adaptateurs LoRA avec des rangs ($r$) différents.

Les solutions existantes pour le LoRA fédéré hétérogène se concentrent principalement sur la phase d'agrégation, employant souvent un remplissage par zéro (zero-padding) pour aligner les adaptateurs de rangs différents avant la moyenne. Cependant, les auteurs identifient une limitation fondamentale de ces approches : le désalignement de représentation. Même lorsqu'une agrégation est mathématiquement définie, la combinaison de modèles entraînés indépendamment avec des rangs différents entraîne des incohérences. Spécifiquement, les clients de rang inférieur ne bénéficient pas des mises à jour des clients de rang supérieur car l'information n'est pas distribuée de manière hiérarchique à travers les dimensions de rang. Le remplissage par zéro dilue également les contributions des clients de rang élevé en introduisant des segments de valeur nulle, ce qui entraîne une utilisation sous-optimale des représentations de bas rang partagées.

2. Méthodologie : PreLort

Les auteurs proposent PreLort (Prefix-Nested LoRA), un cadre qui traite l'hétérogénéité de rang en structurant conjointement la dynamique d'entraînement et les règles d'agrégation. La philosophie centrale consiste à organiser les dimensions de l'adaptateur en une hiérarchie de préfixes, garantissant que les dimensions de rang inférieur encodent des informations pertinentes pour la tâche partagées entre tous les clients, tandis que les dimensions de rang supérieur capturent une capacité incrémentale.

2.1 Structure de bas rang imbriquée

PreLort impose une organisation structurée par préfixe sur les dimensions de l'adaptateur. Pour un client de rang $r_k$, seules les $r_k$ premières dimensions sont actives. Crucialement, pendant l'entraînement local, l'adaptateur n'est pas optimisé uniquement à son rang complet. Au lieu de cela, il est optimisé sous plusieurs configurations de troncature (sous-rangs de préfixe).

• Mécanisme : Un client de rang $r_k$ est entraîné pour être performant non seulement à $r_k$, mais aussi à tous les niveaux de rang inférieurs présents dans la fédération ($r \in \{r \mid r \le r_k\}$).
• Objectif : Cela force le modèle à concentrer les signaux pertinents de la tâche dans les dimensions de tête (bas rang), garantissant que ces dimensions restent informatives et utiles même lorsqu'elles sont tronquées.

2.2 Stratégie d'entraînement local

L'objectif local pour un client $k$ est défini comme la moyenne des pertes de cross-entropie calculées sur tous les sous-rangs de préfixe valides $R_k$ :
$$\mathcal{L}_k = \frac{1}{|R_k|} \sum_{r \in R_k} \mathcal{L}_{CE}(f(\cdot; \theta, r), \mathcal{D}_k)$$
En pratique, chaque étape d'entraînement implique un passage forward/backward complet en rang plein et des passages supplémentaires pour tous les sous-rangs. Cette stratégie garantit que les segments de préfixe sont explicitement optimisés à la même granularité utilisée lors de l'agrégation, encourageant une hiérarchie cohérente où les dimensions inférieures capturent l'information la plus critique.

2.3 Agrégation par segment

Pour éviter la dilution causée par le remplissage par zéro, PreLort introduit une règle d'agrégation par segment opérant dans l'espace du gradient (mise à jour).

• Processus : Au lieu de moyenner tous les clients pour chaque dimension, le serveur agrège les gradients pour chaque segment de rang $[r_{l-1}, r_l)$ uniquement sur le sous-ensemble de clients dont le rang $r_k$ est suffisant pour contribuer à ce segment (c'est-à-dire $r_k \ge r_l$).
• Formule : Pour le $l$-ième segment, le gradient agrégé est :
  $$\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}$$
• Bénéfice : Cela empêche la dilution du signal provenant des clients de rang inférieur remplis de zéros et garantit que les dimensions de rang supérieur ne sont mises à jour que par les clients possédant réellement la capacité de les apprendre, tandis que les dimensions inférieures bénéficient de la population complète.

3. Contributions clés

Le document souligne trois contributions principales :

1. Identification du désalignement de représentation : Les auteurs identifient que l'absence d'alignement structuré à travers les dimensions de rang est une limitation clé des méthodes actuelles de LoRA fédéré hétérogène, menant à des représentations incohérentes.
2. Formulation imbriquée et stratégie unifiée : Ils proposent une formulation de bas rang imbriquée qui impose une structure de préfixe commune. Cela est couplé à une stratégie d'entraînement et d'agrégation unifiée qui permet une agrégation cohérente à travers les rangs hétérogènes sans remplissage par zéro.
3. Validation empirique : La méthode est démontrée comme améliorant systématiquement la qualité et la stabilité du modèle par rapport aux méthodes précédentes de LoRA fédéré hétérogène sur des benchmarks de fine-tuning d'instructions.

4. Résultats expérimentaux

Les auteurs ont évalué PreLort sur deux modèles de base (TinyLlama-1.1B et Qwen2.5-0.5B) à travers trois jeux de données : Alpaca, Databricks-dolly-15, et 20 Newsgroups.

• Métriques de performance : PreLort surpasse systématiquement les bases de référence (ZeroPad, HetLoRA, FLoRA) en termes de Précision (MMLU, 20 Newsgroups) et de scores ROUGE-L.
  • Sur TinyLlama-1.1B, PreLort atteint 33,86 % de précision sur MMLU (Dolly) et 30,76 % sur Alpaca, surpassant les meilleures bases de référence.
  • Sur Qwen2.5-0.5B, il atteint 39,85 % de précision sur MMLU (Dolly).
• Perplexité : PreLort a obtenu la perplexité la plus basse sur Dolly et NewsGroup pour les deux modèles de base, démontrant une supériorité de la vraisemblance au niveau des tokens.
• Robustesse du rang : L'analyse a montré que PreLort maintient une performance stable à travers différents niveaux de troncature de préfixe. En revanche, les bases de référence comme "Nested Aggregation Only" ont montré une instabilité, particulièrement à des rangs élevés, confirmant que l'agrégation seule est insuffisante sans l'objectif d'entraînement imbriqué.
• Importance par dimension : La visualisation de l'importance des dimensions a révélé que PreLort concentre le signal pertinent de la tâche dans les dimensions de tête (une distribution structurée), alors que l'entraînement standard produit une distribution uniforme. Cette hiérarchisation facilite des mises à jour stables à travers les rangs hétérogènes.

5. Signification et affirmations

Le papier affirme que PreLort représente un changement de paradigme dans l'approche de l'hétérogénéité de rang : plutôt que de la traiter comme un problème d'agrégation post-hoc, elle est résolue en structurant le processus d'entraînement lui-même.

• Argument central : Les auteurs soutiennent que traiter l'hétérogénéité de rang nécessite non seulement une agrégation améliorée, mais aussi un alignement structuré des représentations à travers les rangs.
• Efficacité : La méthode permet aux clients de bas rang de bénéficier d'informations plus riches fournies par les clients de rang élevé car les dimensions de préfixe sont apprises et agrégées de manière cohérente.
• Modestie sur les limites : Les auteurs reconnaissent que la stratégie d'entraînement imbriquée introduit environ 40 % de surcharge de calcul supplémentaire en raison des multiples passages forward/backward par batch. Ils notent également que l'efficacité dépend de suffisamment d'étapes d'optimisation locale ; avec des époques locales très limitées, le modèle peut ne pas apprendre pleinement des représentations de préfixe de bas rang bien alignées.

En conclusion, PreLort démontre que l'hétérogénéité de rang dans le LoRA fédéré est mieux gérée en structurant conjointement la manière dont les adaptateurs sont entraînés (via l'optimisation de préfixe imbriquée) et la manière dont leurs mises à jour sont agrégées (via des règles par segment), résultant en des représentations de préfixe alignées et robustes qui surpassent les méthodes d'agrégation hétérogène existantes.