Shaping capabilities with token-level data filtering

# Shaping capabilities with token-level data filtering Neil Rathi 1 2 Alec Radford 3 # Abstract Current approaches to reducing undesired capa-bilities 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 in-tervention of filtering pretraining data is highly effective, robust, and inexpensive at scale. In-spired 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.

Train-ing 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 7000 × 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 methodol-ogy 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. neilrathi/token-filtering # 1. Introduction Frontier language models are pretrained on enormous amounts of text, acquiring a number of diverse capabili-ties (Wei et

, 2022b; Villalobos et, 2024). In turn, an important design goal is capability shaping: selectively reducing undesired capabilities without harming desired ones. For example, we want models to be able to assist with writing quality prose or conducting biology research, but not with running disinformation campaigns or synthesizing bioweapons (Hendrycks et, 2023; Schroeder et, 2026). As models become more generally capable, the associated risks of misuse are increasingly pressing (G ¨otting et, 2025; Ho & Berg, 2025; Xiao et, 2025). > 1 Anthropic 2Stanford 3Independent. Correspondence to: Neil Rathi com >.

Preprint. February 3, 2026. 10 18 10 19 10 20 > 1.0 > 10 > 10 2 > 10 3 > 10 18 10 19 10 20 Pretraining Compute (FLOPs) > Loss-Matched Compute Slowdown Forget (Medical) Retain (Biology) > Document Token (Masking) Token (Removal) Figure 1. Token-level data filtering gets more effective with scale. We plot relative scaling laws that show the effective com-pute 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, e. filtering is more effective for larger models.

For 1.8B parameter models trained on token filtered data, we see a 7000 × compute slowdown on the forget domain (medicine). A standard approach is to apply training or inference-time interventions to an already-pretrained model (Cao & Yang, 2015; Bourtoule et, 2021; Bai et, 2022; Sharma et, 2025). But because these strategies don’t remove unde-sired capabilities from the base model, adversaries can still elicit them via jailbreaks or finetuning (Wei et, 2023; Łucki et, 2024; Chowdhury et, 2025). This creates a perpetual cat-and-mouse game (Rando et, 2025). An alternative is to shape the capabilities of the model dur-ing pretraining itself, for instance by adjusting the data that a model is trained on.

The existing literature is encouraging: data selection can improve targeted downstream capabil-ities as well as decrease undesired attributes like toxicity (Longpre et, 2024; Hojel et, 2025). A natural way of framing the data selection problem is data filtering, e. selectively removing data from the pretraining corpus if it improves undesired capabilities downstream. Classifier-based filtering has shown promise as a way to robustly and effectively reduce dangerous capabilities (O’Brien et, 1 > arXiv:2601.21571v2 LG] 30 Jan 2026 Shaping capabilities with token-level data filtering 2025; Chen et, 2025). Yet beyond this, data filtering has been mostly neglected in the literature.

Here, we aim to improve our understanding of pretraining data filtering as a way of shaping capabilities. The data attribution literature suggests that individual tokens in pretraining can vary in their influence on model capabili-ties (Grosse et, 2023), yet most work on data selection operates at coarser granularity: for example, O’Brien et al. (2025) and Chen et al. (2025) train classifiers to identify doc-uments containing undesired content. We show that filtering tokens is a Pareto improvement over this baseline, achieving equal reduction in undesired capabilities at a lower cost to desired ones (§4.2). Then, training models spanning two orders of magnitude in compute, we find that filtering gets

more effective relative to an unfiltered baseline as we scale pretraining compute: for 1.8B parameter models, token fil-tering reduces compute efficiency 7000 × on the undesired domain (§4.2). Filtering is also 10 × more robust to adver-sarial finetuning attacks than a state-of-the-art unlearning intervention (§4.3). Another concern is that data filtering might make it harder to control model behavior. That is, a model might need to ‘know’ undesired knowledge in order to properly respond to it, for example by refusing (Wu, 2021). Work on detox-ifying language models has shown that while training on proportionally less toxic content reduces toxicity, it also makes it harder to align models on toxic queries (Longpre et

, 2024; Maini et, 2025; Li et, 2025). Surprisingly, we show that this is not the case for capability shaping—in fact, models trained with token filtering generalize to refusal training better than an unfiltered baseline (§4.4). Data filtering also suffers from the fact that generating high quality labels can be expensive, in particular because sam-ple efficient models might learn from just a few mislabeled examples (Welbl et, 2021; Cloud et, 2024; Lee et, 2025; Shilov et, 2025). We develop a weakly-supervised pipeline utilizing sparse autoencoders to label tokens, which beats supervised methods (§5.1, §6.3).

We use this to train token-level classifiers that cost a small fraction of pretrain-ing compute to run (§5.2). We also show that while imper-fect labeling does make filtering less effective, by decreas-ing the classification threshold to trade precision for recall, low-quality classifiers can still be highly effective given enough pretraining compute (§6.2). We also demonstrate that token-level classifiers can bootstrap from weak labels, but document-level classifiers cannot (§6.3). Taken together, our results show empirically that token-level filtering can cost effectively shape model capabilities at scale, and that it can do so both without harming alignment and without requiring perfect labels.

