Test-Time Safety Alignment

1. Planteamiento del Problema

Los Modelos de Lenguaje Grande (LLM) pasan por entrenamiento de seguridad (por ejemplo, RLHF, DPO) para rechazar solicitudes dañinas. Sin embargo, los prompts adversarios (jailbreaks) pueden eludir estas defensas explotando la brecha entre la distribución de entrenamiento del modelo y el espacio de entrada ilimitado. Estos prompts suelen utilizar un encuadre estructural (roleplay, ofuscación) para elicitar contenido dañino.

Las defensas existentes en tiempo de prueba enfrentan limitaciones estructurales:

• Las estrategias a nivel de prompt manipulan tokens discretos, lo cual es frágil frente a ataques semánticos.
• La dirección de activación requiere acceso a estados ocultos y a menudo depende de conjuntos de calibración fijos que pueden no generalizar.
• Los modelos guardianes requieren inferencia separada y a menudo siguen un paradigma de "detectar y luego filtrar" en lugar de dirigir la generación.

La Pregunta Central: ¿Se pueden optimizar los embebidos de palabras de entrada de manera sub-lexica (modificando la representación vectorial continua sin cambiar el token visible) para dirigir modelos alineados hacia salidas seguras, incluso cuando el modelo produce una distribución bimodal de "rechazar o cumplir" en lugar de una continuación suave?

2. Metodología: Alineación de Seguridad en Tiempo de Prueba (TSA)

El método propuesto, TSA, opera completamente en tiempo de prueba sin modificar los pesos del modelo, no requiere entrenamiento auxiliar y trata al LLM como una caja negra. Optimiza directamente la matriz de embebidos del prompt.

A. Función Objetivo

El objetivo es minimizar la nocividad del texto generado $y = f(X)$, donde $f$ es el LLM y $X$ es la matriz de embebidos del prompt.

• Oráculo de Nocividad ($h$): El método utiliza la API de Moderación de OpenAI como un oráculo de caja negra. Puntúa las salidas en 13 categorías de daño (por ejemplo, violencia, autolesión, odio).
• Función de Pérdida: $\Phi(X) = \max_{c} s_c(f(X))$, donde $s_c$ es la puntuación para la categoría $c$. El objetivo es minimizar esta puntuación máxima, asegurando que la respuesta sea segura en todas las dimensiones.

B. Estimación de Gradiente de Orden Cero

Dado que el LLM ($f$) y la API de Moderación ($h$) son cajas negras, la retropropagación es imposible. TSA utiliza estimación de gradiente de orden cero:

1. Perturbación: En cada iteración $k$, la matriz de embebidos actual $X_k$ se perturba añadiendo ruido gaussiano $U$ escalado por $\mu$: $X_k + \mu U$.
2. Diferencia Finita: El gradiente se estima comparando la pérdida de las entradas perturbadas contra la entrada base:
   $$g_k \approx \frac{1}{N} \sum_{i=1}^{N} \frac{\Phi(X_k + \mu U_i) - \Phi(X_k)}{\mu} U_i$$
   donde $N$ es el número de muestras de Monte Carlo.
3. Actualización: Los embebidos se actualizan mediante descenso de gradiente: $X_{k+1} \leftarrow X_k - \eta g_k$.

C. Regularización

Para asegurar que los embebidos optimizados permanezcan válidos y no se desvíen hacia regiones no entrenadas del espacio latente:

• Normalización del Gradiente: El paso de actualización utiliza la dirección del gradiente normalizada ($g_k / \|g_k\|_2$) para controlar la magnitud del paso.
• Restricción de Similitud Coseno: El embebido actualizado se proyecta de nuevo sobre una bola de coseno alrededor del embebido original $X_0$ si la similitud cae por debajo de un umbral $\kappa$. Esto asegura que la perturbación sea sub-lexica: el vector optimizado se decodifica de nuevo al token original, pero la representación continua dentro de la celda de Voronoi de ese token se altera.

