Routing-Aware Expert Calibration for Machine Unlearning in Mixture-of-Experts Language Models

背景

随着大型语言模型（LLM）在社会各领域的广泛应用，机器遗忘（Machine Unlearning）的重要性日益凸显。机器遗忘旨在使模型能够“忘记”特定的训练数据，这通常出于隐私合规（如 GDPR 的“被遗忘权”）、消除偏见或更新知识等需求。

然而，现有的机器遗忘研究主要集中在稠密模型（Dense Models）上，对于近年来广泛采用的混合专家架构（Mixture-of-Experts, MoE）的关注仍然不足。MoE 架构通过在每一层引入一个路由器（Router），将每个 token 分配给稀疏的子集专家（Experts），从而在保持计算效率的同时扩展模型容量。这种稀疏激活机制使得 MoE 模型在遗忘任务中面临着独特的挑战：如何精准定位并调整那些对特定数据敏感的专家，同时避免破坏模型整体的通用能力。

核心内容

本文提出了一种名为 TRACE（Targeted Routing-Aware Calibration of Experts，面向专家的目标路由感知校准）的方法，专门用于解决 MoE 架构下的机器遗忘问题。

1. 问题观察：遗忘-保留路由不匹配

研究人员观察到，在 MoE 模型中，“遗忘数据”（需要被模型忘记的数据）往往不成比例地激活一小部分特定的专家。相比之下，这些专家在“保留数据”（模型需要保留的知识）上的激活强度要弱得多。

这种**遗忘-保留路由不匹配（forget--retain routing mismatch）**导致了一个关键问题：在进行机器遗忘时，那些对遗忘数据至关重要的专家（forget-critical experts）在正则化过程中受到的约束不足（under-regularized）。如果不对这些专家进行针对性处理，模型可能无法有效遗忘目标数据，或者为了遗忘而过度破坏模型的其他能力。

2. TRACE 方法流程

TRACE 方法通过以下两个步骤来解决上述问题：

• 离线检测遗忘关键专家：首先，利用离线激活统计信息，识别出那些在遗忘数据上激活显著高于保留数据的专家。这些专家被标记为“遗忘关键专家”。
• 重加权保留正则化：其次，通过重新加权 token 级别的保留损失（retain losses），对保留数据进行正则化校准。具体而言，该方法调整权重，使得每个选定专家的“保留侧激活频率”更好地匹配其“遗忘侧”的对应情况。

简而言之，TRACE 通过平衡遗忘数据和保留数据在关键专家上的激活分布，确保模型在遗忘特定知识时，能够更精准地调整相关参数，而不影响其他非关键专家或通用能力。

3. 实验验证

研究者在多个 MoE LLM 上进行了实验，评估数据集包括 WMDP（恶意大语言模型基准）和 MUSE-BOOKS（书籍遗忘基准）。

• 遗忘-效用权衡：TRACE 在遗忘质量和模型效用（Utility）之间取得了更优的平衡。
• 性能提升：在可比拟的遗忘质量下，TRACE 相比最强的基线方法，实现了 9% 的相对效用提升。
• 基准表现：在 MUSE-BOOKS 的四个指标中，TRACE 在三个指标上取得了最佳性能。

关键要点

• MoE 遗忘的独特性：MoE 架构中的稀疏专家激活机制导致遗忘数据集中激活少数专家，而保留数据激活较弱，造成正则化偏差。
• TRACE 的核心创新：提出了一种基于路由感知的专家校准方法，通过离线统计识别关键专家，并重加权保留损失以平衡激活频率。
• 解决路由不匹配：该方法直接针对“遗忘-保留路由不匹配”问题，确保遗忘关键专家在训练过程中受到适当的正则化约束。
• 显著的效用提升：实验表明，TRACE 在保持高遗忘质量的同时，显著提升了模型的保留知识能力（9% 相对效用提升），优于现有最强基线。
• 广泛的适用性：该方法在多个 MoE 模型和不同数据集（WMDP, MUSE-BOOKS）上均表现出一致的有效性。

意义与影响

这项研究填补了 MoE 架构下机器遗忘研究的空白。随着 MoE 模型（如 Google 的 Switch Transformer、Mixtral 等）成为主流架构，理解并优化其遗忘机制对于确保 AI 系统的合规性、安全性和可控性至关重要。

TRACE 方法通过精细化的路由感知校准，展示了如何在稀疏架构中实现更精准的参数调整。这不仅提高了机器遗忘的效率，也为未来在更大规模、更复杂的 MoE 模型中实施隐私保护和知识更新提供了可行的技术路径。对于致力于构建负责任 AI 的研究人员和工程师而言，TRACE 提供了一种在保持模型高性能的同时，有效移除敏感或错误知识的实用工具。