MOSAIC: Composable Safety Alignment with Modular Control Tokens

O artigo apresenta o MOSAIC, um framework modular que utiliza tokens de controle aprendíveis para permitir alinhamento de segurança composicional e adaptável em modelos de linguagem, superando as limitações de políticas estáticas e instruções em linguagem natural ao reduzir recusas excessivas e preservar a utilidade do modelo.

O essencial

Imagine que os Grandes Modelos de Linguagem (como o ChatGPT ou o Llama) são como chefes de cozinha extremamente talentosos. Eles sabem cozinhar qualquer prato do mundo, desde uma sopa simples até um banquete complexo. No entanto, para que eles não sirvam veneno ou pratos proibidos (como receitas de drogas ou conselhos perigosos), precisamos "alinhá-los" à segurança.

O problema é que as regras de segurança não são iguais para todos. O que é seguro para um adulto pode ser proibido para uma criança; o que é legal em um país pode ser crime em outro.

Aqui está a explicação do papel MOSAIC usando analogias simples:

O Problema: A "Regra Rígida" vs. O "Convite de Festa"

Atualmente, existem duas formas principais de controlar esses chefs:

1. O "Treinamento Pesado" (Alinhamento de Parâmetros): É como tentar ensinar o chef a nunca cozinhar algo perigoso, mudando a estrutura do cérebro dele. O problema? Se você mudar o cérebro para evitar veneno, ele pode esquecer como fazer um bolo delicioso ou ficar com medo de cozinhar qualquer coisa nova. É difícil mudar essa regra depois sem reescrever todo o livro de receitas.
2. O "Convite de Festa" (Prompts de Texto): É como escrever um bilhete no topo da mesa dizendo: "Por favor, não sirva álcool para menores". O problema é que o chef pode não ler o bilhete com atenção, ou se você tiver que escrever um bilhete gigante com 50 regras diferentes, a conversa fica lenta e confusa.

A Solução: MOSAIC (O "Kit de Botões Mágicos")

O MOSAIC propõe uma terceira via, muito mais inteligente. Em vez de mudar o cérebro do chef ou escrever bilhetes longos, eles criam pequenos "botões mágicos" (tokens de controle) que você pode colar na frente do pedido.

Pense nisso como um controle remoto universal para a segurança:

• Botão Vermelho (Jogos de Azar): Se você quer que o chef recuse qualquer coisa sobre cassinos, você cola o "Botão Vermelho" no pedido. O chef sabe imediatamente: "Ah, regra de cassino ativada! Vou recusar".
• Botão Azul (Álcool): Se o pedido é sobre fazer um coquetel, você cola o "Botão Azul".
• Botão Verde (Violência): Para pedidos sobre violência.

A Mágica da Composição:
O grande trunfo do MOSAIC é que você pode colar vários botões ao mesmo tempo.

• Se uma criança pede uma receita de bolo com álcool, você cola o Botão Azul (Álcool) e o Botão Amarelo (Menores). O chef entende a combinação e diz: "Não posso fazer isso porque envolve álcool e é para menores".
• Se um adulto pede a mesma receita, você só cola o Botão Azul (ou nenhum, dependendo da regra local). O chef faz o bolo.

Isso é chamado de alinhamento modular. Você não precisa reprogramar o chef inteiro; apenas troca os botões que estão ligados.

Como eles ensinaram os botões a funcionarem? (O Segredo da Treinagem)

Criar esses botões não foi fácil. Eles precisavam garantir duas coisas:

1. Que o botão funcionasse sozinho e também quando misturado com outros.
2. Que o botão não fizesse o chef recusar pedidos inofensivos (como "como fazer um mojito" quando o usuário é um adulto e a lei permite).

Para isso, eles usaram duas técnicas criativas:

1. Amostragem de "Misturas" (Order-based Sampling): Em vez de ensinar o botão apenas sozinho, eles treinaram o sistema misturando os botões de todas as formas possíveis (1 botão, 2 botões juntos, 3 botões juntos). É como treinar um músico não só tocando notas soltas, mas tocando acordes complexos, para que ele nunca se perca na música.
2. Distilação de Conhecimento "E se...?" (Counterfactual KD): Esta é a parte mais genial. Para evitar que o chef fique "medroso" e recuse tudo, eles ensinaram o modelo a comparar: "O que eu faria se não tivesse o botão de segurança?" vs. "O que eu faço com o botão?".
   • Se o pedido é seguro (ex: "como fazer um bolo"), o modelo aprende a agir exatamente igual ao que faria sem o botão, mesmo que o botão de segurança esteja lá.
   • Se o pedido é perigoso, o botão muda a resposta.
   • Isso evita o "excesso de recusa" (quando o modelo diz "não" para tudo, por medo).

Por que isso é importante?

• Flexibilidade: Você pode adicionar novas regras de segurança (ex: "proibir discurso de ódio") criando apenas um novo botão, sem precisar reeducar todo o modelo do zero.
• Eficiência: O modelo não fica lento com textos longos de regras.
• Precisão: O modelo entende o contexto. Ele sabe recusar um cassino para um menor, mas pode ajudar um adulto a entender os riscos do jogo.

Resumo Final

O MOSAIC é como dar ao chef de cozinha um kit de adesivos de segurança.

• Você não muda a receita base do chef (o modelo continua inteligente e útil).
• Você não escreve regras gigantes (é rápido e leve).
• Você apenas cola o adesivo certo para a situação (criança, país, contexto).

Isso permite que a inteligência artificial seja segura, mas também útil e adaptável às necessidades específicas de cada pessoa e lugar, sem quebrar nada no processo.

Resumo técnico

Título: MOSAIC: Alinhamento de Segurança Composto com Tokens de Controle Modulares

1. O Problema

O alinhamento de segurança em Grandes Modelos de Linguagem (LLMs) é atualmente implementado de forma estática, embutido nos parâmetros do modelo. Isso cria desafios significativos para a adaptação em cenários do mundo real:

• Falta de Flexibilidade Contextual: As políticas de segurança variam conforme o grupo de usuários (ex: menores vs. adultos), jurisdições (leis diferentes por país) e domínios de aplicação.
• Limitações dos Métodos Atuais:
  • Alinhamento em Nível de Parâmetros (SFT, RLHF): Entrelaça o comportamento de segurança com as capacidades gerais do modelo. Atualizar ou remover regras exige re-treinamento custoso e pode causar interferência catastrófica (perda de utilidade).
  • Métodos Baseados em Prompt: Dependem de instruções em linguagem natural, que são interpretadas probabilisticamente. Isso leva a inconsistências, ineficiência (devido ao aumento de tokens) e falha em impor restrições múltiplas de forma robusta.
• Falta de Representação Explícita: Não existe uma representação reutilizável e composicional das restrições de segurança que permita controle granular sem re-treinamento massivo.

2. Metodologia: MOSAIC

O MOSAIC (Composable Safety Alignment with Modular Control Tokens) reinterpreta o alinhamento de segurança como um problema de aprendizado de representação composicional. Em vez de modificar os parâmetros do modelo base, o framework utiliza tokens de controle aprendíveis em um modelo base congelado (frozen backbone).

Principais Componentes:

1. Tokens de Controle Modulares:

   • Cada categoria de restrição de segurança (ex: "jogos de azar", "álcool") é codificada por um pequeno conjunto de tokens aprendíveis ($z_c$) no espaço de incorporação do modelo.
   • Esses tokens são os únicos parâmetros treináveis.
   • Durante a inferência, os tokens correspondentes às restrições ativas são pré-fixados à entrada, ativando ou desativando comportamentos de recusa de forma modular.

2. Amostragem de Tarefas Combinatórias (Compositional Task Sampling):

   • Para evitar o crescimento exponencial de dados ao treinar combinações de categorias, o método organiza as tarefas por ordem (número de restrições ativas simultâneas).
   • Um orçamento de treinamento fixo é alocado por ordem, garantindo que os tokens aprendam a interagir cooperativamente em combinações complexas (ex: A+B, A+B+C) sem custo computacional proibitivo.