# 2. Motivation and related work Post hoc safeguards One way to shape the capabilities of a deployed model is to steer it into a particular distribution; g. we can teach it to refuse dangerous queries via RLHF (Ouyang et, 2022; Bai et, 2022). But this is easy to bypass by jailbreaking or finetuning (Zou et, 2023; Wei et, 2023; Zhan et, 2023; Qi et, 2023; Anil et, 2024; Andriushchenko et, 2024; Hughes et, 2024). In response, recent work has instead attempted to use ma-chine unlearning to extract capabilities from the pretrain-ing base (Barez et

, 2025; Liu et, 2025). Unlearning approaches are promising because they optimize directly against the model’s representations of dangerous knowledge (Liu et, 2022; Yao et, 2024; Li et, 2024; Sheshadri et, 2024; Rosati et, 2024; Gandikota et, 2024; Zou et, 2024; Tamirisa et, 2025). But current unlearning approaches fail against just a few steps of adversarial fine-tuning (Che et, 2024; Lynch et, 2024; Łucki et, 2024; Zhang et, 2024; Thaker et, 2025; Fan et, 2025; Kaunismaa et, 2026). Models are not organized in a way that naturally lends itself to this kind of surgical post hoc ‘extraction’ of capabilities (Jain et

, 2023; Hu et, 2024; Hong et, 2024; Deeb & Roger, 2025; Lee, 2025). Frontier model developers who maintain API-only access to their models have the additional ability to prevent users from accessing dangerous capabilities using input-output or internals-based classifiers (Sharma et, 2025; OpenAI, 2025b; Anthropic, 2025a; Cunningham et, 2026; Kram ´ar et, 2026). But even these defenses fall to cheap-to-find jailbreaks (elder-plinius, 2025; Chowdhury et, 2025). The unifying thread here is that once a capability exists in a base model, it is extremely hard to remove it (Deeb & Roger, 2025; Lee, 2025).

Large-scale pretraining bestows models with capabilities essentially indiscriminately; posttraining simply elicits these capabilities into a human-usable form (Radford et, 2019; Brown et, 2020; Christiano et, 2021; Wei et, 2021; Ouyang et, 2022; Kirstain et, 2022; Zhou et, 2023; Mallen et, 2023; Toshniwal et, 2024; Raghavendra et, 2024; Hofst ¨atter et, 2025; Donoway et, 2025; Yue et, 2025; Wen et, 2025). Shaping capabilities in pretraining Recent work has instead focused on methods that shape capabilities during pretraining itself. An obvious way to do this is to shape the data the model is trained on: model capabilities directly distill their training corpora.

Prior work (Yu et, 2024; Thrush et, 2024; Hojel et, 2025) has shown that data selection can improve downstream capabilities. Anil et al. (2023), Korbak et al. (2023) and Maini et al. (2025) focus on interventions to pretraining data that encourage aligned behavior, for example by adding control tokens for toxicity or training conditioned on human feedback. Lee et al. (2025) 2Shaping capabilities with token-level data filtering show that pretraining from scratch by distilling from an unlearned model can match the performance of a model trained only on benign data. The simplest manifestation of ‘data shaping’ is data filter-ing.

Much work has shown that data filtering is an effective mitigation for reducing fuzzy characteristics like toxicity (Raffel et, 2019; Gehman et, 2020; Xu et, 2021; Dodge et, 2021; Ngo et, 2021; Welbl et, 2021; Paullada et, 2021; Kreutzer et, 2022; Rauh et, 2022; Birhane et, 2023; Longpre et, 2024; Stranisci & Hardmeier, 2025; Li et, 2025). Most frontier labs use basic data filtering as part of their safety pipeline g. Ope-nAI, 2024; 2025a; Gemma Team, 2025; Google DeepMind, 2025; Grattafiori et, 2024).

Closest to our work, O’Brien et al. (2025) and Chen et al. (2025) show that high quality document-level data filtering is a highly effective and robust intervention for suppression of CBRN-related capabilities; in particular, O’Brien et al. (2025) find that a 6.9B Transformer trained with blocklist-based data filtering is 10 × more robust to adversarial fine-tuning than state-of-the-art posttraining safeguards. On the other hand, Longpre et al. (2024) and Li et al. (2025) both find that decreasing the amount of undesired content in pre-training can make it harder to elicit correct refusal behaviors on that domain.

Relatedly, Cloud et al. (2024) and Shilov et al. (2025) pro-pose gradient routing, which attempts to segment capabili-ties within the model ab initio. Gradient routing and related approaches are akin to posttraining safeguards in that they leverage the representations of the trained model in order to shape its own capabilities, as opposed to using external clas-sifiers. Additionally, they promise robustness to imperfect labeling, since in principle a model would learn to bootstrap classification from weak labels. Token-level data attribution A surprising result from work on early language models was that models would sometimes gain knowledge that was seemingly not present in their training data.

For example, Radford et al. (2019) trained GPT-2 on English documents which occasionally contained small sequences of French tokens g. ‘I’m not the cleverest man in the world, but like they say in French: Je ne suis pas un imbecile’). Despite this, however, they found that basic French capabilities could be elicited from the model in-context. Relatedly, Grosse et al. (2023) estimate influence functions using tokens, rather than documents, as training examples. They find that the influence of individ-ual tokens on model generations within a single document can fluctuate substantially. Work on data cleaning has also found that undesired tokens often appear in otherwise be-nign documents, such that document-level data filtering can

