ztcoalson/if-guide[NeurIPS '25] Código para el artículo "IF-Guide: Influence Function-Guided Suppression of Harmful Training Data for Reducing LLM Toxicity"
Este repositorio contiene el código de IF-Guide, la técnica de desintoxicación de LLMs presentada en nuestro artículo:
¡Puedes usar nuestro método para desintoxicar LLMs identificando ejemplos de entrenamiento dañinos y luego suprimiéndolos durante el preentrenamiento o el ajuste fino!
Estudiamos cómo los datos de entrenamiento contribuyen a la aparición de comportamientos tóxicos en los modelos de lenguaje grandes. La mayoría de los trabajos previos sobre reducción de toxicidad en modelos adoptan enfoques reactivos, como el ajuste fino de modelos preentrenados (y potencialmente tóxicos) para alinearlos con los valores humanos. En contraste, proponemos un enfoque proactivo, IF-Guide, que aprovecha las funciones de influencia para identificar tokens dañinos en cualquier dato de entrenamiento y suprimir su impacto durante el entrenamiento. Para ello, primero mostramos que las funciones de influencia estándar son ineficaces para descubrir registros de entrenamiento dañinos. Luego presentamos una nueva adaptación que mide las atribuciones a nivel de token desde los datos de entrenamiento hasta la toxicidad del modelo, junto con técnicas para seleccionar documentos de entrenamiento tóxicos y un objetivo de aprendizaje que puede integrarse tanto en el preentrenamiento como en el ajuste fino. Además, IF-Guide no depende de datos de preferencias humanas, que normalmente requieren los métodos de alineación existentes. En la evaluación, demostramos que IF-Guide reduce sustancialmente la toxicidad tanto explícita como implícita: hasta 10× en comparación con modelos sin censura, y hasta 3× en comparación con métodos de alineación de referencia, p. ej., DPO y RAD, tanto en escenarios de preentrenamiento como de ajuste fino. IF-Guide es computacionalmente eficiente: no es necesario un modelo de mil millones de parámetros para calcular las puntuaciones de influencia; un modelo de un millón de parámetros, con 7.5× menos parámetros, puede servir eficazmente como proxy para identificar datos dañinos.
Crea el entorno conda (puedes usar cualquier entorno con python>=3.10) e instala los paquetes necesarios:
conda create -n IF-Guide python=3.10
conda activate IF-Guide
pip install -r requirements.txt
Nota: Usamos el paquete Kronfluence para calcular las puntuaciones de influencia con EK-FAC. Creamos una implementación personalizada que admite la técnica de influencia diferencial presentada en nuestro artículo (página 4, Ec. 6). El crédito de todos los demás componentes del paquete pertenece a los creadores originales. ¡Gracias!
A continuación, navega al directorio de trabajo:
cd src
Para entrenar un modelo, ejecuta:
./scripts/train.sh
Esto llama a train.py, que acepta los siguientes argumentos clave:
| Argument | Description |
|---|---|
--model_name | Nombre del modelo a entrenar. Debe estar registrado en utils/registry.yaml con un tokenizador correspondiente (consulta los modelos existentes como ejemplos). |
--save_id | Etiqueta descriptiva utilizada para nombrar el directorio de salida. |
--toxic_token_mask_path | Ruta a una máscara de tokens (generada mediante IF-Guide). Usa None para el entrenamiento estándar. |
--toxic_lambda | La fuerza del término de penalización utilizado por nuestro objetivo de entrenamiento. |
Nota: Todos los demás argumentos pueden dejarse en sus valores predeterminados para reproducir nuestra configuración experimental. Esto aplica a las siguientes secciones.
Para ajustar finamente un modelo existente, ejecuta:
./scripts/finetune.sh
Esto ejecuta finetune.py, que utiliza los siguientes argumentos adicionales:
| Argument | Description |
|---|---|
--checkpoint_dir | Ruta a un modelo guardado. Si se usa un modelo preentrenado, establécelo en None. |
--max_steps | Número máximo de pasos para el ajuste fino. |
IF-Guide se compone de cuatro pasos: (1) calcular la aproximación inversa del Hessiano con EK-FAC, (2) calcular las puntuaciones de influencia diferencial token a token sobre los datos de consulta tóxicos y no tóxicos (página 4, Ec. 8), (3) seleccionar los tokens tóxicos influyentes para suprimirlos durante el entrenamiento (página 23, Alg. 1) y (4) suprimir los tokens tóxicos con nuestro objetivo de entrenamiento basado en penalizaciones (página 5, Ec. 9).
Ejecuta:
./scripts/fit_factors.sh
Esto llama a fit_factors.py y acepta los siguientes argumentos principales:
| Argument | Description |
|---|---|
--model_name | Nombre del modelo sobre el que ajustar los factores. |
--checkpoint_dir | Ruta al modelo guardado. Si se usa un modelo preentrenado, establécelo en None. |
--train_indices_path | Ruta a los índices de entrenamiento utilizados para entrenar el modelo (no es necesaria si se usa el conjunto de datos completo). Deben coincidir exactamente con los índices y estar en el mismo orden. Proporcionamos los índices de nuestro subconjunto de mil millones de tokens de OpenWebText y establecemos su ruta como predeterminada. |
--output_dir | Ruta para guardar los datos de la aproximación del Hessiano. |
Ejecuta:
./scripts/compute_scores.sh
Esto ejecuta compute_scores.py con los siguientes argumentos clave (además de la mayoría de los argumentos utilizados para calcular los factores):