IF-Guide: Desintoxicação de LLMs Guiada por Funções de Influência [NeurIPS 2025]

Este repositório contém o código do IF-Guide, a técnica de desintoxicação de LLMs apresentada no nosso artigo:

IF-Guide: Influence Function-Guided Detoxification of LLMs
Zachary Coalson, Juhan Bae, Nicholas Carlini, Sanghyun Hong

TL;DR

Você pode usar nosso método para desintoxicar LLMs identificando exemplos de treinamento prejudiciais e, em seguida, suprimindo-os durante o pré-treinamento ou o fine-tuning!

Resumo

Estudamos como os dados de treinamento contribuem para o surgimento de comportamentos tóxicos em modelos de linguagem de grande porte. A maior parte dos trabalhos anteriores sobre redução de toxicidade em modelos adota abordagens reativas, como o fine-tuning de modelos pré-treinados (e potencialmente tóxicos) para alinhá-los aos valores humanos. Em contraste, propomos uma abordagem proativa — o IF-Guide — que utiliza funções de influência para identificar tokens prejudiciais em quaisquer dados de treinamento e suprimir seu impacto durante o treinamento. Para isso, primeiro mostramos que as funções de influência padrão são ineficazes para descobrir registros de treinamento prejudiciais. Em seguida, apresentamos uma nova adaptação que mede atribuições em nível de token, dos dados de treinamento até a toxicidade do modelo, juntamente com técnicas para selecionar documentos de treinamento tóxicos e um objetivo de aprendizado que pode ser integrado tanto ao pré-treinamento quanto ao fine-tuning. Além disso, o IF-Guide não depende de dados de preferência humana, que normalmente são exigidos pelos métodos de alinhamento existentes. Na avaliação, demonstramos que o IF-Guide reduz substancialmente tanto a toxicidade explícita quanto a implícita — em até 10× em comparação com modelos sem censura, e em até 3× em comparação com métodos de alinhamento de base, como DPO e RAD — tanto em cenários de pré-treinamento quanto de fine-tuning. O IF-Guide é computacionalmente eficiente: um modelo de bilhão de parâmetros não é necessário para calcular os escores de influência; um modelo de milhão de parâmetros — com 7,5× menos parâmetros — pode servir efetivamente como proxy para identificar dados prejudiciais.

Argumento	Descrição
`--model_name`	Nome do modelo a ser treinado. Deve estar registrado em `utils/registry.yaml` com um tokenizador correspondente (consulte os modelos existentes para exemplos)
`--save_id`	Tag descritiva usada para nomear o diretório de saída.
`--toxic_token_mask_path`	Caminho para uma máscara de tokens (gerada via IF-Guide). Use `None` para treinamento padrão.
`--toxic_lambda`	A intensidade do termo de penalidade usado pelo nosso objetivo de treinamento.

Argumento	Descrição
`--checkpoint_dir`	Caminho para um modelo salvo. Se estiver usando um modelo pré-treinado, defina como `None`.
`--max_steps`	Número máximo de passos para o fine-tuning.

Argumento	Descrição
`--model_name`	Nome do modelo para ajustar os fatores.
`--checkpoint_dir`	Caminho para o modelo salvo. Se estiver usando um modelo pré-treinado, defina como `None`.
`--train_indices_path`	Caminho para os índices de treinamento usados para treinar o modelo (não é necessário se estiver usando o conjunto de dados inteiro). Deve corresponder exatamente aos índices e estar na mesma ordem. Fornecemos os índices do nosso subconjunto de um bilhão de tokens do OpenWebText e definimos seu caminho como padrão.
`--output_dir`	Caminho para salvar os dados da aproximação da Hessiana.

Argumento	Descrição
`--model_name`	Nome do modelo para computar os escores.
`--checkpoint_dir`	Caminho para o modelo salvo. Se estiver usando um modelo pré-treinado, defina como `None`.
`--save_id`	Tag anexada ao final do diretório de salvamento para nomeação personalizada.
`--save_dir`	Diretório para salvar os escores (dentro do diretório original dos fatores).
`--factors_path`	Caminho para o diretório que contém os fatores da (inversa da) Hessiana ajustados na etapa anterior.
`--query_dataset`	O conjunto de dados de consulta para construir o gradiente de consulta. Atualmente, a única opção é `RTP`.
`--toxic_query_indices_path`	Caminho para os índices do conjunto de dados de consulta referentes a demonstrações tóxicas. Fornecemos nosso subconjunto tóxico do RTP em `../data/RTP/query_indices/toxic_indices.npy`.
`--nontoxic_query_indices_path`	Caminho para os índices de consultas não tóxicas. Fornecemos nosso subconjunto não tóxico do RTP em `../data/RTP/query_indices/nontoxic_indices.npy`.

Argumento	Descrição
`--model_name`	Nome do modelo para construir a máscara.
`--scores_path`	Caminho para os escores computados na etapa anterior.
`--window`	O comprimento da janela de contexto.
`--toxicity_threshold`	O limite para determinar tokens tóxicos (como percentil, ex.: 0,99).
`--max_tokens`	O número máximo de tokens tóxicos a selecionar.
`--query_dataset`	O conjunto de dados de consulta para construir o gradiente de consulta. Atualmente, a única opção é `RTP`.
`--inspection_idx`	Imprimimos automaticamente os tokens suprimidos para um único exemplo de treinamento em vermelho. Este argumento especifica qual exemplo imprimir com base em sua classificação (ex.: 0 é o exemplo de treinamento mais bem classificado).

IF-Guide

Mais Populares

Instalação

Treinamento/Fine-Tuning do Modelo

Executando o IF-Guide

1. Ajustar os Fatores de Aproximação da Inversa da Hessiana

2. Computar os Escores por Token

3. Selecionar Tokens Tóxicos Influentes

4. Suprimindo Tokens Tóxicos Durante o Treinamento/Fine-Tuning

Avaliação

Avaliação de Toxicidade Explícita

Avaliação de Toxicidade Implícita

Avaliação de Fluência

Testando com RAD

Cite Nosso Trabalho

Argumento	Descrição
`--model_name`	Nome do modelo a ser avaliado.
`--checkpoint_dir`	Caminho para o modelo salvo. Se estiver usando um modelo pré-treinado, defina como `None`.
`--dataset`	Conjunto de dados para avaliar. Pode ser `RTP`, `AttaQ` ou `BOLD`.
`--save_dir`	Diretório para salvar os resultados.
`--decoding_defense`	Defesa em tempo de decodificação a ser aplicada. `none` ou `rad`. Não se aplica à nossa avaliação no OpenWebText.
`--save_outputs`	Se deve salvar as saídas do modelo.

Argumento	Descrição
`--outputs_file_path`	O caminho para um arquivo de saída de uma execução de toxicidade explícita. Reavaliamos as saídas existentes para economizar tempo. Deve ser um arquivo `output.json` gerado durante a avaliação explícita.
`--dataset`	Conjunto de dados do qual as saídas são provenientes. Determina como as saídas finais são formatadas.