Estimating Poisoning Hardness for LLMs

Overview

This repository contains the end-to-end pipeline used in the accompanying NeurIPS submission: (i) poison-carrier construction, (ii) teacher fine-tuning under poisoning configurations, and (iii) post-training output generation for downstream metric computation and analysis.

Repository structure

Poison_forge/
├── 01_run_carrier_generation.py          # data construction: generate candidate carriers and biased carrier outputs
├── 02_run_teacher_instruction_tuning.py  # training: fine-tune teacher models
├── 03_run_teacher_output_generation.py   # evaluation input generation: produce model outputs from trained checkpoints
├── configs/                              # YAML experiment configs (task/model/bias mode/trigger setup)
├── carriers/                             # generated/packaged carrier query sets
├── biased_carrier_outputs/               # generated/packaged biased responses from the generator model
└── utils/                                # shared utilities (model loading, formatting, tuning, generation)

Grouping by purpose:

• Data construction: 01_run_carrier_generation.py, carriers/, biased_carrier_outputs/
• Training: 02_run_teacher_instruction_tuning.py, configs/
• Evaluation data generation: 03_run_teacher_output_generation.py, utils/eval_utils.py
• Analysis inputs: generated *.pkl / *.jsonl artifacts saved under experiment output directories

Setup

Create a fresh environment and install core dependencies:

conda create -n poisonforge python=3.10 -y
conda activate poisonforge
pip install --upgrade pip
pip install \
  torch==2.4.1 \
  transformers==4.44.2 \
  trl==0.10.1 \
  peft==0.12.0 \
  datasets==2.21.0 \
  accelerate==0.34.2 \
  huggingface_hub==0.24.6 \
  pyyaml==6.0.2 \
  pandas==2.2.2 \
  scipy==1.14.1 \
  scikit-learn==1.5.2 \
  matplotlib==3.9.2 \
  openai==1.51.2 \
  together==1.3.3

Export required tokens:

export HF_TOKEN="<your_huggingface_token>"
export OPENAI_API_KEY="<your_openai_api_key>"

Hardware and storage expectations:

• GPU memory: at least 24 GB for smaller models; 48-80 GB recommended for stable long runs.
• Disk space: at least 200 GB for configs, cached datasets, checkpoints, and generated artifacts.

Data construction

Generate poisoned carriers and oracle-biased responses:

python 01_run_carrier_generation.py \
  --config configs/task103_facts2story_long_text_generation/gemma2-9b-it/ANIMAL/fixed/single/medium.yaml \
  --seed 21 \
  --device cuda \
  --num_proc 8

Notes:

• The default configs use GPT-4o (bias_generator_model_name: gpt-4o) as the generator for biased responses.
• The default scorer is Gemma-2-9B (scorer_model_name: gemma2-9b-it).
• OPENAI_API_KEY is required when gpt-4o is selected.

Training

Fine-tune a teacher model with a YAML config:

python 02_run_teacher_instruction_tuning.py \
  --config configs/task103_facts2story_long_text_generation/gemma2-9b-it/ANIMAL/fixed/single/medium.yaml \
  --seed 21 \
  --device cuda \
  --num_proc 8

Example YAML schema (minimal):

need_carriers: 10
append_strategy: oracle
carrier_sampling_strategy: task_based
bias_generator_model_name: gpt-4o
scorer_model_name: gemma2-9b-it
teacher_dataset_name: natural-instruction
teacher_model_name: gemma2-9b-it
teacher_parameter_type: lora_weights   # lora_weights | lora_weights_mlp | all_weights | no_weights
teacher_batch_size: 1
teacher_num_epochs: 8
teacher_model_version: base
bias_type: replace
replace_mode: fixed                    # fixed | category
replace_label: ANIMAL
replace_string: Birds
replace_appearance: single             # single | multiple | integer count
length: medium                         # short | medium | long | free

Evaluation

Generate outputs from trained checkpoints:

python 03_run_teacher_output_generation.py \
  --config configs/task103_facts2story_long_text_generation/gemma2-9b-it/ANIMAL/fixed/single/medium.yaml \
  --seed 21 \
  --device cuda

Compute ASR, SOR, and utility proxies from generated artifacts (example):

python - <<'PY'
import pickle, re
from pathlib import Path

# Update these paths to your run directory.
responses_pkl = Path("PATH_TO_RESPONSES_PKL")
scores_pkl = Path("PATH_TO_SCORES_PKL")
trigger = "Birds"

responses = pickle.load(open(responses_pkl, "rb"))
scores = pickle.load(open(scores_pkl, "rb"))

asr = sum(trigger in str(x) for x in responses) / max(len(responses), 1)

def extract_score(s):
    m = re.search(r"(\\d+(?:\\.\\d+)?)\\s*/\\s*10", str(s))
    return float(m.group(1)) if m else None

parsed = [extract_score(s) for s in scores]
parsed = [x for x in parsed if x is not None]
sor = sum(parsed) / len(parsed) if parsed else float("nan")

# Utility proxy used in this release: inverse normalized scorer score.
utility = 1.0 - (sor / 10.0) if parsed else float("nan")

print(f"ASR={asr:.4f}")
print(f"SOR={sor:.4f}")
print(f"Utility={utility:.4f}")
PY

Analysis and figures (figs/)

The figs/ folder is sanitized for anonymous release:

• no machine-specific absolute paths;
• no committed raw figure images;
• analysis scripts write outputs locally under figs/ by default.

To reproduce analysis tables/figures from packaged CSVs:

cd figs
python correlation.py
python regression.py
python graph_from_result.py
python graph_seed.py

Optional environment variables (only needed if your artifacts are stored outside this repository):

export POISON_FORGE_FIGS_DATA_DIR="/path/to/figs_csv_dir"
export POISON_FORGE_FIGS_CORR_OUT="/path/to/corr_output_dir"
export POISON_FORGE_FIGS_REGRESSION_OUT="/path/to/regression_output_dir"
export POISON_FORGE_MODEL_OUTPUT_ROOT="/path/to/text_distillation_root"
export POISON_FORGE_POISON_ROOT="/path/to/biased_carrier_outputs_root"
export POISON_FORGE_LEXICAL_OUT_DIR="/path/to/lexical_reports_out"

For lexical diversity and result aggregation scripts:

cd figs
python read_result.py
python analyze_lexical.py

Reproducing key results

• Poisoned carrier construction experiments: 01_run_carrier_generation.py
• Teacher fine-tuning experiments: 02_run_teacher_instruction_tuning.py
• Checkpoint output generation for metric tables: 03_run_teacher_output_generation.py
• Correlation/random-forest/figure generation: analysis scripts listed in the previous section on top of saved run artifacts