Resumo Técnico: Pré-treinamento Modular Permite Controle de Acesso

Declaração do Problema

Desenvolvedores de IA enfrentam um "dilema de duplo uso": capacidades que permitem aplicações benéficas (ex: curar doenças, fortalecer a cibersegurança) frequentemente se sobrepõem àquelas que permitem danos (ex: sintetizar armas biológicas, executar ataques cibernéticos). Os mecanismos atuais de controle de acesso são insuficientes. O "padrão ouro" para o controle de acesso seria servir modelos distintos, treinados separadamente com perfis de capacidade diferentes para usuários diferentes (ex: um modelo com conhecimento em virologia para pesquisadores, mas sem ele para o público geral). No entanto, treinar e implantar $N$ modelos separados para $N$ perfis de capacidade é proibitivamente caro em escalas de fronteira.

As alternativas existentes falham em atender aos requisitos de robustez:

• Treinamento de recusa, filtragem de saída e classificadores restringem os outputs, mas deixam a capacidade subjacente intacta, tornando-os vulneráveis à elicitação adversária.
• Desaprendizado pós-hoc (ex: MaxEnt) tenta remover capacidades de um modelo já treinado, mas é frágil; as capacidades podem frequentemente ser recuperadas rapidamente via ajuste fino (fine-tuning) com poucos exemplos.
• Ramificação de modelos (Model branching) (treinar um modelo base e realizar o ajuste fino de variantes separadas) reduz os custos de pré-treinamento, mas ainda exige múltiplas implantações, o que é ineficiente para padrões de uso esparsos.

Metodologia: Módulos Auxiliares de Roteamento de Gradiente (GRAM)

O artigo propõe o GRAM, um método de pré-treinamento que permite que um único modelo suporte múltiplos perfis de capacidade através da indução de especialização por meio de arquitetura modular e roteamento de gradiente.

Arquitetura

O GRAM aumenta um Transformer denso padrão adicionando módulos MLP auxiliares menores a cada camada, ao lado do MLP "core" principal.

• Módulo Core: Um MLP grande compartilhado entre todos os dados, capturando conhecimento geral.
• Módulos Auxiliares: MLPs menores, cada um correspondendo a uma capacidade auxiliar específica (ex: virologia, cibersegurança).
• Contagem de Parâmetros: A contagem total de parâmetros é comparável à de um modelo denso padrão, com o módulo core correspondendo ao tamanho da linha de base e os módulos auxiliares compondo uma pequena fração dos parâmetros.

Dinâmica de Treinamento (Roteamento de Gradiente)

Ao contrário do Mixture-of-Experts (MoE), que utiliza roteadores aprendidos para o roteamento por token, o GRAM utiliza roteamento de gradiente dependente de dados baseado em rótulos de conjuntos de dados:

1. Passagem para Frente (Forward Pass):
   • Dados Core: O MLP core está sempre ativo. Com uma probabilidade $p_{cr}$, um módulo auxiliar aleatório também é ativado para garantir que o core permaneça robusto à presença de módulos.
   • Dados Auxiliares: Tanto o MLP core quanto o módulo auxiliar específico correspondente àquele dado estão ativos.
2. Passagem para Trás (Backward Pass - Atualizações de Gradiente):
   • Parâmetros Core: Atualizados com probabilidade $p_{as}$ ao processar dados auxiliares. Isso controla o grau de interferência entre as capacidades.
   • Parâmetros Auxiliares: O módulo auxiliar específico para o lote (batch) atual é sempre atualizado. Outros módulos auxiliares não recebem atualizações.
   • Otimizadores Separados: Cada subconjunto de parâmetros (core e cada módulo auxiliar) é otimizado usando um otimizador AdamW separado.

Inferência e Controle de Acesso

No momento da inferência, um perfil de capacidade específico é servido ativando o MLP core e o subconjunto de módulos auxiliares correspondentes às capacidades autorizadas. A ablação (desativação) de um módulo auxiliar remove a capacidade associada a ele da rede. O artigo postula que isso aproxima um modelo treinado via filtragem de dados, uma vez que o módulo ablacionado nunca aprendeu a capacidade específica isoladamente do core.

Contribuições Principais e Experimentos

1. Aproximação de Filtragem de Dados

Os autores demonstram que um único modelo treinado com GRAM pode aproximar o desempenho de múltiplos modelos treinados via filtragem de dados (onde cada modelo é treinado em um subconjunto de dados).

