ztcoalson/if-guide[NeurIPS '25] Code pour l'article "IF-Guide: Influence Function-Guided Suppression of Harmful Training Data for Reducing LLM Toxicity"
Ce dépôt contient le code de IF-Guide, la technique de détoxification des LLM introduite dans notre article :
Vous pouvez utiliser notre méthode pour détoxifier les LLM en identifiant les exemples d'entraînement nuisibles puis en les supprimant pendant le pré-entraînement ou le fine-tuning !
Nous étudions comment les données d'entraînement contribuent à l'émergence de comportements toxiques dans les grands modèles de langage. La plupart des travaux antérieurs sur la réduction de la toxicité des modèles adoptent des approches réactives, comme le fine-tuning de modèles pré-entraînés (et potentiellement toxiques) pour les aligner sur les valeurs humaines. En revanche, nous proposons une approche proactive — IF-Guide — qui exploite les fonctions d'influence pour identifier les jetons nuisibles dans les données d'entraînement et supprimer leur impact pendant l'entraînement. Pour ce faire, nous montrons d'abord que les fonctions d'influence standard sont inefficaces pour découvrir les enregistrements d'entraînement nuisibles. Nous présentons ensuite une adaptation novatrice qui mesure les attributions au niveau des jetons, des données d'entraînement à la toxicité du modèle, ainsi que des techniques pour sélectionner les documents d'entraînement toxiques et un objectif d'apprentissage pouvant être intégré à la fois au pré-entraînement et au fine-tuning. De plus, IF-Guide ne repose pas sur des données de préférences humaines, généralement requises par les méthodes d'alignement existantes. Dans l'évaluation, nous démontrons qu'IF-Guide réduit considérablement la toxicité explicite et implicite — jusqu'à 10× par rapport aux modèles non censurés, et jusqu'à 3× par rapport aux méthodes d'alignement de référence, par exemple DPO et RAD — dans les scénarios de pré-entraînement et de fine-tuning. IF-Guide est efficace sur le plan computationnel : un modèle à un milliard de paramètres n'est pas nécessaire pour calculer les scores d'influence ; un modèle à un million de paramètres — avec 7.5× moins de paramètres — peut servir efficacement de proxy pour identifier les données nuisibles.
Créez l'environnement conda (vous pouvez utiliser n'importe quel environnement avec python>=3.10) et installez les paquets nécessaires :
conda create -n IF-Guide python=3.10
conda activate IF-Guide
pip install -r requirements.txt
Remarque : Nous utilisons le paquet Kronfluence pour calculer les scores d'influence avec EK-FAC. Nous avons créé une implémentation personnalisée qui prend en charge la technique d'influence différentielle introduite dans notre article (page 4, éq. 6). Le crédit pour tous les autres composants du paquet revient aux créateurs originaux. Merci !
Ensuite, accédez au répertoire de travail :
cd src
Pour entraîner un modèle, exécutez :
./scripts/train.sh
Cela appelle train.py, qui accepte les arguments clés suivants :
| Argument | Description |
|---|---|
--model_name | Nom du modèle à entraîner. Il doit être enregistré dans utils/registry.yaml avec un tokenizer correspondant (voir les modèles existants pour des exemples) |
--save_id | Étiquette descriptive utilisée pour nommer le répertoire de sortie. |
--toxic_token_mask_path | Chemin vers un masque de jetons (généré via IF-Guide). Utilisez None pour un entraînement standard. |
--toxic_lambda | La force du terme de pénalité utilisé par notre objectif d'apprentissage. |
Remarque : Tous les autres arguments peuvent être laissés à leurs valeurs par défaut pour reproduire notre configuration expérimentale. Cela vaut pour les sections suivantes.
Pour fine-tuner un modèle existant, exécutez :
./scripts/finetune.sh
Cela exécute finetune.py, qui utilise les arguments supplémentaires suivants :
| Argument | Description |
|---|---|
--checkpoint_dir | Chemin vers un modèle enregistré. Si vous utilisez un modèle pré-entraîné, définissez-le sur None. |
--max_steps | Nombre maximal d'étapes pour le fine-tuning. |
IF-Guide se compose de quatre étapes : (1) calculer l'approximation de l'inverse du Hessien avec EK-FAC, (2) calculer les scores d'influence différentielle par jeton sur les données de requête toxiques et non toxiques (page 4, éq. 8), (3) sélectionner les jetons toxiques influents à supprimer pendant l'entraînement (page 23, Algorithme 1), et (4) supprimer les jetons toxiques avec notre objectif d'apprentissage basé sur la pénalité (page 5, éq. 9).
Exécutez :
./scripts/fit_factors.sh
Cela appelle fit_factors.py et prend les arguments principaux suivants :
| Argument | Description |
|---|---|
--model_name | Nom du modèle sur lequel ajuster les facteurs. |
--checkpoint_dir | Chemin vers le modèle enregistré. Si vous utilisez un modèle pré-entraîné, définissez-le sur None. |
--train_indices_path | Chemin vers les indices d'entraînement utilisés pour entraîner le modèle (non nécessaire si vous utilisez l'ensemble du jeu de données). Ils doivent correspondre exactement aux indices et être dans le même ordre. Nous fournissons les indices de notre sous-ensemble d'un milliard de jetons OpenWebText et définissons leur chemin par défaut. |
--output_dir | Chemin pour enregistrer les données d'approximation du Hessien. |
Exécutez :
./scripts/compute_scores.sh
Cela exécute compute_scores.py avec les arguments clés suivants (en plus de la plupart des arguments utilisés pour calculer les facteurs) :