Resumo Técnico: Destilação de Sistemas LLM Seguros via Soft Prompts para Configurações On-Device

Declaração do Problema

A implantação de Grandes Modelos de Linguagem (LLMs) em dispositivos de borda com recursos limitados (ex: smartphones) apresenta um desafio crítico de segurança. Embora sistemas de modelo duplo — combinando um LLM base com um "modelo de guarda" separado para classificar e filtrar saídas inseguras — forneçam garantias de segurança robustas, eles incorrem em custos proibitivos de memória e computação. A natureza sequencial desta abordagem (gerar uma resposta, depois classificar) dobra as passagens de ida (forward passes), aumentando significamente a latência (tempo para o primeiro token) e a pegada de memória, tornando-a inadequada para implantação em dispositivos de borda. Alternativas existentes, como quantização ou ajuste fino direto, frequentemente comprometem as capacidades de generalização ou falham em equilibrar segurança com utilidade do modelo de forma eficaz.

Metodologia

O artigo propõe o TV-DiSP (Total Variation Distillation via Soft Prompts), uma abordagem de ajuste fino de parâmetros eficientes (PEFT) que destila os comportamentos de segurança de um sistema de modelo duplo para um único LLM equipado com soft prompts aprendíveis.

Estrutura Central

1. Formulação do Sistema: Os autores definem um sistema de LLM seguro $p(r|x)$ que combina um LLM base $p(y|x)$ e um modelo de guarda $p(s|x, y)$. Se o modelo de guarda considerar uma resposta $y$ insegura ($s=0$), o sistema produz uma mensagem de recusa predefinida $y_r$; caso contrário, ele produz $y$.
2. Objetivo de Destilação: O objetivo é aprender um conjunto de embeddings de prompt contínuos $W$ (soft prompts) de tal forma que o modelo destilado $q(r|x, W)$ imite a distribuição de saída do sistema seguro $p(r|x)$.
3. Otimização via Variação Total (TV):
   • Os autores derivam uma garantia teórica (Teorema 3.1) mostrando que minimizar a distância de Variação Total entre o sistema seguro e o modelo destilado limita a lacuna de desempenho em qualquer tarefa de jusante (downstream).
   • O objetivo de otimização (Equação 4) maximiza a log-probabilidade da saída correta (seja a resposta do LLM base para entradas seguras ou a mensagem de recusa para entradas inseguras) ponderada pela probabilidade de segurança do modelo de guarda.
   • O treinamento requer apenas um conjunto de dados de prompts; o modelo de guarda rotula dinamicamente esses prompts como seguros ou inseguros durante o processo.
4. Esquemas Alternativos: O estudo também avalia:
   • KL-DiSP: Minimização da divergência de Kullback-Leibler, que oferece garantias teóricas mais frouxas que a TV.
   • Baselines: Otimização de perplexidade, REINFORCE (gradiente de política) e adaptações do framework de destilação para Low-Rank Adapters (LoRA) e Vetores de Direcionamento (SV).

Configuração Experimental

• Modelos: Os experimentos foram conduzidos em quatro LLMs ajustados por instrução (Llama3-1B/3B, Qwen2-1.5B, Gemma2-2B) quantizados para 4-bits para simular restrições de borda.
• Modelos de Guarda: LlamaGuard3-1B foi utilizado para o treinamento de destilação, e LlamaGuard3-8B para avaliação.
• Datasets: O treinamento utilizou os datasets Beavertails e Toxigen. A avaliação incluiu benchmarks de fora de distribuição (HarmBench, Detect-JailBreak) e benchmarks de utilidade (IFEval, GSM8K, MMLU).
• Hardware: Medições on-device foram realizadas em smartphones equipados com chipsets Qualcomm Snapdragon 8 Elite Gen 5.

Principais Resultados

1. Troca Segurança-Computação (Safety-Compute Trade-off)

• Desempenho: O TV-DiSP destila com sucesso a segurança do sistema de modelo duplo para um único modelo. Por exemplo, no modelo Llama3-1B, a segurança (medida pelo Safety Guard Score, SGS) melhorou de 71% (base) para 99% (sistema duplo), enquanto o TV-DiSP alcançou ganhos de segurança comparáveis com overhead mínimo.
• Eficiência: Comparado ao sistema de modelo duplo, que dobra os requisitos de memória e computação, o TV-DiSP introduz menos de 1% de overhead de memória adicional e menos de 10% de overhead de computação adicional (FLOPs por token).
• Generalização: O método demonstra forte robustez em prompts adversários fora de distribuição (HarmBench, Detect-JailBreak), apesar de ter sido treinado em datasets não adversários.

2. Comparação com Métodos PEFT

• Soft Prompts vs. LoRA/SV: O TV-DiSP (soft prompts) superou consistentemente o TV-DiLoRA e o TV-DiSV (Steering Vectors) através de várias taxas de aprendizado. Os vetores de direcionamento mostraram resultados mistos, às vezes degradando o desempenho em modelos maiores, enquanto o LoRA foi considerado excessivamente intrusivo para configurações de borda.
• Objetivos de Otimização:
  • TV vs. KL: Tanto a destilação TV quanto a KL alcançaram altos escores de segurança. No entanto, o TV-DiSP ofereceu um equilíbrio superior entre segurança e utilidade.
  • Perplexity & REINFORCE: A otimização de perplexidade levou ao sobreajuste (overfitting) e quedas significativas na utilidade. O REINFORCE manteve a utilidade, mas não conseguiu fornecer melhorias substanciais de segurança.

3. Preservação da Utilidade

• Avaliações on-device (IFEval e GSM8K) mostraram que o TV-DiSP e o KL-DiSP mantêm alta utilidade.
• Sobre-recusa (Over-refusal): Um achado crítico foi que o KL-DiSP exibiu uma alta taxa de sobre-recusa (84,6% em prompts seguros), rejeitando incorretamente consultas benignas. Em contraste, o TV-DiSP manteve uma taxa de sobre-recusa (36,1%) comparável à do sistema de modelo duplo, preservando a capacidade do modelo de responder a perguntas seguras.

Significância e Alegações

O artigo estabelece a destilação de soft prompt como a abordagem preferencial para o alinhamento de segurança na implantação de LLMs em dispositivos de borda. Suas principais contribuições são:

1. Trocas Superiores: Demonstra que soft prompts treinados via destilação superam consistentemente o LoRA, vetores de direcionamento e outros objetivos de treinamento de destilação no equilíbrio entre segurança, utilidade e eficiência computacional.
2. Garantias Teóricas: A introdução da destilação baseada em Variação Total fornece garantias probabilísticas no desempenho de jusante, garantindo que o modelo destilado não se desvie significativamente do comportamento do sistema seguro.
3. Viabilidade Prática: Validado através de múltiplas arquiteturas e em hardware real de borda, o método alcança níveis de segurança comparáveis aos sistemas de modelo duplo com overhead de recursos negligenciável (<1% de memória, <10% de computação), tornando-o uma solução prática para implantar LLMs seguros em dispositivos com recursos limitados.

Os autores concluem que, embora nenhum método ofereça imunidade completa contra ataques adversários, o TV-DiSP oferece um mecanismo de defesa robusto, eficiente e generalizável que reduz significativamente a superfície de ataque e a carga computacional em comparação com as arquiteturas tradicionais de modelos de guarda.