lead to overremoval of good data (Dodge et, 2021). These results suggest that models can effectively learn ca-pabilities from short subsequences of tokens within doc-uments. Document-based supervision would require re-moving a large amount of benign tokens in order to catch these small subsequences, sacrificing token-level precision to achieve the same recall. This is particularly important in the limited data regime (Muennighoff et, 2023; Villalo-bos et, 2024; Aschenbrenner, 2024; Kim et, 2025). # 3. Setting and approach Our goal is to study the effectiveness of data filtering as an intervention during pretraining.

We partition capabili-ties into a forget and retain set; we’d like to train models that have near-baseline retain capabilities and as-bad-as-possible forget capabilities. Because we don’t have the resources to train models to sufficient scale to get signal on actual dangerous capabilities, we focus on the representa-tive proxy of preventing models from acquiring medical capabilities while preserving related areas like biology. See 1 for how we define ‘medical’ content. We use model-based classifiers for data filtering, as in O’Brien et al. (2025) and Chen et al. (2025). At a high level, our approach is to (1) label a pretraining corpus using a classifier, (2) filter out data relevant to forget capabilities, (3) train models with varying amounts of pretraining com-pute, and (4) evaluate them on various benchmarks (text perplexity, multiple choice, free-response).

3.1. Data and data filtering We train models on FineWeb-Edu (Penedo et, 2024). We use the Edu split of FineWeb so that models are trained on a sufficient amount of biomedical text to elicit reasonable baseline performance; in early experiments, we found that even 1.8B models trained on the default split of FineWeb performed poorly on relevant benchmarks. We experiment both with document- and token-level data filtering. We go into more detail about how we source ground-truth labels and train classifiers in §5. All results reported below are based on our top performing classifiers, set at the threshold that maximized their F1 score on a held-out subset of FineWeb-Edu (unless otherwise specified).

We chose to set the threshold against F1 in order to most fairly maximize the precision-recall tradeoff; in §6.2 we study the consequences of adjusting this threshold. We consider two strategies for token filtering: loss masking,where we remove gradients computed for forget tokens from the backpass, and removal, where we replace forget tokens with a special <|hidden|> token (and similarly mask the loss on these tokens). In principle loss masking ensures that the model has access to coherent context when predicting retain tokens, but this might consequently allow the model 3Shaping capabilities with token-level data filtering > 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. to develop non-trivial contextual representations for forget tokens (see also Berglund et, 2023; Treutlein et, 2024; Wang et, 2025a). Removal, on the other hand, trades context coherence for complete removal of all forget tokens. 3.2. Model training Pretraining We train compute-optimal Transformers at scales ranging from 61M to 1.8B parameters (Hoffmann et

, 2022). Similar to Jordan et al. (2024a), we use an augmented version of the basic GPT-2 architecture (Radford et, 2019). We optimize using AdamW and scale learn-ing rate with μP (Loshchilov & Hutter, 2017; Yang et, 2022). We train models up to 521M on 2×NVIDIA H200s, and train 1B and 1.8B models on 8×NVIDIA H200s. For complete details on model architecture, hyperparameters, and training, see §A. Instruction tuning While raw cross-entropy loss is a use-ful proxy metric for capability shaping, it is somewhat ‘priv-ileged’ by loss masking, which directly intervenes on the backpass of forget tokens.

Therefore, we also evaluate our largest models 1 (1.8B parameters) on both multiple choice and free-response questions, which more fairly assess if we’ve truly attenuated capabilities. For multiple choice training, we use a custom instruction tuning mix consisting of several standard multiple choice datasets across domains, with consistent formatting for all questions. We used this custom mix instead of more standard ones like Flan (Long-pre et, 2023) or Tulu (Lambert et, 2024) since our primary goal was to elicit high multiple choice accuracy > 1In early experiments, we also tried to evaluate smaller models on these benchmarks, but we found that our baseline models were too weak to get any signal on whether filtering was actually a useful intervention.

on a limited compute budget. For chat training, we used the smol-smoltalk mix (Allal et, 2025). See 3 for further details. 3.3. Evaluation Text perplexity As a proxy for capability, we evaluate small models on their cross-entropy loss on relevant text; this also serves as a sanity check since it’s directly what data filtering intervenes on. We construct three text datasets: medical (PubMed articles), biology (bioRxiv articles; a ca-nary for closely related retain capabilities), and general non-medical (arXiv and PhilPapers articles). We do an additional pass over all datasets with Claude Sonnet 4 (An-thropic, 2025b) to remove non-medical documents from the medical dataset (and vice versa), and a third pass to remove unrelated tokens using the methodology described in §5.1.

Multiple choice For instruction tuned 1.8B models, we also use multiple choice evaluation. We evaluate medical knowledge using MedMCQA (Pal et, 2022), a bench-mark of Indian medical entrance exams, MedQA-USMLE (Jin et, 2020), consisting of clinical-style questions from the S. medical licensing exam, and a medical subset of MMLU (Hendrycks et, 2020). 2 We measure retain performance using various subsets of MMLU (biology, non-biomedical STEM, and non-STEM). Free-response We evaluate our chat trained 1.8B mod-els on free-response answers to HealthSearchQA, a dataset consisting of commonly searched consumer medical ques-tions (Singhal et, 2023).

