Resumen Técnico: HiRoute

Declaración del Problema

Los Modelos de Lenguaje de Gran Escala (LLMs) siguen siendo vulnerables a peticiones dañinas, ataques de jailbreak y consultas de múltiples riesgos. Aunque el Ajuste Fino Supervisado (SFT) y el Aprendizaje por Refuerzo a partir de la Retroalimentación Humana (RLHF) pueden codificar comportamientos de seguridad, incurren en altos costes de entrenamiento y almacenamiento y, a menudo, provocan un compromiso entre la seguridad y la utilidad. Específicamente, el fortalecimiento uniforme de los comportamientos de rechazo puede causar un "sobre-rechazo" (over-refusal) en entradas benignas, degradando las capacidades generales.

Los métodos actuales de alineación de seguridad con eficiencia de parámetros, como el ajuste de prompts (prompt tuning), suelen depender de mecanismos estáticos. Estos enfoques generalmente se dividen en dos categorías con limitaciones complementarias:

1. Prompts Globales de Grano Grueso: Un único prompt de seguridad aplicado a todas las entradas riesgosas establece un límite de seguridad sólido, pero a menudo resulta en rechazos genéricos y poco informativos que no proporcionan orientación específica al riesgo o alternativas seguras.
2. Prompts Dirigidos de Grano Fino: Los métodos que asignan prompts específicos de categoría a diferentes riesgos pueden generar respuestas seguras más dirigidas y útiles. Sin embargo, luchan por mantener un límite de seguridad estable, particularmente cuando las entradas no se alinean precisamente con las categorías de riesgo predefinidas o cuando el emparejamiento de expertos es insuficiente.

El desafío central es cómo equilibrar la constructividad de la respuesta de grano fino con la seguridad trans-categoría mediante la composición de prompts adaptativa al input sin congelar todo el modelo base (backbone).

Metodología: HiRoute

Los autores proponen HiRoute (Hierarchical Routed Prompt Tuning), un marco de trabajo adaptativo al input que separa el control de seguridad agnóstico a la categoría de la guía de respuesta específica de la categoría. El marco opera sobre un LLM ajustado por instrucciones que permanece congelado e involucra tres componentes principales:

1. Enrutamiento de Riesgo Jerárquico

Un enrutador ligero se entrena sobre representaciones extraídas del LLM congelado para realizar dos tareas:

• Detección de Grano Grueso: Determina si una entrada es dañina o benigna ($p_{safe}$ vs. $p_{risk}$).
• Predicción de Riesgo de Grano Fino: Predice puntuaciones de riesgo multietiqueta para categorías específicas (por ejemplo, cibercrimen, violaciones de privacidad, delitos económicos, violencia).
  El enrutador genera una distribución de probabilidad utilizada para ponderar dinámicamente los expertos de prompts de grano fino.

2. Composición de Prompts Jerárquica

HiRoute aprende dos tipos de embeddings de prompt continuos:

• Prompt de Grano Grueso Compartido ($P_c$): Proporciona un límite de seguridad agnóstico a la categoría aplicable a todas las entradas riesgosas.
• Expertos de Prompt de Grano Fino ($\{P_k\}$): Un conjunto de $K$ expertos, cada uno correspondiente a una categoría de riesgo específica.
  Para un input dado $x$, el prompt de grano fino $P_f(x)$ se construye como una suma ponderada de los expertos basada en las puntuaciones de riesgo predichas por el enrutador:
  $$P_f(x) = \sum_{k=1}^{K} \alpha_k(x) P_k$$
  donde $\alpha_k(x)$ son pesos normalizados. El prompt final dependiente del input es $P(x) = [P_f(x); P_c]$.

3. Optimización Desacoplada en Dos Etapas

El entrenamiento procede en dos etapas distintas para evitar que el objetivo de generación altere la frontera de decisión de riesgo:

• Etapa I (Aprendizaje del Enrutador): El enrutador se entrena utilizando Entropía Cruzada (CE) para la clasificación de grano grueso y Entropía Cruzada Binaria (BCE) para la clasificación multietiqueta de grano fino. El modelo base permanece congelado.
• Etapa II (Optimización del Prompt): El modelo base y el enrutador entrenado se congelan. Solo se optimizan los parámetros del prompt ($P_c$ y $\{P_k\}$) utilizando Optimización de Preferencia Directa (DPO) sobre tripletes de preferencia (seguro/útil vs. inseguro/baja calidad). Para reducir la interferencia, los autores emplean actualizaciones de gradiente alternas: cuando se actualiza $P_c$, los gradientes para $P_f(x)$ se enmascaran (stop-gradient), y viceversa.

