Update documentation and some miscellaneous things

.github/workflows/tests.yml

@@ -1,11 +1,11 @@
 name: tests
 
 on:
-  # push:
-  #   paths:
-  #     - "**.py"
-  #     - "requirements.txt"
-  #     - ".github/workflows/*.yml"
+  push:
+    paths:
+      - "**.py"
+      - "requirements.txt"
+      - ".github/workflows/*.yml"
   pull_request:
     paths:
       - "**.py"

README.md

@@ -5,15 +5,12 @@
 <h3><strong>An easily extensible framework unifying LLM unlearning evaluation benchmarks.</strong></h3>
 
   <div style="display: flex; gap: 10px; justify-content: center; align-items: center;">
-      <a href="https://github.com/locuslab/open-unlearning/actions">
-          <img src="https://github.com/locuslab/open-unlearning/actions/workflows/tests.yml/badge.svg" alt="Build Status">
-      </a>
-      <a href="https://huggingface.co/open-unlearning">
-        <img src="https://img.shields.io/badge/Hugging%20Face-white?logo=huggingface" alt="Hugging Face">
-      </a>
-      <a href="https://github.com/locuslab/open-unlearning">
-        <img src="https://img.shields.io/github/stars/locuslab/open-unlearning?style=social" alt="GitHub Repo stars">
-      </a>
+    <a href="https://github.com/locuslab/open-unlearning"><img src="https://img.shields.io/github/stars/locuslab/open-unlearning?style=social" alt="GitHub Repo stars"/></a>
+    <a href="https://github.com/locuslab/open-unlearning/actions"><img src="https://github.com/locuslab/open-unlearning/actions/workflows/tests.yml/badge.svg" alt="Build Status"/></a>
+    <a href="https://huggingface.co/open-unlearning"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Model-blue" alt="HuggingFace 🤗"/></a>
+    <a href="https://github.com/locuslab/open-unlearning"><img src="https://img.shields.io/github/repo-size/locuslab/open-unlearning" alt="GitHub repo size"/></a>
+    <a href="https://github.com/locuslab/open-unlearning"><img src="https://img.shields.io/github/languages/top/locuslab/open-unlearning" alt="GitHub top language"/></a>
+    <a href="https://github.com/locuslab/open-unlearning/blob/main/LICENSE"><img src="https://img.shields.io/badge/License-MIT-blue" alt="License: MIT"/></a>
   </div>
 </div>
 
@@ -30,7 +27,7 @@ We invite the LLM unlearning community to collaborate by adding new benchmarks,
 ### 📢 Updates
 
 #### [Apr 6, 2025]
-⚠️⚠️ **IMPORTANT:** Be sure to run `python setup_data.py` immediately after merging the latest version. This is required to refresh the downloaded eval log files and ensure they're compatible with the latest evaluation metrics.
+🚨🚨 **IMPORTANT:** 🚨🚨 Be sure to run `python setup_data.py` immediately after merging the latest version. This is required to refresh the downloaded eval log files and ensure they're compatible with the latest evaluation metrics.
 - **More Metrics!** Added 6 Membership Inference Attacks (MIA) (LOSS, ZLib, Reference, GradNorm, MinK, and MinK++), along with Extraction Strength (ES) and  Exact Memorization (EM) as additional evaluation metrics.
 - **More TOFU Evaluations!** Now includes a holdout set and supports MIA attack-based evaluation. You can now compute MUSE's privleak on TOFU.
 - **More Documentation!** [`docs/links.md`](docs/links.md) contains resources for each of the implemented features and other useful LLM unlearning resources.
