MOSAIC: Composable Safety Alignment with Modular Control Tokens

El paper presenta MOSAIC, un marco modular que utiliza tokens de control aprendibles sobre un modelo base congelado para habilitar una alineación de seguridad compuesta y adaptable al contexto, logrando un alto rendimiento defensivo con menos rechazos excesivos y sin comprometer la utilidad del modelo.

Lo esencial

Imagina que tienes un chef de cocina muy talentoso (el modelo de lenguaje o LLM) que puede cocinar de todo: desde recetas complejas hasta consejos de viaje. Sin embargo, este chef tiene un problema: a veces, sin querer, prepara platos que no son adecuados para ciertos comensales. Por ejemplo, si un niño pide cómo hacer un cóctel con alcohol, el chef debería decir "no", pero si un adulto lo pide, podría decir "sí".

El problema actual es que la mayoría de los chefs de IA tienen las reglas de seguridad "grabadas a fuego" en su cerebro. Si quieres cambiar una regla (por ejemplo, "prohibido el alcohol para menores"), tienes que reentrenar todo el cerebro del chef, lo cual es costoso, lento y a veces hace que olvide cómo cocinar bien otras cosas. Otras veces, simplemente le gritas instrucciones al chef ("¡No hagas alcohol!"), pero si le gritas muchas cosas a la vez, se confunde y empieza a rechazar hasta recetas inocentes como "cómo hacer un pastel".

Aquí es donde entra MOSAIC.

¿Qué es MOSAIC? (La analogía de los "Botones Mágicos")

MOSAIC es como un panel de control con botones intercambiables que se colocan frente al chef, en lugar de reescribir su cerebro.

1. No tocas el cerebro: El chef sigue siendo el mismo, con todo su talento intacto.
2. Usas "Fichas de Control": Imagina que cada regla de seguridad (como "Prohibido apuestas" o "Prohibido alcohol") es una pequeña ficha magnética o un token de control.
3. Combinación flexible:
   • Si el usuario es un niño, pones la ficha de "Prohibido alcohol" y la de "Prohibido apuestas".
   • Si el usuario es un adulto, quitas esas fichas.
   • Si el usuario es un juego de rol medieval, quizás solo activas la ficha de "Prohibido violencia gráfica".

El sistema MOSAIC aprende a crear estas fichas. Lo genial es que puedes combinarlas libremente. Si activas la ficha de "Alcohol" y la de "Apuestas" al mismo tiempo, el chef sabe exactamente cómo actuar sin confundirse.

¿Cómo aprenden estas fichas a funcionar bien? (El entrenamiento)

El papel explica dos trucos inteligentes para entrenar estas fichas:

1. La "Gimnasia de Combinaciones":
   En lugar de enseñar al chef solo con una regla a la vez, el sistema le muestra mezclas. Le dice: "Oye, hoy vamos a probar la regla de Alcohol + Apuestas + Terror". Esto evita que las fichas se peleen entre sí. Es como entrenar a un equipo de fútbol: no solo practican solos, sino que aprenden a jugar juntos en diferentes formaciones.

2. El "Espejo de Realidad" (Para no ser demasiado estricto):
   A veces, si le pones muchas fichas de seguridad al chef, se vuelve paranoico y rechaza preguntas inocentes (como "¿cómo se hace un mojito?"). Para evitar esto, MOSAIC usa un truco llamado distilación de conocimiento contrafactual.

   • La analogía: Imagina que le preguntas al chef: "¿Qué hubieras dicho si no tuviera estas fichas de seguridad?". El chef responde: "Un mojito fácil". Luego, el sistema le dice: "Muy bien, pero como tengo la ficha de 'Menores', debes decir 'No'".
   • El objetivo es que el chef solo cambie su respuesta cuando sea estrictamente necesario, manteniendo su naturalidad en todo lo demás. Así, no se vuelve un robot gruñón que dice "no" a todo.

¿Por qué es importante esto?

