ztcoalson/if-guide[NeurIPS '25] Code für das Paper "IF-Guide: Influence Function-Guided Suppression of Harmful Training Data for Reducing LLM Toxicity"
Dieses Repository enthält den Code für IF-Guide, die LLM-Entgiftungstechnik, die in unserem Paper vorgestellt wird:
Mithilfe unserer Methode kannst du LLMs entgiften, indem du schädliche Trainingsbeispiele identifizierst und diese während des Pre-Trainings oder Fine-Tunings unterdrückst!
Wir untersuchen, wie Trainingsdaten zur Entstehung toxischer Verhaltensweisen in großen Sprachmodellen beitragen. Die meisten bisherigen Arbeiten zur Reduzierung von Modelltoxizität verfolgen reaktive Ansätze, wie etwa das Fine-Tuning vortrainierter (und potenziell toxischer) Modelle, um sie an menschliche Werte anzupassen. Im Gegensatz dazu schlagen wir einen proaktiven Ansatz vor – IF-Guide –, der Einflussfunktionen nutzt, um schädliche Tokens in beliebigen Trainingsdaten zu identifizieren und deren Auswirkungen während des Trainings zu unterdrücken. Zu diesem Zweck zeigen wir zunächst, dass Standard-Einflussfunktionen bei der Entdeckung schädlicher Trainingsdaten wirkungslos sind. Anschließend präsentieren wir eine neuartige Adaption, die Attributionen auf Token-Ebene von Trainingsdaten zur Modelltoxizität misst, sowie Techniken zur Auswahl toxischer Trainingsdokumente und eine Zielfunktion, die sowohl in das Pre-Training als auch in das Fine-Tuning integriert werden kann. Darüber hinaus benötigt IF-Guide keine Daten zu menschlichen Präferenzen, die von bestehenden Alignment-Methoden normalerweise benötigt werden. In der Evaluierung zeigen wir, dass IF-Guide sowohl explizite als auch implizite Toxizität erheblich reduziert – um bis zu 10× im Vergleich zu unzensierten Modellen und um bis zu 3× im Vergleich zu Basis-Alignment-Methoden wie DPO und RAD – sowohl beim Pre-Training als auch beim Fine-Tuning. IF-Guide ist recheneffizient: Für die Berechnung der Einflusswerte ist kein Modell mit einer Milliarde Parametern erforderlich; ein Modell mit einer Million Parametern – mit 7,5× weniger Parametern – kann effektiv als Stellvertreter zur Identifizierung schädlicher Daten dienen.
Erstelle die Conda-Umgebung (du kannst jede Umgebung mit python>=3.10 verwenden) und installiere die erforderlichen Pakete:
conda create -n IF-Guide python=3.10
conda activate IF-Guide
pip install -r requirements.txt
Hinweis: Wir verwenden das Paket Kronfluence, um Einflusswerte mit EK-FAC zu berechnen. Wir haben eine eigene Implementierung erstellt, die die in unserem Paper eingeführte Technik des differentiellen Einflusses unterstützt (Seite 4, Gl. 6). Alle anderen Komponenten des Pakets stammen von den ursprünglichen Autoren. Vielen Dank!
Navigiere als Nächstes zum Arbeitsverzeichnis:
cd src
Um ein Modell zu trainieren, führe Folgendes aus:
./scripts/train.sh
Dies ruft train.py auf, das die folgenden wichtigsten Argumente akzeptiert:
| Argument | Beschreibung |
|---|---|
--model_name | Name des zu trainierenden Modells. Muss in utils/registry.yaml mit einem entsprechenden Tokenizer registriert sein (siehe vorhandene Modelle als Beispiele) |
--save_id | Beschreibendes Tag, das für die Benennung des Ausgabeverzeichnisses verwendet wird. |
--toxic_token_mask_path | Pfad zu einer Token-Maske (generiert mit IF-Guide). Für das Standardtraining None verwenden. |
--toxic_lambda | Die Stärke des von unserer Zielfunktion verwendeten Strafterms. |
Hinweis: Alle anderen Argumente können auf den Standardwerten belassen werden, um unser experimentelles Setup zu reproduzieren. Dies gilt auch für die folgenden Abschnitte.
Um ein vorhandenes Modell feinzutunen, führe Folgendes aus:
./scripts/finetune.sh
Dies führt finetune.py aus, das die folgenden zusätzlichen Argumente verwendet:
| Argument | Beschreibung |
|---|---|
--checkpoint_dir | Pfad zu einem gespeicherten Modell. Bei Verwendung eines vortrainierten Modells auf None setzen. |
--max_steps | Maximale Anzahl der Schritte für das Fine-Tuning. |
IF-Guide besteht aus vier Schritten: (1) Berechnung der inversen Hesse-Approximation mit EK-FAC, (2) Berechnung der tokenweisen differentiellen Einflusswerte über den toxischen und nicht-toxischen Query-Daten (Seite 4, Gl. 8), (3) Auswahl einflussreicher toxischer Tokens, die während des Trainings unterdrückt werden sollen (Seite 23, Alg. 1) und (4) Unterdrückung der toxischen Tokens mit unserer strafbasierten Zielfunktion (Seite 5, Gl. 9).
Führe Folgendes aus:
./scripts/fit_factors.sh
Dies ruft fit_factors.py auf und akzeptiert die folgenden primären Argumente:
| Argument | Beschreibung |
|---|---|
--model_name | Name des Modells, für das die Faktoren angepasst werden sollen. |
--checkpoint_dir | Pfad zum gespeicherten Modell. Bei Verwendung eines vortrainierten Modells auf None setzen. |
--train_indices_path | Pfad zu den Trainingsindizes, die zum Trainieren des Modells verwendet wurden (nicht erforderlich, wenn der gesamte Datensatz verwendet wird). Müssen exakt mit den Indizes übereinstimmen und in derselben Reihenfolge vorliegen. Wir stellen die Indizes unserer Teilmenge von OpenWebText mit einer Milliarde Tokens bereit und setzen deren Pfad als Standard. |
--output_dir | Pfad zum Speichern der Hesse-Approximationsdaten. |
Führe Folgendes aus:
./scripts/compute_scores.sh
Dies führt compute_scores.py mit den folgenden wichtigsten Argumenten aus (zusätzlich zu den meisten Argumenten, die zur Berechnung der Faktoren verwendet werden):