We use Claude Sonnet 4 as a judge along three criteria: (1) relevance to the question, (2) coherence and (3) correctness of the response (§E). As a control, we also evaluate models on Alpaca, a free-response instruction following dataset (Taori et, 2023). 3 # 4. Token-level data filtering works and scales In §4.1, we show that token filtering, compared to document filtering, can achieve an equal hit to forget capabilities at a lower cost to retain capabilities. We then demonstrate that both kinds of filtering are effective across all three kinds of benchmarks, and that they get more effective with scale.

We also show that filtering is robust to elicitation of forget capabilities under adversarial finetuning (§4.3). Finally, in §4.4 we show that models trained with token filtering can still be aligned on the forget domain. > 2We use the college medicine, professional medicine, medical genetics, anatomy, virology, and clinical knowledge categories. > 3Note that we use Alpaca, rather than AlpacaEval and its as-sociated eval harness (Li et, 2023). We chose Alpaca as it is syntactically quite similar to HealthSearchQA. We additionally filter out medical questions using Claude Sonnet 4. 4Shaping capabilities with token-level data filtering 3.5 4 4.5 > 3 > 4 > 5 > better > Tokens Filtered > 50% > 25% > 12% > 6% > 3% > Biology Loss > Medical Loss > Baseline Document Token (Masking) Token (Removal) 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 pa-rameter models. We observe that token filtering can consistently achieve the same recall e. equal medical loss) at higher precision e. lower biology loss) than document filtering. 4.1. Token filtering Pareto dominates document filtering Our motivation for token filtering is that we can achieve equal recall with higher precision compared to document fil-tering. To test this empirically, we sweep across the decision boundary of our token- and document-level classifiers. We set the threshold based on the proportion of tokens filtered, filtering between 3% and 50% of all tokens from pretraining.

We then train 521M parameter models on the filtered data for each classification threshold, evaluating them on text perplexity. Figure 3 shows that token filtering is a Pareto improvement over document filtering, in that it can achieve lower retain loss at equal forget loss. 4.2. Filtering works, and filtering scales Text perplexity In Figure 4 we plot the forget and retain loss of each model series; we see that capabilities scale predictably under data filtering and that token filtering is close to the frontier of high forget loss and low retain loss. To more concretely understand scaling behavior, in Figure 1 we plot, for each model size, the proportion of pretrain-ing compute required to train a model on unfiltered data to matched loss (see Held et

, 2025; Shilov et, 2025). We compute this value by linearly interpolating the log-log compute-to-loss plot of the baseline model (see Figure 16 and 1). We find that (1) token-level filtering is more effective than document filtering at all scales of pretraining compute and (2) both kinds of data filtering get more effec-tive as we scale pretraining compute. In other words, the gap between models trained on filtered and unfiltered data 3.5 4 4.5 5 > Non-Medical Loss > 3 > 4 > 5 > better > Parameters > 61M > 113M > 224M > 521M > 1.0B > 1.8B > 345 > Biology Loss > better > Medical Loss > Baseline Document Token (Masking) Token (Removal) 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). gets larger with scale. Another way of interpreting this is that models trained with data filtering have lower magnitude scaling exponents on the forget domain. For the largest models we trained, token removal obtains over a 7000 × effective compute slowdown, compared to around 30 × for document filtering.

Multiple choice On multiple choice evaluations, we see that models trained with data filtering are substantially worse than the baseline on forget benchmarks, perform-ing around chance on MedMCQA and MedQA-USMLE (Figure 5). We see no noticeable degradation on the retain sets. We also evaluate using cloze-style selection, which bears out similar distinctions (see 2). Free response In Figure 6, we see that models trained with token-level filtering are substantially worse at respond-ing to medical -related queries: they are 4× less coherent and relevant, and 10 × less correct. Meanwhile, document-level filtering has a more muted effect.

On the other hand, we see no major performance hit on Alpaca (Figure 17). Amongst models trained with data filtering, we find consid-erable qualitative variance in their responses. While models do generate medical tokens when conditioned on them, they almost always fail to use them correctly. Sometimes model outputs show no relevance to the question (‘ A red eye is a serious condition that can be caused by a combination of factors, including a combination of factors such as a red eye ’) or fall into repet-itive cycles (‘ Bone cysts are a type of bacteria that [· · ·] caused by various factors such as bacteria, bacteria, bacteria, bacteria [· · ·]’).

5Shaping capabilities with token-level data filtering MedMCQA MedQA-USMLE MMLU Medicine > 0 > 10 > 20 > 30 > 40 > 50 Forget (↓) > MMLU Bio MMLU STEM MMLU Non-STEM Retain (↑) > Accuracy (%) > Baseline Document Token (Masking) Token (Removal) 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.

In other instances, models output mostly coherent yet totally false answers (‘ Dry lips can indeed be a symptom of various conditions, including cancer, heart disease, or other medical conditions ’). See §D for more examples. Correctness Relevance Coherence > 0 > 25 > 50 > 75 > 100 > Proportion of Generations (%) > Baseline Document Token (Masking) Token (Removal) Figure 6. Token filtering decreases free response quality in the forget domain. Responses to open-ended questions from the forget domain (HealthSearchQA) are judged by Claude Sonnet 4. Comparing different filtering methods, we see that token filtering decreases correctness up to 20 ×, and relevance and coherence 3×,relative to the baseline.

