Shaping capabilities with token-level data filtering

TL;DR

This work investigates shaping model capabilities by filtering pretraining data at the token level, focusing on a medical forget domain as a proxy for undesired capabilities. It demonstrates that token-level filtering outperforms document-level filtering in both effectiveness and efficiency, and that its benefits scale with model size, producing large compute savings for forget-related capabilities. The authors introduce a weakly supervised pipeline using sparse autoencoders to label forget tokens and train cheap token-level classifiers, showing that even imperfect labels can yield strong filtering when scaled. Additionally, token-level filtering preserves alignment and enables refined control such as refusal training, contrasting with unlearning and other posthoc safeguards. The results advocate for token-level data filtering as a scalable, robust, and safer approach to shaping capabilities during pretraining, while highlighting avenues for improved evaluation and broader safeguards against misuse.

Abstract

Current approaches to reducing undesired capabilities in language models are largely post hoc, and can thus be easily bypassed by adversaries. A natural alternative is to shape capabilities during pretraining itself. On the proxy task of removing medical capabilities, we show that the simple intervention of filtering pretraining data is highly effective, robust, and inexpensive at scale. Inspired by work on data attribution, we show that filtering tokens is more effective than filtering documents, achieving the same hit to undesired capabilities at a lower cost to benign ones. Training models spanning two orders of magnitude, we then demonstrate that filtering gets more effective with scale: for our largest models, token filtering leads to a 7000x compute slowdown on the forget domain. We also show that models trained with token filtering can still be aligned on the forget domain. Along the way, we introduce a methodology for labeling tokens with sparse autoencoders and distilling cheap, high-quality classifiers. We also demonstrate that filtering can be robust to noisy labels with sufficient pretraining compute.

Figures (28)

• Figure 1: Token-level data filtering gets more effective with scale. We plot relative scaling laws that show the effective compute required to train a Transformer on filtered data that matches the loss on a baseline trained on completely unfiltered data. Larger models require proportionally more compute, i.e. filtering is more effective for larger models. For 1.8B parameter models trained on token filtered data, we see a $7000\times$ compute slowdown on the forget domain (medicine).
• Figure 2: Operationalizing token filtering. After labeling our pretraining set using a model-based classifier, we remove forget tokens from the Transformer backpass. When loss masking, we allow models to see forget tokens during the forwards pass. We also experiment with removal, where we additionally replace forget tokens with <|hidden|> tokens.
• Figure 3: Token filtering Pareto dominates document filtering. We sweep across classifier boundaries for both our token- and document-level classifiers to filter pretraining data for 521M parameter models. We observe that token filtering can consistently achieve the same recall (i.e. equal medical loss) at higher precision (i.e. lower biology loss) than document filtering.
• Figure 4: Token filtering scales better than document filtering. We plot forget vs. retain loss for all model series; each point is a model. We observe that token filtering is close to the 'frontier,' achieving high forget loss for any given level of retain loss (top left of the plot).
• Figure 5: Data filtering decreases MCQ performance on the forget domain without substantial damage to the retain domain. On MedMCQA and MedQA-USMLE, models trained with data filtering score near chance. Token filtering slightly reduces capabilities near the classification boundary (biology) but has no effect outside (STEM, non-STEM). The models trained with token filtering are weaker than the one trained with document filtering on MedQA-USMLE and MMLU Medicine, but equivalent on retain evaluations.