Few Tokens, Big Leverage: Preserving Safety Alignment by Constraining Safety Tokens during Fine-tuning

O artigo apresenta o framework PACT, que preserva o alinhamento de segurança durante o ajuste fino de modelos de linguagem ao restringir especificamente a confiança em tokens relacionados à segurança, evitando assim o desvio de alinhamento sem comprometer a utilidade do modelo em tarefas downstream.

O essencial

Imagine que você tem um assistente de IA muito inteligente e educado, que foi treinado para nunca fazer nada perigoso ou ofensivo. Ele sabe dizer "não" para pedidos como "como fazer uma bomba?".

Agora, imagine que você quer ensinar esse assistente a fazer algo específico, como resolver problemas de matemática ou escrever poemas. Para isso, você precisa fazer um "treinamento especial" (chamado de fine-tuning) com novos dados.

O Problema: O Treinamento que "Esquece" as Regras
O problema é que, durante esse treinamento especial, mesmo que você tente usar apenas dados bons, o assistente pode começar a "esquecer" suas regras de segurança. Pior ainda: se houver um pouquinho de dados ruins misturados no treinamento, ele pode começar a obedecer a pedidos perigosos, esquecendo completamente que não deveria fazer isso. É como se, ao aprender a fazer matemática, ele decidisse que as regras de segurança eram apenas sugestões e não mais leis.

Métodos antigos para consertar isso eram como colocar uma "mordaça" no cérebro inteiro do assistente. Eles tentavam travar grandes partes do modelo para que ele não mudasse. O resultado? O assistente ficava seguro, mas também ficava burro e ruim em fazer a tarefa nova (a matemática, por exemplo).

A Solução: PACT (O "Filtro de Palavras" Inteligente)
Os autores deste artigo propuseram uma solução chamada PACT. A ideia deles é baseada em uma descoberta interessante: a segurança da IA não depende de todas as palavras que ela conhece, mas sim de um pequeno grupo de palavras-chave.

Pense assim:

• Quando o assistente vai dizer "não", ele usa certas palavras específicas, como "não", "não posso", "não ajudo", "perigo".
• O PACT descobriu que, se você garantir que o assistente mantenha a mesma confiança ao usar apenas essas poucas palavras, ele continuará seguro, mesmo enquanto aprende coisas novas.

A Analogia do Maestro e os Violinos
Imagine que o modelo de IA é uma orquestra gigante.

• O Treinamento Antigo (Métodos Globais): Era como pedir para o maestro (o modelo) parar de tocar qualquer nota que não fosse da música nova. Isso fazia a orquestra tocar a música nova, mas sem emoção e sem a "alma" original.
• O Método PACT: É como se o maestro dissesse: "Pessoal, continuem tocando a música nova com toda a liberdade! Mas, atenção: quando chegarmos na parte onde alguém pede algo perigoso, os violinos (que representam as palavras de segurança) devem tocar exatamente na mesma nota forte e clara que tocavam antes. O resto da orquestra pode fazer o que quiser."

Como funciona na prática?

1. Identificação: O sistema primeiro descobre quais são essas "palavras de segurança" (o grupo de violinos). São apenas cerca de 50 palavras em todo o vocabulário gigante.
2. Proteção Seletiva: Durante o treinamento, o sistema vigia apenas essas palavras. Se o assistente começar a hesitar em dizer "não", o sistema o corrige imediatamente, forçando-o a manter a confiança nessas palavras específicas.
3. Liberdade para o Resto: Todas as outras palavras (para matemática, poesia, etc.) podem mudar livremente para que o assistente aprenda a nova tarefa com perfeição.

O Resultado?
Com o PACT, o assistente aprende a nova tarefa (fica excelente em matemática) e continua sendo super seguro (recusa pedidos perigosos com a mesma firmeza de antes). É como se você pudesse ensinar um guarda-costas a ser um ótimo chef de cozinha sem que ele perca a habilidade de proteger você de perigos.

