MOSAIC: Composable Safety Alignment with Modular Control Tokens

• URL: http://arxiv.org/abs/2603.16210v1
• Date: Tue, 17 Mar 2026 07:40:43 GMT
• Title: MOSAIC: Composable Safety Alignment with Modular Control Tokens
• Abstract summary: Real-world deployments often require context-dependent safety rules that vary across users, regions, and applications.<n>We propose MOSAIC, a modular framework that enables compositional safety alignment through learnable control tokens optimized over a frozen backbone model.<n>Experiments show that MOSAIC achieves strong defense performance with substantially lower over-refusal while preserving model utility.
• Score: 22.486612919976977
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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.

Related papers

• Palette: A Modular, Controllable, and Efficient Framework for On-demand Authorized Safety Alignment Relaxation in LLMs [39.416807881947875]
  Current safety alignment of foundation models largely follows a emphone-size-fits-all paradigm.<n>We propose textscPalette, a modular, controllable, and efficient framework that selectively relaxes refusal behavior on authorized target domains.<n>Our method identifies a refusal direction via multi-objective search and internalizes it into the model through lightweight adaptation.
  arXiv  Detail & Related papers  (2026-05-22T19:22:17Z)
• Selective Safety Steering via Value-Filtered Decoding [54.87935112120107]
  Large language models (LLMs) are trained to align with human values, but their generations may still violate safety constraints.<n>Existing decoding-time steering methods often intervene unnecessarily, modifying generations that would have been safe under the base model.<n>We propose a new test-time steering method designed to reduce such unnecessary interventions while improving the safety of unsafe responses.
  arXiv  Detail & Related papers  (2026-05-14T12:13:08Z)
• Guardrails in Logit Space: Safety Token Regularization for LLM Alignment [9.180430520680611]
  Fine-tuning well-aligned large language models (LLMs) on new domains often degrades their safety alignment.<n>Existing safety alignment techniques primarily focus on pretraining, leaving fine-tuned models vulnerable to behavioral shifts.<n>We introduce safety token regularization (STR), a lightweight method designed to preserve safety properties during fine-tuning.
  arXiv  Detail & Related papers  (2026-04-19T02:41:04Z)
• Robust Multimodal Safety via Conditional Decoding [52.92816441364308]
  Multimodal large-language models (MLLMs) often experience degraded safety alignment when harmful queries exploit cross-modal interactions.<n>We propose a simple conditional decoding strategy, CASA, that utilizes internal representations of MLLMs to predict a binary safety token before response generation.
  arXiv  Detail & Related papers  (2026-03-31T23:19:50Z)
• Few Tokens, Big Leverage: Preserving Safety Alignment by Constraining Safety Tokens during Fine-tuning [4.839980912290382]
  Large language models (LLMs) often require fine-tuning (FT) to perform well on downstream tasks.<n>FT can induce safety-alignment drift even when the training dataset contains only benign data.<n>We propose a fine-tuning framework called Preserving Safety Alignment via Constrained Tokens (PACT), which stabilizes the model's confidence on safety tokens.
  arXiv  Detail & Related papers  (2026-03-08T03:42:55Z)
• BarrierSteer: LLM Safety via Learning Barrier Steering [83.12893815611052]
  BarrierSteer is a novel framework that formalizes safety by embedding learned non-linear safety constraints directly into the model's latent representation space.<n>We show that BarrierSteer substantially reduces adversarial success rates, decreases unsafe generations, and outperforms existing methods.
  arXiv  Detail & Related papers  (2026-02-23T18:19:46Z)
• PoSafeNet: Safe Learning with Poset-Structured Neural Nets [49.854863600271614]
  existing approaches often enforce multiple safety constraints uniformly or via fixed priority orders, leading to infeasibility and brittle behavior.<n>We formalize this setting as poset-structured safety, modeling safety constraints as a partially ordered set and treating safety composition as a structural property of the policy class.<n>Building on this formulation, we propose PoSafeNet, a differentiable neural safety layer that enforces safety via sequential closed-form projection.
  arXiv  Detail & Related papers  (2026-01-29T22:03:32Z)
• Taxonomy-Adaptive Moderation Model with Robust Guardrails for Large Language Models [3.710103086278309]
  Large Language Models (LLMs) are typically aligned for safety during the post-training phase.<n>They may still generate inappropriate outputs that could potentially pose risks to users.<n>This challenge underscores the need for robust safeguards that operate across both model inputs and outputs.
  arXiv  Detail & Related papers  (2025-12-05T00:43:55Z)
• Efficient Switchable Safety Control in LLMs via Magic-Token-Guided Co-Training [1.5349686675266894]
  Current methods for content safety in Large Language Models (LLMs) rely on multi-stage training pipelines.<n>We propose a unified co-training framework that efficiently integrates multiple safety behaviors.<n>We show that our method matches the safety alignment quality of SFT+DPO, with our 8B model notably surpassing DeepSeek-R1 (671B) in safety performance.
  arXiv  Detail & Related papers  (2025-08-12T02:39:33Z)
• Shape it Up! Restoring LLM Safety during Finetuning [65.75757313781104]
  Finetuning large language models (LLMs) enables user-specific customization but introduces critical safety risks.<n>We propose dynamic safety shaping (DSS), a framework that uses fine-grained safety signals to reinforce learning from safe segments of a response while suppressing unsafe content.<n>We present STAR-DSS, guided by STAR scores, that robustly mitigates finetuning risks and delivers substantial safety improvements across diverse threats, datasets, and model families.
  arXiv  Detail & Related papers  (2025-05-22T18:05:16Z)
• Controllable Safety Alignment: Inference-Time Adaptation to Diverse Safety Requirements [46.79887158348167]
  The current paradigm for safety alignment of large language models (LLMs) follows a one-size-fits-all approach.<n>We propose Controllable Safety Alignment (CoSA), a framework designed to adapt models to diverse safety requirements without re-training.
  arXiv  Detail & Related papers  (2024-10-11T16:38:01Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.