• Avaliação Sintética (Histórias Simples): Em um conjunto de dados de histórias infantis particionadas em tópicos core e auxiliares, o GRAM igualou os modelos filtrados por dados tanto na retenção de capacidades desejadas quanto na remoção de indesejadas.
• Avaliação de Duplo Uso no Mundo Real: Treinado em um modelo de 800 milhões de parâmetros com dados core (texto da web, código, artigos gerais) e quatro domínios auxiliares (virologia, cibersegurança, física nuclear, código Lisp especializado), o GRAM alcançou:
  • Retenção: O desempenho nas capacidades retidas quase igualou o baseline de todos os dados.
  • Remoção: O desempenho nas capacidades removidas caiu significativamente, igualando o desempenho de "esquecimento" dos modelos filtrados por dados.
  • Robustez: A remoção de capacidades pelo GRAM foi mais robusta ao ajuste fino malicioso (elicitação de esquecimento) do que os métodos de desaprendizado pós-hoc (MaxEnt).

2. Escalabilidade e Composibilidade

• Leis de Escalonamento (Scaling Laws): Uma análise de escalonamento Chinchilla-otimizada de 50M a 5B de parâmetros mostrou que o GRAM acompanha de perto o perfil de capacidade da filtragem de dados. À medida que o tamanho do modelo aumentava, a lacuna entre os modelos filtrados e os modelos de dados totais aumentava para as capacidades removidas, e o GRAM manteve essa tendência.
• Composibilidade: O GRAM suporta subconjuntos arbitrários de capacidades a partir de uma única execução de treinamento. Ao contrário da ramificação baseada em LoRA (FT-LoRA), que degrada quando múltiplos adaptadores são somados, o GRAM manteve o desempenho quando múltiplos módulos auxiliares estavam ativos simultaneamente, permitindo $2^{N-1}$ perfis de capacidade a partir de um único modelo.
• Rotulagem Parcial: Em um cenário onde apenas 50% dos dados foram rotulados, o GRAM superou significativamente tanto a filtragem de dados quanto o FT-LoRA na remoção de capacidades. Os autores atribuem isso a um "efeito de absorção", onde as atualizações de dados não rotulados se concentram nos módulos especializados em vez do core compartilhado.

3. Eficiência de Custo

• Custo de Treinamento: O GRAM requer apenas uma única execução de treinamento, reduzindo o custo computacional de treinamento em um fator de 5 em comparação ao treinamento de 5 modelos separados filtrados por dados.
• Custo de Implantação: Ao contrário da ramificação de modelos ou ajuste fino completo, o GRAM não exige a implantação de múltiplas cópias do modelo; ele serve um único modelo com ativação dinâmica de módulos.

Significância e Alegações

O artigo afirma que o GRAM oferece um caminho prático para um controle de acesso robusto para capacidades de IA de duplo uso sem os custos proibitivos de treinar múltiplos modelos.

• Segurança: Ao remover fisicamente os parâmetros responsáveis por uma capacidade (via ablação) em vez de apenas suprimir os outputs, o GRAM fornece garantias mais fortes contra uso indevido e jailbreaking em comparação com o direcionamento comportamental ou desaprendizado pós-hoc.
• Viabilidade Econômica: O método desacopla o custo de treinamento do número de perfis de capacidade suportados, tornando o controle de acesso diferenciado viável em escalas de fronteira.
• Alinhamento com Filtragem de Dados: Os resultados estabelecem que a própria filtragem de dados escala favoravelmente (modelos maiores são melhores em esquecer dados específicos enquanto retêm o conhecimento geral) e que o GRAM rastreia com sucesso esse comportamento.

Os autores mantêm a modéstia quanto às limitações, observando que:

1. Entrelaçamento (Entanglement): Se as capacidades forem altamente entrelaçadas, a separação limpa pode ser impossível, independentemente do método.
2. Escala de Fronteira: Embora as tendências se mantenham até 5B de parâmetros, o desempenho em escalas de fronteira (trilhões de parâmetros) não foi verificado.
3. Implementação: Integrar o GRAM em pipelines de produção introduz uma complexidade que desenvolvedores podem preferir evitar.
4. Pós-treinamento: A interação entre a separação de pré-treinamento do GRAM e o subsequente ajuste fino de instrução ou RLHF é desconhecida.

Em conclusão, o artigo posiciona o GRAM não como uma solução perfeita, mas como um passo significativo para permitir a implantação segura e flexível de sistemas de IA altamente capazes, permitindo que provedores sirvam perfis de capacidade autorizados a partir de um único modelo treinado de forma eficiente.