social-data-attribution-w2/notes/hparam_rationale.md

NGDiff Unlearning — Hyperparameter Rationale

Last updated: 2026-03-18
Experiment: Topic-bin machine unlearning on OLMo-3 7B using Normalized Gradient Difference (Bu & Xu, NAACL 2025).

---

1. Model

Choice	Value
Base model	allenai/OLMo-3-1025-7B
Precision	bfloat16
Attention	sdpa (scaled-dot-product, PyTorch 2.0+)

Why OLMo-3 7B:
Open weights with fully documented training provenance (Dolma corpus). The training data is publicly described topic-by-topic, making it possible to construct a validated forget set — documents that are known to be in the model's training mix. This is a prerequisite for meaningful unlearning experiments. Closed-source models (GPT-4, Gemini) cannot be used because their training data is unknown.

---

2. LoRA Configuration

Hyperparameter	Value	Rationale
r (rank)	8	Standard rank for 7B-scale selective unlearning. Enough capacity to encode forget-direction gradients; low enough that only ~0.1% of parameters are trainable, preventing catastrophic forgetting.
lora_alpha	16	Set to 2×r by convention — gives a scaling factor of 2 at initialization, which keeps gradient magnitudes stable regardless of rank choice.
lora_dropout	0.05	Light regularization to prevent overfitting to the small forget set (1000 docs). Higher dropout (e.g. 0.1) degraded retain-set PPL in preliminary trials.
target_modules	q_proj, k_proj, v_proj, o_proj	Attention projections only. The attention heads mediate topic-specific associative recall; updating FFN layers risks altering factual representations needed for the retain set. This is consistent with the Bu & Xu NAACL 2025 experiments.
bias	none	Standard LoRA setting; biases are not fine-tuned.

---

3. Training Setup

Hyperparameter	Value	Rationale
per_device_train_batch_size	4	Maximum that fits in ~40–45 GB VRAM with OLMo-3 7B + LoRA + gradient checkpointing. A100 80G and H100/H200 allow this comfortably; V100 32G does not.
gradient_accumulation_steps	4	Effective batch = 16 sequences. Each NGDiff step processes 1 forget + 1 retain mini-batch (2 forward/backward passes). Grad accum of 4 → 1 optimizer step = 4 NGDiff steps = 16 total sequences. Matches the batch size used in the Bu & Xu paper.
learning_rate	1e-5	Conservative. Too high (>5e-5) causes rapid retain-set PPL degradation after ~500 steps even with NGDiff normalization. Too low (<5e-6) produces negligible forget-set PPL increase within 5000 steps. 1e-5 reliably triggers the PPL stopping criterion at 500–2000 optimizer steps for single-topic runs.
lr_scheduler_type	constant	No decay. The AutoLR mechanism (quadratic probe every 10 optimizer steps) dynamically adjusts the learning rate, so a separate scheduler would conflict.
warmup_steps	0	Warmup is irrelevant here: we want maximum unlearning signal from the first step. Adding warmup delays the forget-loss gradient from flowing.
weight_decay	0.0	No regularization beyond LoRA dropout. Weight decay on LoRA matrices would counteract the forget gradient unnecessarily.
max_grad_norm	1.0	Standard gradient clipping to prevent occasional large NGDiff update spikes.
bf16	true	A100/L40S/H100/H200 all support BF16 natively. FP16 risks loss spikes with the NGDiff normalization step.
gradient_checkpointing	true	Reduces peak VRAM from ~55 GB to ~40–45 GB by recomputing activations during backward pass at the cost of ~20% slower throughput. Required for the 4-batch setting.

---

4. NGDiff-Specific

Hyperparameter	Value	Rationale
auto_lr	true	Quadratic probe on retain loss every 10 optimizer steps. Automatically keeps LR near the local minimum of retain loss, improving forget/retain trade-off over fixed LR.
lr_delta	1e-5	Perturbation step for the AutoLR probe (equal to the base LR). Small enough that the quadratic approximation is locally accurate; large enough that numerical noise is dominated by the signal.

---

5. Data Sampling

Hyperparameter	Value	Rationale
max_forget_docs	1000	A forget set of 1000 documents provides strong unlearning signal without excessively long tokenization (1000 × avg 800 tokens ≈ 800K tokens, fits in memory). Larger sets (5000+) did not change the PPL trajectory significantly given that the stopping criterion fires early.
max_retain_docs	9000	9:1 retain:forget ratio. The retain objective anchors general capabilities; higher ratios (>9:1) slow down unlearning; lower ratios (<3:1) cause rapid MMLU degradation. 9000 docs is also small enough that data loading is fast (few seconds).
max_length	2048	OLMo-3's native positional encoding supports 2048 tokens. Truncating to this avoids position-ID out-of-range errors and controls memory usage.
seed	42	Fixed for reproducibility of the random forget/retain doc sample.

---

6. Stopping Criteria

Criterion	Value	Rationale
max_steps	5000	Hard ceiling ≈ 23 hours on H100. Prevents unbounded runs if PPL stopping never fires (e.g. null-bin control).
PPL stopping threshold	Dynamic (per-run)	Computed as: base model's perplexity on chunk-shuffled forget-set tokens, measured before LoRA is applied. Shuffling destroys local syntax but preserves token statistics, so this PPL represents the floor of "random gibberish." Training stops when the fine-tuned model's forget-set PPL meets or exceeds this threshold — i.e., when the model can no longer predict the forget documents better than shuffled noise. This is more principled than a fixed threshold because it adapts to the actual forget set's difficulty.
MMLU safety threshold	0.90 × baseline	If in-training MMLU drops below 90% of the pre-training baseline, training stops. This guards against catastrophic forgetting of general world knowledge. MMLU is checked every 50 NGDiff steps using 100 pre-loaded validation questions.
Wall-time guard	930 min (PACE ICE only)	SLURM 16-hr jobs stop 30 min early to ensure a clean checkpoint is saved. Disabled on RunPod (pass trainer.max_walltime_minutes=null) since there is no SLURM preemption deadline.

---

7. PPL Logging

Parameter	Value	Rationale
PPL_CHECK_INTERVAL	1000 NGDiff steps (= 250 optimizer steps)	Checking every 250 optimizer steps is frequent enough to catch early stopping without significant eval overhead (~2 extra forward passes per check).
PPL_LOG_INTERVAL	2000 NGDiff steps (= 500 optimizer steps)	Forget-set PPL is written to forget_ppl_log.json every 500 optimizer steps. For grouped-bin experiments where unlearning is faster and stronger than single-topic, 500-step resolution is needed to capture the full PPL trajectory before the stopping criterion fires.

---

8. Chain Job / Resume Design (PACE ICE)

• 3 SLURM jobs chained per run (--dependency=afterok)
• Each job is 16 hrs; 3 × 16 hr = 48 hr wall-clock ceiling
• At job start, train.py checks for an existing ppl_stopping_threshold.json and forget_ppl_log.json; if the log already shows PPL ≥ threshold, the job exits immediately without reloading the model
• This avoids wasting GPU-hours when the stopping criterion fired in an earlier job of the chain

---

References

• Bu, J., & Xu, W. (2025). Forget What You Know: Machine Unlearning via Normalized Gradient Difference. NAACL 2025.
• Hu, E. J., et al. (2022). LoRA: Low-Rank Adaptation of Large Language Models. ICLR 2022.
• Groeneveld, D., et al. (2024). OLMo: Accelerating the Science of Language Models. ACL 2024.