Résumé technique : LoRIF – Fonctions d'influence à faible rang pour l'attribution de données d'entraînement évolutives

Énoncé du problème

L'attribution des données d'entraînement (TDA) vise à identifier quels exemples d'entraînement spécifiques ont le plus influencé la prédiction d'un modèle sur une requête donnée. Bien que les fonctions d'influence (IF) offrent une approche théoriquement fondée pour la TDA en approximant le réentraînement sans un exemple via des gradients et des matrices hessiennes inverses, elles sont restées pratiquement inaccessibles pour les modèles à l'échelle moderne (milliards de paramètres) entraînés sur des millions d'exemples.

Les variantes évolutives existantes, telles que TRAK, LoGRA et TrackStar, résolvent les goulots d'étranglement computationnels en stockant les gradients projetés par exemple et en utilisant une approximation de Hessien de Gauss-Newton amortie. Cependant, ces méthodes font face à deux goulots d'étranglement critiques en matière d'évolutivité qui créent un compromis entre la qualité de l'attribution et la faisabilité computationnelle :

1. Stockage des gradients et E/S : Le stockage des gradients projetés par exemple pour $N$ exemples d'entraînement nécessite un espace de $O(ND)$, où$D$ est la dimension de projection effective. Au moment de la requête, le chargement de ces gradients domine la latence (environ 96 % du temps de requête).
2. Mise à l'échelle de l'inverse du Hessien : La formation et le stockage de l'approximation de l'inverse du Hessien $D \times D$ nécessitent une mémoire de $O(D^2)$.

Augmenter la dimension de projection $D$ est nécessaire pour améliorer la qualité de l'attribution, mais cela exacerbe les deux goulots d'étranglement de stockage/E/S et de mémoire, rendant l'attribution de haute qualité impraticable pour les modèles de pointe.

Méthodologie : LoRIF

