通过Token级数据过滤塑造模型能力 Shaping capabilities with token-level data filtering
在预训练阶段用token级过滤移除不想要的能力,比文档级过滤更有效且随规模提升
前置知识
数据过滤 (Data Filtering)
在预训练数据集中选择性地移除包含特定内容的数据,以防止模型学习到不想要的能力。这与后训练干预(如RLHF)不同,它从源头上控制模型能学到什么。数据过滤可以在文档级别(移除整篇文档)或token级别(只移除特定词汇单元)进行。文档级过滤简单但粗糙,可能误删有用数据;token级过滤精确但需要更细粒度的标注。
本文的核心贡献就是证明token级数据过滤优于文档级过滤,理解这个概念是把握全文的基础。
稀疏自编码器 (Sparse Autoencoder, SAE)
一种用于分解模型激活的无监督学习方法。SAE将模型的高维激活向量分解为多个稀疏的'特征',每个特征对应一个可解释的概念(如'医疗信息'、'编程代码'等)。通过分析SAE的特征,可以理解模型内部在处理什么信息。本文使用Gemma 2 9B的SAE来识别与医疗相关的token。
本文使用SAE来生成token级标注,这是方法的关键创新之一——利用可解释性工具来做数据标注,而非人工标注。
线性探针 (Linear Probe)
在预训练模型的某一层表示之上训练的简单线性分类器。由于只使用线性层,探针的性能反映了模型表示中是否已经编码了目标任务的信息。线性探针常用于评估模型学到了什么,但本文将其用于训练token级分类器来过滤数据。使用L-BFGS优化器而非全模型微分可以提高对噪声标签的鲁棒性。
本文训练线性探针作为高效的token分类器,这是将SAE标注扩展到大规模预训练数据集的关键技术。
机器遗忘 (Machine Unlearning)
从已训练模型中移除特定知识或能力的技术。现有的遗忘方法(如RMU)通过修改模型内部表示来'擦除'特定能力,但这些方法容易被对抗性微分绕过——只需几步微分就能恢复被遗忘的能力。遗忘方法的脆弱性促使本文探索从预训练阶段就避免学习不想要的能力。
理解遗忘方法的局限性是本文研究动机的核心——后训练干预不够鲁棒,所以需要在预训练阶段塑造能力。
缩放定律 (Scaling Laws)
描述模型性能如何随计算量、数据量或参数量变化的数学关系。通常表现为幂律关系:$L \propto C^{-\alpha}$,其中$L$是损失,$C$是计算量,$\alpha$是缩放指数。本文发现数据过滤的有效性随规模增加而提高——这是因为他们观察到过滤后模型的缩放指数$\alpha$在遗忘领域更小,意味着增加计算量带来的收益更少。
本文的核心发现之一是token过滤随规模变得更有效,理解缩放定律有助于理解这个结论的数学含义。
研究动机
当前大语言模型的安全防护主要依赖后训练干预,如RLHF教模型拒绝危险请求、机器遗忘从已训练模型中移除能力。然而这些方法存在根本性缺陷:RLHF教模型'不回答'危险问题,但底层能力依然存在,攻击者可以通过越狱(jailbreak)或微分(finetuning)轻松绕过。实验表明,对抗性微分只需几百步就能让经过遗忘处理的模型恢复被移除的能力。论文引用了大量研究表明,一旦能力在预训练中被'学进去',就极难通过后训练完全移除——这不是技术实现问题,而是根本性的架构限制。例如,RMU等遗忘方法虽然在静态评估中表现良好,但面对对手只有少量计算资源时就会失效,形成了一场永无止境的'猫鼠游戏'。
本文的目标是本文的直接目标是在预训练阶段通过数据过滤来移除不想要的能力,而不是在模型训练完成后再试图'擦除'。具体来说,他们以医疗能力作为代理任务(因为真正的危险能力如生物武器无法在研究中测试),目标是:(1) 训练出在医疗领域表现接近随机、但在相关领域(如生物学)保持基线水平的模型;(2) 证明这种方法随规模增大而更有效;(3) 展示过滤后的模型仍然可以被对齐训练控制。他们使用模型规模从61M到1.8B参数,跨越两个数量级来研究缩放行为。
与已有工作不同的是,本文的独特视角在于将数据归因(data attribution)领域的洞察引入数据过滤。以往的数据过滤研究主要在文档级别进行,但数据归因文献表明,单个token对模型能力的影响差异巨大——一个文档中的几个关键token可能决定了模型是否学到某种能力。本文抓住了这个被忽视的点:文档级过滤为了移除几个关键token会误删大量有用的保留领域数据,而token级过滤可以在相同召回率下实现更高精确率。此外,本文还探索了用可解释性工具(SAE)来生成token级标注,这是一种弱监督方法,避免了昂贵的人工标注。
核心方法
本文的方法可以类比为'精准医疗'与'全身化疗'的区别。文档级过滤就像全身化疗——为了杀死癌细胞(移除医疗能力),必须对整个身体(文档)进行处理,这会误伤健康细胞(保留领域的有用token)。Token级过滤则像精准医疗——只针对癌细胞进行治疗,保留健康组织。技术路线分为四个阶段:首先,使用稀疏自编码器(SAE)识别预训练语料中与医疗相关的token并生成标注;其次,训练一个高效的双向语言模型分类器来对token进行分类;然后,使用分类器过滤预训练数据(通过损失掩码或token替换);最后,在过滤后的数据上训练不同规模的模型并评估效果。
本文最核心的创新点是token级过滤相比文档级过滤是一个帕累托改进。具体来说,在相同的'遗忘损失'(即医疗能力的退化程度)下,token级过滤对'保留损失'(即生物学等相关领域能力的影响)更小。这是因为医疗相关内容通常只占文档的一小部分——一篇生物学论文可能只有几句话提到医学概念。文档级过滤为了移除这几句话会丢弃整篇文档中有价值的生物学内容,而token级过滤可以只移除那几个关键token,保留其余内容供模型学习。实验表明,对于521M参数模型,token过滤在相同召回率下可以达到更高精确率,这意味着更少的'误伤'。
方法步骤详情
方法分为以下步骤:(1) 数据标注:使用Gemma 2 9B的SAE(16k宽度,第31层)识别医疗相关特征。首先用Claude 3.5 Haiku为每个SAE特征生成自然语言解释,然后用Claude Sonnet 4分类解释是否与医疗相关,并使用嵌入评分过滤低质量解释,最终得到600个医疗相关特征。Token标注规则为:如果token在至少2个医疗特征上的激活值高于均值4个标准差,或在至少1个医疗特征上有正激活且与已标注医疗token相邻,则标注为医疗。(2) 分类器训练:在SAE标注的基础上训练双向语言模型分类器。训练了一个224M参数的双向语言模型(由两个113M的自回归模型组成,分别从左到右和从右到左),在FineWeb-Edu上训练4.8B token。然后在8.2M token的标注数据上训练线性探针,使用L-BFGS优化器。(3) 数据过滤:将分类器应用于FineWeb-Edu数据集,移除被标记为医疗的token。考虑两种策略:损失掩码(保留token但在反向传播时忽略其梯度)和token替换(将医疗token替换为特殊<|hidden|>标记)。(4) 模型训练与评估:在过滤后的数据上训练61M到1.8B参数的Transformer模型,使用Chinchilla最优配置,评估文本困惑度、多项选择和自由回答三个维度。
技术新颖性
本文的技术新颖性体现在三个方面。首先,将数据过滤的粒度从文档级推进到token级,这是对数据归因研究的直接应用——Grosse et al. (2023) 的工作表明单个token的影响力差异很大,本文首次将这个洞察用于能力塑造。其次,创新性地使用可解释性工具(SAE)进行数据标注,这是一种弱监督方法,避免了昂贵的人工标注,且可以扩展到没有现成标注数据的领域。最后,本文系统地研究了过滤的缩放行为,发现过滤的有效性随规模增大而提高——这是首次在两个数量级的计算范围内展示这种现象,为大规模应用提供了信心。此外,本文还发现了token过滤的一个意外优势:过滤后的模型更容易被对齐训练控制,这与直觉相反——人们可能认为'不知道危险知识就无法拒绝',但实际上模型学会了区分'学过的'和'没学过的',从而更好地泛化到拒绝训练。
实验结果
本文的核心实验结果可以分为四个维度。首先是帕累托优势:在521M参数模型上,token过滤在相同召回率下(即相同医疗损失)达到了更高精确率(即更低生物学损失),实现了帕累托改进。其次是缩放有效性:token过滤的有效性随规模增大而提高。对于1.8B参数模型,token替换过滤实现了7000倍的有效计算减速——这意味着一个在过滤数据上训练的模型要达到基线在未过滤数据上的医疗损失,需要7000倍的计算量。相比之下,文档级过滤只有约30倍减速。在多项选择评估中,过滤后的模型在MedMCQA和MedQA-USMLE上接近随机水平(准确率约25%),但在保留领域的MMLU子集上没有明显退化。在自由回答评估中,token过滤使HealthSearchQA的正确性下降10倍,相关性和连贯性下降3-4倍。第三是鲁棒性:token过滤比RMU等遗忘方法更抗对抗性微分。对于1.8B模型,RMU所需的对抗性微分token数量比token替换少13倍,比token损失掩码少3倍。随着规模增大,RMU的鲁棒性曲线比数据过滤更陡峭,意味着差距会继续扩大。最后是对齐兼容性:token过滤后的模型更容易学习拒绝。在拒绝训练实验中,token替换模型在HealthSearchQA上的拒绝率比基线高2倍,而文档级过滤模型则难以泛化,对Alpaca也产生大量拒绝。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 医疗知识 (MedMCQA) | 准确率 | ~25% (接近随机) | ~35-40% | 降低约30-40%相对性能 |
| 医疗知识 (MedQA-USMLE) | 准确率 | ~25% (接近随机) | ~35-40% | 降低约30-40%相对性能 |
| 有效计算减速 (1.8B模型) | 计算减速倍数 | 7000x (token替换) | 30x (文档过滤) | 233倍更有效 |
| 对抗性微分鲁棒性 (1.8B) | 所需微分token比例 | 13x多于RMU (token替换) | RMU基线 | 鲁棒性提升13倍 |
| 拒绝训练泛化 (HealthSearchQA) | 拒绝率 | ~60% (token替换) | ~30% (基线) | 2倍提升 |
| 分类器性能 | F1分数 | 0.894 (测试集) | 0.812 (ModernBERT-large) | 10%相对提升 |
局限与改进
本文存在几个重要局限。首先是代理任务的局限:使用医疗能力作为危险能力的代理,但医疗能力与真正的危险能力(如生物武器合成)在知识结构和获取难度上可能有本质区别。医疗知识广泛存在于互联网,而真正的危险知识可能更稀缺,过滤行为本身可能更明显。其次是规模局限:最大模型只有1.8B参数,而当前前沿模型已达到数百B参数。虽然缩放趋势表明过滤在更大规模上应该更有效,但也存在'U型缩放'的可能——足够大的模型可能从极少量漏网样本中学到能力。第三是数据集局限:使用FineWeb-Edu作为预训练数据集,这是一个教育内容丰富的数据集,医疗内容比例可能高于一般网络数据。在医疗内容更稀疏的数据集上,token级过滤的优势可能不那么明显。最后是分类器的局限:SAE标注是基于Gemma 2 9B的特征,这可能引入特定于该模型的偏差。此外,分类器只能识别与医疗'相关'的token,而非真正'影响'医疗能力的token——某些不直接包含医疗信息但对推理能力至关重要的token可能被误删。
独立分析的弱点
本文的一个显著弱点是缺乏对过滤'副作用'的系统性评估。虽然作者评估了生物学和STEM等保留领域,但没有评估更广泛的能力,如编程、数学推理或常识。过滤医疗相关token可能会意外移除对这些能力重要的上下文信息,特别是在医疗内容与其他领域交织的文本中。改进方向是开发更精细的'能力图谱',追踪不同领域知识之间的依赖关系,避免过度过滤。另一个弱点是分类器的泛化性不足。本文的分类器是针对'医疗'这个特定领域训练的,对于其他'不想要的能力'(如虚假信息、黑客技术),需要重新训练分类器。未来可以探索多领域通用的分类器架构,或者开发可以快速适应新领域的少样本分类方法。此外,损失掩码和token替换两种策略各有优劣,但本文没有深入研究混合策略或自适应策略——例如,对某些关键token使用替换,对其他使用掩码。
未来方向
作者提出的研究方向包括:(1) 探索使用模型自身的表示来指导过滤,而非外部分类器——例如基于数据归因方法直接计算每个token对能力的影响;(2) 研究是否可以让模型在预训练中学会'按能力自我组织',使其内部表示自然地将不同能力分离;(3) 推动缩放定律研究到7B以上规模,验证是否存在'U型缩放';(4) 将过滤应用于对齐风险(如减少危险倾向、塑造角色先验),而不仅仅是能力移除。基于本文成果可延伸的方向包括:(1) 开发'能力编辑'而非'能力移除'——用类似技术替换模型的知识,而非简单删除;(2) 研究过滤与课程学习的结合——在训练的不同阶段使用不同的过滤策略;(3) 探索多模态过滤——将类似方法应用于图像、代码等非文本数据;(4) 开发实时过滤系统,在预训练过程中动态调整过滤策略。
复现评估
本文的复现难度中等。代码已开源在GitHub(neilrathi/token-filtering),包含完整的训练和评估流程。数据方面,使用FineWeb-Edu(公开可用)和CommonPile的PubMed部分,都可以获取。SAE模型使用的是Gemma Scope的公开权重。主要的复现挑战在于计算资源:训练224M参数的双向语言模型分类器需要约4倍Chinchilla的计算量(约4.8B token),训练1.8B参数的主模型需要8张NVIDIA H200。对于学术实验室来说,完整复现所有实验(61M到1.8B的模型系列)可能需要数万美元的计算成本。然而,核心创新点(token级过滤优于文档级过滤)可以在小规模(如521M参数)上验证,这大大降低了复现门槛。此外,分类器训练数据集(128k文档,8.2M token)的规模也相对可控。总体而言,如果只验证核心结论而非完整缩放研究,复现是可行的。
论文图表