Résumé Technique : Le pré-entraînement modulaire permet le contrôle d'accès

Problématique

Les développeurs d'IA font face à un « dilemme du double usage » : les capacités qui permettent des applications bénéfiques (ex. : guérir des maladies, renforcer la cybersécurité) chevauchent souvent celles qui permettent des dommages (ex. : synthétiser des armes biologiques, exécuter des cyberattaques). Les mécanismes de contrôle d'accès actuels sont insuffisants. Le « standard d'or » pour le contrôle d'accès consisterait à servir des modèles distincts, entraînés séparément avec des profils de capacités différents à différents utilisateurs (ex. : un modèle doté de connaissances en virologie pour les chercheurs, mais sans celles-ci pour le grand public). Cependant, l'entraînement et le déploiement de $N$ modèles distincts pour $N$ profils de capacités est prohibitif aux échelles de pointe.

Les alternatives existantes ne répondent pas aux exigences de robustesse :

• L'entraînement au refus, le filtrage de sortie et les classificateurs restreignent les sorties mais laissent la capacité sous-jacente intacte, ce qui les rend vulnérables à l'élicitation adversaire.
• Le désapprentissage post-hoc (ex. : MaxEnt) tente de supprimer les capacités d'un modèle entraîné mais est fragile ; les capacités peuvent souvent être rapidement récupérées via un ajustement fin (fine-tuning) sur quelques exemples.
• Le branchement de modèles (entraîner un modèle de base et des variantes de fine-tuning séparées) réduit les coûts de pré-entraînement mais nécessite toujours plusieurs déploiements, ce qui est inefficace pour des schémas d'utilisation épars.

Méthodologie : Modules Auxiliaires par Routage de Gradient (GRAM)

Le papier propose GRAM, une méthode de pré-entraînement qui permet à un seul modèle de supporter plusieurs profils de capacités en induisant une spécialisation par une architecture modulaire et un routage de gradient.

Architecture

GRAM augmente un Transformer dense standard en ajoutant de plus petits modules MLP auxiliaires à chaque couche, aux côtés du MLP « central » principal.

• Module Central (Core) : Un grand MLP partagé pour toutes les données, capturant les connaissances générales.
• Modules Auxiliaires : De plus petits MLP, chacun correspondant à une capacité auxiliaire spécifique (ex. : virologie, cybersécurité).
• Nombre de paramètres : Le nombre total de paramètres est comparable à celui d'un modèle dense standard, le module central correspondant à la taille de référence et les modules auxiliaires ne représentant qu'une petite fraction des paramètres.

Dynamique d'entraînement (Routage de Gradient)

Contrairement au mélange d'experts (MoE) qui utilise des routeurs appris pour le routage par jeton (token), GRAM utilise un routage de gradient dépendant des données basé sur les étiquettes de l'ensemble de données :

1. Passe Avant (Forward Pass) :
   • Données Centrales : Le MLP central est toujours actif. Avec une probabilité $p_{cr}$, un module auxiliaire aléatoire est également activé pour garantir que le cœur reste robuste à la présence de modules.
   • Données Auxiliaires : Le MLP central et le module auxiliaire spécifique correspondant à ces données sont tous deux actifs.
2. Passe Arrière (Backward Pass - Mise à jour des gradients) :
   • Paramètres Centraux : Mis à jour avec une probabilité $p_{as}$ lors du traitement de données auxiliaires. Cela contrôle le degré d'interférence entre les capacités.
   • Paramètres Auxiliaires : Le module auxiliaire spécifique pour le lot (batch) actuel est toujours mis à jour. Les autres modules auxiliaires ne reçoivent aucune mise à jour.
   • Optimiseurs Séparés : Chaque sous-ensemble de paramètres (cœur et chaque module auxiliaire) est optimisé à l'aide d'un optimiseur AdamW distinct.

Inférence et Contrôle d'Accès

Au moment de l'inférence, un profil de capacité spécifique est servi en activant le MLP central et le sous-ensemble de modules auxiliaires correspondant aux capacités autorisées. L'ablation (désactivation) d'un module auxiliaire supprime la capacité associée du réseau. Le papier postule que cela s'approche d'un modèle entraîné sur des données filtrées, car le module abléré n'a jamais appris la capacité spécifique de manière isolée du cœur.

Contributions Clés et Expérimentations

1. Approximation du Filtrage de Données

Les auteurs démontrent qu'un modèle entraîné avec GRAM peut approximer les performances de plusieurs modèles entraînés via le filtrage de données (où chaque modèle est entraîné sur un sous-ensemble de données).

