Résumé technique : Sondage causal en temps d'inférence dans les LLM (HDMI)

1. Énoncé du problème

L'article aborde les limites des méthodes actuelles de sondage causal utilisées pour comprendre comment les grands modèles de langage (LLM) exploitent les représentations internes pour la prédiction.

• Limitation du sondage corrélationnel : Les méthodes traditionnelles entraînent des classificateurs auxiliaires (« sondes ») sur les états cachés pour détecter des propriétés linguistiques (par exemple, le nombre du sujet). Bien qu'elles montrent ce qui est encodé, elles ne prouvent pas comment le modèle utilise cette information pour générer le token suivant.
• Limitation du sondage causal : Les approches causales existantes tentent d'intervenir sur les états cachés pour modifier des propriétés (par exemple, changer le singulier en pluriel) et observer les changements dans la sortie. Cependant, ces méthodes reposent généralement sur l'entraînement d'une sonde distincte pour déterminer la direction de l'intervention (par exemple, utiliser le gradient de la sonde pour inverser une étiquette).
  • Inconvénients : Cela introduit des coûts d'entraînement supplémentaires, nécessite une supervision spécifique à la propriété et risque un désalignement. La frontière de classification de la sonde peut ne pas coïncider avec la « géométrie prédictive » interne du modèle, conduisant à des interventions qui ne dirigent pas efficacement la génération du modèle ou altèrent involontairement des propriétés sans rapport.

Les auteurs proposent le Sondage causal en temps d'inférence, un cadre où les interventions sont appliquées entièrement au moment de l'inférence, sans réentraîner le modèle ni entraîner de sondes auxiliaires.

2. Méthodologie : Intervention de marge pilotée par l'état caché (HDMI)

La contribution centrale est HDMI, une technique sans sonde et basée sur le gradient qui exploite directement la tête de sortie native du modèle génératif pour piloter les états cachés.

2.1 Mécanisme central

HDMI fonctionne en définissant un objectif de marge entre un token « source » (la continuation préférée originale du modèle) et un token « cible » (la continuation contrefactuelle souhaitée).

• Objectif : Maximiser la marge de logit $L(x) = \phi(x)_\tau - \phi(x)_\sigma$, où $\phi(x)$ est le vecteur de logits, $\tau$ est l'index du token cible et $\sigma$ est l'index du token source.
• Intervention : La méthode effectue une ascension de gradient sur l'état caché $h_\ell(x)$ à une couche spécifique $\ell$ pour augmenter la probabilité du token cible tout en diminuant celle du token source.
• Efficacité computationnelle : Pour la dernière couche, le gradient est un produit matrice-vecteur sous forme fermée : $\nabla_{h_L} L(x) = W_U^\top (e_\tau - e_\sigma)$. Cela évite la nécessité d'un entraînement itératif de sonde.
• Extension multi-token : La méthode prend en charge des ensembles de tokens (par exemple, plusieurs verbes au pluriel valides) en sommant les logits sur les ensembles cibles ($T^+$) et les ensembles sources ($T^-$), définis comme $L_{set}(x) = \sum \phi(x)_{\tau_i} - \sum \phi(x)_{\sigma_j}$.

2.2 LA-HDMI (Lookahead HDMI) pour l'édition de texte

Pour répondre à l'édition de texte où un utilisateur fournit une séquence modifiée, les auteurs introduisent LA-HDMI.

• Défi : La génération autoregressive standard ne permet pas aux états cachés à l'étape $t$ d'être conscients des futures modifications à $t' > t$.
• Solution : LA-HDMI introduit un mécanisme de « regard en avant » qui rétropropage les gradients à travers les embeddings de tokens attendus.
  • Il découple le passage avant (utilisant une température faible $\beta_f$ pour maintenir la fluidité) du chemin de gradient (utilisant une température élevée $\beta_g$ pour calculer les embeddings attendus).
  • La fonction objectif $J_t(h_t)$ accumule les objectifs de marge pour les futures positions de modification (sur $S_{max}$ étapes), ajustant l'état caché actuel $h_t$ pour orienter les générations futures vers les modifications de l'utilisateur tout en préservant la fluidité dans les régions non modifiées.

3. Cadre d'évaluation

L'article évalue les interventions en utilisant deux métriques principales, suivant le cadre de Canby et al. [2024] :

1. Complétude : Mesure si l'intervention déplace avec succès la propriété ciblée (par exemple, changer le nombre du sujet) dans l'état caché, vérifié par une sonde de validation distincte.
2. Sélectivité : Mesure si des propriétés sans rapport (variables parasites, par exemple le nombre de noms non-sujets) restent stables après l'intervention.
3. Fiabilité : La moyenne harmonique de la Complétude et de la Sélectivité, servant de mesure globale de la qualité de l'intervention.

4. Résultats clés

Les auteurs ont évalué HDMI sur Meta-Llama-3-8B-Instruct et Pythia-70M en utilisant le corpus d'accord LGD et le benchmark CausalGym.

• Performance : HDMI atteint systématiquement une Fiabilité plus élevée que les bases de référence pilotées par sonde (AlterRep, FGSM, PGD) sur la plupart des tâches.
  • Dans les suites d'accords (par exemple, accord sujet-verbe), HDMI atteint souvent une Complétude parfaite et une Sélectivité élevée.
  • Contrairement aux bases de référence, la performance de HDMI est indépendante de la qualité d'une sonde interventionnelle, car elle n'en utilise pas. Des bases de référence comme AlterRep et PGD ont montré des baisses de performance significatives sur les modèles plus petits (Pythia-70M) lorsque la précision de la sonde était inférieure.
• Robustesse : HDMI maintient une fiabilité cohérente à travers différentes tailles de modèles et tâches, tandis que les méthodes basées sur des sondes ont montré des résultats mitigés, sensibles à la propriété spécifique et au modèle.
• Édition de texte : LA-HDMI a démontré avec succès la capacité à orienter la génération vers des modifications spécifiées par l'utilisateur (par exemple, changer « girl » en « owl » ou « was » en « were ») tout en maintenant la fluidité grammaticale et en ajustant le contexte nécessaire (par exemple, changer « a » en « an »).

5. Signification et affirmations

L'article affirme que HDMI représente un passage du sondage causal dépendant de la sonde au sondage causal natif du modèle.

• Alignement : En utilisant la propre tête de sortie du modèle pour définir la direction de l'intervention, HDMI garantit que la perturbation s'aligne sur la géométrie prédictive réelle du modèle, évitant ainsi les risques de désalignement des sondes externes.
• Efficacité : La méthode est légère sur le plan computationnel, nécessitant uniquement une ascension de gradient sur les états cachés sans la surcharge de l'entraînement de classificateurs auxiliaires.
• Généralité : L'approche s'étend naturellement aux propriétés multi-token et offre un mécanisme pour l'édition fluide de texte via des gradients de regard en avant.

Les auteurs reconnaissent des limites, notamment la sensibilité aux hyperparamètres (taille de l'étape, température, horizon) et l'atténuation potentielle des gradients dans les tâches d'édition de texte sur de longues distances. Ils suggèrent un travail futur sur la sélection adaptative de couches et la planification de la taille des étapes pour améliorer davantage la sélectivité sans sacrifier la complétude.