Two to Tango: Coupled Task-Reference Selection for Safe LLM Fine-tuning

Resumo Técnico: Two to Tango: Seleção Acoplada de Tarefa–Referência para Ajuste Fino Seguro de LLMs

Declaração do Problema

O ajuste fino (fine-tuning) de Large Language Models (LLMs) alinhados para segurança em dados de tarefas subsequentes frequentemente melhora a adaptação à tarefa, mas corre o risco de erodir os comportamentos de segurança aprendidos durante a fase inicial de alinhamento. As estratégias de mitigação existentes baseiam-se tipicamente em uma de três abordagens: filtragem de dados de tarefa, mistura de exemplos de segurança fixos ou imposição de restrições globais de segurança. No entanto, os diagnósticos apresentados no artigo revelam uma limitação crítica nesses métodos: eles tratam as referências de segurança como sinais estáticos ou globais. Na realidade, diferentes direções de atualização de tarefas ativam restrições de segurança distintas. Consequentemente, um subconjunto fixo de referências de segurança pode falhar em expor as restrições específicas relevantes para uma determinada atualização de tarefa, levando à degradação da segurança mesmo quando dados benignos são utilizados.

Metodologia: DualSelect

Os autores propõem o DualSelect, um framework acoplado que seleciona conjuntamente amostras de tarefas e referências de segurança com base na geometria da atualização do modelo. O insight central é que a preservação da segurança requer o emparelhamento das atualizações de tarefa com restrições de segurança ativas específicas para essa atualização, em vez de aplicar um filtro de segurança global.

1. Formulação Min-Max Acoplada
O DualSelect formula o ajuste fino seguro como um problema de seleção min-max de estilo bilevel:

• Maximização Interna (Seleção de Referência): Identifica referências de segurança que exibem alta perda de preservação e alto conflito (oposição de gradiente) com a direção atual da tarefa. Estas referências representam as restrições ativas de "pior caso" para a atualização atual.
• Minimização Externa (Seleção de Tarefa): Seleciona amostras de tarefa que são compatíveis com a direção do gradiente induzida pelas referências selecionadas.

2. Substitutos de Pontuação com Regularização de Entropia
Para tornar o problema de seleção discreta tratável, os autores derivam substitutos de pontuação regularizados por entropia.

• Pontuação de Referência: As referências são pontuadas com base em sua perda supervisionada e seu conflito de cosseno com o gradiente da tarefa atual. Referências de alta perda e alto conflito recebem pontuações mais altas.
• Pontuação de Tarefa: As amostras de tarefa são pontuadas com base em sua perda supervisionada e sua compatibilidade de cosseno com o gradiente da referência selecionada. Amostras de baixa perda e alta compatibilidade são preferidas.
• Seleção Rígida (Hard Selection): Embora os substitutos forneçam probabilidades suaves (soft), o algoritmo implementa seleção rígida (Top-K para referências, seleção de amostra com orçamento correspondente para tarefas) para manter a compatibilidade com pipelines padrão de ajuste fino.

3. Atualização Preguiçosa (Lazy Refresh) e Correção de Gradiente

• Atualização Preguiçosa (Lazy Refresh): A seleção de referência é atualizada por época (ou preguiçosamente) em vez de por batch para amortizar os custos de pontuação, enquanto a seleção de tarefa ocorre por mini-batch.
• Correção de Gradiente de Referência: O gradiente de atualização final é corrigido adicionando um componente escalonado do gradiente da referência selecionada ($g_{final} = g_{new} + \rho g_{ref}$). Isso direciona a atualização de parâmetros para a direção da referência selecionada, garantindo que a atualização realizada se alinhe com as restrições de segurança identificadas.

4. Análise Teórica
O artigo fornece uma análise geométrica local mostrando que, sob força de correção suficiente ($\rho$), a direção de descida corrigida se alinha com o gradiente da referência selecionada, garantindo uma diminuição de primeira ordem na perda da referência selecionada.

Principais Contribuições

1. Mudança de Paradigma: Os autores identificam e demonamstram uma mudança de uso de pools de referências de segurança estáticas para a seleção de referência condicionada à tarefa. Eles mostram que as referências selecionadas são estruturadas e preservam a segurança retida significativamente melhor do que subconjuntos estáticos.
2. Framework Acoplado: Eles introduzem o DualSelect, que acopla a seleção de amostras de tarefa e referências de segurança como variáveis interdependentes em vez de alvos de pontuação separados, formulado como uma otimização min-max sobre referências de alta perda e conflito com a tarefa e amostras de tarefa compatíveis.
3. Derivação Algorítmica: O artigo deriva substitutos regularizados por entropia que induzem regras de seleção rígida, incorporando atualização preguiçosa de referência, correção de gradiente de referência e análise de alinhamento de atualização ao nível de direção.
4. Validação Empírica: Experimentos extensos em LLMs de 1B a 8B parâmetros (Gemma, Qwen, Llama) nos datasets REDORCA e GSM8K validam a abordagem.

Resultados Experimentais

• Preservação de Segurança: O DualSelect supera consistentemente os fortes baselines (incluindo TOSS, SQSD, SEAL e LARF) na preservação da segurança. No dataset REDORCA, ele melhora a pontuação de Média de Segurança em pelo menos 5,10 pontos sobre o baseline mais forte em modelos de 1B–8B.
• Retenção de Utilidade: O método mantém a utilidade da tarefa comparável ou superior aos baselines. No GSM8K, ele alcança a maior Média de Segurança enquanto mantém a precisão próxima aos níveis de SFT Padrão.
• Eficiência: O DualSelect introduz um overhead moderado (aproximadamente 1,34× as horas de GPU do SFT Padrão) e uso de memória comparável ao ajuste fino LoRA, sendo significativamente mais eficiente que métodos de otimização bilevel como o SEAL.
• Estudos de Ablação: Remover qualquer componente (seleção de referência, pontuação de tarefa, compatibilidade ou correção final) leva a uma queda na preservação da segurança, confirmando a necessidade do mecanismo acoplado.
• Aprendizado Contínuo: O framework estende-se ao aprendizado contínuo ao tratar exemplos de tarefas antigas como referências de retenção, reduzindo com sucesso o esquecimento catastrófico enquanto preserva capacidades anteriores.

Significância e Alegações

O artigo alega que a principal contribuição não é meramente um novo heurístico de seleção de dados, mas um reenquadramento fundamental da preservação de segurança como um problema condicionado à atualização. Os autores argumentam que os dados de segurança não são uniformemente úteis; sua relevância depende da direção específica da tarefa que está sendo aprendida. Ao expor dinamicamente as restrições ativas através de referências condicionadas à tarefa e filtrar tarefas compatíveis, o DualSelect alcança um trade-off superior entre segurança e utilidade.

Os autores observam modestamente que o método assume o acesso a um sinal de referência (respostas seguras ou alvos de tarefas antigas) e que o desempenho depende da cobertura da referência. Eles também reconhecem que a avaliação de segurança depende parcialmente de juízes automáticos, embora utilizem análise de juízes cruzados e métricas absolutas para mitigar artefatos. O trabalho posiciona-se como uma solução prática e escalável para ajuste fino seguro que evita os altos custos de reponderação global ou edição pós-hoc complexa.