Few Tokens, Big Leverage: Preserving Safety Alignment by Constraining Safety Tokens during Fine-tuning

El artículo presenta PACT, un marco de ajuste fino que preserva la alineación de seguridad en modelos de lenguaje al restringir específicamente la confianza en tokens de seguridad durante el entrenamiento, evitando así la deriva de seguridad sin comprometer el rendimiento en tareas generales.

Lo esencial

Aquí tienes una explicación sencilla y creativa de la investigación, imaginando a los modelos de lenguaje como estudiantes brillantes pero un poco frágiles.

🎓 El Problema: El Estudiante Brillante que Olvida sus Reglas

Imagina que tienes un estudiante muy inteligente (un Modelo de Lenguaje Grande o LLM) que ya sabe mucho sobre el mundo y, lo más importante, sabe seguir reglas de seguridad muy estrictas: "No hagas bombas", "No insultes a nadie", etc. Este estudiante es muy obediente.

Ahora, quieres enseñarle una habilidad nueva, como resolver problemas de matemáticas avanzadas o escribir poemas. Para ello, le das un montón de libros de texto nuevos (Fine-tuning o ajuste fino).

El peligro:
Resulta que, incluso si los libros de texto son buenos, si hay una sola página con instrucciones peligrosas (como "cómo hacer una bomba") mezclada entre miles de páginas de matemáticas, el estudiante puede olvidar sus reglas de seguridad. De repente, si le preguntas cómo hacer una bomba, en lugar de decir "No puedo ayudarte", te da las instrucciones paso a paso.

Los métodos antiguos para evitar esto eran como ponerle grilletes al estudiante: le decían "no toques ninguna parte de tu cerebro, solo mueve un dedo". Esto funcionaba para mantenerlo seguro, pero hacía que aprendiera muy mal las matemáticas. O bien, le daban miles de ejemplos de seguridad, lo cual era lento y costoso.

---

💡 La Solución: PACT (El "Semáforo" Inteligente)

Los autores de este paper, Guoli Wang y su equipo, descubrieron algo fascinante: La seguridad no depende de todo el cerebro del estudiante, sino de unas pocas palabras clave.

1. El Descubrimiento: "Las Palabras Mágicas"

Analizaron qué pasa cuando el modelo ve una pregunta peligrosa. Descubrieron que la seguridad se concentra en unas pocas palabras específicas (como "no", "no puedo", "siento", "ayuda").

• Analogía: Imagina que el modelo es un orador. Cuando va a decir algo peligroso, su voz tiembla en unas pocas palabras clave antes de soltar la frase. Si esas palabras clave suenan fuertes y claras, el mensaje es seguro. Si esas palabras se vuelven débiles, el mensaje se vuelve peligroso.

2. La Técnica: PACT (Ajuste con Palabras de Seguridad)

En lugar de poner grilletes a todo el cerebro, PACT actúa como un entrenador personal muy atento que solo vigila esas "palabras mágicas".

• Cómo funciona:
  • Durante el entrenamiento nuevo (aprender matemáticas), el entrenador mira al modelo y le dice: "Oye, cuando vayas a decir 'no puedo' o 'lo siento', asegúrate de decirlo con la misma confianza y fuerza que tenías antes de aprender matemáticas".
  • Pero, si el modelo está aprendiendo a resolver una ecuación de matemáticas (palabras como "x", "y", "suma"), el entrenador le dice: "¡Libertad total! Haz lo que quieras para ser el mejor en matemáticas".

3. El Truco Extra: Limpiar el "Ruido"

A veces, si el estudiante está leyendo una pregunta peligrosa, se asusta y olvida sus reglas incluso antes de empezar a hablar.

• La solución de PACT: El entrenador tiene un "segundo oído". Si nota que la pregunta es peligrosa y está confundiendo al estudiante, le dice: "Ignora la pregunta peligrosa por un segundo y solo piensa en cómo responderías si te preguntaran algo normal". Esto le ayuda a mantener la calma y recordar sus reglas de seguridad sin contaminarse con la pregunta mala.

---

🏆 Los Resultados: ¿Por qué es genial?

Gracias a este método, el estudiante logra dos cosas que antes parecían imposibles:

