# MOSAIC: Composable Safety Alignment with Modular Control Tokens

arXiv: https://arxiv.org/abs/2603.16210 · HTML: https://mlsanity.com/paper/2603.16210/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Safety alignment and jailbreak attacks, one of 10 directions in LLM alignment / RLHF / preference optimisation.
- Builds on 19 of its 35 references that are in the map.
- First posted on arXiv 17 Mar 2026.

## Placed in

- LLM alignment / RLHF / preference optimisation / [Safety alignment and jailbreak attacks](https://mlsanity.com/alignment/safety-alignment-and-jailbreak-attacks/)
- LLM safety and security / [Safety alignment, refusal and guardrails](https://mlsanity.com/safety/safety-alignment-refusal-and-guardrails/)

## Builds on (most cited first, 19 of 19)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Lost in the Middle: How Language Models Use Long Contexts](https://mlsanity.com/paper/2307.03172/) (2023)
- [Red Teaming Language Models with Language Models](https://mlsanity.com/paper/2202.03286/) (2022)
- [SimPO: Simple Preference Optimization with a Reference-Free Reward](https://mlsanity.com/paper/2405.14734/) (2024)
- [BeaverTails: Towards Improved Safety Alignment of LLM via a Human-Preference Dataset](https://mlsanity.com/paper/2307.04657/) (2023)
- [ORPO: Monolithic Preference Optimization without Reference Model](https://mlsanity.com/paper/2403.07691/) (2024)
- [Safety Alignment Should Be Made More Than Just a Few Tokens Deep](https://mlsanity.com/paper/2406.05946/) (2024)
- [Understanding the Effects of RLHF on LLM Generalisation and Diversity](https://mlsanity.com/paper/2310.06452/) (2023)
- [ProSA: Assessing and Understanding the Prompt Sensitivity of LLMs](https://mlsanity.com/paper/2410.12405/) (2024)
- [Modular Pluralism: Pluralistic Alignment via Multi-LLM Collaboration](https://mlsanity.com/paper/2406.15951/) (2024)
- [Evaluating Cultural and Social Awareness of LLM Web Agents](https://mlsanity.com/paper/2410.23252/) (2024)
- [The Alignment Ceiling: Objective Mismatch in Reinforcement Learning from Human Feedback](https://mlsanity.com/paper/2311.00168/) (2023)
- [Controllable Safety Alignment: Inference-Time Adaptation to Diverse Safety Requirements](https://mlsanity.com/paper/2410.08968/) (2024)
- [From Distributional to Overton Pluralism: Investigating Large Language Model Alignment](https://mlsanity.com/paper/2406.17692/) (2024)
- [Logic Jailbreak: Efficiently Unlocking LLM Safety Restrictions Through Formal Logical Expression](https://mlsanity.com/paper/2505.13527/) (2025)
- [GRP-Obliteration: Unaligning LLMs With a Single Unlabeled Prompt](https://mlsanity.com/paper/2602.06258/) (2026)
- [Turning the Spell Around: Lightweight Alignment Amplification via Rank-One Safety Injection](https://mlsanity.com/paper/2508.20766/) (2025)
- [PICACO: Pluralistic In-Context Value Alignment of LLMs via Total Correlation Optimization](https://mlsanity.com/paper/2507.16679/) (2025)
- [NeST: Neuron Selective Tuning for LLM Safety](https://mlsanity.com/paper/2602.16835/) (2026)

## Built on by (most cited first, 1 of 1)

- [HiRoute: Hierarchical Routed Prompt Tuning for Safety Alignment of Large Language Models](https://mlsanity.com/paper/2608.12821/) (2026)

## Abstract (arXiv)

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.