Guiding Giants: Lightweight Controllers for Weighted Activation Steering in LLMs

1. O Problema

Os Grandes Modelos de Linguagem (LLMs) apresentam riscos significativos ao gerar conteúdo malicioso, enviesado ou factualmente incorreto. As abordagens atuais para garantir a segurança desses modelos geralmente envolvem:

• Ajuste Fino (Fine-tuning): Custo computacional elevado, risco de "esquecimento catastrófico" (perda de capacidades gerais) e dificuldade de adaptação a novas ameaças de segurança sem retreinamento.
• Métodos de Intervenção em Tempo de Inferência (Steering): Técnicas existentes, como Activation Steering, muitas vezes aplicam alterações fixas e uniformes em todas as camadas do modelo. Isso carece de granularidade: se a intervenção for muito forte, o modelo pode recusar indiscriminadamente até prompts benignos (perdendo utilidade); se for muito fraca, falha em bloquear solicitações prejudiciais.

O desafio central é desenvolver um mecanismo de controle em tempo de inferência que seja adaptativo, leve e capaz de aplicar intervenções finas e específicas sem alterar os parâmetros originais do modelo base.

2. Metodologia: Weighted Activation Steering (WAS)

Os autores propõem o Weighted Activation Steering (WAS), uma nova abordagem que utiliza uma rede controladora leve e treinável para modular dinamicamente as ativações do LLM durante a geração.

Arquitetura e Funcionamento

O processo consiste em quatro etapas principais:

1. Cache de Ativações: Um LLM congelado processa prompts maliciosos e benignos, armazenando as ativações intermediárias (hidden states).
2. Treinamento do Controlador ($f_c$): Uma rede neural leve (MLP) é treinada de forma discriminativa. Ela recebe as ativações concatenadas de uma subconjunto de camadas (geralmente as últimas 1/3 do modelo) e prevê dois parâmetros:
   • Um fator de escala escalar global ($s$): Determina a intensidade geral da intervenção (alto para prompts nocivos, baixo para benignos).
   • Um vetor de pesos por camada ($w$): Determina a importância relativa de cada camada específica para a intervenção naquele prompt.
3. Cálculo do Vetor de Direção ($d_{steer}$): Um vetor de "direção de recusa" é pré-calculado subtraindo a média dos embeddings de tokens de resposta ("Sim", "Claro") dos embeddings de tokens de recusa ("Não posso", "Desculpe").
4. Aplicação Dinâmica: Durante a inferência, o controlador observa as ativações, prevê $(s, w)$ e aplica um "patch" de ativação modificado em todas as camadas:
   $$\Delta h_{l} = s \cdot w_l \cdot \alpha_{global} \cdot d_{steer}$$
   Onde $\alpha_{global}$ é um hiperparâmetro fixo.

Vantagens da Abordagem

• Leveza: O controlador é uma MLP pequena, adicionando latência mínima.
• Adaptabilidade: Aprende a ajustar a força da intervenção e em quais camadas focar dependendo do contexto do prompt.
• Não Invasivo: Não requer alteração nos pesos do modelo base.

3. Contribuições Principais

• Introdução do WAS: Um mecanismo de controle em tempo de inferência que combina magnitude global e ponderação específica por camada.
• Treinamento Discriminativo Indireto: O controlador é treinado apenas para prever a magnitude global ($s$) correta (1.0 para nocivo, 0.0 para benigno), mas, devido à arquitetura compartilhada, os pesos por camada ($w$) aprendem implicitamente padrões estruturados e interpretáveis, identificando "hotspots" de segurança em camadas específicas.
• Implementação Eficiente: Demonstração prática usando hooks do PyTorch para injetar os patches sem modificar o código do modelo base.

4. Resultados Experimentais

Os experimentos foram realizados nos modelos Llama-3.1-8B, Llama-3.2-1B e Mistral-7B.

• Segurança (Benchmarks ToxicChat, Jailbreak, AdvBench):
  • O WAS superou consistentemente métodos existentes (CAST, ACT, ITI).
  • No ToxicChat, o WAS alcançou taxas de recusa de 93.0% (Llama-3.1-8B) e 95.0% (Mistral-7B), enquanto os concorrentes ficaram abaixo de 50% na maioria dos casos.
  • No AdvBench, atingiu até 98.8% de recusa.
• Utilidade (Benchmarks Gerais - AlpacaEval, MMLU, GSM8K):
  • O WAS manteve a capacidade do modelo em tarefas benignas. A taxa de vitória no AlpacaEval contra o modelo base foi estatisticamente indistinguível de 50% (ex: 49.82% para Llama-3.1-8B).
  • Houve impacto mínimo ou nulo em benchmarks de raciocínio (GSM8K) e conhecimento geral (MMLU, HellaSwag).
• Eficiência: O WAS oferece um equilíbrio superior entre segurança e latência, sendo significativamente mais rápido que o ITI e mais seguro que o CAST.
• Estudos de Ablação:
  • A remoção dos pesos por camada (usando apenas o escalar $s$) reduziu a taxa de recusa em mais de 20%, provando que a ponderação específica por camada é crucial.
  • A aplicação em todas as camadas foi superior à aplicação apenas em camadas tardias, embora as camadas tardias tenham contribuído mais significativamente.

5. Significância e Conclusão

O trabalho demonstra que é possível alinhar e proteger LLMs de grande escala de forma eficiente, sem o custo massivo do ajuste fino. O WAS oferece uma camada de segurança "plug-and-play" que:

1. Maximiza a segurança ao recusar conteúdo tóxico com alta precisão.
2. Preserva a utilidade ao não degradar o desempenho em tarefas legítimas.
3. Fornece interpretabilidade, revelando que diferentes tipos de ameaças (ex: assédio vs. desinformação) ativam padrões de intervenção em camadas distintas do modelo.

Limitações: O método depende da qualidade do vetor de direção pré-calculado e pode ser vulnerável a ataques adversariais projetados especificamente para enganar o controlador (fazendo-o prever uma magnitude baixa para um prompt nocivo). Além disso, a generalização para novos tipos de ameaças não vistos no treinamento pode ser um desafio.

Em suma, o WAS representa um avanço significativo na área de Activation Engineering, oferecendo uma solução prática e escalável para a segurança de LLMs em produção.