SafeMERGE: Preserving Safety Alignment in Fine-Tuned Large Language Models via Selective Layer-Wise Model Merging

Resumen Técnico: SafeMERGE

1. El Problema

El ajuste fino (fine-tuning) de Modelos de Lenguaje Grandes (LLMs) es una práctica estándar para adaptar modelos generales a dominios especializados (como matemáticas, código o biomedicina). Sin embargo, estudios recientes han demostrado que este proceso puede erosionar la alineación de seguridad del modelo.

• Fragilidad de la Seguridad: Incluso el ajuste fino con datos benignos puede degradar inadvertidamente la capacidad del modelo para rechazar solicitudes dañinas o poco éticas.
• Limitaciones de las Defensas Actuales: Las soluciones existentes se dividen en tres categorías (durante la alineación inicial, durante el ajuste fino o post-ajuste). Sin embargo, muchas requieren algoritmos de entrenamiento personalizados complejos, son difíciles de implementar en bibliotecas de código abierto estándar, o sacrifican el rendimiento en la tarea principal (utility) para recuperar la seguridad.

2. Metodología: SafeMERGE

Los autores proponen SafeMERGE, un marco ligero y posterior al ajuste fino (post-fine-tuning) diseñado para restaurar la seguridad sin reentrenar el modelo ni comprometer su utilidad.

Mecanismo Central:
SafeMERGE no fusiona todo el modelo, sino que realiza una fusión selectiva a nivel de capas (específicamente en adaptadores LoRA). El proceso sigue estos pasos:

1. Construcción de Modelos:

   • Se entrena un modelo ajustado para la tarea específica (task-specific).
   • Se entrena un modelo "seguro" (safe model) utilizando datos de alineación de seguridad (pares de prompts dañinos y respuestas seguras). Este modelo es agnóstico a la tarea y se reutiliza.

2. Detección de Capas Inseguras (Criterio de Similitud Coseno):

   • Se define un subespacio de alineación de seguridad ($V_i$) basado en la diferencia de pesos entre un modelo base y un modelo alineado (instruct/chat).
   • Para cada capa $i$ del modelo ajustado, se calcula la similitud del coseno ($\rho_i$) entre la actualización LoRA del modelo ajustado ($\Delta W^f_i$) y su proyección sobre el subespacio de seguridad ($C_i \Delta W^f_i$).
   • Umbral de Seguridad ($\tau$): Si $\rho_i < \tau$, la capa se considera que se ha desviado del comportamiento seguro y es candidata a fusión.

3. Fusión Selectiva:

   • Solo las capas que fallan el criterio de similitud se fusionan con las capas correspondientes del modelo seguro.
   • La fusión se realiza mediante una combinación lineal: $\Delta W^m_i = \alpha \Delta W^f_i + (1 - \alpha) \Delta W^s_i$.
   • Las capas que mantienen una alta similitud con el subespacio de seguridad se dejan intactas para preservar la utilidad de la tarea.

Ventajas Técnicas:

• Sin Reentrenamiento: No requiere retropropagación ni bucles de entrenamiento adicionales.
• Eficiencia Computacional: Opera solo sobre los parámetros LoRA (aprox. 0.5-1% de los parámetros totales), permitiendo ejecutarse completamente en CPU.
• Interpretabilidad: El umbral $\tau$ tiene una interpretación geométrica clara (ángulo de desviación respecto a la dirección de seguridad).

3. Contribuciones Clave

• Enfoque Selectivo: A diferencia de métodos que proyectan globalmente todas las actualizaciones (como SafeLoRA) o fusionan indiscriminadamente todo el modelo (como RESTA), SafeMERGE fusiona solo las capas degradadas, preservando el conocimiento útil en las capas seguras.
• Marco Ligero y Estándar: Utiliza bibliotecas de código abierto estándar y no requiere algoritmos de entrenamiento personalizados, facilitando su adopción por parte de la industria.
• Validación Rigurosa: Se evalúa en cuatro modelos LLM diferentes (Llama-2, Llama-3.1, Qwen-2, Qwen-2.5) y múltiples tareas (GSM8K, PubMedQA, y tareas de telecomunicaciones).

4. Resultados Experimentales

Los experimentos demuestran que SafeMERGE logra un equilibrio superior entre seguridad y utilidad en comparación con los baselines (SafeInstruct, RESTA, RESTA-Instruct, SafeLoRA).

• Reducción de Daños: SafeMERGE reduce consistentemente las salidas dañinas (medidas en benchmarks DirectHarm y HexPhi) a niveles cercanos o inferiores a los del modelo original alineado, superando a otros métodos de defensa.
  • Ejemplo: En Llama-3.1 ajustado en GSM8K, redujo la toxicidad de DirectHarm del 28.30% (modelo ajustado) al 8.80%, mientras que otros métodos oscilaron entre 11% y 15%.
• Preservación de Utilidad: A diferencia de RESTA, que a menudo sufre de "olvido catastrófico" (pérdida de rendimiento en la tarea), SafeMERGE mantiene o incluso mejora la precisión en la tarea objetivo.
  • Ejemplo: En Llama-3.1 (GSM8K), SafeMERGE alcanzó una precisión del 78.50%, superando al modelo ajustado original (78.24%).
• Análisis de Ablación:
  • Similitud Coseno vs. Distancia Euclidiana: El uso de similitud coseno es crucial. La distancia euclidiana (basada en magnitud) penaliza actualizaciones grandes pero seguras, degradando la utilidad drásticamente. La similitud coseno (basada en dirección) preserva la utilidad mientras corrige la dirección insegura.
  • Umbral Óptimo: Se identificó que un umbral $\tau$ entre 0.6 y 0.8 ofrece el mejor equilibrio, fusionando típicamente entre 28 y 34 capas LoRA (de un total de ~50-60).

5. Significado e Impacto

SafeMERGE representa un avance significativo en la seguridad de LLMs por las siguientes razones:

1. Solución Práctica: Resuelve el dilema de "seguridad vs. utilidad" sin requerir infraestructura de entrenamiento compleja, haciendo viable la protección de modelos ajustados en entornos de producción.
2. Robustez: Demuestra que la alineación de seguridad no es monolítica; es posible identificar y reparar selectivamente las partes del modelo que han "roto" la seguridad sin tocar las partes que funcionan bien.
3. Generalización: Su efectividad en dominios altamente especializados (como telecomunicaciones) sugiere que es una solución escalable para la adaptación de modelos en industrias críticas.

En conclusión, SafeMERGE establece un nuevo estándar para las defensas post-ajuste fino, ofreciendo un mecanismo simple, interpretable y altamente efectivo para garantizar que los LLMs especializados permanezcan seguros y alineados con los valores humanos.