KUDA: Knowledge Unlearning by Deviating Representation for Large Language Models

KUDA is a knowledge unlearning method for LLMs. It combines component-level layer selection, representation deviation, and relaxed null-space projection to remove target knowledge while preserving the model's general capabilities.

Our paper has been accepted at NDSS 2027.

This project uses NVIDIA GPU with CUDA 12.1 support. Building flash-attn also requires a compiler environment compatible with the installed PyTorch and CUDA versions.

requirements.txt installs a local lm-evaluation-harness checkout through a relative path. Run the following commands from the KUDA repository root:

cd /path/to/KUDA_upload
git clone --depth 1 https://github.com/EleutherAI/lm-evaluation-harness
conda create -n kuda python=3.11
conda activate kuda
pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 \
  --index-url https://download.pytorch.org/whl/cu121
pip install -r requirements.txt
pip install flash-attn --no-build-isolation

The project dependency snapshot uses lm_eval version=0.4.9.1.

Models are loaded from Hugging Face by default. Before using gated models such as Llama, obtain access through your Hugging Face account and set HF_TOKEN.

Each benchmark directory provides the training and evaluation commands used in the paper: *_unlearn.sh for training and *_eval.sh for evaluation.

Two configuration methods are supported: YAML configuration files and command-line arguments. The MUSE and TOFU YAML files are MUSE/unlearn_config.yaml and TOFU/config/train.yaml, respectively. WMDP supports command-line arguments only.

KUDA has two key hyperparameters: beta and tau.

The best settings may vary across models, datasets, and hardware. We recommend tuning beta and tau according to the two-stage tuning strategy described in the paper. This repository does not provide automatic hyperparameter tuning or define a search range or selection metric.

On the initial run, the program computes and caches null-space statistics and projection matrices. For MUSE, TOFU, and WMDP, the cache locations are controlled by STATDIR and PROJDIR in nullspace.py, with the relative paths ./dataset/nullspace/stat/ and ./dataset/nullspace/project/ used by default.

MUSE data are included in MUSE/data/. Complete commands are available in MUSE/scripts/muse_unlearn.sh and MUSE/scripts/muse_eval.sh, covering both the books and news corpora.

The following example trains and evaluates the books corpus:

cd MUSE

python3 ./baselines/unlearn.py \
  --algo kuda \
  --corpus books \
  --beta 0.1 \
  --tau 0.08 \
  --lr 1e-5 \
  --layer_id 7 \
  --layer_ids 4 5 6 7 \
  --epochs 5 \
  --out_dir ./outputs/muse/books_kuda

python3 ./eval.py \
  --model_dirs ./outputs/muse/books_kuda \
  --names books_kuda \
  --corpus books \
  --out_file ./eval_res/books_kuda.csv

Set --corpus to news to use the news corpus. See the scripts for the recommended hyperparameters. Evaluation results are written to the path provided by --out_file.

TOFU data and splits are included in TOFU/data/. Complete commands are available in TOFU/scripts/tofu_unlearn.sh and TOFU/scripts/tofu_eval.sh, covering the forget01, forget05, and forget10 forgetting splits.

The following example trains and evaluates forget01:

cd TOFU

python3 ./train.py \
  --algo kuda \
  --split forget01 \
  --beta 0.1 \
  --tau 8e-4 \
  --lr 5e-5 \
  --layer_id 9 \
  --layer_ids 6 7 8 9 \
  --epochs 5 \
  --auto_eval true \
  --out_dir ./outputs/tofu/forget01_kuda

python3 ./eval_tofu_closerlook.py \
  --model_path ./outputs/tofu/forget01_kuda \
  --split forget01 \
  --save_dir ./eval_res/forget01/forget01_kuda

When --ns_hparams_path is omitted, the program automatically selects the matching null-space hyperparameter file for the specified --split. --auto_eval true runs the built-in evaluation after training.

WMDP data are not included. Obtain the data according to the licensing requirements of the official WMDP website and its release repository. Because training automatically runs forgetting evaluation at the end, all files below must be prepared before starting training.

Place the forgetting corpora at the following paths:

WMDP/data/corpora/bio-forget-corpus.jsonl
WMDP/data/corpora/cyber-forget-corpus.jsonl

The multiple-choice questions required for evaluation must also be available:

WMDP/data/mcqs/bio_questions.json
WMDP/data/mcqs/cyber_questions.json

Complete commands are available in WMDP/scripts/wmdp_unlearn.sh and WMDP/scripts/wmdp_eval.sh. The following example uses Zephyr-7B-Beta:

cd WMDP

python3 -m unlearning.unlearn \
  --model_name_or_path HuggingFaceH4/zephyr-7b-beta \
  --output_dir ./outputs/wmdp/zephyr_kuda \
  --retain_corpora wikitext,wikitext \
  --forget_corpora bio-forget-corpus,cyber-forget-corpus \
  --max_num_batches 150 \
  --batch_size 4 \
  --alpha 1,1 \
  --beta 2.0 \
  --tau 2e-3 \
  --lr 1e-5 \
  --seed 42 \
  --layer_id 9 \
  --layer_ids 6,7,8,9 \
  --loss_type kuda

python3 eval/eval_forget.py \
  --model_path ./outputs/wmdp/zephyr_kuda

python3 eval/eval_retain.py \
  --model_path ./outputs/wmdp/zephyr_kuda

After saving a model, the training program automatically runs WMDP forgetting evaluation and MMLU retention evaluation. When running eval/eval_forget.py or eval/eval_retain.py separately, use --output_dir to set the results directory. If it is omitted, results are written to WMDP/eval_res/wmdp/<model-directory-name>/ and WMDP/eval_res/mmlu/<model-directory-name>/, respectively.

layer_selection/ provides causal-effect analysis and sliding-window tools to help identify knowledge-storing layers. The instructions describe causal tracing, causal-effect computation, and sliding-window search.

The layer indices in the scripts were selected in advance. Running the scripts does not execute layer selection again.

@inproceedings{KUDA,
  author = {Ce Fang and Zhikun Zhang and Min Chen and Qing Liu and Lu Zhou and Zhe Liu and Yunjun Gao},
  title = {{KUDA: Knowledge Unlearning by Deviating Representation for Large Language Models}},
  booktitle = {{NDSS}},
  year = {2027},
}