PaperlayerView exact revision on arXiv

Source-linked AI summary

Shaping capabilities with token-level data filtering

Neil Rathi, Alec Radford

arXiv:2601.21571v2cs.LGcs.AIcs.CL

TL;DR

Existing post hoc methods can leave undesired capabilities vulnerable to adversarial elicitation, so the paper studies shaping capabilities during pretraining by filtering data. On a medical-capability removal proxy, token filtering is effective, scales favorably, preserves alignment, and remains useful with imperfect labels.

  • Problem

    Post hoc interventions can leave undesired capabilities in the base model, motivating evidence on pretraining-time data filtering for capability shaping.

  • Method

    The paper labels pretraining data with classifiers and compares document- and token-level filtering across model scales, benchmarks, and alignment evaluations.

  • Results

    Token filtering Pareto-dominates document filtering, becomes more effective with scale, and yields a 7000× forget-domain compute slowdown for 1.8B models.

  • Takeaways & Limitations

    Pretraining token filtering can shape capabilities at scale without harming alignment and without requiring perfect labels.

  • Takeaways & Limitations

    Filtering struggles with in-domain classification, limiting fine-grained control across multiple forget domains, and classifier-based filtering is difficult for dual-use information.

Abstract

from arXiv · show

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.

1. Introduction

Post hoc interventions can leave undesired capabilities available to jailbreaks or finetuning, motivating capability shaping during pretraining. This paper studies filtering pretraining data, finding token-level filtering effective, scalable, alignment-compatible, and robust to imperfect labels.

  • Post hoc interventions can be bypassed because they do not remove undesired capabilities from the base model.
  • The paper frames data filtering as a pretraining intervention for selectively reducing undesired capabilities while preserving desired ones.
  • Token filtering achieves equal reduction in undesired capabilities at a lower cost to desired capabilities than document filtering.
  • For 1.8B parameter models, token filtering produces a 7000× compute slowdown on the forget domain.
  • Models trained with token filtering generalize to refusal training better than an unfiltered baseline.
  • A weakly supervised sparse-autoencoder pipeline labels tokens and supports cheap classifiers, while filtering remains possible despite imperfect labels.

2. Motivation and related work

Prior work primarily uses posttraining safeguards, unlearning, or coarse document filtering to shape capabilities. The paper motivates token-level pretraining filtering because capabilities can be learned from short token subsequences and existing safeguards remain bypassable.

  • Posttraining safeguards steer model behavior but remain vulnerable to jailbreaking or finetuning.
  • Once a capability exists in a base model, removing it is extremely difficult, motivating interventions during pretraining.
  • Prior document-level filtering suppresses CBRN-related capabilities and can improve robustness to adversarial finetuning, but reducing undesired pretraining content may hinder refusal elicitation.
  • Gradient-routing approaches shape capabilities using model representations and may bootstrap classification from imperfect labels.
  • Models can learn capabilities from short token subsequences, so document filtering may remove many benign tokens to achieve comparable recall.

3. Setting and approach

The study uses medical capability removal as a proxy for pretraining-time capability shaping while retaining related biology and other capabilities. It compares document- and token-level filtering across model scales and evaluates language-model, instruction-following, and free-response behavior.

  • 3. Setting and approach: The study partitions capabilities into forget and retain sets, targeting medical capability removal while preserving related biology capabilities.
  • 3. Setting and approach: The filtering pipeline labels a pretraining corpus, removes forget-relevant data, trains models at varied compute levels, and evaluates multiple benchmarks.
  • 3.1. Data and data filtering: The experiments compare document-level filtering with token-level loss masking and removal.
  • 3.1. Data and data filtering: Loss masking removes forget-token gradients while preserving their visibility during the forward pass; removal additionally replaces them with <|hidden|> tokens.
  • 3.2. Model training: Models range from 61M to 1.8B parameters and are trained as compute-optimal Transformers.
  • 3.2. Model training: The largest models receive instruction tuning and are evaluated with multiple-choice and free-response tasks, alongside text perplexity evaluations for smaller models.

4. Token-level data filtering works and scales

