MASCing: Configurable Mixture-of-Experts Behavior via Activation Steering Masks

1. Declaración del Problema

Los Modelos de Lenguaje Grandes (LLMs) que utilizan arquitecturas de Mezcla de Expertos (MoE) se han vuelto dominantes debido a su capacidad para escalar la cantidad de parámetros manteniendo bajos costos de inferencia mediante la activación dispersa. Sin embargo, este paradigma de activación dispersa introduce desafíos únicos de seguridad:

• Comportamiento Acoplado: El comportamiento del modelo está estrechamente acoplado a las decisiones de enrutamiento dinámicas. Dado que solo un subconjunto de "expertos" se activa por token, los mecanismos de seguridad pueden ser eludidos si un adversario manipula el enrutador para evitar expertos alineados con la seguridad.
• Inflexibilidad de las Defensas Actuales: La alineación de seguridad tradicional depende del ajuste fino completo o del reentrenamiento. Estos métodos son computacionalmente prohibitivos para los grandes modelos MoE y carecen de la agilidad para adaptarse a políticas de seguridad que cambian rápidamente (por ejemplo, nuevas estrategias de jailbreak o directrices de contenido en evolución).
• La Brecha: Existe una falta de mecanismos ligeros y sin entrenamiento para reconfigurar flexiblemente el comportamiento de los MoE para escenarios de seguridad específicos (por ejemplo, endurecer las defensas contra jailbreaks o relajar restricciones para dominios específicos) sin reentrenar los pesos del modelo.

2. Metodología: El Marco MASCing

MASCing (Configuración de Dirección de Activación de MoE) es un marco en tres fases, sin entrenamiento, que interviene a nivel de puerta de enrutamiento para dirigir el comportamiento del modelo. Evita modificar los pesos del modelo en su totalidad.

Fase 1: Modelado Secuencial del Comportamiento

• Modelo Sustituto: En lugar de optimizar sobre selecciones discretas de los $k$ mejores expertos (que son no diferenciables), MASCing entrena un modelo sustituto LSTM ligero sobre logits de enrutamiento continuos y no normalizados.
• Entrada: El LSTM procesa los logits de enrutamiento de todos los expertos a través de todas las capas para una secuencia de tokens.
• Objetivo: El LSTM aprende a mapear estos patrones de enrutamiento a un comportamiento objetivo (por ejemplo, "rechazo" vs. "cumplimiento"). Al utilizar logits continuos, el modelo preserva información detallada sobre los expertos que no son seleccionados pero que aún influyen en la distribución de enrutamiento.
• Entrenamiento: El LSTM se entrena para clasificar si una secuencia de logits conduce a un comportamiento específico (clasificación binaria).

Fase 2: Creación de la Máscara de Dirección

• Identificación del Circuito: Utilizando el LSTM entrenado como un proxy diferenciable, el marco optimiza una matriz de dirección ($S$) para identificar qué expertos son responsables del comportamiento objetivo.
• Optimización: La matriz se optimiza para minimizar una función de pérdida compuesta que consiste en:
  1. Entropía Cruzada Binaria (BCE): Para impulsar la predicción del sustituto hacia la clase objetivo.
  2. Regularización $L_1$: Para imponer dispersión, forzando los pesos no importantes a cero.
• Escalado Adaptativo: Para manejar la variación entre capas, los valores de dirección se escalan por la desviación estándar de los logits no dirigidos en cada capa.
• Poda: Se aplica una puerta de magnitud simétrica (umbral $\tau$) a la matriz optimizada para crear una máscara de dirección dispersa final. Esta máscara identifica expertos específicos para potenciar (valores positivos) o suprimir (valores negativos).

Fase 3: Aplicación de la Máscara de Dirección

• Intervención en Inferencia: Durante la inferencia, la máscara de dirección estática se aplica directamente a los logits de enrutamiento del modelo MoE objetivo mediante ganchos pre-forward.
• Mecanismo: La intervención añade un sesgo escalado a los logits: $g'_{l,t} = g_{l,t} + \alpha(\sigma_l \cdot \hat{S}_l)$.
• Resultado: Esto desplaza la distribución de probabilidad del enrutador, anulando efectivamente la selección de los $k$ mejores para activar o desactivar circuitos de expertos específicos. El modelo mantiene su utilidad general de lenguaje porque la intervención es dispersa y dirigida.

3. Contribuciones Clave