LoRIF (Fonctions d'influence à faible rang) résout ces goulots d'étranglement en exploitant la structure de faible rang inhérente aux gradients des réseaux de neurones. La méthode suit un pipeline en trois étapes :

1. Stockage de gradients à faible rang (Résolution du goulot d'étranglement des E/S)

Au lieu de stocker la matrice complète des gradients projetés par exemple $\tilde{G}_i \in \mathbb{R}^{d_1 \times d_2}$ (où $D = d_1 d_2$), LoRIF stocke une factorisation de rang-$c$ :
$$\tilde{G}_i \approx u_i v_i^\top$$
où $u_i \in \mathbb{R}^{d_1 \times c}$ et $v_i \in \mathbb{R}^{d_2 \times c}$.

• Mécanisme : Cette factorisation est calculée via des itérations de puissance par blocs.
• Impact : Le stockage et les E/S au moment de la requête par couche par échantillon sont réduits de $O(D)$ à $O(c(\sqrt{D}))$. Les résultats empiriques suggèrent que même avec $c=1$, l'approximation conserve une qualité d'attribution significative, et privilégier un $D$ plus grand plutôt qu'un $c$ plus élevé offre une meilleure efficacité coût-performance.

2. Approximation de courbure par SVD tronquée (Résolution du goulot d'étranglement de la mémoire)

Pour éviter de former l'inverse du Hessien dense $D \times D$ $(G^\top G + \lambda I)^{-1}$, LoRIF calcule une décomposition en valeurs singulières (SVD) tronquée de rang-$r$ de la matrice de gradients agrégée $G$ :
$$G \approx U_r \Sigma_r V_r^\top$$
En utilisant l'identité de Woodbury, le terme d'inverse du Hessien est approximé dans un sous-espace de dimension $r$ :
$$(V_r \Sigma_r^2 V_r^\top + \lambda I)^{-1} = \frac{1}{\lambda}I - \frac{1}{\lambda^2} V_r \left( \Sigma_r^{-2} + \frac{1}{\lambda}I_r \right)^{-1} V_r^\top$$

• Mécanisme : Seuls les $r$ vecteurs singuliers principaux ($V_r$) et les valeurs singulières ($\Sigma_r$) sont stockés.
• Impact : La complexité mémoire pour le terme de courbure est réduite de $O(D^2)$ à $O(Dr)$. Cela permet une dimension de projection effective $D$ significativement plus grande dans un budget mémoire fixe.

3. Calcul du score d'influence

Au moment de la requête, les scores d'influence sont calculés en projetant le gradient de la requête dans le sous-espace de dimension $r$ et en utilisant les facteurs de faible rang stockés pour les gradients d'entraînement. Le score d'influence $I(x_{tr}, x_{te})$ est approximé comme suit :
$$I(x_{tr}, x_{te}) \approx \frac{1}{\lambda} g_{te}^\top g_{tr} - \frac{1}{\lambda^2} (g'_{te})^\top \left( \Sigma_r^{-2} + \frac{1}{\lambda}I_r \right)^{-1} g'_{tr}$$
où $g'$ représente les gradients projetés sur les vecteurs singuliers $V_r$. Le premier terme est calculé efficacement en utilisant les facteurs de rang-$c$, et le second terme n'implique que $O(r)$ opérations en raison de la nature diagonale de la matrice centrale.

Contributions clés

1. Identification des goulots d'étranglement : L'article identifie que la dimension de projection $D$ est la contrainte principale dans la TDA de style IF évolutif, créant un compromis qualité-évolutivité où l'attribution précise (grand $D$) est actuellement impraticable.
2. Factorisation de rang-$c$ : La proposition de stocker les gradients projetés sous forme de facteurs de rang-$c$, réduisant le stockage par exemple par couche de $O(D)$ à $O(c\sqrt{D})$.
3. Approximation de courbure de rang-$r$ : La proposition d'utiliser la SVD tronquée et l'identité de Woodbury pour approximer le terme de courbure, réduisant la mémoire par couche de $O(D^2)$ à $O(Dr)$.
4. Démonstration d'évolutivité : La première démonstration d'une attribution de données efficace sur des modèles allant jusqu'à 70 milliards de paramètres (Apertus-70B) avec des millions d'exemples d'entraînement, tout en maintenant un coût et une latence pratiques.

Résultats expérimentaux

LoRIF a été évalué sur des modèles allant de 124M à 70B de paramètres (GPT2-small, OLMo-3-7B, Apertus-70B) avec des jeux de données contenant jusqu'à 3,8 millions d'exemples.

• Stockage et latence : Par rapport à LoGRA, LoRIF atteint un stockage de gradients 2,3–20 fois inférieur et une latence de requête 1,3–20 fois inférieure tout en égalant ou dépassant la qualité d'attribution.
  • Sur GPT2-small, LoRIF a réduit le stockage jusqu'à 9 fois et la latence de 5,6 fois dans des régimes de stockage moyen.
  • Sur OLMo-3-7B, LoRIF a obtenu des scores de tail-patch presque identiques à LoGRA avec un stockage 20,3 fois inférieur et une latence 22,0 fois inférieure.
  • Sur Apertus-70B, LoRIF a amélioré les scores de tail-patch de 1,5 à 2,2 fois par rapport à LoGRA avec un stockage et une latence inférieurs.
• Qualité d'attribution : Mesurée via le score de modélisation de données linéaire (LDS) et le score de tail-patch, LoRIF a constamment amélioré la frontière de Pareto. En permettant des dimensions de projection effectives $D$ plus grandes (qui étaient auparavant irréalisables en raison des contraintes de mémoire), LoRIF a atteint une qualité supérieure aux bases contraintes à un $D$ plus petit.
• Évaluation qualitative : Une évaluation par un LLM en tant que juge sur les exemples récupérés en top-1 a montré que LoRIF récupère des exemples avec une pertinence sémantique plus élevée (score moyen 3,04 contre 2,49 sur OLMo-3-7B) et moins de récupérations complètement non pertinentes par rapport à LoGRA.

Importance et revendications

L'article revendique que LoRIF rend l'attribution de données d'entraînement basée sur les gradients pratique à l'échelle de pointe. En exploitant la structure de faible rang des gradients, LoRIF résout le compromis qualité-évolutivité qui a entravé l'application des fonctions d'influence aux grands modèles de langage.

Les auteurs soulignent que LoRIF ne repose pas sur de nouveaux primitifs de faible rang, mais démontre plutôt comment les structures de faible rang existantes (observées à la fois dans les gradients par exemple et les gradients agrégés) peuvent être systématiquement exploitées dans les méthodes d'influence par gradient projeté. Les résultats suggèrent que la TDA basée sur les gradients peut désormais être utilisée pour l'audit, le débogage et les vérifications de sécurité des LLM à grande échelle lorsque l'index d'attribution est réutilisé sur de nombreuses requêtes, sans nécessiter les coûts de stockage ou d'E/S prohibitifs des méthodes précédentes.

L'article reconnaît des limites, notant que la construction de l'index d'attribution nécessite toujours un passage de gradient sur le corpus (coût de prétraitement) et que les expériences à grande échelle actuelles se concentrent sur les données de fine-tuning supervisé plutôt que sur les corpus d'entraînement complet. Cependant, ce travail établit une nouvelle référence pour l'attribution évolutive, permettant l'analyse de modèles précédemment considérés hors de portée pour de telles méthodes.