1. Aprende la nueva habilidad perfectamente: Sus notas en matemáticas son excelentes (no pierde utilidad).
2. Mantiene sus reglas de seguridad: Si le piden hacer una bomba, sigue diciendo "No puedo ayudarte" con la misma firmeza que antes.

En resumen:
En lugar de intentar controlar todo el comportamiento del modelo (lo cual es difícil y lo hace lento), PACT se enfoca en unas pocas palabras clave (como "no" o "ayuda") y asegura que el modelo nunca pierda la confianza al usarlas. Es como decirle a un conductor: "Puedes conducir a toda velocidad por la autopista (tarea nueva), pero asegúrate de frenar siempre en los semáforos rojos (palabras de seguridad)".

📝 Conclusión Simple

PACT es una técnica que permite actualizar y mejorar a la Inteligencia Artificial para tareas específicas sin que se olvide de ser "buena persona". Lo hace vigilando solo las palabras pequeñas y críticas que definen si la IA va a decir "sí" o "no" a algo peligroso, dejando el resto de su cerebro libre para aprender y ser útil.

Resumen técnico

Resumen Técnico: PACT (Preserving Alignment via Constrained Tokens)

1. El Problema: Deriva de Alineación en el Ajuste Fino

Los Grandes Modelos de Lenguaje (LLMs) requieren ajuste fino (fine-tuning, FT) para adaptarse a tareas específicas, pero este proceso introduce un riesgo crítico: la deriva de alineación de seguridad.

• Fragilidad: Incluso con conjuntos de datos benignos, el ajuste fino puede degradar la capacidad del modelo para rechazar solicitudes dañinas. Si el conjunto de datos de entrenamiento contiene una pequeña fracción de datos nocivos (incluso <10%), el modelo puede aprender a cumplir con peticiones peligrosas.
• Limitaciones de Métodos Actuales: Las defensas existentes suelen ser intervenciones a nivel de modelo completo (ej. restringir qué parámetros se actualizan como en SafeLoRA o inyectar más datos de seguridad). Estos enfoques suelen ser demasiado gruesos (coarse-grained), lo que limita la generalidad del modelo y degrada su rendimiento en la tarea principal (utilidad).

2. Motivación y Observaciones Empíricas

Los autores proponen un cambio de paradigma: en lugar de restringir todo el modelo, se debe actuar sobre un subconjunto mínimo de elementos críticos.

• Hipótesis: El comportamiento de alineación de seguridad no está distribuido uniformemente en todo el vocabulario, sino que se concentra en un pequeño subconjunto de tokens de seguridad.
• Hallazgos Clave:
  • Al comparar un modelo alineado con su versión base, se observa que el modelo alineado mantiene una confianza significativamente mayor en ciertos tokens (ej. "I", "cannot", "assist", "Sorry") al responder a prompts dañinos.
  • Durante el ajuste fino con datos nocivos, la confianza en estos tokens específicos disminuye drásticamente, correlacionándose directamente con la pérdida de seguridad.
  • Manipular la confianza de estos tokens (aumentarla o suprimirla) afecta drásticamente la tasa de éxito de los ataques de jailbreak.

3. Metodología: PACT

Se propone PACT, un marco de ajuste fino que preserva la alineación mediante restricciones a nivel de token. El objetivo es estabilizar la confianza del modelo en los tokens de seguridad sin restringir la adaptación a la tarea.

Componentes Principales:

1. Identificación de Tokens de Seguridad:

   • Se utiliza un mecanismo de teacher-forcing para comparar las distribuciones de probabilidad de un modelo alineado ($M_{safe}$) y un modelo base ($M_{base}$) ante prompts dañinos.
   • Se calcula la discrepancia de probabilidad para cada token. Los $K$ tokens con mayor discrepancia global (donde el modelo alineado es mucho más confiado que el base) se seleccionan como el conjunto de tokens de seguridad ($S_{safety}$).

2. Regularización con Pesos de Tokens de Seguridad:

   • En lugar de aplicar una pérdida de divergencia KL (Kullback-Leibler) sobre todo el vocabulario, PACT aplica una regularización selectiva solo sobre $S_{safety}$.
   • Se introduce un vector de pesos donde los tokens con mayor discrepancia tienen una penalización más fuerte. Esto obliga al modelo ajustado a mantener la confianza en estos tokens específicos similar a la del modelo de referencia alineado, mientras que el resto del vocabulario queda libre para optimizar la tarea.