Token-level filtering reduces undesired medical capabilities while preserving retain capabilities, outperforms document filtering, becomes more effective with scale, and remains robust to adversarial finetuning and alignment.

  • Filtering works: Token filtering achieves equal reduction in forget capabilities at lower retain-capability cost than document filtering.It is a Pareto improvement over document filtering in the paper’s comparisons.
  • Filtering works: Token-level filtering remains close to the frontier of high forget loss and low retain loss across model series.The forget–retain loss plot places token filtering near the preferred top-left frontier.
  • Filtering scales: 7000× effective compute slowdown is achieved on the forget domain by token removal in the largest models, compared with around 30× for document filtering.Both filtering methods become more effective as pretraining compute scales.
  • Filtering works: Data filtering drives forget-domain multiple-choice performance near chance while causing no noticeable degradation on retain evaluations.This pattern appears on MedMCQA and MedQA-USMLE, with similar distinctions under cloze-style selection.
  • Filtering works: Token-level filtering reduces medical free-response quality more strongly than document filtering while leaving Alpaca performance largely unaffected.Token filtering makes responses 4× less coherent and relevant and 10× less correct in the reported evaluation.
  • Robustness: Both token and document filtering are more robust to adversarial finetuning than RMU, with relative robustness increasing with model scale up to 10× for 1.8B models.At 1.8B parameters, RMU requires fewer tokens than filtering baselines but has steeper scaling.
  • Alignment: Filtered models can distinguish forget from retain tokens, and the classification gap relative to the baseline closes with scale.However, filtering struggles to distinguish subdomains within the forget distribution, limiting fine-grained control across multiple forget domains.
  • Alignment: Models trained with token filtering can still be aligned on the forget domain, with token-level removal producing refusals at twice the baseline rate on HealthSearchQA.The increased refusal rate does not appear on Alpaca, while document-level filtering struggles to generalize to the task.

5. How to train your classifier

The classifier pipeline labels tokens relevant to forget-domain knowledge, using SAE-derived weak supervision and bidirectional linear probes to produce cheap, accurate token-level filters.

  • Classifier objective: The classifier targets tokens relevant to forget-domain knowledge as a proxy for tokens influential to forget capabilities.The proxy is imperfect: some identified datapoints may not influence capabilities, while some influential datapoints may not directly contain forget knowledge.
  • SAE-based labeling: SAEs identify forget-domain latents, which label tokens using high activation on multiple features and iterative propagation to adjacent tokens.The procedure extends token labels across spans so phrases such as ‘insert the catheter’ are classified as medical rather than only individual keywords.
  • SAE-based labeling: SAEs label a subset of data, which is then used to distill token-level classifiers because corpus-wide SAE inference is prohibitively expensive.The paper also notes that SAEs can lag behind simple linear probes for classification.
  • Classifier training: Bidirectional linear probes provide context-sensitive token classification, reducing ambiguity when the same token can belong to different domains.The probes use L-BFGS rather than full finetuning to improve robustness to spurious correlations under noisy labels.
  • Classifier training: 0.894 F1 on the test set is achieved by the final 224M-parameter biLM classifier, with accuracy increasing as classifier scale increases.The same classifier achieves 0.856 F1 on the validation set.
  • Classifier training: Small, task-specific base models outperform large, general-purpose models for token-level classification, and higher classification performance correlates with more effective filtering.Domain-specific pretraining makes classification-relevant features more salient.

6. How bad are bad labels?

The paper finds that imperfect labels need not prevent effective filtering: scaling and aggressive decision boundaries recover suppression, while token-level probes generalize better than document-level probes.

  • 6.1. Noise hurts: Random label noise makes filtering substantially worse and produces power-law scaling in compute slowdown.The simulated labels are randomly flipped between forget and retain, with error measured against SAE-generated ground truth.
  • 6.2. ...but good things come to those who scale: Bad classifiers can still become effective filters by shifting the decision boundary toward very high recall and scaling model size.The paper frames this as effective suppression in unbounded-compute regimes.
  • 6.2. ...but good things come to those who scale: Filtering proportionally more tokens moves models closer to the low-forget/high-retain loss frontier when sufficient scale is available.More aggressive filtering also incurs a larger hit to retain loss.
  • 6.3. Token-level classifiers generalize from weak labels: Coarse document- or sentence-level labels yield token classifiers only marginally worse than classifiers trained on fine-grained labels.Each token inherits the label of its containing document or sentence before probe training.
  • 6.3. Token-level classifiers generalize from weak labels: Token-level probes show weak-to-strong generalization from low-quality labels, whereas document-level probes do not.The comparison trains weak classifiers on a 13M-parameter biLM and evaluates strong probes on a 224M-parameter biLM.

