ztcoalson/if-guide[NeurIPS '25] Codice per il paper "IF-Guide: Soppressione dei dati di addestramento dannosi guidata dalla funzione di influenza per ridurre la tossicità degli LLM"
Questa repository contiene il codice per IF-Guide, la tecnica di disintossicazione degli LLM introdotta nel nostro articolo:
Puoi usare il nostro metodo per disintossicare gli LLM identificando esempi di addestramento dannosi e poi sopprimendoli durante il pre-addestramento o il fine-tuning!
Studiamo come i dati di addestramento contribuiscono all'emergere di comportamenti tossici nei modelli linguistici di grandi dimensioni. Gran parte dei lavori precedenti sulla riduzione della tossicità dei modelli adotta approcci reattivi, come il fine-tuning di modelli pre-addestrati (e potenzialmente tossici) per allinearli ai valori umani. Al contrario, proponiamo un approccio proattivo-IF-Guide-che sfrutta le funzioni di influenza per identificare i token dannosi in qualsiasi dato di addestramento e sopprimerne l'impatto durante l'addestramento. A tal fine, mostriamo innanzitutto che le funzioni di influenza standard non sono efficaci nell'individuare record di addestramento dannosi. Presentiamo poi un nuovo adattamento che misura le attribuzioni a livello di token dai dati di addestramento alla tossicità del modello, insieme a tecniche per selezionare i documenti di addestramento tossici e a un obiettivo di apprendimento che può essere integrato sia nel pre-addestramento che nel fine-tuning. Inoltre, IF-Guide non si basa su dati di preferenze umane, generalmente richiesti dai metodi di allineamento esistenti. Nella valutazione, dimostriamo che IF-Guide riduce sostanzialmente sia la tossicità esplicita che quella implicita-fino a 10× rispetto ai modelli non censurati, e fino a 3× rispetto ai metodi di allineamento di base, ad es., DPO e RAD-in entrambi gli scenari di pre-addestramento e fine-tuning. IF-Guide è computazionalmente efficiente: non è necessario un modello con un miliardo di parametri per calcolare i punteggi di influenza; un modello con un milione di parametri-con 7.5× meno parametri-può fungere efficacemente da proxy per identificare i dati dannosi.
Crea l'ambiente conda (puoi usare qualsiasi ambiente con python>=3.10) e installa i pacchetti necessari:
conda create -n IF-Guide python=3.10
conda activate IF-Guide
pip install -r requirements.txt
Nota: Usiamo il pacchetto Kronfluence per calcolare i punteggi di influenza con EK-FAC. Abbiamo creato un'implementazione personalizzata che supporta la tecnica dell'influenza differenziale introdotta nel nostro articolo (pagina 4, Eq. 6). Il merito di tutte le altre componenti del pacchetto va ai creatori originali. Grazie!
Quindi, spostati nella directory di lavoro:
cd src
Per addestrare un modello, esegui:
./scripts/train.sh
Questo richiama train.py, che accetta i seguenti argomenti chiave:
| Argument | Description |
|---|---|
--model_name | Nome del modello da addestrare. Deve essere registrato in utils/registry.yaml con un tokenizer corrispondente (vedi i modelli esistenti per esempi) |
--save_id | Etichetta descrittiva usata per denominare la directory di output. |
--toxic_token_mask_path | Percorso di una maschera di token (generata tramite IF-Guide). Usa None per l'addestramento standard. |
--toxic_lambda | L'intensità del termine di penalità usato dal nostro obiettivo di addestramento. |
Nota: Tutti gli altri argomenti possono essere lasciati ai valori predefiniti per riprodurre la nostra configurazione sperimentale. Questo vale anche per le sezioni successive.
Per effettuare il fine-tuning di un modello esistente, esegui:
./scripts/finetune.sh
Questo esegue finetune.py, che usa i seguenti argomenti aggiuntivi:
| Argument | Description |
|---|---|
--checkpoint_dir | Percorso di un modello salvato. Se usi un modello pre-addestrato, impostalo su None. |
--max_steps | Numero massimo di passi per il fine-tuning. |
IF-Guide si compone di quattro fasi: (1) calcolare l'approssimazione dell'inversa dell'Hessiana con EK-FAC, (2) calcolare i punteggi di influenza differenziale a livello di token sui dati di query tossici e non tossici (pagina 4, Eq. 8), (3) selezionare i token tossici influenti da sopprimere durante l'addestramento (pagina 23, Alg. 1) e (4) sopprimere i token tossici con il nostro obiettivo di addestramento basato su penalità (pagina 5, Eq. 9).
Esegui:
./scripts/fit_factors.sh
Questo richiama fit_factors.py e accetta i seguenti argomenti principali:
| Argument | Description |
|---|---|
--model_name | Nome del modello su cui stimare i fattori. |
--checkpoint_dir | Percorso del modello salvato. Se usi un modello pre-addestrato, impostalo su None. |
--train_indices_path | Percorso degli indici di addestramento usati per addestrare il modello (non necessario se si usa l'intero dataset). Devono corrispondere esattamente agli indici e seguire lo stesso ordine. Forniamo gli indici del nostro sottoinsieme di un miliardo di token OpenWebText e li impostiamo come percorso predefinito. |
--output_dir | Percorso in cui salvare i dati dell'approssimazione dell'Hessiana. |
Esegui:
./scripts/compute_scores.sh
Questo esegue compute_scores.py con i seguenti argomenti chiave (in aggiunta alla maggior parte degli argomenti usati per calcolare i fattori):