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

1. O Problema

O ajuste fino (fine-tuning) de Grandes Modelos de Linguagem (LLMs) é uma prática comum para adaptar modelos generalistas a domínios especializados (como matemática, medicina ou telecomunicações). No entanto, estudos recentes demonstram que esse processo pode erodir o alinhamento de segurança (safety alignment).

• Fragilidade da Segurança: Mesmo o ajuste fino com dados benignos pode, inadvertidamente, degradar a capacidade do modelo de recusar prompts maliciosos ou antiéticos.
• Limitações das Defesas Atuais: Métodos existentes para realinhar a segurança geralmente caem em duas categorias problemáticas:
  1. Algoritmos Complexos: Exigem treinamento personalizado, integração difícil com bibliotecas de código aberto e expertise especializada.
  2. Compromisso Utilidade-Segurança: Defesas mais simples frequentemente sacrificam o desempenho na tarefa principal (utilidade) em favor da segurança, ou não restauram a segurança completamente.

O desafio central é: Como recuperar a segurança de um modelo após o ajuste fino sem perder a utilidade da tarefa e sem depender de algoritmos de treinamento personalizados?

2. Metodologia: SafeMERGE

Os autores propõem o SafeMERGE, um framework leve e pós-ajuste fino (post-fine-tuning) que restaura a segurança através de uma fusão seletiva de camadas (selective layer-wise merging).

Princípios Fundamentais:

1. Dois Modelos Complementares:

   • Modelo Ajustado (Fine-tuned): Treinado em dados específicos da tarefa (ex: GSM8K, PubMedQA).
   • Modelo Seguro (Safe Model): Treinado em dados de alinhamento de segurança (pares de prompts prejudiciais e respostas seguras). Este modelo é agnóstico à tarefa e pode ser reutilizado.

2. Identificação de Camadas Inseguras:

   • O SafeMERGE não funde o modelo inteiro. Em vez disso, ele analisa camada por camada.
   • Utiliza um subespaço de alinhamento de segurança ($V_i$), calculado como a diferença entre os pesos de um modelo alinhado (ex: Instruct) e um modelo base não alinhado.
   • Calcula a similaridade de cosseno ($\rho_i$) entre as atualizações do LoRA do modelo ajustado e a projeção dessas atualizações no subespaço seguro.
   • Critério de Decisão: Se a similaridade de cosseno for inferior a um limiar ($\tau$), a camada é considerada "desviada" do comportamento seguro e, portanto, insegura.

3. Fusão Seletiva:

   • Apenas as camadas identificadas como inseguras são fundidas com as camadas correspondentes do modelo seguro.
   • Camadas que mantêm alta similaridade com o subespaço seguro (preservando a utilidade da tarefa) permanecem intactas.
   • A fusão é realizada via uma estratégia linear simples: $\Delta W_{merge} = \alpha \Delta W_{f} + (1 - \alpha) \Delta W_{s}$.

Vantagens Técnicas:

• Sem Retreinamento: Opera apenas nos pesos após o ajuste fino, sem necessidade de backpropagation ou loops de treinamento adicionais.
• Interpretabilidade: O limiar $\tau$ tem uma interpretação geométrica clara (ângulo de desvio em relação à direção de segurança).
• Eficiência Computacional: A complexidade é linear em relação ao número de parâmetros LoRA, podendo ser executada inteiramente na CPU.

3. Contribuições Chave

1. Framework Pós-Ajuste Fino: Introduz uma solução que não requer modificação no pipeline de treinamento existente, sendo fácil de integrar em fluxos de trabalho padrão.
2. Estratégia Seletiva vs. Global: Diferente de métodos que projetam todas as atualizações (como SafeLoRA) ou fundem modelos inteiros (como RESTA), o SafeMERGE funde apenas as camadas degradadas, preservando a utilidade da tarefa.
3. Validação Abrangente: O método foi testado em quatro LLMs diferentes (Llama-2, Llama-3.1, Qwen-2, Qwen-2.5) e múltiplas tarefas (GSM8K, PubMedQA e domínios de telecomunicações).
4. Análise de Ablação: Demonstra que a similaridade de cosseno é uma métrica superior à distância euclidiana para detectar desvios de segurança, pois ignora a magnitude da mudança e foca na direção.

4. Resultados Experimentais

Os experimentos compararam o SafeMERGE com baselines como SafeInstruct, RESTA, RESTA-Instruct e SafeLoRA.

• Desempenho de Segurança: O SafeMERGE reduziu consistentemente as saídas prejudiciais (avaliadas nos benchmarks DirectHarm e HexPhi) para níveis próximos ou inferiores aos dos modelos originais alinhados, superando todas as outras defesas.
  • Exemplo: No Llama-3.1 ajustado para GSM8K, o SafeMERGE reduziu a taxa de harmfulness de 28,30% para 8,80% (DirectHarm), enquanto outros métodos mantiveram taxas mais altas.
• Preservação de Utilidade: Ao contrário de métodos que fundem indiscriminadamente (que causam esquecimento catastrófico), o SafeMERGE manteve ou até melhorou a precisão na tarefa principal.
  • Exemplo: No Llama-3.1 (GSM8K), a precisão subiu de 78,24% para 78,50% após a aplicação do SafeMERGE, ao mesmo tempo que a segurança foi restaurada.
• Robustez em Domínios Especializados: O método demonstrou eficácia em tarefas fora da distribuição (telecomunicações), provando sua generalização.
• Eficiência: O método não requer recursos computacionais significativos adicionais, pois evita o treinamento e opera apenas na fusão de pesos.

5. Significado e Conclusão

O SafeMERGE estabelece um novo paradigma para a manutenção da segurança em LLMs ajustados:

• Simplicidade e Eficácia: Demonstra que soluções complexas de treinamento não são necessárias para preservar a segurança; uma fusão inteligente e seletiva de pesos é suficiente.
• Equilíbrio Otimizado: Resolve o dilema clássico entre segurança e utilidade, oferecendo o melhor compromisso (trade-off) entre os dois, superando métodos que priorizam um em detrimento do outro.
• Adoção Prática: Por ser leve, não exigir re-treinamento e funcionar com bibliotecas padrão (como HuggingFace PEFT), o SafeMERGE é altamente viável para implementação em cenários do mundo real onde modelos são frequentemente ajustados para nichos específicos.

Em suma, o trabalho prova que a segurança em LLMs ajustados pode ser restaurada de forma robusta através da fusão seletiva de camadas, protegendo contra a perda inadvertida de alinhamento sem sacrificar a capacidade do modelo de realizar tarefas úteis.