7. Wrapping up

Token filtering is presented as an effective pretraining intervention for shaping capabilities, with benefits that extend across scale, robustness, alignment, labeling cost, and deployment strategy. The discussion also identifies unresolved limits involving classifier precision, scaling, evaluation, and dual-use information.

  • Wrapping up: Token filtering is a Pareto improvement over document filtering, becomes more effective with scale, and remains robust to adversarial finetuning without harming alignment.The authors also report that filtering can be performed cheaply and without perfect labels.
  • Shaping capabilities in pretraining: External classifiers remove knowledge related to dangerous capabilities rather than exactly identifying tokens that directly improve those capabilities.This makes pretraining filtering a blunt and potentially imprecise intervention.
  • Weak-to-strong generalization: Reliable external classification becomes harder as model size increases because weak supervision depends on annotators or representations sufficiently capable of judging data relevance.The paper identifies scalable oversight and reducing dependence on noisy labels as open questions.
  • Scaling further: Filtering improves in effectiveness with scale, but the paper notes that models beyond 7B could potentially learn dangerous capabilities from small amounts of data that evade filtering.The authors therefore treat further scaling as an unresolved question rather than assuming monotonic improvement indefinitely.
  • Better evaluations for capability shaping: Current capability-shaping evaluations are limited because they often target unlearning or capabilities that emerge only at large scales.The medical-capability proxy still required reasonably large models to produce evaluation signal.
  • Building effective safeguards against misuse: Pretraining filtering should complement rather than replace deployment safeguards, and the authors advocate defense in depth.Results on refusal training suggest that pretraining and posttraining safeguards can compound.
  • Building effective safeguards against misuse: Classifier-based filtering is difficult for dual-use information because it targets underlying capabilities rather than only the behavior exposed to end users.The paper suggests filtered public models and trusted access to fully capable models as one possible strategy.
  • Filtering for alignment: Filtering for dangerous capabilities is related to, but distinct from, filtering for misalignment risk and character priors.The paper identifies alignment filtering as a separate direction with multiple possible forms.

A.1. Architecture

The appendix describes the model architectures, optimization and instruction-tuning setup, and evaluations used to measure filtering effects. It also reports scaling-law and compute-matching analyses for the forget domain.

  • Architecture: Medical-filtering experiments use a modified GPT-2-style architecture with RoPE, ReLU2, pre-RMSNorm, and a constant width-to-depth ratio of 64.Pretraining models use block size 2048, while classifier models use block size 1024.
  • Architecture: RoBERTa classifiers use a six-layer, 65M-parameter architecture and train for 100k iterations at effective batch size 491,520.The setup uses the default RoBERTa-base architecture with the layer count reduced from 12 to 6.
  • Optimization: All experiments use AdamW; initial Muon experiments were abandoned after scaling revealed undertraining.The training setup also uses µP for hyperparameter transfer, with cosine learning-rate decay and a 10% linear warmup.
  • Instruction tuning: Instruction tuning combines ARC, BIG-Bench, BoolQ, MCTest, OpenBookQA, PIQA, RACE, and an auxiliary MMLU train set.Models train for one pass through 122k examples and reach 0.66 held-out accuracy versus 0.23 before instruction tuning.
  • Chat training: Chat training uses one pass through 460k smol-smoltalk examples; training on the full 1.1M-example dataset degraded coherence on Alpaca and HealthSearchQA.The reported chat-training learning rate is 10^-5 after hyperparameter search.
  • Evaluation: Free-response performance is generally comparable across models, although filtering causes slight degradation in a single-random-seed evaluation.The evaluation uses a 3k-question Alpaca subset judged by Claude Sonnet 4.
  • Scaling analysis: Filtering models have a smaller compute-to-loss exponent on the forget domain, meaning they scale worse there than baseline models.Figure 16 reports consistently higher medical loss for token filtering than document filtering and baseline.
  • Scaling analysis: The compute-matching analysis interpolates baseline compute-to-loss curves to estimate the compute required for an unfiltered model to reach a filtered model’s forget-domain loss.The relative compute slowdown is obtained by comparing the matched baseline compute with the filtered model’s compute.

