RASA: Routing-Aware Safety Alignment for Mixture-of-Experts Models

Resumen Técnico: RASA (Routing-Aware Safety Alignment)

1. El Problema: Atajos de Alineación en Modelos MoE

Los modelos de lenguaje basados en Mezcla de Expertos (MoE) presentan desafíos únicos para la alineación de seguridad debido a sus mecanismos de enrutamiento dispersos. Los autores identifican un fallo crítico en los enfoques tradicionales de fine-tuning (ajuste fino) de parámetros completos:

• Atajos de Alineación (Alignment Shortcuts): Al aplicar un ajuste fino estándar de parámetros completos a modelos MoE, la mejora en la seguridad a menudo no proviene de reparar los expertos que generan comportamientos inseguros. En su lugar, el modelo aprende a "evitar" el problema modificando el enrutador (router) o haciendo que los expertos ya seguros dominen la respuesta.
• Consecuencias:
  • Los Expertos Críticos de Seguridad (Safety-Critical Experts - SCEs), que son los responsables reales de generar contenido dañino, permanecen sin corregir.
  • Estos expertos pueden ser reactivados fácilmente por ataques de jailbreak adaptativos o no vistos, rompiendo la seguridad.
  • Los expertos restantes, al volverse excesivamente dominantes para compensar, provocan una sobre-denegación (over-refusal), degradando la utilidad general del modelo.
• Evidencia: Experimentos preliminares muestran que si se restaura el patrón de enrutamiento original en un modelo ajustado completamente, la seguridad colapsa, demostrando que la protección era frágil y dependía del enrutamiento, no de una reparación real de los parámetros inseguros.

2. Metodología: RASA (Routing-Aware Safety Alignment)

Para abordar esto, los autores proponen RASA, un marco de alineación a nivel de experto que es consciente del enrutamiento. El objetivo es reparar explícitamente los SCEs sin depender de atajos de enrutamiento.

El método se formula como un problema de optimización de dos niveles (bi-level optimization) y se resuelve mediante una estrategia de optimización alterna:

• Paso 1: Identificación de Expertos Críticos (Safety-Critical Experts):

  • Se construyen lotes mixtos con pares de prompts: Ancla (maliciosos pero rechazados por el modelo) y Adversarios (variantes de jailbreak que logran eludir la seguridad).
  • Se calcula la Discrepancia de Activación Adversaria (AAD) para cada experto: la diferencia en la frecuencia de activación entre los prompts adversarios y los de ancla.
  • Los expertos con una AAD alta se identifican como SCEs (son los que activan desproporcionadamente ante ataques exitosos).

• Paso 2: Optimización Interna (Ajuste Fino Selectivo de Expertos - SCE-FT):

  • El enrutador se congela.
  • Solo se actualizan los parámetros de los SCEs identificados.
  • Se entrena a estos expertos para generar respuestas de rechazo (refusal) ante los prompts adversarios, utilizando una función de pérdida de entropía cruzada con respuestas de referencia diversas.
  • Los expertos benignos permanecen congelados para preservar la capacidad general.

• Paso 3: Optimización Externa (Consistencia del Enrutador):

  • Los parámetros de los expertos se congelan (ya reparados).
  • Se optimiza el enrutador para que la distribución de enrutamiento de los inputs adversarios sea consistente con la de los inputs seguros (ancla).
  • Se utiliza una pérdida de divergencia KL (Forward KL) para asegurar que los inputs maliciosos sigan el mismo camino de expertos seguros, evitando que el enrutador aprenda a desviar el tráfico para evadir la seguridad.

3. Contribuciones Clave

1. Identificación de un Fallo de Diseño: Demostración empírica y causal de que el ajuste fino de parámetros completos en MoE a menudo satisface objetivos de seguridad mediante atajos de enrutamiento en lugar de reparar los expertos defectuosos.
2. Propuesta de RASA: Un marco de alineación a nivel de experto que desacopla la corrección de expertos del control del enrutador, reparando explícitamente los SCEs y luego estabilizando el enrutamiento.
3. Eficacia y Eficiencia: El método logra una defensa casi perfecta contra ataques de jailbreak, generaliza bien a ataques no vistos, mantiene las capacidades generales del modelo y controla la tasa de sobre-denegación, todo ello con una fracción de los datos de entrenamiento necesarios.

4. Resultados Experimentales

Los autores evaluaron RASA en dos arquitecturas MoE representativas (Qwen3-30B-A3B y OLMoE-1B-7B-0125) frente a una variedad de ataques de jailbreak (FlipAttack, DeepInception, Persuasion, etc.).

• Robustez: RASA alcanzó tasas de inocuidad (Harmlessness) cercanas al 100% en ataques alineados y mostró una fuerte generalización cruzada a ataques no vistos, superando significativamente a modelos ajustados completamente y a métodos baselines como SteerMoE o SafeMoE-Fine-tuning.
• Preservación de Capacidades: A diferencia del ajuste fino completo, que a menudo degrada el rendimiento en benchmarks generales (MMLU, GSM8K, TruthfulQA), RASA mantuvo o incluso mejoró ligeramente estas métricas.
• Control de Sobre-Denegación: El método redujo drásticamente la tasa de rechazo excesivo en consultas legítimas (medido en XsTest), un problema común en los enfoques de seguridad agresivos.
• Validación Causal: Los experimentos de restauración de enrutamiento confirmaron que la seguridad de RASA es robusta (no colapsa al restaurar el enrutamiento original), mientras que la de los métodos de parámetros completos sí colapsa, validando que RASA repara realmente los expertos.
• Eficiencia de Datos: Se demostró que es suficiente con un 25-50% de las muestras de ataque disponibles para lograr una alineación efectiva.

5. Significado e Impacto

Este trabajo es fundamental para el futuro de la seguridad en modelos de lenguaje escalables (MoE):

• Cambio de Paradigma: Sugiere que la alineación de seguridad en arquitecturas dispersas no debe tratarse como un problema global de parámetros, sino como un problema de reparación localizada de componentes específicos (expertos).
• Prevención de Vulnerabilidades Latentes: Al corregir directamente los SCEs en lugar de depender de atajos de enrutamiento, RASA elimina vulnerabilidades que podrían ser explotadas por adversarios adaptativos en el futuro.
• Arquitectura Preservadora: RASA no requiere modificar la arquitectura del modelo, añadir expertos auxiliares ni realizar intervenciones en tiempo de inferencia, lo que lo hace una solución práctica y escalable para la industria.
• Dirección Futura: Establece una nueva dirección para la alineación consciente de la arquitectura, donde la comprensión de la dinámica interna (enrutamiento y especialización de expertos) es crucial para lograr defensas robustas y generalizables.

En resumen, RASA demuestra que la verdadera seguridad en modelos MoE requiere ir más allá de la superficie del enrutamiento y reparar quirúrgicamente los componentes internos responsables del comportamiento inseguro.