Test-Time Safety Alignment

1. Énoncé du problème

Les modèles de langage de grande taille (LLM) subissent un entraînement de sécurité (par exemple, RLHF, DPO) pour refuser les demandes nuisibles. Cependant, des invites adverses (jailbreaks) peuvent contourner ces défenses en exploitant l'écart entre la distribution d'entraînement du modèle et l'espace d'entrée illimité. Ces invites utilisent souvent un cadrage structurel (jeu de rôle, obfuscation) pour provoquer la génération de contenu nuisible.

Les défenses existantes au moment du test font face à des limitations structurelles :

• Les stratégies au niveau de l'invite manipulent des tokens discrets, ce qui les rend fragiles face aux attaques sémantiques.
• Le pilotage des activations nécessite l'accès aux états cachés et repose souvent sur des ensembles de calibration fixes qui peuvent ne pas généraliser.
• Les modèles de garde nécessitent une inférence séparée et suivent souvent un paradigme « détecter puis filtrer » plutôt que de piloter la génération.

La question centrale : Peut-on optimiser les plongements de mots d'entrée de manière sub-lexicale (en modifiant la représentation vectorielle continue sans changer le token visible) pour orienter les modèles alignés vers des sorties sûres, même lorsque le modèle produit une distribution bimodale « refuser ou se conformer » plutôt qu'une continuation lisse ?

2. Méthodologie : Alignement de sécurité au moment du test (TSA)

La méthode proposée, TSA, opère entièrement au moment du test sans modifier les poids du modèle, ne nécessitant aucun entraînement auxiliaire et traitant le LLM comme une boîte noire. Elle optimise directement la matrice de plongement de l'invite.

A. Fonction objectif

L'objectif est de minimiser la nocivité du texte généré $y = f(X)$, où $f$ est le LLM et $X$ est la matrice de plongement de l'invite.

• Oracle de nocivité ($h$) : La méthode utilise l'API de modération OpenAI comme oracle boîte noire. Elle note les sorties selon 13 catégories de nuisances (par exemple, violence, automutilation, haine).
• Fonction de perte : $\Phi(X) = \max_{c} s_c(f(X))$, où $s_c$ est le score pour la catégorie $c$. L'objectif est de minimiser ce score maximal, garantissant que la réponse est sûre dans toutes les dimensions.

B. Estimation du gradient d'ordre zéro

Puisque le LLM ($f$) et l'API de modération ($h$) sont des boîtes noires, la rétropropagation est impossible. TSA utilise une estimation de gradient d'ordre zéro :

1. Perturbation : À chaque itération $k$, la matrice de plongement actuelle $X_k$ est perturbée en ajoutant un bruit gaussien $U$ mis à l'échelle par $\mu$ : $X_k + \mu U$.
2. Différence finie : Le gradient est estimé en comparant la perte des entrées perturbées à celle de l'entrée de base :
   $$g_k \approx \frac{1}{N} \sum_{i=1}^{N} \frac{\Phi(X_k + \mu U_i) - \Phi(X_k)}{\mu} U_i$$
   où $N$ est le nombre d'échantillons de Monte Carlo.
3. Mise à jour : Les plongements sont mis à jour par descente de gradient : $X_{k+1} \leftarrow X_k - \eta g_k$.

C. Régularisation

Pour garantir que les plongements optimisés restent valides et ne dérivent pas vers des régions non entraînées de l'espace latent :

• Normalisation du gradient : L'étape de mise à jour utilise la direction de gradient normalisée ($g_k / \|g_k\|_2$) pour contrôler l'amplitude de l'étape.
• Contrainte de similarité cosinus : Le plongement mis à jour est projeté sur une boule de cosinus autour du plongement original $X_0$ si la similarité tombe en dessous d'un seuil $\kappa$. Cela garantit que la perturbation est sub-lexicale : le vecteur optimisé se décode de nouveau en le token original, mais la représentation continue à l'intérieur de la cellule de Voronoï de ce token est modifiée.

