MOSAIC: Composable Safety Alignment with Modular Control Tokens

Jingyu Peng; Hongyu Chen; Jiancheng Dong; Maolin Wang; Wenxi Li; Yuchen Li; Kai Zhang; Xiangyu Zhao

cs.AIMar 17, 2026

MOSAIC: Composable Safety Alignment with Modular Control Tokens

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

Organizations: University of Science and Technology of China · City University of Hong Kong · Baidu Inc · Minzu University of China.

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.

Figures & tables

Figure 1: An illustration comparing two existing safety alignment paradigms and the proposed MOSAIC framework at inference time.

Figure 2: Overview of the proposed MOSAIC approach, showing the sampling strategy in the upper panel and the training objective in the lower panel.

Model	Method	# Params $\downarrow$	1-order		2-order		3-order		4-order
Model	Method	# Params $\downarrow$	DSR $\uparrow$	OR $\downarrow$	DSR $\uparrow$	OR $\downarrow$	DSR $\uparrow$	OR $\downarrow$	DSR $\uparrow$	OR $\downarrow$
Llama-3.1-8B	In-context	0.4M $/C$	76.7	10.2	62.0	13.6	51.7	11.2	44.5	13.9
	ORPO	33.6M	79.8	29.1	75.3	28.9	78.1	30.2	76.4	28.7
	SFT	33.6M	99.4	7.3	99.5	6.6	98.3	6.3	98.9	6.1
	MOSAIC-2	8.2k $/C$	99.8	7.1	99.6	7.5	99.8	6.2	100.0	5.9
	MOSAIC-5	20.5k $/C$	99.8	4.3	99.8	4.3	100.0	2.0	100.0	1.8

Table 1: Performance comparison of various safety-alignment methods. # Params denotes the number of trainable parameters for optimization-based methods (SFT, ORPO) or the additional parameter overhead per category for prompt-based approaches (In-context, MOSAIC). The suffix

/C

indicates that the overhead scales linearly with the number of target categories (

C

). MOSAIC-N indicates that the number of control tokens per category is set to N. K-order refers to the optimization and usage of control token combinations across K categories.

Model	Method	1-order	2-order	3-order
Llama-3.1-8B	Raw	0.552	0.552	0.552
	In-context	0.534	0.529	0.527
	SFT	0.547	0.544	0.547
	ORPO	0.549	0.542	0.549
	MOSAIC	0.551	0.552	0.551
Llama-3.2-3B	Raw	0.507	0.507	0.507

Table 2: Comparison of performance across task orders on MMLU.

Figure 3: Performance on Llama-3.1-8B under different numbers of control tokens per category.

Figure 4: Ablation Study. N-token uses only the control tokens of N subset categories during inference.

Figure 5: Performance on Llama-3.1-8B under different negative-to-positive ratio.

Model	Setting	DSR (%)	OR (%)
Llama-8B	Initial (2 categories)	99.8	4.3
	+1 Category	99.4	2.1
	+2 Categories	99.6	2.3
	+1 +1 Categories	99.8	2.8
Llama-3B	Initial (2 categories)	99.3	3.3
	+1 Category	99.5	3.4

Table 3: Incremental category expansion results. “+1” and “+2” indicate adding one or two new categories to the initial set, while “+1+1” denotes sequentially adding one category at a time.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Performance of MOSAIC on Llama3.1-8B under different

\lambda

Figure 7: Prompt for judgement

Explore similar work

Few Tokens, Big Leverage: Preserving Safety Alignment by Constraining Safety Tokens during Fine-tuning
Mar 8, 2026Guoli Wang, Haonan Shi, Tu Ouyang +1LLM Safety Alignment LLM Fine-Tuning
Case Western Reserve University, Cleveland, OH, USA
PolicyAlign: Direct Policy-Based Safety Alignment for Large Language Models
Jun 24, 2026Chang Wu, Junfeng Fang, Houcheng Jiang +5LLM Alignment LLM Safety Alignment
Qwen Large Model Application Team, Alibaba · National University of Singapore · Zhongguancun Academy
MESA: Improving MoE Safety Alignment via Decentralized Expertise
May 30, 2026Yitong Sun, Yao Huang, Teng Li +5Mixture-of-Experts Language Models LLM Safety Alignment
Institute of Artificial Intelligence, State Key Laboratory of Virtual Reality Technology and Systems, Beihang University, Beijing, 100191, China · College of AI, Tsinghua University, Beijing, 100083, China · School of Computer Science, Fudan University, Shanghai, 200438, China +2

Mar 8, 2026J/K move · Enter open · S save

Few Tokens, Big Leverage: Preserving Safety Alignment by Constraining Safety Tokens during Fine-tuning

Guoli Wang, Haonan Shi, Tu Ouyang, An Wang

Case Western Reserve University, Cleveland, OH, USA

LLM Safety Alignment LLM Fine-Tuning

Open paper AskPDF arXiv

Large language models (LLMs) often require fine-tuning (FT) to perform well on downstream tasks, but FT can induce safety-alignment drift even when the training dataset contains only benign data. Prior work shows that introducing a small fraction of harmful data can substantially compromise LLM refusal behavior, causing LLMs to comply with harmful requests. Existing defense methods often rely on model-wide interventions, such as restricting which parameters are updated or injecting additional safety data, which can limit generality and degrade downstream task performance. To address these limitations, we propose a fine-tuning framework called Preserving Safety Alignment via Constrained Tokens (PACT), which stabilizes the model's confidence on safety tokens. Our approach is motivated by the empirical observation that safety-aligned behavior is reflected in the model's token-level output confidence and is often concentrated on a small subset of safety-related tokens. During downstream fine-tuning, we regularize the fine-tuned model to match the aligned reference model's confidence on safety-related tokens at each response step, while leaving non-safety tokens largely unconstrained to allow effective task adaptation. This targeted constraint prevents alignment drift without imposing global restrictions that typically trade off with model utility. Our code is available at {https://github.com/Glresearch1/PACT}.

Abstract

Explore similar work

Few Tokens, Big Leverage: Preserving Safety Alignment by Constraining Safety Tokens during Fine-tuning

PolicyAlign: Direct Policy-Based Safety Alignment for Large Language Models

MESA: Improving MoE Safety Alignment via Decentralized Expertise