Resumo em uma frase:
O PACT é uma técnica inteligente que protege a segurança da IA focando apenas em um punhado de palavras críticas, permitindo que o resto do cérebro da máquina aprenda e se adapte livremente sem "esquecer" suas regras de ouro.

Resumo técnico

Aqui está um resumo técnico detalhado do artigo "Few Tokens, Big Leverage: Preserving Safety Alignment by Constraining Safety Tokens during Fine-tuning" (Poucos Tokens, Grande Alavancagem: Preservando o Alinhamento de Segurança ao Restringir Tokens de Segurança durante o Ajuste Fino), traduzido e estruturado em português.

---

1. Problema: Deriva de Alinhamento de Segurança no Ajuste Fino

Os Grandes Modelos de Linguagem (LLMs) frequentemente requerem ajuste fino (Fine-Tuning - FT) para se adaptarem a tarefas específicas. No entanto, o processo de ajuste fino introduz um risco crítico: a deriva de alinhamento de segurança (safety-alignment drift).

• O Fenômeno: Mesmo quando o conjunto de dados de treinamento contém apenas dados benignos, o ajuste fino pode degradar a capacidade do modelo de recusar solicitações prejudiciais. Se houver uma pequena fração de dados nocivos no corpus de treinamento, o modelo pode aprender a cumprir solicitações perigosas que antes rejeitaria.
• Limitações das Defesas Existentes: Métodos anteriores tentam mitigar esse problema através de intervenções em nível de modelo ou parâmetro (ex: restringir quais parâmetros são atualizados, como no SafeLoRA, ou injetar dados de segurança adicionais).
  • Desvantagem: Essas abordagens são frequentemente "grossas" (coarse-grained), restringindo globalmente a atualização dos pesos. Isso limita a generalidade do modelo e degrada o desempenho nas tarefas downstream (utilidade), criando um trade-off negativo entre segurança e eficácia.

2. Metodologia: Framework PACT

Os autores propõem o PACT (Preserves safety Alignment via Constrained Tokens), um framework de ajuste fino que opera em nível de token. A premissa central é que o comportamento de segurança não é distribuído uniformemente por todo o vocabulário, mas sim concentrado em um pequeno subconjunto de "tokens de segurança".

O método consiste em três etapas principais:

A. Identificação de Tokens de Segurança

Os autores analisam a discrepância de probabilidade entre um modelo alinhado à segurança ($M_{safe}$) e seu modelo base ($M_{base}$) ao responder a prompts prejudiciais.

• Mecanismo: Utilizando teacher forcing, eles comparam as distribuições de probabilidade de próxima palavra em cada passo de geração.
• Resultado: Identificam um conjunto pequeno de tokens (ex: os 50 com maior discrepância) onde o modelo alinhado mantém uma confiança significativamente maior do que o modelo base. Exemplos incluem tokens como "I'm", "cannot", "but", "assistant".
• Hipótese: Manter a confiança do modelo ajustado nesses tokens específicos é suficiente para preservar o comportamento de recusa.

B. Regularização com Pesos de Tokens de Segurança

Em vez de aplicar uma perda de divergência KL (Kullback-Leibler) global sobre todo o vocabulário, o PACT aplica restrições seletivas:

• Vetor de Pesos: Cria-se um vetor de pesos esparsos onde apenas os tokens de segurança identificados têm pesos baseados em sua importância (discrepância).
• Função de Perda: A regularização KL força o modelo ajustado a corresponder à confiança do modelo de referência apenas nos tokens de segurança, permitindo que os demais tokens sejam otimizados livremente para a tarefa downstream. Isso preserva a utilidade da tarefa enquanto mantém a segurança.

C. Calibração do Sinal de Segurança (Mitigação de Prefixo Nocivo)

Um desafio é que, durante o ajuste fino com dados nocivos, o modelo de referência (congelado) pode ser forçado a condicionar-se em prefixos inseguros, o que pode suprimir sua própria confiança nos tokens de segurança.