@@ -89,13 +86,13 @@ We provide several variants for each of the components in the unlearning pipelin
 ## ⚡ Quickstart
 
 ```bash
-# environment setup
+# Environment setup
 conda create -n unlearning python=3.11
 conda activate unlearning
 pip install .
 pip install --no-build-isolation flash-attn==2.6.3
 
-# data setup
+# Data setup
 python setup_data.py  # saves/eval now contains evaluation results of the uploaded models
 # Downloads log files with metric eval results (incl retain model logs) from the models 
 # used in the supported benchmarks.
@@ -175,7 +172,7 @@ For more in-depth information on specific aspects of the framework, refer to the
 | [`docs/contributing.md`](docs/contributing.md)       | Instructions on how to add new methods, benchmarks, components such as trainers, benchmarks, metrics, models, datasets, etc.              |
 | [`docs/evaluation.md`](docs/evaluation.md)       | Detailed instructions on creating and running evaluation metrics and benchmarks.                                     |
 | [`docs/experiments.md`](docs/experiments.md)     | Guide on running experiments in various configurations and settings, including distributed training, fine-tuning, and overriding arguments. |
-| [`docs/hydra.md`](docs/hydra.md)                 | Explanation of the Hydra features used in configuration management for experiments.                                  |
+| [`docs/hydra.md`](docs/hydra.md)                 | A short tutorial on Hydra features, Hydra is the configuration management package we use extensively.                                  |
 | [`community/leaderboard.md`](community/leaderboard.md)             | Reference results from various unlearning methods run using this framework on TOFU and MUSE benchmarks.              |
 | [`docs/links.md`](docs/links.md)             | List of all links to the research papers or other sources the implemented features are sourced from.              |
 | [`docs/repro.md`](docs/repro.md)            | Results are provided solely for reproducibility purposes, without any parameter tuning.             |
@@ -193,26 +190,25 @@ If you use OpenUnlearning in your research, please cite OpenUnlearning and the b
 
 ```bibtex
 @misc{openunlearning2025,
-  title={OpenUnlearning: A Unified Framework for LLM Unlearning Benchmarks},
+  title={{OpenUnlearning}: A Unified Framework for LLM Unlearning Benchmarks},
   author={Dorna, Vineeth and Mekala, Anmol and Zhao, Wenlong and McCallum, Andrew and Kolter, J Zico and Maini, Pratyush},
   year={2025},
   howpublished={\url{https://github.com/locuslab/open-unlearning}},
   note={Accessed: February 27, 2025}
 }
 @inproceedings{maini2024tofu,
-  title={TOFU: A Task of Fictitious Unlearning for LLMs},
+  title={{TOFU}: A Task of Fictitious Unlearning for LLMs},
   author={Maini, Pratyush and Feng, Zhili and Schwarzschild, Avi and Lipton, Zachary Chase and Kolter, J Zico},
   booktitle={First Conference on Language Modeling},
   year={2024}
 }
-@article{shi2024muse,
-  title={MUSE: Machine Unlearning Six-Way Evaluation for Language Models},
+@inproceedings{
+  shi2025muse,
+  title={{MUSE}: Machine Unlearning Six-Way Evaluation for Language Models},
   author={Weijia Shi and Jaechan Lee and Yangsibo Huang and Sadhika Malladi and Jieyu Zhao and Ari Holtzman and Daogao Liu and Luke Zettlemoyer and Noah A. Smith and Chiyuan Zhang},
-  year={2024},
-  eprint={2407.06460},
-  archivePrefix={arXiv},
-  primaryClass={cs.CL},
-  url={https://arxiv.org/abs/2407.06460},
+  booktitle={The Thirteenth International Conference on Learning Representations},
+  year={2025},
+  url={https://openreview.net/forum?id=TArmA033BU}
 }
 ```
 </details>
@@ -231,6 +227,4 @@ This project is licensed under the MIT License. See the [`LICENSE`](LICENSE) fil
 
 ---
 
-### Star History
-
 [![Star History Chart](https://api.star-history.com/svg?repos=locuslab/open-unlearning&type=Date)](https://www.star-history.com/#locuslab/open-unlearning&Date)

community/benchmarks/template/README.md

@@ -26,7 +26,7 @@ Please include the experimental setup for the baselines
 
 - [ ] **Hyperparameters & Search Space:** Specify key hyperparameters, their search ranges, number of trials etc.
 - [ ] **Computational Setup:** Mention the type and number of GPUs used.
-- [ ] **DeepSpeed Configuration:** If any modifications were made to the default DeepSpeed config, specify them here. (You may include the config as a code block.)
+- [ ] **DeepSpeed Configuration** (if used): If any modifications were made to the default DeepSpeed config, specify them here. (You may include the config as a code block.)
 - [ ] **Other Details:** Any additional setup details crucial for reproducing your method.
 
 To replicate your results, provide a `run.sh` script that contains all necessary commands to reproduce the final results. Ensure the script is well-documented.

community/methods/template/README.md

@@ -11,7 +11,7 @@ Please include the experimental setup such as
 
 - [ ] **Hyperparameters & Search Space:** Specify key hyperparameters, their search ranges, number of trials etc.
 - [ ] **Computational Setup:** Mention the type and number of GPUs used.
-- [ ] **DeepSpeed Configuration:** If any modifications were made to the default DeepSpeed config, specify them here. (You may include the config as a code block.)
+- [ ] **DeepSpeed Configuration** (if used): If any modifications were made to the default DeepSpeed config, specify them here. (You may include the config as a code block.)
 - [ ] **Other Details:** Any additional setup details crucial for reproducing your method.
 
 # Results

configs/experiment/examples/muse_unlearn.yaml

@@ -22,13 +22,15 @@ trainer:
     per_device_train_batch_size: 4
     per_device_eval_batch_size: 16
     gradient_accumulation_steps: 8
-    learning_rate: 1.0e-05
+    learning_rate: 3.0e-05
     bf16: true
     bf16_full_eval: true
     logging_steps: 5
     output_dir: ${paths.output_dir}
     logging_dir: ${trainer.args.output_dir}/logs
     report_to: tensorboard
+    ddp_find_unused_parameters: None
+    gradient_checkpointing: false
     optim: paged_adamw_32bit
     save_strategy: 'no'
     save_only_model: true
@@ -53,22 +55,20 @@ data:
       args:
         hf_args:
           path: muse-bench/MUSE-News
-          name: train
+          name: raw
           split: ${forget_split}
         text_key: text
-        max_length: 128
-        insert_space: true
+        max_length: 2048
   retain:
     MUSE_retain:
       handler: PretrainingDataset
       args:
         hf_args:
           path: muse-bench/MUSE-News
-          name: train
+          name: raw
           split: ${retain_split}
         text_key: text
-        max_length: 128
-        insert_space: true
+        max_length: 2048
   anchor: forget
 collator:
   DataCollatorForSupervisedDataset:
@@ -119,64 +119,144 @@ eval:
         handler: rouge
         rouge_type: rougeL_f1
         batch_size: 16
+      retain_knowmem_ROUGE:
+        datasets:
+          MUSE_retain_knowmem:
+            handler: QADataset
+            args:
+              hf_args:
+                path: muse-bench/MUSE-${eval.muse.data_split}
+                name: knowmem
+                split: retain_qa
+              few_shot_dataset_hf_args:
+                path: muse-bench/MUSE-${eval.muse.data_split}
+                name: knowmem
+                split: retain_qa_icl
+              question_key: question
+              answer_key: answer
+              max_length: 512
+              predict_with_generate: true
+        collators:
+          DataCollatorForSupervisedDataset:
+            handler: DataCollatorForSupervisedDataset
+            args:
+              padding_side: left
+              index: index
+        generation_args:
+          do_sample: false
+          top_p: null
+          temperature: null
+          max_new_tokens: 32
+          use_cache: true
+          stopwords:
+          - '
+
+
+            '
+          - '
+
+            Question'
+          - 'Question:'
+        handler: rouge
+        rouge_type: rougeL_f1
+        batch_size: 16
+      forget_verbmem_ROUGE:
+        datasets:
+          MUSE_forget_verbmem:
+            handler: CompletionDataset
+            args:
+              hf_args:
+                path: muse-bench/MUSE-${eval.muse.data_split}
+                name: verbmem
+                split: forget
+              prefix_key: prompt
+              text_key: gt
+              max_length: 2048
+              insert_space: true
+              predict_with_generate: true
+        collators:
+          DataCollatorForSupervisedDataset:
+            handler: DataCollatorForSupervisedDataset
+            args:
+              padding_side: left
+              index: index
+        generation_args:
+          do_sample: false
+          top_p: null
+          temperature: null
+          max_new_tokens: 128
+          use_cache: true
+        handler: rouge
+        rouge_type: rougeL_f1
+        batch_size: 8
       privleak:
         pre_compute:
-          forget_minKpc_neg_logprob:
+          mia_min_k:
             datasets:
-              MUSE_forget_privleak:
-                handler: PretrainingDataset
+              MUSE_MIA_holdout:
+                access_key: holdout
+                handler: CompletionDataset
                 args:
                   hf_args:
                     path: muse-bench/MUSE-${eval.muse.data_split}
                     name: privleak
-                    split: forget
+                    split: holdout
                   prefix_key: prompt
                   text_key: text
-            collators:
-              DataCollatorForSupervisedDataset:
-                handler: DataCollatorForSupervisedDataset
-                args:
-                  padding_side: right
-                  index: index
-            handler: minKpc_negative_logprob
-            batch_size: 8
-            k: 0.4
-            access_key: forget
-          holdout_minKpc_neg_logprob:
-            datasets:
-              MUSE_holdout_privleak:
-                handler: PretrainingDataset
+                  max_length: 2048
+              MUSE_MIA_forget:
+                access_key: forget
+                handler: CompletionDataset
                 args:
                   hf_args:
                     path: muse-bench/MUSE-${eval.muse.data_split}
                     name: privleak
-                    split: holdout
+                    split: forget
                   prefix_key: prompt
                   text_key: text
+                  max_length: 2048
             collators:
               DataCollatorForSupervisedDataset:
                 handler: DataCollatorForSupervisedDataset
                 args:
                   padding_side: right
                   index: index
-            handler: minKpc_negative_logprob
             batch_size: 8
+            handler: mia_min_k
             k: 0.4
-            access_key: holdout
+            access_key: forget
         reference_logs:
           retain_model_logs:
             path: ${eval.muse.retain_logs_path}
             include:
-              forget_minKpc_neg_logprob:
+              mia_min_k:
                 access_key: retain
-              holdout_minKpc_neg_logprob:
-                access_key: holdout
         handler: privleak
         ref_value: 0.5
+      extraction_strength:
+        datasets:
+          MUSE_forget_verbmem:
+            handler: CompletionDataset
+            args:
+              hf_args:
+                path: muse-bench/MUSE-${eval.muse.data_split}
+                name: verbmem
+                split: forget
+              prefix_key: prompt
+              text_key: gt
+              max_length: 2048
+              insert_space: true
+        collators:
+          DataCollatorForSupervisedDataset:
+            handler: DataCollatorForSupervisedDataset
+            args:
+              padding_side: right
+              index: index
+        handler: extraction_strength
+        batch_size: 8
     handler: MUSEEvaluator
-    device: cuda
     output_dir: ${paths.output_dir}
-    overwrite: false
+    overwrite: true
     data_split: ${data_split}
     retain_logs_path: ${retain_logs_path}
 paths:
@@ -188,6 +268,6 @@ paths:
 data_split: News
 forget_split: forget
 retain_split: retain1
-retain_logs_path: saves/eval/muse_news_retain/MUSE_EVAL.json
-task_name: llama2_news_NPO
+retain_logs_path: saves/eval/muse_Llama-2-7b-hf_News_retrain/MUSE_EVAL.json
+task_name: muse_npo_unlearn
 mode: unlearn