Technical Summary: Shaping Capabilities with Token-Level Data Filtering

Problem Statement

Current approaches to reducing undesired capabilities in large language models (LLMs) are largely post hoc, relying on interventions after pretraining (e.g., RLHF, machine unlearning, or input-output classifiers). The paper argues that these methods are easily bypassed by adversaries through jailbreaking or adversarial finetuning because they do not remove the underlying capabilities from the base model's representations. While shaping capabilities during pretraining via data selection is a promising alternative, existing literature has largely focused on coarse-grained document-level filtering. The authors posit that this granularity is suboptimal: models can learn dangerous capabilities from short subsequences of tokens within otherwise benign documents, leading to either the over-removal of useful data (low precision) or the retention of harmful knowledge (low recall).

Methodology

The paper proposes token-level data filtering as a more precise intervention during pretraining. The methodology involves four main stages:

1. Labeling via Sparse Autoencoders (SAEs): To generate ground-truth token-level labels without expensive manual annotation, the authors utilize sparse autoencoders trained on a 9B parameter model (Gemma 2). They identify SAE latents corresponding to the "forget" domain (medical knowledge) and use a stronger LLM (Claude Sonnet 4) to classify these latent descriptions. Tokens are labeled as "medical" if they exhibit high activation on these specific medical latents. This process creates a weakly supervised dataset of token-level labels.
2. Classifier Training: The authors train task-specific, bidirectional language models (biLMs) to act as token-level classifiers. They demonstrate that small, domain-specific models (e.g., 224M parameters) trained on relevant corpora outperform larger, general-purpose models (e.g., ModernBERT-large) for this specific classification task.
3. Filtering Strategies: Two token-level filtering strategies are implemented during pretraining:
   • Loss Masking: The model sees the "forget" tokens during the forward pass, but gradients are masked (zeroed out) during the backward pass for those tokens.
   • Removal: "Forget" tokens are replaced with a special <|hidden|> token, and their loss is masked.
4. Scaling Experiments: The authors train Transformer models ranging from 61M to 1.8B parameters on the FineWeb-Edu dataset, applying the filtering strategies. They evaluate performance on "forget" domains (medical capabilities) and "retain" domains (biology, general STEM) using text perplexity, multiple-choice questions (MedMCQA, MedQA-USMLE), and free-response generation.

Key Contributions and Results

1. Token Filtering Pareto Dominates Document Filtering

The paper demonstrates that token-level filtering achieves a Pareto improvement over document-level filtering. By filtering at the token level, the authors achieve the same reduction in "forget" capabilities (medical knowledge) with significantly less degradation in "retain" capabilities (biology). Document-level filtering often discards large amounts of benign data to catch small sequences of harmful tokens, whereas token filtering preserves the context of benign data while removing specific harmful tokens.

2. Effectiveness Scales with Model Size

A central finding is that the efficacy of data filtering increases with model scale.

• Compute Slowdown: For the largest models (1.8B parameters), token filtering leads to a 7,000× compute slowdown on the "forget" domain compared to a baseline trained on unfiltered data. This means an adversary would need 7,000 times more compute to re-learn the filtered medical capabilities via adversarial finetuning compared to a model trained on unfiltered data.
• Scaling Laws: The gap in performance between filtered and unfiltered models widens as pretraining compute increases. Token filtering results in lower magnitude scaling exponents on the forget domain, indicating that larger models are more effectively "shaped" by the filtering.

3. Robustness to Adversarial Attacks

Token filtering is significantly more robust to adversarial finetuning than state-of-the-art machine unlearning methods (specifically RMU).

• While RMU performance degrades rapidly with scale (requiring fewer adversarial tokens to re-learn capabilities), data filtering becomes relatively more robust as models scale.
• For 1.8B models, RMU requires 13× fewer tokens than token removal to achieve parity with the baseline, highlighting the superior tamper-resistance of pretraining filtering.

4. Alignment and Refusal Generalization

Contrary to concerns that removing knowledge might hinder a model's ability to refuse dangerous queries, the authors find that token-filtered models generalize better to refusal training than unfiltered baselines.

• Models trained with token removal generate refusals for medical queries at a rate 2× higher than the baseline while maintaining normal behavior on non-medical queries.
• The authors hypothesize that this is because the model learns to distinguish between "trained" and "untrained" tokens, a simpler task than classifying toxicity, allowing for effective alignment even without the underlying dangerous knowledge.

5. Robustness to Label Noise

The paper investigates the impact of imperfect labels.

• Weak-to-Strong Generalization: Token-level classifiers can generalize from weak or coarse labels (e.g., sentence or document-level labels) to produce effective token-level filters. Document-level classifiers, conversely, fail to generalize from weak labels.
• Scaling Mitigates Noise: Even with noisy labels, scaling up the model size and adjusting the classification threshold (prioritizing recall over precision) allows the filtering to remain effective.

Significance and Claims

The paper claims that token-level data filtering is a cost-effective, robust, and scalable method for shaping model capabilities during pretraining. Its primary significance lies in:

• Preemptive Safety: It addresses the "cat-and-mouse" game of post hoc safeguards by removing capabilities at the source (pretraining) rather than trying to extract them later.
• Efficiency: It achieves high precision in removing specific capabilities without the massive data loss associated with document-level filtering.
• Scalability: The method becomes more effective as models grow larger, suggesting it is a viable long-term strategy for frontier models.
• Alignment Compatibility: It demonstrates that removing dangerous knowledge does not necessarily compromise a model's ability to be aligned or to refuse harmful requests; in some cases, it may improve alignment generalization.

The authors conclude that while pretraining filtering is a "blunt instrument" compared to ideal influence-based removal, it is currently the most effective and robust intervention available for preventing frontier models from acquiring undesired capabilities, particularly when combined with a defense-in-depth approach including post-training safeguards.