• Évaluation Synthétique (Histoires Simples) : Sur un ensemble de données d'histoires pour enfants partitionnées en thèmes centraux et auxiliaires, GRAM a égalé les modèles filtrés par données, tant dans la rétention des capacités souhaitées que dans la suppression des capacités indésirables.
• Évaluation de Double Usage Réel : Entraîné sur un modèle de 800 millions de paramètres avec des données centrales (texte web, code, articles généraux) et quatre domaines auxiliaires (virologie, cybersécurité, physique nucléaire, code Lisp spécialisé), GRAM a obtenu :
  • Rétention : La performance sur les capacités retenues était presque identique à celle de la référence (all-data baseline).
  • Suppression : La performance sur les capacités supprimées a chuté de manière significative, rejoignant la performance d'« oubli » des modèles filtrés par données.
  • Robustesse : La suppression des capacités par GRAM était plus robuste au fine-tuning malveillant (élicitation d'oubli) que les méthodes de désapprentissage post-hoc (MaxEnt).

2. Scalabilité et Composabilité

• Lois d'Échelle (Scaling Laws) : Une analyse d'échelle optimale selon Chinchilla, de 50M à 5B de paramètres, a montré que GRAM suit de près le profil de capacité du filtrage de données. À mesure que la taille du modèle augmentait, l'écart entre les modèles filtrés et les modèles de données complètes s'élargissait pour les capacités supprimées, et GRAM maintenait cette tendance.
• Composabilité : GRAM supporte des sous-ensemts arbitraires de capacités à partir d'un seul entraînement. Contra-rément au branchement par LoRA (FT-LoRA), qui se dégrade lorsque plusieurs adaptateurs sont additionnés, GRAM a maintenu ses performances lorsque plusieurs modules auxiliaires étaient actifs simultanément, permettant de générer $2^{N-1}$ profils de capacités à partir d'un seul modèle.
• Étiquetage Partiel : Dans un scénario où seulement 50 % des données étaient étiquetées, GRAM a nettement surpassé le filtrage de données et le FT-LoRA en termes de suppression de capacités. Les auteurs attribuent cela à un « effet d'absorption », où les mises à jour des données non étiquetées se concentrent dans les modules spécialisés plutôt que dans le cœur partagé.

3. Efficacité des Coûts

• Coût d'Entraînement : GRAM ne nécessite qu'un seul cycle d'entraînement, réduisant le coût de calcul de l'entraînement d'un facteur 5 par rapport à l'entraînement de 5 modèles séparés filtrés par données.
• Coût de Déploiement : Contra-rément au branchement de modèles ou au fine-tuning complet, GRAM ne nécessite pas le déploiement de multiples copies du modèle ; il sert un seul modèle avec une activation dynamique des modules.

Signification et Revendications

Le papier affirme que GRAM offre une voie pratique vers un contrôle d'accès robuste pour les capacités d'IA à double usage, sans les coûts prohibitifs de l'entraînement de plusieurs modèles.

• Sécurité : En supprimant physiquement les paramètres responsables d'une capacité (via l'ablation) plutôt qu'en supprimant simplement les sorties, GRAM offre des garanties plus fortes contre les usages malveillants et le jailbreaking par rapport au pilotage comportemental ou au désapprentissage post-hoc.
• Viabilité Économique : La méthode découple le coût d'entraînement du nombre de profils de capacités supportés, rendant le contrôle d'accès différencié réalisable aux échelles de pointe.
• Alignement avec le Filtrage de Données : Les résultats établissent que le filtrage de données lui-même passe favorablement à l'échelle (les modèles plus grands sont meilleurs pour oublier des données spécifiques tout en retenant les connaissances générales) et que GRAM suit avec succès ce comportement.

Les auteurs restent modestes quant aux limites, notant que :

1. Enchevêtrement (Entanglement) : Si les capacités sont fortement enchevêtrées, une séparation nette peut être impossible, quelle que soit la méthode.
2. Échelle de Pointe : Bien que les tendances se confirment jusqu'à 5 milliards de paramètres, l'efficacité aux échelles de pointe (billions de paramètres) n'est pas vérifiée.
3. Implémentation : L'intégration de GRAM dans les pipelines de production introduit une complexité que les développeurs pourraient vouloir éviter.
4. Post-Entraînement : L'interaction entre la séparation par GRAM lors du pré-entraînement et l'ajustement d'instructions ultérieur ou le RLHF est inconnue.

En conclusion, le papier positionne GRAM non pas comme une solution parfaite, mais comme une étape significative vers le déploiement sûr et flexible de systèmes d'IA hautement capables, en permettant aux fournisseurs de servir des profils de capacités autorisés à partir d'un seul modèle entraîné efficacement.