3. Objetivo de Alinhamento com Distilação de Conhecimento Contrafactual (Counterfactual KD):

   • Para mitigar o excesso de recusa (over-refusal — onde o modelo recusa consultas benignas apenas porque tokens de segurança estão presentes), o MOSAIC introduz um objetivo de distilação.
   • Compara a distribuição de saída do modelo com tokens de controle ($p_{ctrl}$) contra a distribuição do modelo base sem tokens ($p_{base}$) para entradas benignas.
   • Minimiza a divergência KL entre as duas distribuições, forçando os tokens a intervir apenas quando estritamente necessário, preservando o comportamento original do modelo em consultas seguras.

3. Contribuições Chave

• Reconceitualização do Alinhamento: Propõe tratar o controle de segurança como ativação de restrições modulares e compostas, em vez de modificação monolítica de parâmetros.
• Framework MOSAIC: Um sistema que permite a expansão incremental de categorias de segurança (adicionar novos tokens sem re-treinar o modelo base) e a composição dinâmica de políticas.
• Benchmarks Realistas: Construção de um novo conjunto de dados de avaliação com 1.500 solicitações cobrindo 5 categorias de segurança. Diferente de benchmarks tradicionais, estes dados são projetados para não serem rejeitados por LLMs alinhados padrão, permitindo testar a ativação seletiva e condicional de segurança.

4. Resultados Experimentais

Os experimentos foram realizados nos modelos Llama-3.1-8B e Llama-3.2-3B, comparando o MOSAIC com métodos In-context, ORPO e SFT (Fine-tuning Supervisionado).

• Desempenho de Defesa (DSR - Defense Success Rate):
  • O MOSAIC alcançou taxas de sucesso de recusa superiores a 99% em todas as ordens de tarefas (1 a 4 categorias simultâneas), superando consistentemente métodos baseados em prompt e ORPO.
  • O SFT também teve alto DSR, mas com problemas de utilidade.
• Redução de Excesso de Recusa (OR - Over-Refusal Rate):
  • O MOSAIC reduziu drasticamente o excesso de recusa. Enquanto o SFT manteve uma OR de ~6%, o MOSAIC-5 (5 tokens por categoria) reduziu a OR para 1.8% em tarefas de alta ordem no modelo de 8B.
  • A técnica de Counterfactual KD foi crucial para distinguir entre consultas maliciosas e benignas quando os tokens de segurança estavam ativos.
• Preservação de Utilidade:
  • O MOSAIC preservou a capacidade de modelagem de linguagem geral, com degradação negligenciável em benchmarks como MMLU, ao contrário do SFT que pode degradar o desempenho em tarefas gerais.
• Expansão Incremental:
  • O framework demonstrou capacidade de adicionar novas categorias de segurança sem re-treinar as existentes, com degradação mínima de desempenho, validando sua escalabilidade.

5. Significado e Impacto

O MOSAIC representa um avanço significativo na segurança de IA ao oferecer:

• Adaptabilidade Dinâmica: Permite que provedores de LLMs ajustem políticas de segurança em tempo de inferência para diferentes regiões ou grupos demográficos sem re-treinamento.
• Eficiência: Elimina o custo de re-treinamento de modelos grandes para cada nova regra de segurança.
• Precisão: Resolve o dilema entre segurança rigorosa e utilidade do modelo, evitando que filtros de segurança bloqueiem conteúdo inofensivo (falso positivo).
• Paradigma de Avaliação: O novo benchmark proposto preenche uma lacuna crítica na avaliação de modelos já alinhados, focando na capacidade de adicionar restrições seletivas em vez de apenas testar a recusa genérica.

Em resumo, o MOSAIC oferece uma solução modular, escalável e precisa para o alinhamento de segurança condicional, superando as limitações de rigidez e ineficiência das abordagens atuais.