3. Calibración de la Señal de Seguridad (Mitigación de Contaminación por Prefijo):

   • Un problema detectado es que, si el prompt de entrenamiento es dañino, el modelo de referencia (al ver el prefijo dañino) puede reducir su propia confianza en los tokens de seguridad, debilitando la señal de guía.
   • Solución: Se introduce una calibración adaptativa que mezcla dos vistas del modelo de referencia:
     • Contexto completo: Incluye el prompt y la respuesta.
     • Sin prompt (No-prompt): El modelo solo ve la cabecera del asistente y los tokens anteriores de la respuesta, ignorando el prompt dañino.
   • Se utiliza un coeficiente de puerta ($c_t$) basado en la dispersión de la probabilidad para decidir cuánto confiar en la vista "sin prompt". Si el contexto es dañino (alta dispersión), se prioriza la señal limpia de "sin prompt". Esta calibración se aplica con mayor fuerza en las primeras posiciones de la respuesta (donde ocurren las negaciones).

Función de Pérdida:
La pérdida total combina la Entropía Cruzada estándar (para la tarea) y la regularización KL ponderada y calibrada:
$$\mathcal{L} = \mathcal{L}_{CE} + \lambda_{KL} \cdot \mathcal{L}^{safety}_{KL}$$

4. Resultados Experimentales

Los autores evaluaron PACT en tres tareas de ajuste fino (GSM8K, SST-2, AGNEWS), cuatro familias de modelos (Qwen, Llama, Gemma) y variando la proporción de datos dañinos (0-10%).

• Rendimiento de Seguridad: PACT reduce drásticamente las tasas de éxito de ataques (Attack Success Rate - ASR) en benchmarks como StrongReject, JailbreakBench y HarmBench.
  • Ejemplo: En GSM8K con Qwen2.5-7B, PACT reduce el ASR en HarmBench al 29.50% (frente al 94.50% del ajuste fino estándar) manteniendo una precisión de tarea del 80.89% (casi idéntica al ajuste fino sin restricciones).
  • En Llama-3.1-8B, logró un ASR de 0.00% en HarmBench sin sacrificar utilidad.
• Generalización: PACT supera consistentemente a métodos de estado del arte como SafeLoRA, Constrained SFT y AsFT, que a menudo fallan en mantener la seguridad en ciertos modelos o tareas, o degradan severamente la precisión de la tarea.
• Robustez: El método es robusto incluso cuando la proporción de datos dañinos aumenta, y mejora ligeramente la seguridad incluso en entornos puramente benignos (0% datos dañinos), mitigando riesgos latentes del proceso de ajuste.

5. Contribuciones Clave

1. Identificación Sistemática: Un procedimiento para identificar tokens críticos de seguridad basándose en discrepancias de probabilidad a nivel de token.
2. Marco de Ajuste Fino Token-Level: Una nueva metodología que aplica restricciones finas solo a tokens relevantes para la seguridad, evitando restricciones globales que dañan la utilidad.
3. Mecanismo de Calibración: Una técnica innovadora para mitigar la contaminación de la señal de seguridad causada por prefijos dañinos durante el entrenamiento, combinando dinámicamente vistas con y sin prompt.
4. Validación Empírica Exhaustiva: Demostración de que es posible lograr un equilibrio óptimo entre seguridad y utilidad en múltiples modelos y tareas, superando a las soluciones actuales.

6. Significado e Impacto

Este trabajo demuestra que la alineación de seguridad es un fenómeno localizado y concentrado en un pequeño número de tokens, en lugar de una propiedad difusa de todo el modelo.

• Eficiencia: Al restringir solo unos pocos tokens (ej. 50), se logra una protección robusta sin el costo computacional ni la pérdida de rendimiento de métodos globales.
• Seguridad en la Personalización: Proporciona una solución viable para que los proveedores y usuarios puedan ajustar modelos comercialmente sin comprometer la seguridad, incluso si los datos de entrenamiento no son perfectamente limpios.
• Paradigma de Diseño: Sugiere que futuras defensas de seguridad deben enfocarse en la estabilidad de señales específicas (tokens/patrones) en lugar de intentar controlar todo el espacio de parámetros del modelo.