Resumen Técnico: Destilación de Sistemas LLM Seguros mediante Soft Prompts para Entornos en Dispositivos

Planteamiento del Problema

El despliegue de Modelos de Lenguaje de Gran Escala (LLMs) en dispositivos de borde con recursos limitados (por ejemplo, teléfonos inteligentes) presenta un desafío crítico de seguridad. Si bien los sistemas de modelo dual —que combinan un LLM base con un "modelo de guardia" separado para clasificar y filtrar salidas inseguras— proporcionan garantías de seguridad robustas, incurren en costos prohibitivos de memoria y computación. La naturaleza secuencial de este enfoque (generar una respuesta, luego clasificarla) duplica las pasadas hacia adelante (forward passes), lo que aumenta significamente la latencia (tiempo hasta el primer token) y la huella de memoria, resultando inadecuado para el despliegue en dispositivos de borde. Las alternativas existentes, como la cuantización o el ajuste fino directo, a menudo comprometen las capacidades de generalización o no logran equilibrar la seguridad con la utilidad del modelo de manera efectiva.

Metodología

El artículo propone TV-DiSP (Destilación de Variación Total mediante Soft Prompts), un enfoque de ajuste fino de parámetros eficientes (PEFT) que destila los comportamientos de seguridad de un sistema de modelo dual en un único LLM equipado con soft prompts aprendibles.

Marco Central

1. Formulación del Sistema: Los autores definen un sistema de LLM seguro $p(r|x)$ que combina un LLM base $p(y|x)$ y un modelo de guardia $p(s|x, y)$. Si el modelo de guardia considera que una respuesta $y$ es insegura ($s=0$), el sistema emite un mensaje de rechazo predefinido $y_r$; de lo contrario, emite $y$.
2. Objetivo de Destilación: El objetivo es aprender un conjunto de incrustaciones de prompts continuos $W$ (soft prompts) tales que el modelo destilado $q(r|x, W)$ imite la distribución de salida del sistema seguro $p(r|x)$.
3. Optimización mediante Variación Total (TV):
   • Los autores derivan una garantía teórica (Teorema 3.1) que muestra que minimizar la distancia de Variación Total entre el sistema seguro y el modelo destilado acota la brecha de rendimiento en cualquier tarea descendente.
   • El objetivo de optimización (Ecuación 4) maximiza la log-probabilidad de la salida correcta (ya sea la respuesta del LLM base para entradas seguras o el mensaje de rechazo para entradas inseguras) ponderada por la probabilidad de seguridad del modelo de guardia.
   • El entrenamiento requiere únicamente un conjunto de datos de prompts; el modelo de guardia etiqueta dinámicamente estos como seguros o inseguros durante el proceso.
4. Esquemas Alternativos: El estudio también evalúa:
   • KL-DiSP: Minimización de la divergencia de Kullback-Leibler, que ofrece garantías teóricas más laxas que la TV.
   • Baselines: Optimización de perplejidad, REINFORCE (gradiente de política) y adaptaciones del marco de destilación a Low-Rank Adapters (LoRA) y Vectores de Dirección (Steering Vectors, SV).

Configuración Experimental

• Modelos: Los experimentos se realizaron en cuatro LLMs ajustados por instrucciones (Llama3-1B/3B, Qwen2-1.5B, Gemma2-2B) cuantizados a 4 bits para simular restricciones de dispositivos de borde.
• Modelos de Guardia: Se utilizó LlamaGuard3-1B para el entrenamiento de la destilación, y LlamaGuard3-8B para la evaluación.
• Conjuntos de Datos: El entrenamiento utilizó los conjuntos de datos Beavertails y Toxigen. La evaluación incluyó benchmarks fuera de la distribución (HarmBench, Detect-JailBreak) y benchmarks de utilidad (IFEval, GSM8K, MMLU).
• Hardware: Las mediciones en dispositivos se realizaron en teléfonos inteligentes equipados con chips Qualcomm Snapdragon 8 Elite Gen 5.

Resultados Clave

1. Compromiso Seguridad-Computación

• Rendimiento: TV-DiSP destila con éxito la seguridad del sistema de modelo dual en un solo modelo. Por ejemplo, en el modelo Llama3-1B, la seguridad (medida por el Safety Guard Score, SGS) mejoró del 71% (base) al 99% (sistema dual), mientras que TV-DiSP logró ganancias de seguridad comparables con un overhead mínimo.
• Eficiencia: En comparación con el sistema de modelo dual, que duplica los requisitos de memoria y computación, TV-DiSP introduce menos del 1% de overhead de memoria adicional y menos del 10% de overhead de computación adicional (FLOPs por token).
• Generalización: El método demuestra una fuerte robustez ante prompts adversarios fuera de la distribución (HarmBench, Detect-JailBreak) a pesar de haber sido entrenado con conjuntos de datos no adversarios.

2. Comparación con Métodos PEFT

• Soft Prompts vs. LoRA/SV: TV-DiSP (soft prompts) superó consistentemente a TV-DiLoRA y TV-DiSV (Steering Vectors) a través de varias tasas de aprendizaje. Los vectores de dirección mostraron resultados mixtos, degradando a veces el rendimiento en modelos más grandes, mientras que LoRA se consideró demasiado intrusivo para entornos de borde.
• Objetivos de Optimización:
  • TV vs. KL: Tanto la destilación TV como la KL lograron puntuaciones de seguridad altas. Sin embargo, TV-DiSP ofreció un equilibrio superior entre seguridad y utilidad.
  • Perplejidad y REINFORCE: La optimización de la perplejidad condujo al sobreajuste y caídas significativas en la utilidad. REINFORCE mantuvo la utilidad pero no proporcionó mejoras sustanciales en la seguridad.

3. Preservación de la Utilidad

• Las evaluaciones en dispositivos (IFEval y GSM8K) mostraron que TV-DiSP y KL-DiSP mantienen una alta utilidad.
• Sobre-rechazo (Over-refusal): Un hallazgo crítico fue que KL-DiSP exhibió una alta tasa de sobre-rechazo (84.6% en prompts seguros), rechazando incorrectamente consultas benignas. En contraste, TV-DiSP mantuvo una tasa de sobre-rechazo (36.1%) comparable a la del sistema de modelo dual, preservando la capacidad del modelo para responder preguntas seguras.

Significación y Reivindicaciones

El artículo establece la destilación de soft prompts como el enfoque preferido para el alineamiento de seguridad en el despliegue de LLMs en dispositivos. Sus principales contribuciones son:

1. Compromisos Superiores: Demuestra que los soft prompts entrenados mediante destilación superan consistentemente a LoRA, los vectores de dirección y otros objetivos de entrenamiento en el equilibrio entre seguridad, utilidad y eficiencia computacional.
2. Garantías Teóricas: La introducción de la destilación basada en Variación Total proporciona garantías probabilísticas sobre el rendimiento descendente, asegurando que el modelo destilado no se desvíe significativamente del comportamiento del sistema seguro.
3. Viabilidad Práctica: Validado a través de múltiples arquitecturas y en hardware real de borde, el método logra niveles de seguridad comparables a los sistemas de modelo dual con un overhead de recursos insignificante (<1% de memoria, <10% de computación), lo que lo convierte en una solución práctica para desplegar LLMs seguros en dispositivos con recursos limitados.

Los autores concluyen que, si bien ningún método ofrece inmunidad total contra ataques adversarios, TV-DiSP ofrece un mecanismo de defensa robusto, eficiente y generalizable que reduce significativamente la superficie de ataque y la carga computacional en comparación con las arquitecturas tradicionales de modelos de guardia.