Document filtering also degrades response quality, but to a lesser extent. 4.3. Filtering is more robust than unlearning We consider the setting where an adversary has open-weight access to a model and wishes to train-in dangerous capabil-ities. We show that token and document filtering are both substantially more robust to adversarial finetuning attacks that a state-of-the-art unlearning safeguard, and that the rel-ative strength of this robustness increases with model scale (up to 10 × for 1.8B parameter models). Experimental setup We finetune models on medical text and evaluate their in-domain loss. We use the PubMed section of the Common Pile (Kandpal et

, 2025). For each model, we select the learning rate that enables finetuning to parity with the baseline in the fewest steps; see §A for detailed hyperparameters. Unlearning baseline We use RMU as an example of a state-of-the-art unlearning safeguard (Li et, 2024). RMU is a representation-based method that finetunes a model against an objective that encourages (1) preservation of retain representations and (2) stochasticity of forget repre-sentations (by aligning these representations to a random vector). RMU is at, or close to, the Pareto frontier of effec-tiveness and robustness amongst unlearning methods (Che et

, 2024). We use PubMed documents as the forget set and text from Project Gutenberg as the retain set. See §A for hyperparameters. Results We are interested in the amount of finetuning com-pute required to achieve parity with the unfiltered baseline. Figure 7 shows how this changes with scale. We notice that RMU exhibits substantially steeper scaling than all of our 6Shaping capabilities with token-level data filtering 10 18 10 19 10 20 > 0.01 > 0.1 > Pretraining Compute (FLOPs) > Tokens of Adversarial Finetuning Required > to Elicit Baseline (% of Pretraining) > Document Token (Masking) Token (Removal) RMU Figure 7.

Data filtering scales more robustly than unlearn-ing. Larger models need fewer adversarial finetuning samples to achieve baseline performance (as a proportion of pretraining compute), but the RMU curve is steeper; in other words, as pre-training compute scales, the robustness gap between RMU and data filtering will greaten. filtering baselines. That is, RMU gets less robust with scale at a rate faster than data filtering; for the 1.8B parameter models, RMU requires 1.5× fewer tokens than document filtering, 3× fewer than token loss masking, and 13 × fewer than token removal. This is notable especially given that

RMU has a substantially higher initial loss on the test set. Figure 26 shows that finetuning an RMU-tuned model re-sults in a steep decrease in loss almost immediately, while models trained with data filtering are more gradual. 4.4. Token-level filtering makes alignment easier Prior work has shown that models trained on proportionally more toxic data can be better at identifying when data is toxic, and are therefore more robustly ‘alignable’ (Longpre et, 2024; Li et, 2025; Maini et, 2025; Geng et, 2025; Wichers et, 2025; Tan et, 2025; Azarbal et

, 2025). In the context of capabilities shaping, while we’d like to remove unsafe knowledge, we’d still like to be able to control model behavior in these domains as opposed to having completely unpredictable outputs. Intuitively, it seems as though filtering data would be less effective than teaching the model the dangerous material and then teaching it how to respond to it (Wu, 2021). Here, we show that a surprising advantage of token filtering over document filtering is that it still allows us to control models in the forget distribution. Classifying forget tokens A simple version of this prob-lem is identification: can models trained on filtered data still distinguish the forget domain?

We fit a linear probe on 61M 113M 224M 521M 1030M 1816M > 0.70 > 0.75 > 0.80 > 0.85 > Model Size > Forget vs. Retain Classification F1 > Document Token (Masking) Token (Removal) Baseline Figure 8. Models trained with data filtering can reliably distin-guish the forget domain. We fit a linear probe to each model to classify forget vs. retain tokens using the same setup as §5. Though small models trained with token filtering are worse at clas-sification, the gap closes with scale. We include the performance of the pretraining filter (trained on 4× as many tokens) as a baseline.

top of each model to classify tokens as medical vs. non-medical, using a 2.05M-token subset of our classifier train-ing corpus and sweeping across layers. We find that models trained with data filtering are only marginally worse than the baseline, and that this gap closes with scale (Figure 8). Refusal training A more realistic setting is refusal train-ing: say we remove dangerous biology knowledge from pretraining. We’d still want to control the model’s behavior on dangerous biology-related queries, g. to have it gen-erate a refusal. To simulate this setting, we finetune our already-chat trained 1.8B parameter models on questions from HealthSearchQA and Alpaca.

On HealthSearchQA, we train the model to generate single-sentence refusals; on Alpaca, we use normal completions. We then evaluate on a held-out subset of both datasets, using Claude Sonnet 4 to classify refusals. Models that learn the correct generaliza-tion would generate refusals to HealthSearchQA questions and normal responses to questions from Alpaca. We repeat refusal training across three random seeds and use the same hyperparameters as we do for chat training 3). Surprisingly, we find that token-level data filtering actually improves control in this setting, while document-level fil-tering is less corrigible (Figure 9). Models trained with

token-level removal generate refusals at a rate 2× higher than the baseline on HealthSearchQA, while showing no notable increase on Alpaca. Models trained with token-level loss masking generate slightly fewer refusals than the base-line on HealthSearchQA but similarly do not output refusals on Alpaca. Meanwhile, models trained with document-level filtering struggle to generalize to the task, refusing Alpaca 7Shaping capabilities with token-level data filtering Alpaca (↓) HealthSearchQA (↑) > 0 > 10 > 20 > 30 > 40 > 50 > Refusal Rate (%) > Baseline Document Token (Masking) Token (Removal) > Figure 9. Token-level removal makes forget set alignment easier.

