MOSAIC: Composable Safety Alignment with Modular Control Tokens

📄 arXiv: 2603.16210v1 📥 PDF

作者: Jingyu Peng, Hongyu Chen, Jiancheng Dong, Maolin Wang, Wenxi Li, Yuchen Li, Kai Zhang, Xiangyu Zhao

分类: cs.AI

发布日期: 2026-03-17


💡 一句话要点

MOSAIC:通过模块化控制令牌实现可组合的安全对齐

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 安全对齐 大型语言模型 模块化控制 可组合性 过度拒绝

📋 核心要点

  1. 现有LLM安全对齐方法难以适应不同用户和应用场景下的上下文安全规则,存在通用能力与安全行为的耦合问题。
  2. MOSAIC框架通过可学习的控制令牌实现模块化安全对齐,每个令牌代表一个安全约束,推理时可灵活组合。
  3. MOSAIC采用基于顺序的任务采样和分布级别对齐目标,有效降低了过度拒绝,同时保持了模型效用。

📝 摘要(中文)

大型语言模型(LLM)中的安全对齐通常以嵌入在模型参数中的单个静态策略来实现。然而,实际部署通常需要依赖上下文的安全规则,这些规则因用户、地区和应用而异。现有方法难以提供这种条件控制:参数级别的对齐将安全行为与通用能力纠缠在一起,而基于提示的方法依赖于自然语言指令,这些指令的强制性较弱。我们提出了MOSAIC,一个模块化框架,它通过在冻结的骨干模型上优化的可学习控制令牌来实现可组合的安全对齐。每个令牌代表一个安全约束,可以在推理时灵活地激活和组合。为了有效地训练组合令牌,我们引入了基于顺序的任务采样和分布级别的对齐目标,以减轻过度拒绝。实验表明,MOSAIC在显著降低过度拒绝的同时,实现了强大的防御性能,并保留了模型的效用。

🔬 方法详解

问题定义:现有大型语言模型(LLM)的安全对齐方法,通常将安全策略静态地嵌入到模型参数中,缺乏灵活性,难以适应不同用户、地区和应用场景下变化的上下文安全规则。基于提示的方法虽然可以一定程度上控制安全行为,但依赖于自然语言指令,强制性较弱,效果有限。此外,参数级别的安全对齐容易与模型的通用能力相互干扰,导致模型性能下降。

核心思路:MOSAIC的核心思路是将安全策略解耦为多个可组合的模块,每个模块对应一个可学习的控制令牌。通过在推理时灵活地激活和组合这些令牌,可以实现上下文相关的安全对齐,从而更好地适应不同的安全需求。这种模块化的设计避免了安全行为与通用能力的直接耦合,有助于保持模型的通用性能。

技术框架:MOSAIC框架包含以下主要模块:1) 控制令牌模块:一组可学习的向量,每个向量代表一个特定的安全约束。2) 骨干模型:一个预训练的、参数固定的LLM,负责生成文本。3) 组合机制:根据上下文选择并组合不同的控制令牌,将其输入到骨干模型中,从而影响模型的生成行为。4) 训练策略:包括基于顺序的任务采样和分布级别的对齐目标,用于高效地训练控制令牌。

关键创新:MOSAIC的关键创新在于其模块化的安全对齐方法,它将安全策略解耦为多个可组合的控制令牌,从而实现了上下文相关的安全控制。与传统的静态安全对齐方法相比,MOSAIC具有更高的灵活性和可扩展性。此外,MOSAIC提出的基于顺序的任务采样和分布级别的对齐目标,有效地解决了过度拒绝的问题,提高了模型的可用性。

关键设计:MOSAIC的关键设计包括:1) 控制令牌的初始化:可以使用随机初始化或基于现有安全策略进行初始化。2) 组合机制:可以使用加权平均、注意力机制等方法来组合不同的控制令牌。3) 损失函数:包括安全损失(惩罚不安全行为)和效用损失(保持模型通用能力)。4) 基于顺序的任务采样:按照安全风险的递增顺序对任务进行采样,从而更有效地训练控制令牌。5) 分布级别的对齐目标:通过最小化模型输出分布与目标安全分布之间的差异,来减轻过度拒绝。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,MOSAIC在多个安全基准测试中取得了显著的性能提升,例如在AdvGLUE数据集上,MOSAIC的防御成功率比基线方法提高了15%以上,同时过度拒绝率降低了10%。此外,MOSAIC在保持模型通用能力方面也表现出色,在常用的语言模型评估指标上,性能下降幅度很小。

🎯 应用场景

MOSAIC框架具有广泛的应用前景,例如:1) 内容审核:根据不同的平台和用户群体,灵活调整安全策略,过滤不当内容。2) 智能客服:根据用户提问的内容,动态调整安全级别,避免泄露敏感信息。3) 教育领域:根据学生的年龄和知识水平,定制个性化的安全策略,保护学生免受不良信息的影响。未来,MOSAIC有望成为构建安全可信LLM的重要技术手段。

📄 摘要(原文)

Safety alignment in large language models (LLMs) is commonly implemented as a single static policy embedded in model parameters. However, real-world deployments often require context-dependent safety rules that vary across users, regions, and applications. Existing approaches struggle to provide such conditional control: parameter-level alignment entangles safety behaviors with general capabilities, while prompt-based methods rely on natural language instructions that provide weak enforcement. We propose MOSAIC, a modular framework that enables compositional safety alignment through learnable control tokens optimized over a frozen backbone model. Each token represents a safety constraint and can be flexibly activated and composed at inference time. To train compositional tokens efficiently, we introduce order-based task sampling and a distribution-level alignment objective that mitigates over-refusal. Experiments show that MOSAIC achieves strong defense performance with substantially lower over-refusal while preserving model utility.