4. Inferencia Adaptativa al Input

En la inferencia, el enrutador primero computa la probabilidad de seguridad $p_{safe}(x)$.

• Inputs Benignos: Si $p_{safe}(x) > \tau_s$ (un umbral de seguridad), el sistema evita todos los prompts de seguridad y genera directamente desde el backbone congelado, preservando la utilidad general.
• Inputs Riesgosos: Si $p_{safe}(x) \leq \tau_s$, el sistema activa la rama de seguridad, construyendo el prompt jerárquico $P(x)$ para condicionar la generación.

Contribuciones Clave

1. Validación Empírica de Compromisos: Los autores demuestran que los prompts compartidos de grano grueso ofrecen una seguridad más fuerte pero menor utilidad, mientras que las mezclas de prompts dirigidos de grano fino mejoran la utilidad pero producen menores ganancias en seguridad.
2. Marco de Trabajo HiRoute: Proponen una arquitectura novedosa que combina un prompt de grano grueso compartido (para un límite de seguridad estable) con expertos dirigidos de grano fino (para una guía específica al riesgo). Este diseño evita la intervención innecesaria en inputs benignos mediante una puerta de seguridad.
3. Evaluación Exhaustiva: El método se valida en tres modelos ajustados por instrucciones (Mistral-7B, Vicuna-7B, Zephyr-7B) y múltiples benchmarks, mostrando que mantiene una alta seguridad y utilidad mientras preserva la utilidad general.

Resultados Experimentales

Los experimentos se realizaron en tres modelos contra líneas base que incluyen modelos Base, RPO (Robust Prompt Optimization), DRO (Dynamic Risk Optimization) y ACD (Adversarial Contrastive Decoding).

• Seguridad y Utilidad: HiRoute logró las tasas de seguridad promedio más altas en todos los modelos (por ejemplo, 93.2% en Mistral, 97.7% en Vicina, 94.8% en Zephyr) mientras alcanzaba simultáneamente las puntuaciones de utilidad de respuesta segura más altas o cercanas a las más altas. Superó a las líneas base más fuertes por 1.3–2.9 puntos porcentuales en seguridad.
• Utilidad General: A diferencia de los métodos que degradan las capacidades generales, HiRoute preservó el rendimiento en GSM8K (razonamiento matemático), MT-Bench (seguimiento de instrucciones) y TruthfulQA. Por ejemplo, en Mistral, aumentó la seguridad del 53.2% al 93.2% manteniendo un 48.0% de precisión en GSM8K.
• Sobre-rechazo: HiRoute redujo significativamente el sobre-rechazo en inputs benignos (XSTest). En Mistral, logró una tasa de sobre-rechazo del 2.5%, una mejora masiva respecto a RPO (40.5%) y DRO (18.0%).
• Estudios de Ablación:
  • Estrategia de Entrenamiento: El entrenamiento alterno conjunto (actualizando los prompts en un contexto compuesto consistente) superó significativamente al entrenamiento secuencial ("Coarse después de fine"), mejorando la seguridad promedio del 74.8% al 93.2%.
  • Longitud de Prompt: Una división de 15/5 entre la longitud del prompt de grano grueso y de grano fino (total 20) ofreció el mejor equilibrio entre seguridad y utilidad.
  • Umbral: Aumentar el umbral de la puerta de seguridad $\tau_s$ fortaleció el control de seguridad con un impacto mínimo en la capacidad general.

Significado y Reivindicaciones

El artículo afirma que HiRoute aborda eficazmente la tensión entre la seguridad y la utilidad en la alineación de LLMs. Al desacoplar la identificación del riesgo de la generación de respuestas y emplear una estructura de prompt jerárquica y adaptativa al input, el método:

• Establece un límite de seguridad estable y agnóstico a la categoría sin depender de rechazos agresivos y uniformes.
• Proporciona una guía constructiva y específica al riesgo, así como alternativas seguras para peticiones dañinas.
• Minimiza el sobre-rechazo permitiendo que los inputs benignos eviten por completo el prompting de seguridad.
• Logra estos objetivos con eficiencia de parámetros, congelando el modelo base y optimizando solo un pequeño número de parámetros de prompt y un enrutador ligero.

Los autores concluyen que el prompting jerárquico dependiente del input ofrece una solución robusta para la alineación de seguridad que equilibra la seguridad, la calidad de la respuesta, la eficiencia de parámetros y la utilidad general. Reconocen limitaciones respecto a la precisión del enrutador en riesgos compuestos fuera de la distribución y el enfoque actual en interacciones de texto de un solo turno, sugiriendo trabajos futuros en el reconocimiento de riesgos de conjunto abierto y entornos multimodales.