> We train models to refuse queries from HealthSearchQA, but not queries from Alpaca. We observe that models trained with token filtering generalize as well as or better than the baseline, while the model trained with document filtering generalizes poorly. queries at a rate only slightly lower than HealthSearchQA. In 4 we show similar results when training models to generate a single refusal token rather than a prose refusal. What’s going on? Previous work has shown that decreas-ing the proportion of toxic data seen in pretraining makes models worse at classifying whether new data is toxic (Li et

, 2025; Longpre et, 2024). We claim that this does not, as it might seem, contradict our results. In the case of filtering a capability like medicine, refusal training essen-tially asks a model to discriminate between tokens it has seen and tokens it has not; this is a much simpler task than classifying whether a piece of text is toxic or not, because the model will have seen ‘toxic’ tokens in pretraining, just not in the toxic context. In other words, it seems like the mechanism is something more akin to the model learning to separate ‘trained’ versus ‘untrained’ tokens.

To study this further, we analyze whether models trained on filtered data can discriminate on in-domain classification, e. between subdomains. We fit linear probes on top of each model to classify tokens sourced from the medRxiv sections on neurology and infectious disease. We find that though filtering achieves parity with the baseline on forget-retain classification, it struggles on in-domain classification, consistent with our hypothesis (Figure 21). A consequence of this is that filtering does not allow for fine-grained control on multiple forget domains. But this is sufficient for refusal training: we simply need the model to refuse when asked a question it does not have an answer to.

# 5. How to train your classifier In this section, we describe various engineering improve-ments that allow us to train a cheap and accurate token-level classifier. Our approach is to train a classifier to determine whether a token is relevant to forget domain knowledge,with the idea that this approximates whether a token is influ-ential for forget domain capabilities. Note that the objective we train our classifiers on is really a proxy for what we actually want to remove: datapoints that lead to downstream improvements on forget capabilities. Not all identified datapoints will be necessarily influential for capabilities, and not

all influential datapoints will be identified by the classifier; some datapoints influence forget capabilities without directly containing forget knowledge (Grosse et, 2023). We return to this distinction in §7, but our results in §4.2 confirm that this proxy objective is generally well-aligned with the true objective at scale. 5.1. Sourcing ground-truth labels Training a classifier requires annotated data. While labeled documents are relatively plentiful (or at the very least easy to generate synthetically), it’s not immediately obvious how we’d get token-level annotations in an unsupervised or weakly supervised way. Recent work in mechanistic interpretability has made sub-stantial progress on decomposing and interpreting model activations using sparse dictionary learning with sparse au-toencoders (Olshausen & Field, 1997; Cunningham et

, 2023; Bills et, 2023; Paulo et, 2024). Here, rather than using SAEs to understand model activations, we consider SAE latents (and their corresponding explanations) as a set of natural language descriptions of tokens (Movva et, 2025; Jiang et, 2025; Nguyen et, 2025). Our approach is simple: 1. Collect forget-domain latents from a pretrained SAE. 2. Label tokens as medical if they have high activations on a certain number of these latents. 3. Iteratively label adjacent tokens as medical if they have positive activations on at least one of these latents.

The first step essentially identifies which features are rel-evant for our task. We then need to determine if a given token actually belongs to the forget domain: does it have high activation on any of these features? We require that a token activate multiple latents because of feature split-ting (Bricken et, 2023) and high variance in autointerp quality. For example, Gemma Scope’s Gemma 2 9B SAE has features ranging from ‘references to health and medi-cal information’ to ‘pharmaceutical and medical research data related to ’ Many tokens would activate gen-eral health or medical related latents without actually being 8Shaping capabilities with token-level data filtering MethodsWe extended a Susceptible-Infected-Recovered-Deceased (SIR-D) model with vaccination, ran extensive simulations with different settings, and compared the level of infection attack rate (IAR) under different reach ratios between two vaccine types under different resource allocation decisions.

ResultsWe found that when there were limited resources, allocating resources to a vaccine with high efficacy that became available earlier than a vaccine with lower efficacy did not always lead to a lower IAR There has been an increase in cases of leishmaniasis in Syria, according to Dr. Glen Thomas of the World Health Organization, (WHO). Leishmaniasis is a disease transmitted by the bite of the sand fly and characterized by skin sores. Dr. Thomas says the disease is spreading because of poor waste and hygiene management caused by the continuing conflict in the country.

Even though we are not in an area that is susceptible to drought, we still think that it's a great demonstration of a technology that is very water efficient." The hydroponic growing environment is sterile as well, which eliminates the risk of pathogens-particularly important in light of the increase in foodborne illnesses, such as E. coli and salmonella, from fresh vegetables. Puri and his partners have seen to the frugal use of energy too. Figure 10. Ground-truth labels for three randomly selected classifier training documents. Highlighted tokens are labeled as forget, unhighlighted tokens are retain.

