Shaping capabilities with token-level data filtering

TL;DR

A short paper from Neil Rathi (Anthropic/Stanford) and Alec Radford that argues for shaping a model’s capabilities at the pretraining-data level rather than via post-hoc unlearning or RLHF. Using removal of medical knowledge as a proxy task, they show that filtering individual tokens (mask the loss on tokens flagged by a lightweight classifier) Pareto-dominates filtering entire documents — same hit to the forget domain at lower cost to adjacent benign domains. The intervention scales: at 1.8B parameters they report a 7000× effective-compute slowdown for a model trained on filtered data to recover the filtered capability, and 10× more robustness than RMU-style unlearning under adversarial fine-tuning attacks.

Key claims

Post-hoc capability-reduction methods (unlearning, RLHF safety tuning) are easily bypassed by adversaries, motivating a shift to interventions that act during pretraining [Abstract, §1].
Filtering pretraining data on the proxy task of removing medical capabilities is highly effective, robust, and inexpensive at scale [Abstract].
Token-level filtering (loss masking on flagged tokens) Pareto-dominates document-level filtering — equal reduction in the forget domain at lower cost to a related retain domain (biology) [Abstract, Fig. 1].
The intervention gets more effective with scale: across compute-optimal Transformers from 61M to 1.8B parameters trained on filtered data, the effective-compute slowdown on the forget domain grows, reaching 7000× at 1.8B [§1, Fig. 1].
On free-response medical questions at 1.8B, token filtering reduces answer correctness by up to 20× and relevance/coherence by 3× vs. an unfiltered baseline; multiple-choice scores collapse to near chance on MedMCQA and MedQA-USMLE while retain-domain benchmarks are unaffected [HF summary, §4].
Token filtering is ~10× more robust than state-of-the-art unlearning (RMU) under adversarial fine-tuning attacks at the 1.8B scale; recovering the filtered capability from the token-filtered base requires ~13× more tokens than recovering it from an RMU-unlearned base [HF summary].
The recipe tolerates noisy classifier labels given sufficient compute, so the token-filter classifier does not have to be perfectly precise [HF summary].

Method

The pipeline: (1) train a lightweight classifier on (token, context) pairs to flag tokens belonging to a target “forget” domain (here: medical content); (2) during pretraining, mask the cross-entropy loss on flagged tokens so they do not contribute to the backward pass, while non-flagged tokens in the same document continue to train normally. This is contrasted against the more common document-level filter that drops any document containing flagged content.

Models are compute-optimal Transformers from 61M to 1.8B parameters, an augmented GPT-2 architecture similar to Jordan et al. (2024a), trained with AdamW on 8× NVIDIA H200s [§ method / §A]. Evaluation uses (a) raw cross-entropy on held-out forget vs. retain text, (b) multiple-choice benchmarks (MedMCQA, MedQA-USMLE for forget; general benchmarks for retain), (c) free-response medical Q&A scored for correctness and relevance/coherence, and (d) adversarial fine-tuning attacks that attempt to recover the filtered capability with a small amount of medical data, measured as tokens-to-recovery.

The “effective compute” framing (Fig. 1) plots how much extra compute a model trained on filtered data needs to match the loss of an unfiltered baseline on the forget domain — i.e. how much the filter has “moved” the model along the scaling-law curve in the forget direction. Larger models trained on the same filtered corpus require proportionally more compute to recover, which is the empirical statement of “filtering scales”.

Results

7000× effective-compute slowdown on the forget domain at 1.8B parameters under token filtering [§1, Fig. 1].
Multiple choice: near-chance on MedMCQA and MedQA-USMLE; retain-domain MCQ scores ~unchanged [HF summary].
Free response at 1.8B: up to 20× reduction in medical-answer correctness, 3× reduction in relevance/coherence vs. an unfiltered baseline [HF summary].
Adversarial robustness: ~10× more robust than RMU unlearning at 1.8B; RMU-unlearned models recover the capability with ~13× fewer adversarial tokens than token-filtered models [HF summary].
Token-level filtering Pareto-dominates document-level filtering on the (forget-loss, retain-loss) frontier across the scales tested [Abstract, Fig. 1].

Why it’s interesting

This is the inverse twin of the “synthetic / curated pretraining data improves capabilities” thread the wiki is already tracking via SmolLM3, CLIMB, and Cosmopedia: same operating point (curation during pretraining), opposite sign (remove a capability instead of inject one). The token-level granularity is the load-bearing claim — it implies the curation gate that several recently filed papers rely on can be made sub-document, and that the benefit grows with scale rather than washing out. The reported ~7000× slowdown and the 10× robustness gap over RMU also reframe the unlearning literature: if you can pay the cost up-front during pretraining, post-hoc methods look very expensive in robustness-per-FLOP. For a generative-models team, the relevant translation is whether the same token/region-level masking idea transfers to image/video pretraining — e.g. “filter at the patch level rather than dropping the whole shot.”