Test-Time Safety Alignment

1. Declaração do Problema

Os Grandes Modelos de Linguagem (LLMs) passam por treinamento de segurança (por exemplo, RLHF, DPO) para recusar solicitações prejudiciais. No entanto, prompts adversariais (jailbreaks) podem contornar essas defesas explorando a lacuna entre a distribuição de treinamento do modelo e o espaço de entrada ilimitado. Esses prompts frequentemente utilizam enquadramento estrutural (roleplay, ofuscação) para elicitar conteúdo prejudicial.

As defesas existentes em tempo de teste enfrentam limitações estruturais:

• Estratégias no nível do prompt manipulam tokens discretos, o que é frágil contra ataques semânticos.
• Direcionamento de ativação requer acesso a estados ocultos e frequentemente depende de conjuntos de calibração fixos que podem não generalizar.
• Modelos de guarda exigem inferência separada e frequentemente seguem um paradigma de "detectar e depois filtrar" em vez de direcionar a geração.

A Questão Central: É possível otimizar os embeddings de palavras de entrada de maneira sub-lexical (modificando a representação vetorial contínua sem alterar o token visível) para direcionar modelos alinhados a saídas seguras, mesmo quando o modelo produz uma distribuição bimodal de "recusar ou cumprir" em vez de uma continuação suave?

2. Metodologia: Alinhamento de Segurança em Tempo de Teste (TSA)

O método proposto, TSA, opera inteiramente em tempo de teste sem modificar os pesos do modelo, não requer treinamento auxiliar e trata o LLM como uma caixa preta. Ele otimiza diretamente a matriz de embeddings do prompt.

A. Função Objetivo

O objetivo é minimizar a nocividade do texto gerado $y = f(X)$, onde $f$ é o LLM e $X$ é a matriz de embeddings do prompt.

• Oráculo de Nocividade ($h$): O método utiliza a API de Moderação da OpenAI como um oráculo de caixa preta. Ela pontua as saídas em 13 categorias de dano (por exemplo, violência, automutilação, ódio).
• Função de Perda: $\Phi(X) = \max_{c} s_c(f(X))$, onde $s_c$ é a pontuação para a categoria $c$. O objetivo é minimizar essa pontuação máxima, garantindo que a resposta seja segura em todas as dimensões.

B. Estimação de Gradiente de Ordem Zero

Como o LLM ($f$) e a API de Moderação ($h$) são caixas pretas, a retropropagação é impossível. O TSA utiliza estimação de gradiente de ordem zero:

1. Perturbação: Em cada iteração $k$, a matriz de embeddings atual $X_k$ é perturbada adicionando ruído Gaussiano $U$ escalado por $\mu$: $X_k + \mu U$.
2. Diferença Finita: O gradiente é estimado comparando a perda das entradas perturbadas contra a entrada base:
   $$g_k \approx \frac{1}{N} \sum_{i=1}^{N} \frac{\Phi(X_k + \mu U_i) - \Phi(X_k)}{\mu} U_i$$
   onde $N$ é o número de amostras de Monte Carlo.
3. Atualização: Os embeddings são atualizados via descida de gradiente: $X_{k+1} \leftarrow X_k - \eta g_k$.

C. Regularização

Para garantir que os embeddings otimizados permaneçam válidos e não se desviem para regiões não treinadas do espaço latente:

• Normalização do Gradiente: O passo de atualização usa a direção do gradiente normalizada ($g_k / \|g_k\|_2$) para controlar a magnitude do passo.
• Restrição de Similaridade Cosseno: O embedding atualizado é projetado de volta para uma bola de cosseno ao redor do embedding original $X_0$ se a similaridade cair abaixo de um limiar $\kappa$. Isso garante que a perturbação seja sub-lexical: o vetor otimizado decodifica de volta para o token original, mas a representação contínua dentro da célula de Voronoi desse token é alterada.

