Token-level Data Selection for Safe LLM Fine-tuning¶
日期: 2026-03-01
arXiv: 2603.01185
代码: GitHub
领域: LLM效率 / AI安全
关键词: safe fine-tuning, token-level selection, safety alignment, data curation, LoRA
一句话总结¶
提出 TOSS/TOSS-Pro，将 LLM 微调中的安全退化问题从"样本级"细化到"token 级"，通过训练安全退化/效用两个参考模型对每个 token 打分并遮蔽高风险 token，在 Llama-3-8B 上平均安全 win rate 比标准 SFT 高 33%，比最强 baseline SEAL 高 22%，同时保持效用不降。
研究背景与动机¶
- 领域现状：在自定义数据集上微调 LLM 会破坏预训练阶段建立的安全对齐，即使数据中只含少量有害样本也会导致安全能力显著下降。
- 现有方法的局限：
- 样本级防御（DSIR、SEAL）：丢弃整个被判定有害的样本，但同一样本中有害和有用信号共存，丢弃样本会损失有价值的效用信息。
- 后处理防御（SafetyLock、Circuit Breakers）：在推理时添加安全过滤，但对抗性微调可以轻松绕过。
- 核心洞察：安全退化从根本上是一个 token 级问题——有害信号集中在响应的初始 token 中，但也散布在整个序列里。需要在 token 粒度做精细化的数据筛选。
- 本文切入角度：训练两个"探针"参考模型（安全退化模型 + 效用模型），对比两者在每个 token 上的损失差异来识别和遮蔽危险 token，同时保留有用 token 用于训练。
方法详解¶
整体框架¶
TOSS 分三阶段：(1) 训练两个 LoRA 参考模型（安全退化 + 效用导向）；(2) 对自定义数据集每个 token 计算安全风险分数并全局排序遮蔽；(3) 用遮蔽后的 token 级 loss 微调目标模型。TOSS-Pro 在此基础上迭代优化安全退化参考模型。
关键设计¶
-
双参考模型训练
- 做什么：分别训练安全退化模型 \(f_{\theta^h}\) 和效用导向模型 \(f_{\theta^u}\)
- 核心思路：安全退化模型用 Anthropic HH-RLHF Red Teaming 子集（犯罪/暴力/辱骂等，~10% 数据），效用模型用 OpenOrca（~3% 数据），均用 LoRA（rank 16, Q/K 投影层）微调。两个模型设计为"极端化"——一个最大化学习有害模式，另一个最大化学习有用模式。
- 设计动机：通过对比两个极端，可以精确定位哪些 token 更像"有害分布"而不像"有用分布"。消融证实只用一个参考模型效果会打折。
-
Token 级风险评分
- 做什么：为自定义数据集中每个 token 计算安全风险分数
- 核心思路：分数定义为两个参考模型在该 token 上的损失差异： $\(\mathcal{S}(y_{i,j}) = \underbrace{-\log P(y_{i,j} | \mathbf{x}_i, \mathbf{y}_{i,:j-1}; \theta^u)}_{\text{效用损失（高=不像有用信息）}} + \underbrace{\log P(y_{i,j} | \mathbf{x}_i, \mathbf{y}_{i,:j-1}; \theta^h)}_{\text{安全损失（高=像有害信息）}}\)$ 高分 token = 在安全退化分布下概率高 + 在效用分布下概率低 → 有害风险大。
-
设计动机：单纯用安全退化模型会误伤高效用 token，加入效用模型的"反向约束"可以区分"真正有害"和"只是安全模型不认识"。
-
全局排序 + Token 遮蔽
- 做什么：将全数据集所有 token 按风险分数全局排序，遮蔽 top-\(d\)（默认 10%）
- 核心思路：遮蔽的 token 在微调 loss 中乘以 \(m_{i,j}=0\)，不参与梯度更新： $\(\mathcal{L}^{\text{cus}} = \frac{1}{\sum_{i=1}^{N} L_i} \sum_{i=1}^{N} \sum_{j=1}^{L_i} -m_{i,j} \log P(y_{i,j} | \mathbf{x}_i, \mathbf{y}_{j-1}; \theta)\)$
-
设计动机：全局排序比逐样本（局部）排序更好——有害样本中危险 token 密度远高于良性样本，局部排序会在良性样本中误遮蔽有用 token。
-
TOSS-Pro：渐进式参考模型优化
- 做什么：迭代改进安全退化参考模型，捕捉更多自定义数据中的隐含有害模式
- 核心思路：\(T\) 轮迭代——每轮用当前参考模型打分 → 选出含最高风险 token 的 top-\(k\) 样本 → 加入有害数据集重新训练参考模型。最终用最后一轮的参考模型做 token 评估。
- 设计动机：初始有害数据集可能不覆盖自定义数据中的全部有害模式。渐进优化使检测器逐步适应目标分布。消融证实 2 轮迭代即可收敛，且对初始参考模型质量不敏感。
损失函数 / 训练策略¶
- 所有模型用 LoRA（rank 16, Q/K 层），冻结其余参数
- 自定义微调：LR \(1 \times 10^{-5}\), batch 64, 2 epoch
- 安全退化模型：LR \(1 \times 10^{-4}\), batch 64, 2 epoch
- 效用模型：LR \(1 \times 10^{-5}\), batch 64, 2 epoch
- Token 遮蔽比例 \(d = 0.1\)（全局 top 10%）
实验关键数据¶
主实验¶
Win Rate（vs 标准 SFT, 由 GPT-4o 评判）：
| 方法 | Llama-3-8B (HH/HEx-PHI/SLIMORCA/AVG) | Llama-2-7B (AVG) |
|---|---|---|
| Standard SFT | 50/50/50/50 | 50 |
| Random | 52.5/65.0/48.0/55.2 | 50.0 |
| SafeInstr | 51.5/64.6/50.5/55.5 | 50.9 |
| DSIR | 67.4/60.8/53.8/60.7 | 57.6 |
| SEAL | 58.2/68.8/57.4/61.5 | 53.8 |
| TOSS | 88.8/87.5/68.4/81.6 | 70.1 |
| TOSS-Pro | 88.9/93.8/68.9/83.8 | 74.0 |
TOSS 平均 +31.6%（Llama-3），+20.1%（Llama-2）。TOSS-Pro 再提升 2-4%。
消融实验¶
| 配置 | HH test | HEx-PHI | SLIMORCA |
|---|---|---|---|
| TOSS (full) | 88.82 | 87.54 | 68.37 |
| 全局→局部排序 | ↓显著 | ↓显著 | ↓显著 |
| Token→Sample 级选择 | ↓↓ | ↓↓ | ↓↓ |
| 仅安全退化参考模型 | 安全↑ 效用↓↓ | - | - |
| 仅效用参考模型 | 安全不变 效用↑ | - | - |
| TOSS-Pro (2 iter) | 88.85 | 93.79 | 68.85 |
其他重要结果¶
- 通用能力保持：MMLU 63.07（vs SFT 62.99），HellaSwag 59.03（vs 57.95），TruthfulQA 45.78（vs SFT 15.91!）
- 跨模型迁移：在 Llama-3-8B 上选出的 token mask 直接用于 Llama-3.2-1B/3B，仍然有效（AVG 63.9/68.1 vs SFT 50）
- 后处理防御失败：SafetyLock（50.4）和 Circuit Breakers（40.3）在对抗性微调下几乎无效，TOSS（81.6）从源头阻止
- ASR 下降：Llama-3-8B HEx-PHI ASR 从 86.96%（SFT）→54.54%（TOSS）→43.63%（TOSS-Pro）
关键发现¶
- Token 级是正确粒度：安全退化确实是由特定 token 驱动的，不是整个样本。有害样本中危险 token 密度远高于良性样本（约 10x）。
- 全局排序优于局部排序：因为有害样本中几乎所有 token 都有高风险分数，局部排序无法正确处理这种不均匀分布。
- 两个参考模型缺一不可：只用安全退化模型会同时遮蔽效用 token（安全好但效用差），只用效用模型无法识别有害 token（效用好但安全不变）。
- TruthfulQA 强烈信号：标准 SFT 后 TruthfulQA 从 47.12 暴跌到 15.91，TOSS 恢复到 45.78——说明安全退化和事实幻觉高度相关。
亮点与洞察¶
- Token 级视角的核心洞察极具说服力：同一样本中有害和有用信号共存，样本级方法必然丢信息。这个观察可能改变安全微调领域的研究范式。
- 双参考模型的"对比探测"设计简洁优雅：用两个极端模型的损失差做定位，比单一探测器更精准。
- 跨模型迁移很实用：只需在大模型上做一次 token 评估，mask 可以直接复用给同 tokenizer 的小模型，大幅降低成本。
局限性 / 可改进方向¶
- 额外计算开销：需要训练两个 LoRA 参考模型，比无需训练的后处理方法成本高。
- 依赖有害参考数据：需要获取少量有害数据来训练安全退化模型，在某些合规场景下可能受限。
- 遮蔽比例固定：\(d=0.1\) 全局统一，不同数据集/领域的最优比例可能不同。
- OOD 泛化：仅在 privacy 类别上测试了 OOD 检测，更多有害类别的覆盖未充分探索。
相关工作与启发¶
- vs SEAL: 最强 baseline，做样本级安全数据选择。TOSS 在安全 win rate 上高 20 个百分点（81.6 vs 61.5），因为 SEAL 丢弃了有害样本中的有用信息。
- vs SafetyLock / Circuit Breakers: 后处理防御在对抗性微调下完全失效（win rate ~50），TOSS 从数据源头阻止，鲁棒性本质不同。
- 启发：Token 级粒度的思路可以扩展到其他数据质量问题（如去噪、去偏见），"哪些 token 驱动了某种行为"是一个通用的分析框架。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ Token 级安全分析视角非常新颖，双参考模型对比探测设计精巧
- 实验充分度: ⭐⭐⭐⭐⭐ 2 模型 × 3 基准 × 多种 baseline + 迁移/泛化/通用能力 + ASR/Harmfulness
- 写作质量: ⭐⭐⭐⭐ 方法描述清晰，实验全面
- 价值: ⭐⭐⭐⭐⭐ Token 级视角可能影响安全微调领域的研究方向