3. Contribuciones Clave

1. Extensión a Modelos Alineados: Demuestra que el control a nivel de embebido, previamente mostrado para desintoxicar modelos base, funciona efectivamente en modelos instruidos alineados en seguridad, los cuales tienen una distribución de salida fundamentalmente diferente (bimodal).
2. Control Sub-Lexico: Prueba que optimizar el vector de embebido continuo dentro de la celda del vecino más cercano de un token (sin cambiar el token en sí) es suficiente para redirigir el comportamiento del modelo de dañino a seguro. Esto desafía la visión de los embebidos como meras búsquedas discretas.
3. Optimización de Caja Negra: El método no requiere acceso a los internos del modelo (gradientes, estados ocultos o logits). Solo requiere acceso a los embebidos de entrada y a una API de moderación de caja negra.
4. Aplicación Incondicional: El método se aplica a todos los prompts independientemente de su intención. No fuerza un rechazo; dirige al modelo a generar una respuesta sustantiva y segura si es posible, o un rechazo seguro si es necesario.

4. Resultados Experimentales

El método fue evaluado en cinco modelos instruidos (que van desde 1B hasta 20B parámetros: Gemma 3, Phi-3.5, Llama 3.1, Qwen3, GPT-OSS) en dos benchmarks:

• WildJailbreak: Prompts adversarios (dañinos y benignos).
• HarmBench: Consultas dañinas directas.

Hallazgos Clave:

• Eliminación de Seguridad: TSA redujo el número de respuestas marcadas (dañinas) a cero para la mayoría de los modelos en benchmarks estándar. Por ejemplo, en Llama 3.1-8B, redujo los prompts dañinos marcados de 218 a 0.
• Convergencia: La optimización típicamente converge dentro de 1 a 2 pasos de gradiente.
• Preservación Benigna: En prompts benignos adversarios, TSA no aumentó las puntuaciones de moderación, confirmando que no induce un sesgo de "rechazo generalizado" ni degrada la utilidad.
• Comparación con Líneas Base: TSA superó significativamente a líneas base como SmoothLLM (perturbación aleatoria de tokens), AdaSteer (dirección de activación) y RESTA (ruido de embebido). Las líneas base a menudo fallaron en jailbreaks semánticos o aumentaron las tasas de marcados en consultas dañinas directas.
• Eficiencia: El principal costo de latencia es la llamada a la API del oráculo de moderación, no la inferencia del modelo.

5. Significado e Implicaciones

• Nueva Superficie de Control: El artículo revela que la capa de embebidos ofrece una superficie de control más rica de lo que se pensaba. Pequeñas perturbaciones sub-lexicas pueden navegar las complejas fronteras de seguridad de los modelos alineados.
• Viabilidad de Despliegue: Como una capa de seguridad "plug-and-play", TSA no requiere reentrenamiento ni ajuste fino del modelo. Puede desplegarse en cualquier modelo donde los embebidos de entrada sean accesibles (por ejemplo, modelos de pesos abiertos) combinados con una API de moderación pública.
• Limitaciones:
  • Latencia: La optimización iterativa añade latencia por prompt (segundos), aunque esto está dominado por los viajes de ida y vuelta de la API.
  • Dependencia del Modelo: La efectividad escala con la calidad del entrenamiento de seguridad existente del modelo. Modelos con alineación débil (por ejemplo, Gemma 3-1B) mostraron un éxito reducido, ya que el optimizador no pudo encontrar una dirección segura cercana.
  • Acceso a API: Requiere acceso a embebidos de entrada (no disponible para modelos cerrados puramente protegidos por API) y una API de moderación.

En conclusión, Alineación de Seguridad en Tiempo de Prueba proporciona un mecanismo robusto y agnóstico al modelo para neutralizar ataques adversarios optimizando la representación de entrada continua, cerrando efectivamente la brecha entre la seguridad teórica y el despliegue práctico sin reentrenamiento.