1. Marco Novel: Introducción de MASCing, el primer marco que permite la reconfiguración flexible y bidireccional del comportamiento de seguridad de MoE (tanto endurecimiento como relajación) sin reentrenamiento.
2. Análisis de Logits Continuos: Un cambio metodológico de analizar selecciones discretas de los $k$ mejores a analizar logits de enrutamiento continuos. Esto elude el cuello de botella no diferenciable y captura información de seguridad latente en expertos no seleccionados.
3. Descubrimiento de Circuitos Guiado por Sustituto: El uso de un sustituto LSTM para mapear dependencias de enrutamiento complejas y cruzadas entre capas a circuitos de comportamiento específicos, permitiendo la identificación de intervenciones mínimas y de alto impacto en expertos.
4. Flexibilidad Bidireccional: Capacidad demostrada para tanto mejorar la seguridad (defensa contra jailbreaks) como relajar restricciones (cumplimiento de políticas específicas de dominio) en la misma arquitectura de modelo.

4. Resultados Experimentales

El marco fue evaluado en siete modelos MoE de código abierto (incluyendo DeepSeek-MoE, Mixtral, Qwen y GPT-OSS) en dos escenarios contrastantes:

A. Defensa contra Jailbreak Multivuelta

• Objetivo: Fortalecer las defensas contra prompts adversarios que dirigen gradualmente al modelo hacia contenido dañino a lo largo de varias vueltas.
• Rendimiento:
  • Línea Base: Tasa promedio de éxito de defensa del 52.5%.
  • MASCing: Mejoró la tasa promedio de éxito al 83.9%.
  • Rendimiento Máximo: Hasta 89.2% (Qwen3-30B-A3B-Instruct-2507).
• Comparación: Superó significativamente al estado del arte SteerMoE (58.4% promedio), que depende de la selección discreta de expertos e ignora el contexto del historial de conversación.

B. Generación de Contenido para Adultos (Relajación de Políticas)

• Objetivo: Permitir que los modelos generen contenido para adultos en contextos apropiados (simulando un cambio de política), reduciendo el "rechazo excesivo".
• Rendimiento:
  • Línea Base: Tasa promedio de éxito de generación del 52.6%.
  • MASCing: Mejoró la tasa promedio de éxito al 82.0%.
  • Rendimiento Máximo: Hasta 93.0% (Phi-3.5-MoE-Instruct).

C. Preservación de Utilidad

• Punto de Referencia: Evaluado en MMLU (conocimiento factual) y GSM8K (razonamiento matemático).
• Resultado: MASCing causó una disminución promedio negligible en la utilidad (~4.1%). Los modelos mantuvieron su capacidad para realizar razonamiento complejo y tareas generales, demostrando que la dirección no causa olvido catastrófico.

D. Eficiencia

• Entrenamiento: El sustituto LSTM puede entrenarse en ~5 minutos en una sola GPU NVIDIA H100.
• Inferencia: Añade latencia negligible (suma elemento a elemento a los logits) y no requiere memoria adicional para los pesos del modelo.

5. Significado e Impacto

• Practicidad para Despliegue: MASCing ofrece una solución "plug-and-play" para la configuración de seguridad. Los desarrolladores pueden adaptar rápidamente los modelos a nuevas políticas o amenazas emergentes sin el costo y el tiempo del reentrenamiento completo.
• Insight Mecanístico: El artículo demuestra que la seguridad en los modelos MoE no es monolítica, sino que está distribuida a través de circuitos de expertos específicos que pueden aislarse y manipularse mediante logits de enrutamiento.
• Control Bidireccional: Cuestiona la visión binaria de la seguridad (seguro vs. inseguro) al mostrar que las restricciones de seguridad pueden ajustarse dinámicamente: endurecidas para la seguridad o relajadas para casos de uso específicos, manteniendo la integridad del modelo.
• Superioridad sobre Métodos Discretos: Los resultados confirman que aprovechar la información de enrutamiento continua (logits) es superior al análisis discreto de los $k$ mejores para comprender y controlar comportamientos complejos y multivuelta en arquitecturas dispersas.

En conclusión, MASCing establece un nuevo paradigma para la seguridad configurable en modelos MoE, proporcionando un mecanismo ligero, eficiente y altamente efectivo para gestionar la compensación entre la utilidad del modelo y la alineación de seguridad en entornos de despliegue dinámicos.