Token labels are mostly good at identifying related tokens and ignoring benign ones, but there is still some noise. ‘medical’ under our classification g. biochemistry tokens). The final step is important because crucially, our goal is not only to classify keywords but rather spans of tokens. For example, we’d like the entire phrase ‘insert the catheter’ to be classified as medical, not just ’ This also helps further reduce noise from various steps of the pipeline. We frame classifier training as a kind of weak-to-strong generalization problem (Burns et, 2023). Token labels, despite our best efforts, are noisy in systematic ways (Fig-ure 10).

Our goal is to create a dataset that is hill climbable, and upon which hill climbing leads to improvements in ef-fectiveness. But a ‘good’ classifier will not achieve perfect accuracy on this set; rather, we want a classifier that gener-alizes from noisy labels to learn the ‘correct’ ground truth direction. In §6.3 we describe other annotation approaches. Technical details We use Lieberum et al. (2024)’s pre-trained SAEs for Gemma 2 9B (Gemma Team, 2024). We use the 16k width SAE at layer 31. 4 We first use Claude 3.5 Haiku to generate an explanation for each latent using the Neuronpedia API (Anthropic, 2024b; Bills et

, 2023; Lin & Bloom, 2023). We then classify each explanation as medical or non-medical with Claude Sonnet 4 (full prompt in Appendix). We additionally score all explanations using Paulo et al. (2024)’s embedding scoring, and discard latents with scores lower than 0.9. This leaves us with 600 latents. > 4 Later layers tended to have better latents for labeling. We suspect this is because the medical/bio distinction is likely clearer later in the forward pass of a model. Tokens are labeled as medical if they are at least 4SD above the mean activation on at least two medical latents, or if they have positive activation on at least one medical latent and are adjacent to a token already classified as medical (we repeat this process iteratively until convergence).

We select these hyperparameters mostly by inspection. While we use SAEs to generate ground-truth labels, we do not use them to label the entire pretraining corpus. One reason is simply that running 9B SAE inference over an entire pretraining corpus is prohibitively expensive. Fur-ther, recent work has shown that SAEs—while useful for un supervised concept detection—lag behind simple linear probes for classification (Wu et, 2025; Kantamneni et, 2025). Our core methodology is thus to use SAEs to label a subset of data, which we use to distill a much smaller probe. Training data We annotate a mix of academic papers and web documents for classifier training; the split is roughly 75-25.

We use academic papers from PubMed, bioRxiv, medRxiv, chemRxiv, arXiv, Project Gutenberg, and the Stan-ford Encyclopedia of Philosophy, with an equal distribution between them. For web documents, we use FineWeb-Edu, which we label using Claude Sonnet 4. In total, our dataset consists of 128k documents. All classifiers are trained on 8.2M tokens sampled from these documents, with an even split of forget and retain tokens. We evaluate on a held out val set of 1.64M tokens (from the train distribution) and a test set of 0.82M tokens (consisting solely of FineWeb-Edu documents). Because our pretraining experiments used a different tokenizer than Gemma, we retokenize and rela-bel the dataset after applying the SAE pipeline to generate labels for Gemma tokens.

We relabel tokens such that if a Gemma forget token maps to a partial token of the new tokenizer, the whole token is labeled as forget. 5.2. A good representation is hard to find We now move to actually training a classifier. Our first claim is that using bidirectional context for classification will offer significant performance gains: whether a token like ‘virus’ is relevant to virology or computer security de-pends entirely on context (Wittgenstein, 1953). Our method is therefore to fit linear probes to bidirectional models. 5 We choose to fit linear probes using L-BFGS rather than doing full finetuning in order to improve robustness to spu-rious correlations (Pimentel et

, 2020; Kumar et, 2022; Kirichenko et, 2022), especially given that our ground-truth labels are already somewhat noisy. Here, we show that small task-specific base models can beat larger general ones for token-level classification for a fraction of the cost. As a baseline, we find that ModernBERT-large (Warner > 5 We sweep across layers. All results reported are for the highest performing probe. 9Shaping capabilities with token-level data filtering model f1 (val) f1 (test) ModernBERT-large 0.794 0.812 base → RoBERTa 0.808 0.834 base → biLM 0.830 0.880 upsample PubMed 0.834 0.877 61M → 113M 0.844 0.885 113M → 224M (final) 0.856 0.894 Table 1.

Small, task-specific base models outperform large, general-purpose ones. Our ModernBERT-large baseline is out-performed on medical classification by changing base model archi-tecture, training objective, and pretraining corpus. We can scale up a working recipe to achieve additional gains. et, 2024), a 395M parameter BERT-like model, does rea-sonably well out-of-the-box, reaching an F1 score of 0.794 on our val set. 6 But this is a big (and therefore expensive) model, and we’d like to push performance more if we can. As a first stab, we pretrain a 65M parameter RoBERTa-like model on FineWeb-Edu with a masked language modeling objective.

This leads to a modest improvement on our val set (0.808 F1) at a fraction of the cost. However, we believed this could be improved upon. Masked language modeling induces a number of strange artifacts which can make frozen-representation probes weaker (Clark et, 2020; Meng et, 2024). Autoregressive models also benefit from significantly more updated training and inference infrastructure. Inspired by earlier work, we exper-iment with training bidirectional models by jointly training separate left-to-right and right-to-left autoregressive models (Graves & Schmidhuber, 2005; McCann et, 2017; Peters et, 2018). 7 For classification, we simply fit the probe to the concatenated representations of the two models.

