Compositional Steering of Large Language Models with Steering Tokens

1. Problématique

Le déploiement des Grands Modèles de Langage (LLM) dans des applications réelles nécessite un contrôle précis de leur sortie pour satisfaire simultanément plusieurs critères (comportements). Bien que le "steering" (pilotage) d'un LLM vers un comportement unique soit bien documenté, le pilotage compositionnel (contrôler plusieurs comportements à la fois, par exemple : "répondre en français" ET "en 10 à 50 mots" ET "avec un ton formel") reste un défi majeur.

Les approches existantes présentent des limites :

• Le fine-tuning est coûteux en calcul et risque d'entraîner un oubli des connaissances pré-entraînées ou des interférences négatives. De plus, pour $N$ comportements, il faudrait théoriquement $2^N$ procédures de fine-tuning pour couvrir toutes les combinaisons.
• Les instructions en langage naturel (Prompting) sont fragiles : des prompts sémantiquement équivalents peuvent produire des comportements incohérents.
• Le pilotage par activation (Activation Steering) et les tokens Gist fonctionnent bien pour des comportements individuels mais échouent souvent à généraliser efficacement lorsqu'ils sont combinés (composition), car ils ne capturent pas bien le concept de "composition" lui-même.

2. Méthodologie : Les Tokens de Pilotage Compositionnels

Les auteurs proposent une nouvelle approche basée sur des tokens de pilotage (steering tokens) qui opèrent dans l'espace des tokens d'entrée (embeddings) plutôt que dans l'espace des activations internes du modèle. Le modèle de base reste entièrement figé (frozen).

La méthode se déroule en deux étapes principales via une auto-distillation :

A. Apprentissage des Tokens de Comportement Individuels

Pour chaque comportement $b$ (exprimé par une instruction naturelle), un token trainable $\langle b \rangle$ (un vecteur d'embedding) est appris.

• Mécanisme : Un modèle "enseignant" (frozen) reçoit le prompt avec l'instruction textuelle complète. Un modèle "élève" (frozen) reçoit le même prompt mais avec l'instruction remplacée par le token $\langle b \rangle$.
• Objectif : Minimiser la divergence KL entre les distributions de sortie des deux modèles. Le token $\langle b \rangle$ apprend à reproduire le comportement induit par l'instruction textuelle.

B. Apprentissage du Token de Composition $\langle and \rangle$

C'est l'innovation centrale. Pour combiner deux comportements $b_i$ et $b_j$, les auteurs introduisent un token dédié $\langle and \rangle$.

• Architecture d'entrée : Pour un prompt $x$ avec deux contraintes, l'entrée devient : $[x, \langle b_i \rangle, \langle and \rangle, \langle b_j \rangle]$.
• Entraînement : Le token $\langle and \rangle$ est entraîné sur des paires de comportements. L'enseignant reçoit les deux instructions textuelles, tandis que l'élève reçoit les deux tokens de comportement séparés par le token $\langle and \rangle$.
• Contrainte d'orthogonalité : Pour éviter que le token $\langle and \rangle$ ne se confonde avec les tokens de comportement existants, une régularisation d'orthogonalité est appliquée, forçant $\langle and \rangle$ à être orthogonal aux embeddings des comportements appris.
• Généralisation Zero-Shot : Le token $\langle and \rangle$ apprend la fonction abstraite de "composition" et non pas simplement une fusion spécifique de deux comportements. Cela lui permet de généraliser à des combinaisons jamais vues (nouveaux comportements, ou un nombre différent de comportements, comme 3 ou 4).

3. Contributions Clés

1. Espace d'entrée vs Espace d'activation : Contrairement aux méthodes de pilotage par activation, cette approche opère dans l'espace des tokens d'entrée, permettant une composition plus efficace et zero-shot.
2. Opérateur de composition explicite : L'introduction d'un token $\langle and \rangle$ appris spécifiquement pour gérer la notion de combinaison, ce qui permet une généralisation robuste.
3. Évaluation rigoureuse à grande échelle : Contrairement à des travaux antérieurs basés sur des preuves anecdotiques, les auteurs évaluent leur méthode sur des contraintes vérifiables automatiquement (langue, longueur, format, structure) sur plus d'un million d'évaluations.
4. Approche hybride : Ils démontrent que combiner les tokens de pilotage avec des instructions textuelles naturelles offre les meilleures performances.

4. Résultats Expérimentaux

Les expériences ont été menées sur plusieurs architectures (Qwen, Llama, SmolLM, OLMo, Gemma) et tailles (de 3B à 14B paramètres).

• Supériorité sur la composition : Les tokens de pilotage surpassent les instructions textuelles pures, en particulier sur les compositions non vues (zero-shot). Par exemple, pour des combinaisons de 3 comportements non vues, les tokens atteignent 59,5 % de précision contre 54,0 % pour les instructions (sur Qwen-8B).
• Généralisation aux combinaisons inédites : Le modèle généralise bien à :
  • Des combinaisons de comportements vus mais jamais combinés ensemble.
  • Des comportements totalement nouveaux (jamais vus durant l'entraînement du token $\langle and \rangle$).
  • Des nombres de comportements différents (passer de 2 à 3 ou 4 comportements sans réentraînement).
• Robustesse à la taille du modèle : Les performances s'améliorent avec la taille du modèle (jusqu'à 14B). L'entraînement uniquement sur des paires (2 comportements) suffit pour généraliser à 3 ou 4 comportements sur les grands modèles.
• Approche Hybride : La combinaison des tokens appris et des instructions textuelles (Hybrid) offre le meilleur compromis précision/robustesse, éliminant les modes d'échec spécifiques à chaque méthode.
• Comparaison avec l'état de l'art : La méthode surpasse nettement le pilotage par activation (CAA, LM-Steer) et le fine-tuning avec fusion de LoRA (LoRA DARE), qui échouent souvent à composer correctement plusieurs contraintes.

5. Signification et Impact

Ce travail propose un mécanisme efficace et économe en calcul pour le contrôle compositionnel des LLMs.

• Efficacité : Il ne modifie pas les paramètres internes du modèle (seuls quelques vecteurs d'embedding sont appris), évitant ainsi le coût du fine-tuning et les risques d'oubli catastrophique.
• Modularité : La capacité à composer des comportements de manière zero-shot ouvre la voie à des systèmes d'IA plus flexibles capables de s'adapter dynamiquement à des contraintes multiples et changeantes sans réentraînement.
• Complémentarité : La démonstration que les tokens appris et le langage naturel sont complémentaires suggère que les futures interfaces de contrôle des LLMs devraient probablement intégrer les deux paradigmes.

En résumé, cette recherche établit les tokens de pilotage comme une solution robuste et généralisable pour le contrôle multi-comportements des LLMs, résolvant un problème fondamental de l'alignement et du contrôle des modèles génératifs.