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:

```text
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](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:

```bash
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.