• Ahorro de tiempo y dinero: No necesitas reentrenar al modelo gigante cada vez que quieras cambiar una regla para un país o una edad específica. Solo cambias las fichas.
• Menos errores: Evita que el modelo sea demasiado estricto y rechace cosas útiles (como explicar cómo funciona una ruleta en un contexto educativo, en lugar de enseñar a ganar).
• Adaptabilidad: Puedes crear un "chef" que se adapte a un abuelo, a un niño, a un abogado o a un médico, simplemente cambiando las fichas que le pones delante.

En resumen

MOSAIC es como darle a una IA un set de gafas de realidad aumentada. Dependiendo de quién sea el usuario, le pones unas gafas diferentes que le muestran el mundo con ciertas reglas de seguridad activadas, sin tener que cambiar la forma en que la IA ve el mundo en su interior. Es más rápido, más flexible y hace que la IA sea más inteligente y menos torpe al decir "no".

Resumen técnico

Resumen Técnico: MOSAIC

1. El Problema

La alineación de seguridad en los Modelos de Lenguaje Grandes (LLM) actuales suele implementarse como una política estática incrustada en los parámetros del modelo. Este enfoque presenta limitaciones críticas en escenarios de despliegue real:

• Falta de adaptabilidad: Las reglas de seguridad varían según la región, la edad del usuario, el rol profesional o el dominio de aplicación. Un modelo con una política fija no puede adaptarse dinámicamente a estos contextos sin un costoso reentrenamiento.
• Entrelazamiento de capacidades: Los métodos basados en parámetros (como el Fine-Tuning supervisado - SFT o RLHF) entrelazan el comportamiento de seguridad con las capacidades generales del modelo. Actualizar una política de seguridad a menudo degrada el rendimiento en tareas útiles o requiere reentrenar todo el modelo.
• Ineficiencia de los métodos basados en prompts: Los enfoques que usan instrucciones en lenguaje natural (prompts) son frágiles, inconsistentes y generan una sobrecarga de tokens, especialmente cuando se combinan múltiples restricciones.
• Sobre-rechazo (Over-refusal): Los modelos alineados tienden a rechazar consultas benignas cuando se activan restricciones de seguridad, perdiendo utilidad.

El objetivo es lograr un control de seguridad condicional y composicional, donde diferentes subconjuntos de reglas se habiliten dinámicamente sin alterar el núcleo del modelo.

2. Metodología: MOSAIC

Los autores proponen MOSAIC (Composable Safety Alignment with Modular Control Tokens), un marco que trata la alineación de seguridad como un problema de aprendizaje de representaciones modulares.

Concepto Central:
En lugar de modificar los parámetros del modelo base (que permanece congelado), cada restricción de seguridad se representa mediante un pequeño conjunto de tokens de control aprendibles en el espacio de incrustación (embedding space).

Componentes Clave:

• Tokens de Control Modulares: Cada categoría de seguridad (ej. apuestas, alcohol, adicción) se asocia a un conjunto de $m$ tokens ($z_c$). Estos son los únicos parámetros entrenables.
• Activación Condicional: En la inferencia, se prefienden los tokens correspondientes a las categorías activas para el usuario específico. Por ejemplo, para un menor, se activan tokens de "alcohol" y "apuestas"; para un adulto, quizás solo "apuestas".
• Muestreo de Tareas Combinatorias (Compositional Task Sampling):
  • Entrenar para todas las combinaciones posibles de categorías es exponencialmente costoso ($2^K - 1$).
  • MOSAIC introduce una estrategia de muestreo basado en orden. Se asigna un presupuesto de entrenamiento fijo por "orden" (número de categorías activas simultáneamente), en lugar de por cada subconjunto específico. Esto permite exponer al modelo a composiciones diversas sin crecimiento exponencial de datos.