3. Contribuições Principais

1. Extensão para Modelos Alinhados: Demonstra que o controle no nível de embedding, anteriormente mostrado para detoxificar modelos base, funciona efetivamente em modelos instruídos alinhados à segurança, que possuem uma distribuição de saída fundamentalmente diferente (bimodal).
2. Controle Sub-Lexical: Prova que otimizar o vetor de embedding contínuo dentro da célula do vizinho mais próximo de um token (sem alterar o token em si) é suficiente para redirecionar o comportamento do modelo de prejudicial para seguro. Isso desafia a visão de embeddings como meras consultas discretas.
3. Otimização de Caixa Preta: O método requer nenhum acesso aos internos do modelo (gradientes, estados ocultos ou logits). Requer apenas acesso aos embeddings de entrada e a uma API de moderação de caixa preta.
4. Aplicação Incondicional: O método aplica-se a todos os prompts, independentemente da intenção. Não força uma recusa; direciona o modelo a gerar uma resposta substantiva e segura se possível, ou uma recusa segura se necessário.

4. Resultados Experimentais

O método foi avaliado em cinco modelos instruídos (variando de 1B a 20B parâmetros: Gemma 3, Phi-3.5, Llama 3.1, Qwen3, GPT-OSS) em duas benchmarks:

• WildJailbreak: Prompts adversariais (prejudiciais e benignos).
• HarmBench: Consultas prejudiciais diretas.

Principais Descobertas:

• Eliminação de Segurança: O TSA reduziu o número de respostas sinalizadas (prejudiciais) para zero na maioria dos modelos em benchmarks padrão. Por exemplo, no Llama 3.1-8B, reduziu prompts prejudiciais sinalizados de 218 para 0.
• Convergência: A otimização geralmente converge dentro de 1 a 2 passos de gradiente.
• Preservação Benigna: Em prompts benignos adversariais, o TSA não aumentou as pontuações de moderação, confirmando que não induz um viés de "recusa abrangente" ou degrada a utilidade.
• Comparação com Baselines: O TSA superou significativamente baselines como SmoothLLM (perturbação aleatória de tokens), AdaSteer (direcionamento de ativação) e RESTA (ruído de embedding). As baselines frequentemente falharam em jailbreaks semânticos ou aumentaram as taxas de sinalização em consultas prejudiciais diretas.
• Eficiência: O principal custo de latência é a chamada de API ao oráculo de moderação, não a inferência do modelo.

5. Significado e Implicações

• Nova Superfície de Controle: O artigo revela que a camada de embedding oferece uma superfície de controle mais rica do que se pensava anteriormente. Pequenas perturbações sub-leksicais podem navegar nas complexas fronteiras de segurança de modelos alinhados.
• Viabilidade de Implantação: Como uma camada de segurança "plug-and-play", o TSA não requer retreinamento ou ajuste fino do modelo. Pode ser implantado em qualquer modelo onde os embeddings de entrada sejam acessíveis (por exemplo, modelos de pesos abertos) combinados com uma API de moderação pública.
• Limitações:
  • Latência: A otimização iterativa adiciona latência por prompt (segundos), embora isso seja dominado pelas viagens de ida e volta da API.
  • Dependência do Modelo: A eficácia escala com a qualidade do treinamento de segurança existente do modelo. Modelos com alinhamento tênue (por exemplo, Gemma 3-1B) mostraram sucesso reduzido, pois o otimizador não pôde encontrar uma direção segura próxima.
  • Acesso à API: Requer acesso aos embeddings de entrada (não disponível para modelos fechados puramente protegidos por API) e uma API de moderação.

Em conclusão, o Alinhamento de Segurança em Tempo de Teste fornece um mecanismo robusto e agnóstico ao modelo para neutralizar ataques adversariais otimizando a representação de entrada contínua, efetivamente fechando a lacuna entre segurança teórica e implantação prática sem retreinamento.