We train two 61M parameter models (so, 122M altogether) on FineWeb-Edu, each for 4.8B tokens (4 × Chinchilla). This again leads to a slight improvement (0.830 F1). One of our hypotheses for why our from-scratch RoBERTa slightly outperformed the much larger ModernBERT-large is that training on FineWeb-Edu gave it representations that were more salient for medical classification (compared to a default web text split). To push this further, we re-run biLM pretraining on a domain-upsampled corpus, where 50% of tokens were sourced from the PubMed section of the CommonPile (Kandpal et, 2025) and 50% were sourced from FineWeb-Edu.

And again, we see another incremental improvement: 0.834 F1. > 6 We also tried a number of other off-the-shelf pretrained friends of BERT: BERT, RoBERTa, DeBERTa, SciBERT, BioLinkBERT (Devlin et, 2019; Liu et, 2019; He et, 2021; Beltagy et, 2019; Yasunaga et, 2022). They were all worse. > 7 See §A for architecture details. In the future, genetically manufactured crops may be used to produce edible vaccines against infectious diseases. The production of Norwalk virus and the Hepatitis B virus can be produced in genetically modified sweet potatoes. This will in turn ensure the cost of production of vaccines is reduced availing it to the poor in

developing nations and therefore ensuring widespread vaccination. Hallucinogenic drugs or hallucinogens, as they are sometimes called, are drugs that effect users’ sensory perceptions. There are many drugs that fall into this broad category, but they are all characterized by their ability to change one’s subjective perception of reality. In the last 15 years, there has been a significant revival of interest in these drugs in particular amongst young Americans. Bloomy rinds are white and soft, sometimes even fuzzy. Cheesemakers spray a solution containing edible mold spores (Penicillium candidum, camemberti or glaucum) on the cheese. Humidity in the room where the cheese is ripened encourages this mold to grow, or bloom, and form a rind.

Figure 11. Classifier predictions for three randomly selected FineWeb-Edu documents. Annotations are from the classifier trained atop the 224M biLM, representing p(medical) ranging from low to high based on the F1-maximizing threshold. We also test whether scaling the size of these biLMs im-proves performance by training models at 113M and 224M parameters (again at 4 × Chinchilla). Table 1 shows the core result: as classifier scale increases, accuracy incrementally increases as well. Our final 224M parameter biLM classifier achieves 0.856 F1 on the val set and 0.894 F1 on the test set. These results are summarized in Table 1.

The upshot is that small, task-specific base models outperform large, general-purpose ones for token-level classification. Do-main specific pretraining helps models build representations where classification-relevant features are more salient. In 3 we show that higher classification performance indeed correlates with more effective filtering. 5.3. Document-level classification For document-level classification we mostly use the same approach, training a probe on top of the 224M biLM. We train on the same dataset as we do for the token-level classi-fier, but use Claude Sonnet 4 for labels; we use the same set of 128k documents for probe training. Our document-level classifier achieves 0.922 val and 0.941 test F1.

# 6. How bad are bad labels? A common critique of data filtering is that it is hard to get high quality labels, both for determining what to filter during pretraining and for actually training classifiers (Welbl et, 2021; Cloud et, 2024; Lee et, 2025; Shilov et, 2025). Here, we empirically study how much this matters. We show that while filtering is highly sensitive to label noise, even bad classifiers can be made into good filters, 10 Shaping capabilities with token-level data filtering 10% 30% 50% > 1 > 10 > 100 > Classifier Error Rate > Loss-Matched Compute Slowdown Model Size 61M 113M 224M 521M Figure 12.

Artificially noising labels makes filtering substan-tially worse. We simulate classifier error by randomly flipping labels (forget ↔ retain) with a given probability. For clas-sifier accuracy a = 0.89 and flip rate r, we plot error rate 1 − a(1 − r) − r(1 − a). Note that the error rate is in terms of SAE-generated ground truth labels, so our best performing clas-sifier still has an error rate of 11%. simply by shifting the decision boundary to be very high recall and scaling up model size. We also show that (1) token-level probes can be trained on coarse labels and (2) token-level probes easily generalize from low quality labels, while document-level probes do not.

6.1. They’re pretty bad... In some settings, it might be difficult to push classifier accu-racy beyond a certain level—compute scaling might plateau, labels might be too noisy, or the domain might just be too difficult. How bad is this? We simulate the noisy-label setting by randomly perturbing the labels generated by our gold-standard 224M biLM classifier. For each noise level, we train a series of models up to 521M parameters. Fig-ure 12 shows that this noising leads to power law scaling in compute slowdown: in the low error regime, increasing the error rate even a small amount leads to significantly less effective filtering, but this saturates in the high error regime.

but good things come to those who scale In cases like this, we still want to be able to effectively sup-press capabilities. Here, we show that in unbound compute regimes, bad classifiers can still be effective filters. To be precise: setting the decision boundary of our classifier to be extremely high recall at the cost of low precision, if we can scale models indefinitely, we can get models close to the frontier of low forget / high retain performance. Intuitively, this is because ‘aggressive’ classifiers are likely to remove proportionally more forget content than retain content;

, 3.75 4 4.25 4.5 4.75 5 > Non-Medical Loss > 3 > 4 > 5 > better > Parameters > 61M > 113M > 224M >