B.2. Multiple choice evaluations

Multiple-choice evaluations measure answer selection from conditional answer-string losses and show the same capability-shaping pattern as other evaluations. Filtering reduces forget-domain performance while generally preserving retain-domain performance.

  • Multiple choice evaluations: Filtering consistently decreases forget-domain cloze accuracy, with token filtering outperforming document filtering.For each question, the model selects the answer string with the lowest loss conditioned on the question.
  • Multiple choice evaluations: Filtered models score around chance on forget evaluations while generally matching the baseline on retain questions.These trends are reported for base-model MCQ cloze accuracy.

B.3. Robustness

The robustness appendix specifies RMU and adversarial-finetuning procedures and describes compute-to-loss matching for baseline comparisons. These procedures support evaluation of forgetting and recovery under intervention.

  • Robustness: RMU uses AdamW with learning rate 1 × 10^-4, weight decay 0.01, batch size 8192, α = 100.0, and c = 20.0.Optimization runs for 1,000 steps on middle and preceding MLP layers, where forget loss has begun to plateau.
  • Robustness: The compute-matching procedure interpolates baseline compute-to-loss curves to estimate the compute needed to match filtered-model loss.Figure 19 presents this baseline-compute estimate.
  • Robustness: Adversarial finetuning sweeps learning rates from 1 × 10^-5 to 1 × 10^-3 and weight decay values of 0.01 and 0.1.Hyperparameters are selected by the fewest steps needed to reach parity with baseline loss.

B.4. Training to generate refusal tokens

Refusal-token training preserves the pattern observed with free-response refusal: token removal substantially increases refusal on medical questions while retaining baseline-level behavior on benign instruction following.

  • Token removal produces substantially higher refusal rates on HealthSearchQA than the baseline model.
  • Token masking is on par with the baseline on the refusal-token evaluation.
  • Document filtering is also compared in the refusal-token experiment, with its benign Alpaca behavior described separately in Figure 20.

B.5. Training dynamics

The experiments show that token filtering benefits from early application, precise token-level labels, and larger models, while weakening within-domain discrimination and remaining vulnerable to classifier quality and adversarial finetuning.

  • Training dynamics: Delaying filtering substantially degrades its effectiveness, making early loss masking important throughout pretraining.The study changes when loss masking begins and reports substantial degradation when filtering starts later.
  • Domain classification: Token-filtered models match the baseline on forget-versus-retain classification but struggle to distinguish subdomains within the forget domain.This limits fine-grained control across multiple forget domains.
  • Token-level filtering: Only around 23% of documents contain zero medical tokens, while 37% contain more than 10% medical tokens, supporting token-level filtering.Document filtering would need to discard many partially medical documents or tolerate leakage to match token-level precision and recall.
  • Classifier construction: Classifiers trained with coarse sentence- or document-level labels generalize well to token-level ground truth, performing only marginally worse than token-labeled classifiers.
  • Adversarial finetuning: Filtered models change more gradually than RMU under adversarial finetuning, although RMU begins with a test loss 3× higher than token removal and improves steeply within a few steps.
  • Classifier quality: Better classifiers generally produce better filters, with higher AUROC associated with higher normalized AUC relative to the baseline.

D. Example responses to free-response medical questions

Free-response examples show that data-filtered models may emit medical tokens but generally fail to use them correctly, producing irrelevant, repetitive, or medically distorted answers across several questions.

  • Botulism: For botulism, token masking attributes the illness to bullying, peer pressure, and social media instead of its medical causes.
  • Lung cancer: The lung-cancer examples include repetition, malformed explanations, and inaccurate treatment or disease descriptions after filtering.
  • Cold versus Covid: On the cold-versus-Covid question, the token-masking response gives contradictory temperature definitions for both conditions.
  • Cold versus Covid: The token-removal response to the same question replaces medical distinctions with contradictory statements about heat and cold.
  • Diabetic retinopathy: For diabetic retinopathy, the document-filtered response incorrectly describes a pancreatic insulin condition rather than the retinal disease.
Loading 2601.21571v2…