• Objetivo de Entrenamiento y Destilación de Conocimiento Contrafactual:
  • Muestras Positivas: Se usa pérdida de entropía cruzada para forzar la respuesta de rechazo cuando los tokens de seguridad están activos.
  • Destilación de Conocimiento Contrafactual (Counterfactual KD): Para mitigar el sobre-rechazo, se compara la distribución de salida del modelo con los tokens de control frente a la distribución del modelo base sin ellos (en consultas benignas). Se minimiza la divergencia KL entre ambas. Esto obliga a los tokens a intervenir solo cuando es estrictamente necesario, preservando el comportamiento original del modelo en consultas seguras.

3. Contribuciones Clave

1. Reconceptualización: Plantear la alineación de seguridad como un problema de aprendizaje de representaciones composicionales, donde el control es modular y desvinculado del modelo base.
2. Marco MOSAIC: Un sistema que utiliza tokens de control aprendibles sobre un modelo congelado, permitiendo la expansión incremental de categorías y la activación flexible de políticas sin reentrenamiento.
3. Técnicas de Optimización: Introducción de muestreo de tareas basado en orden y un objetivo de destilación contrafactual para reducir el sobre-rechazo y mejorar la precisión en la activación de restricciones.
4. Nuevo Benchmark: Construcción de un conjunto de datos de evaluación realista (3,000 solicitudes) diseñado específicamente para probar la activación condicional en modelos ya alineados, evitando la familiaridad con datos de entrenamiento previos.

4. Resultados Experimentales

Los experimentos se realizaron en modelos Llama-3.1-8B y Llama-3.2-3B, comparando MOSAIC con In-context prompting, ORPO y SFT.

• Tasa de Éxito de Defensa (DSR): MOSAIC logra tasas de rechazo correcto superiores al 99% en todas las configuraciones (1, 2 y 3 categorías activas), superando o igualando al SFT.
• Tasa de Sobre-rechazo (OR): MOSAIC reduce drásticamente el rechazo de consultas benignas.
  • Mientras que el SFT mantiene un OR alrededor del 6%, MOSAIC logra reducirlo a ~1.8% en tareas de alto orden (3 categorías) en el modelo de 8B.
  • Esto demuestra que la destilación contrafactual y el entrenamiento composicional permiten una frontera de rechazo más precisa.
• Utilidad del Modelo: La evaluación en MMLU muestra que MOSAIC preserva las capacidades generales de lenguaje del modelo base con una degradación insignificante, a diferencia de los métodos de ajuste de parámetros que pueden causar interferencia catastrófica.
• Escalabilidad e Incremento: El método permite agregar nuevas categorías de seguridad (ej. +1 o +2 categorías) sin reentrenar las existentes, manteniendo un rendimiento estable y una degradación mínima.
• Eficiencia: Requiere muy pocos parámetros entrenables (ej. 2 tokens por categoría), lo que es significativamente más eficiente que el ajuste fino completo (SFT) o el uso de LoRA masivo.

5. Significado e Impacto

MOSAIC representa un avance significativo hacia la alineación pluralista y contextual de los LLMs:

• Flexibilidad Operativa: Permite a los desarrolladores adaptar las políticas de seguridad en tiempo de inferencia según el contexto (país, edad, aplicación) sin tocar el modelo base.
• Eficiencia de Recursos: Elimina la necesidad de reentrenar modelos costosos para cada nueva regulación o requisito de seguridad.
• Equilibrio Seguridad-Utilidad: Resuelve el dilema tradicional entre seguridad estricta y utilidad, logrando un rechazo preciso de contenido dañino sin sacrificar la capacidad de respuesta en consultas legítimas.
• Paradigma de Control: Cambia el enfoque de "políticas monolíticas" a "controles modulares", ofreciendo una solución escalable para la gestión dinámica de riesgos en IA.

En conclusión, MOSAIC demuestra que es posible lograr una seguridad condicional robusta y adaptable mediante el uso de tokens de control aprendibles y técnicas de entrenamiento composicional, superando las limitaciones de los enfoques actuales basados en parámetros o prompts.