3. Contributions clés

1. Extension aux modèles alignés : Démontre que le contrôle au niveau des plongements, précédemment montré pour la désintoxication des modèles de base, fonctionne efficacement sur les modèles instructifs alignés sur la sécurité, qui possèdent une distribution de sortie fondamentalement différente (bimodale).
2. Contrôle sub-lexicale : Prouve que l'optimisation du vecteur de plongement continu à l'intérieur de la cellule du plus proche voisin d'un token (sans changer le token lui-même) suffit à rediriger le comportement du modèle de nuisible à sûr. Cela remet en question la vision des plongements comme de simples recherches discrètes.
3. Optimisation boîte noire : La méthode ne nécessite aucun accès aux internals du modèle (gradients, états cachés ou logits). Elle nécessite uniquement l'accès aux plongements d'entrée et à une API de modération boîte noire.
4. Application inconditionnelle : La méthode s'applique à toutes les invites, indépendamment de leur intention. Elle ne force pas un refus ; elle oriente le modèle à générer une réponse substantielle et sûre si possible, ou un refus sûr si nécessaire.

4. Résultats expérimentaux

La méthode a été évaluée sur cinq modèles instructifs (allant de 1 à 20 milliards de paramètres : Gemma 3, Phi-3.5, Llama 3.1, Qwen3, GPT-OSS) sur deux benchmarks :

• WildJailbreak : Invites adverses (nuisibles et bénignes).
• HarmBench : Requêtes nuisibles directes.

Constats principaux :

• Élimination de la sécurité : TSA a réduit le nombre de réponses signalées (nuisibles) à zéro pour la plupart des modèles sur les benchmarks standards. Par exemple, sur Llama 3.1-8B, elle a réduit les invites nuisibles signalées de 218 à 0.
• Convergence : L'optimisation converge généralement en 1 à 2 étapes de gradient.
• Préservation du bénin : Sur les invites bénignes adverses, TSA n'a pas augmenté les scores de modération, confirmant qu'elle n'induit pas de biais de « refus généralisé » ni ne dégrade l'utilité.
• Comparaison avec les bases de référence : TSA a nettement surpassé des bases de référence telles que SmoothLLM (perturbation aléatoire de tokens), AdaSteer (pilotage des activations) et RESTA (bruit de plongement). Les bases de référence ont souvent échoué sur les jailbreaks sémantiques ou ont augmenté les taux de signalement sur les requêtes nuisibles directes.
• Efficacité : Le principal coût de latence est l'appel API vers l'oracle de modération, et non l'inférence du modèle.

5. Signification et implications

• Nouvelle surface de contrôle : L'article révèle que la couche de plongement offre une surface de contrôle plus riche que prévu. De petites perturbations sub-lexicales peuvent naviguer dans les frontières de sécurité complexes des modèles alignés.
• Faisabilité du déploiement : En tant que couche de sécurité « plug-and-play », TSA ne nécessite ni réentraînement ni fine-tuning du modèle. Elle peut être déployée sur n'importe quel modèle dont les plongements d'entrée sont accessibles (par exemple, modèles à poids ouverts) combiné à une API de modération publique.
• Limites :
  • Latence : L'optimisation itérative ajoute une latence par invite (secondes), bien que celle-ci soit dominée par les allers-retours API.
  • Dépendance au modèle : L'efficacité évolue avec la qualité de l'entraînement de sécurité existant du modèle. Les modèles avec un alignement faible (par exemple, Gemma 3-1B) ont montré un succès réduit, car l'optimiseur ne pouvait pas trouver de direction sûre à proximité.
  • Accès API : Nécessite l'accès aux plongements d'entrée (non disponibles pour les modèles fermés purement gérés par API) et une API de modération.

En conclusion, Test-Time Safety Alignment fournit un mécanisme robuste et agnostique au modèle pour neutraliser les attaques adverses en optimisant la représentation d'entrée continue, comblant efficacement le fossé entre la sécurité théorique et le déploiement pratique sans réentraînement.