• Solução: O PACT introduz um mecanismo de calibração adaptativa.
  • Gera-se duas visões de referência: Full-context (com o prompt) e No-prompt (apenas com o cabeçalho do assistente e tokens anteriores, sem o prompt nocivo).
  • Coeficiente de Portão ($c_t$): Calcula-se uma métrica de dispersão de probabilidade para determinar se o contexto está "contaminado". Se a contaminação for alta (o modelo está inseguro devido ao prefixo nocivo), o sistema aumenta o peso da visão "No-prompt" (mais segura) na distribuição de referência.
  • Decaimento Posicional: A calibração é mais forte no início da resposta (onde as recusas ocorrem) e decai suavemente para tokens subsequentes.

3. Contribuições Principais

1. Identificação e Análise de Tokens de Segurança: Demonstração empírica de que o alinhamento de segurança é governado por um conjunto pequeno e crítico de tokens, identificáveis através de discrepâncias de confiança entre modelos.
2. Framework de Ajuste Fino Preservador de Segurança: Proposta do PACT, que utiliza restrições em nível de token e calibração adaptativa para evitar a deriva de alinhamento sem sacrificar o desempenho da tarefa.
3. Validação Empírica Abrangente: Testes extensivos em múltiplos modelos (Qwen, Llama, Gemma), tarefas (GSM8K, SST-2, AGNEWS) e proporções de dados nocivos (0-10%).

4. Resultados Experimentais

Os autores compararam o PACT com baselines de última geração (SFT padrão, SafeLoRA, Constrained SFT, AsFT) usando métricas de Utilidade (Acurácia) e Segurança (Taxa de Sucesso de Ataque - ASR).

• Desempenho Geral: O PACT consistentemente alcançou o melhor equilíbrio entre utilidade e segurança.
  • Redução de Ataques: Reduziu as taxas de sucesso de ataques (ASR) para 5.75% - 9.27% no StrongReject e 13.50% - 29.50% no HarmBench, comparado a taxas muito mais altas (frequentemente >50% ou >90%) em outros métodos quando expostos a dados nocivos.
  • Preservação de Utilidade: A acurácia nas tarefas downstream (ex: GSM8K) permaneceu comparável ao ajuste fino padrão (Vanilla SFT), superando métodos que sacrificam utilidade para manter segurança.
• Generalização: O método funcionou bem em diferentes arquiteturas de modelos (Llama, Gemma, Qwen) e tamanhos, sem necessidade de ajuste específico para cada arquitetura.
• Robustez: O PACT manteve a segurança mesmo com proporções crescentes de dados nocivos (até 10%), enquanto outros métodos colapsaram rapidamente.
• Análise de Componentes: Estudos de ablação mostraram que cada componente (identificação de tokens, pesos, calibração de sinal e decaimento posicional) contribuiu significativamente para o desempenho final. A restrição apenas aos tokens de segurança foi mais eficaz do que restrições globais.

5. Significado e Conclusão

O trabalho demonstra que a segurança em LLMs não requer a preservação de todo o espaço de parâmetros ou vocabulário. Ao focar em "poucos tokens" críticos, é possível obter uma "grande alavancagem" para manter o alinhamento de segurança.

• Inovação: O PACT muda o paradigma de restrições globais (que prejudicam a utilidade) para restrições localizadas e inteligentes.
• Impacto Prático: Oferece uma solução viável para provedores de serviços de nuvem e usuários que desejam ajustar modelos comercialmente disponíveis para tarefas específicas sem comprometer a segurança inerente do modelo base, mesmo na presença acidental ou maliciosa de dados nocivos no conjunto de treinamento.

Em resumo, o PACT prova que é possível desacoplar a adaptação de tarefas da deriva de segurança, garantindo que modelos ajustados continuem a recusar solicitações perigosas com alta confiança, mantendo ao mesmo tempo sua eficácia nas tarefas para as quais foram treinados.