Before the Last Token: Diagnosing Final-Token Safety Probe Failures

Resumen Técnico: Antes del Último Token: Diagnóstico de Fallos en las Sondas de Seguridad del Último Token

Enunciado del Problema
Los sistemas actuales de monitoreo de seguridad para Modelos de Lenguaje Grandes (LLM), como SafeSwitch, a menudo dependen de "sondas del último token". Estos sistemas entrenan clasificadores sobre el estado oculto del último token en la fase de prellenado de un prompt para detectar contenido inseguro antes de que comience la generación. El artículo identifica un modo de fallo crítico en esta arquitectura: mientras que estas sondas logran una alta recuperación en prompts dañinos "limpios", a menudo fallan al detectar prompts de "jailbreak" (evasión). Los jailbreaks incrustan solicitudes dañinas dentro de marcos adversarios, juegos de roles u ofuscación. Los autores argumentan que estos envoltorios desplazan la representación del último token lejos de la dirección específica de contraste "dañino vs. benigno" aprendida durante el entrenamiento, lo que hace que la sonda pase por alto la solicitud insegura a pesar de que el contenido dañino está presente en el prompt. Este fallo no es meramente un problema de capacidad, ya que aumentar el ancho del cuello de botella de la sonda no resuelve de manera fiable la brecha de detección.

Metodología
El estudio evalúa tres modelos ajustados con instrucciones: Llama-3.1-8B-Instruct, Mistral-7B-Instruct-v0.1 y OLMo3-7B-Instruct.

1. Entrenamiento de la Sonda: Se entrenaron sondas al estilo de SafeSwitch exclusivamente con prompts dañinos y benignos "limpios" del conjunto de entrenamiento original de SafeSwitch. Estas sondas leen el estado oculto del último token del prompt, lo proyectan a través de un cuello de botella (inicialmente con ancho $w=64$) y predicen la seguridad.
2. Evaluación: Las sondas se probaron en:
   • Prompts dañinos limpios retenidos (SorryBench).
   • Los mismos prompts dañinos envueltos con plantillas de jailbreak de HarmBench.
   • Prompts benignos adyacentes a la seguridad (XSTest) para medir los falsos positivos.
   • Benchmarks dañinos directos (AdvBench) para descartar desplazamientos de distribución genéricos.
3. Análisis Geométrico: Los autores analizaron el subespacio representacional de las sondas. Computaron vectores de diferencia de medias para varios conjuntos (por ejemplo, jailbreaks fallidos frente a prompts benignos) y midieron su "energía visible para la sonda": la fracción de la norma de una dirección capturada por el espacio de filas de la sonda (mediante SVD).
4. Análisis de Trayectoria: En lugar de depender únicamente del último token, los autores analizaron las trayectorias de estados ocultos token a token durante la fase de prellenado. Investigaron si la evidencia de inseguridad aparecía antes en la secuencia pero colapsaba en el último token.
5. Modelado Diagnóstico: Para probar si la información de la trayectoria podía recuperar las detecciones fallidas, los autores ajustaron un PCA-HMM simple (Análisis de Componentes Principales - Modelo Oculto de Markov) en los datos de entrenamiento limpios. Este modelo aprendió la dinámica de trayectoria de prompts dañinos limpios frente a benignos y se aplicó a las trayectorias de jailbreak sin re-ajustar en datos de jailbreak.

Resultados Clave

• Fallo de la Sonda en Jailbreaks: Las sondas del último token mantuvieron una alta recuperación (~95%) en prompts dañinos limpios, pero cayeron significativamente en jailbreaks (oscilando entre ~63% y ~87% dependiendo del modelo). Por el contrario, produjeron falsos positivos no triviales (23%–36%) en prompts benignos adyacentes a la seguridad.
• La Capacidad No es el Cuello de Botella: Aumentar el ancho del cuello de botella de la sonda de 64 a 1024 no mejoró de manera fiable la detección de jailbreaks. Para OLMo3, el rendimiento incluso se degradó. Esto sugiere que el fallo es geométrico (dónde reside la señal) en lugar de una falta de capacidad representacional.
• Desalineación Geométrica: La dirección que separa los jailbreaks fallidos de los prompts benignos ($d_{miss}$) tuvo la menor energía visible para la sonda (aprox. 2.4%–3.5%) en comparación con el contraste de entrenamiento limpio ($d_{harm}$). El espacio de filas de la sonda es efectivamente ortogonal a las características específicas que distinguen los jailbreaks fallidos de las entradas benignas.
• Fallo Posicional: La evidencia de inseguridad en los jailbreaks fallidos a menudo es visible en posiciones de token intermedias (específicamente cerca de la solicitud dañina incrustada) pero desaparece o "colapsa" para el último token.
• Fallo del Agrupamiento Ingenuo: Simplemente tomar la puntuación máxima a través de todas las posiciones de token no es una solución viable, ya que se activa en el 100% de los prompts benignos de XSTest (falsos positivos), lo que indica que las puntuaciones intermedias altas no son semánticamente fiables por sí solas.
• Recuperación de Trayectoria: El diagnóstico PCA-HMM, entrenado solo en datos limpios, recuperó con éxito la gran mayoría de los fallos del último token (tasa de recuperación del 97%–98%) mientras mantenía una tasa de falsos positivos significativamente menor en XSTest (6%–15%) en comparación con el agrupamiento ingenuo. La dirección de jailbreak fallido mostró una energía significativamente mayor en el subespacio PCA de la trayectoria que en el espacio de filas de la sonda del último token.

Significado y Afirmaciones
El artículo se posiciona como un estudio diagnóstico en lugar de un mecanismo de defensa desplegado. Su contribución principal es identificar por qué fallan las sondas del último token: dependen de una lectura "tipo atajo" que es sensible a la distribución de entrenamiento limpia pero ciega a los desplazamientos geométricos específicos inducidos por los envoltorios de jailbreak.

Los autores afirman que:

1. Las sondas del último token son insuficientes para detectar solicitudes dañinas envueltas porque el desplazamiento inducido por el envoltorio mueve la representación fuera del subespacio sensible de la sonda.
2. La información relevante para la seguridad a menudo existe en la trayectoria del estado oculto antes del último token, pero se pierde para el momento en que ocurre la lectura final.
3. Los modelos simples conscientes de la trayectoria (como el PCA-HMM), entrenados solo en datos limpios, pueden servir como complementos diagnósticos efectivos para las sondas del último token. Pueden recuperar detecciones fallidas analizando la evolución de las señales de seguridad a través del prompt, ofreciendo una visión más robusta del estado interno del prompt sin las tasas catastróficas de falsos positivos del agrupamiento de tokens ingenuo.

El trabajo motiva un cambio en los diagnósticos de seguridad, pasando de lecturas de un solo punto a análisis conscientes de la trayectoria para comprender y mitigar mejor las limitaciones de los sistemas de seguridad actuales basados en sondas.