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.