social-data-attribution-w2/docs/unlearning.md

Unlearning Pipeline

The unlearning package trains LoRA adapters with NGDiff to forget documents from target WebOrganizer topic bins while preserving performance on retain data. The default method uses stratified forget and retain sampling, retain-pool resampling, length-normalized loss, segment-shuffled PPL stopping, retain-only cooldown support, and MMLU macro aggregation. The unlearning evaluation surface excludes the previously used math-reasoning benchmark.

No experiment results are generated by these commands. Run outputs, generated manifests, scheduler wrappers, and paper figures remain outside this core package surface.

Method Defaults

Area	Default	Meaning
Forget documents	2000	documents sampled per target topic
Retain documents	1000 per non-target topic	stratified retain set, about 23K docs for one target
Minimum document length	512 estimated tokens	short documents are filtered before sampling
Retain resampling	every 200 optimizer steps	a fresh stratified retain set is drawn deterministically
Checkpoint cadence	every 200 optimizer steps	intermediate adapter checkpoints are retained
Loss	length-normalized	mean token loss per document, then mean across batch
PPL shuffle	segments	segment-shuffled forget text sets the stopping threshold
Cooldown	disabled	optional retain-only phases are available but off by default

The top-level Hydra config logs to the social-data-attribution W&B project and the hcai-lab entity by default.

Data Flow

The package reads from a local 6T-filtered Arrow cache. Set DOLMA_CACHE to the cache location before running training.

1. Load the filtered Arrow cache.
2. Apply the minimum-token filter.
3. Sample the forget set from the target topic, or load it from data.forget_manifest_path.
4. Sample retain documents stratified across non-target topics.
5. Write sampling_manifest.json in the run output directory.
6. Tokenize forget and retain documents and build paired dataloaders.

The sampling manifest records selected document IDs, seed, target topics, and sampling config so the exact sampled run can be audited later.

Training Behavior

Each NGDiff step computes one forget gradient and one retain gradient:

g_update = g_retain / ||g_retain|| - g_forget / ||g_forget||

The trainer also supports:

• GeN AutoLR every 10 optimizer steps.
• MMLU safety checks on a small fixed validation subset.
• Forget and retain PPL checks every 200 NGDiff steps.
• Resume detection from Hugging Face checkpoints.
• Optional retain-only cooldown phases through trainer.cooldown_enabled.

Intervals in the config are optimizer steps unless stated otherwise. With the default gradient_accumulation_steps: 4, one optimizer step equals four NGDiff steps.

Evaluation

Use src/unlearning/eval_harness.py after training. See docs/eval_harness.md for benchmark details and output schema.

The retained unlearning evaluation surface is:

• SocialIQA accuracy.
• MMLU Social Sciences macro-average accuracy.
• MMLU STEM macro-average accuracy.
• Wikitext-2 word perplexity.

The previously used math-reasoning benchmark is intentionally excluded from this unlearning cleanup. Other repository areas may still support historical or non-unlearning flows.

Core Configs

Data config:

Parameter	Default	Description
max_forget_docs	2000	forget documents per target topic
max_retain_docs	9000	fallback if docs_per_retain_bin is unset
docs_per_retain_bin	1000	retain documents per non-target topic
min_tokens	512	minimum estimated token count
resample_interval	200	optimizer steps between retain resamples
forget_manifest_path	null	optional manifest for fixed forget sets

Trainer config:

Parameter	Default	Description
max_steps	5000	hard ceiling on optimizer steps
per_device_train_batch_size	4	batch size per device
gradient_accumulation_steps	4	effective batch multiplier
learning_rate	1e-5	initial learning rate
auto_lr	true	enable GeN quadratic learning-rate adjustment
save_steps	200	checkpoint interval
max_walltime_minutes	930	graceful wall-time guard
length_normalized_loss	true	per-document loss normalization
ppl_shuffle_mode	segments	segment or chunk shuffle for PPL threshold
cooldown_enabled	false	retain-only cooldown phases

Local Smoke Shape

The minimal local shape is:

PYTHONPATH=src .venv/bin/python -m unlearning.train \
    topic_bin=entertainment \
    trainer.max_steps=10 \
    data.max_forget_docs=50 \
    data.docs_per_retain_bin=10

This command still loads the configured model and dataset. It is a shape check, not a cheap unit test. The unit tests under tests/unlearning/ cover sampling, retain resampling, cooldown, shuffle helpers, and length-normalized loss.