We study how to allocate a fixed supervised fine-tuning budget when three objectives must be balanced at once: multi-turn safety alignment, low over-refusal on benign boundary queries, and instruction following under verifiable constraints. We propose MOSAIC (Multi-Objective Slice-Aware Iterative Curation for Alignment), a multi-objective framework for closed-loop data mixture search built on a unified L1-L3 evaluation interface. MOSAIC turns slice-level failure profiles into executable data actions, including dataset-level mixture ratios, bucket-level weights, and focus criteria. Under a fixed 1M-token budget and five rounds of independent fine-tuning from the same base model, MOSAIC improves internal XGuard from 2.76 to 4.67 while keeping OrBench at 4.41 and IFEval at 3.65. The final Pareto solution also generalizes better than a random static LoRA baseline on independent attack, over-refusal, and capability tests, suggesting that structured failure diagnosis can serve as a practical control signal for budgeted data construction. Code is available at https://github.com/douyipu/mosaic.
Figures & tables
Figure 1: Overview of MOSAIC. The agent updates the next mixture from slice-level failure profiles under a fixed token budget.
Round
XGuard
OrBench
IFEval
Dataset mixture (x/o/i)
Base
2.7600
4.6667
3.4300
no fine-tuning
0
3.2267
3.6433
3.5300
0.50 / 0.30 / 0.20
1
3.3867
3.8033
3.5700
0.40 / 0.40 / 0.20
2
4.4567
4.3300
3.7033
0.35 / 0.45 / 0.20
3
3.9700
4.2967
3.5767
0.35 / 0.45 / 0.20
4
4.6700
4.4067
3.6533
0.35 / 0.45 / 0.20
Table 1: MOSAIC search history under a fixed 1M-token budget.
Figure 2: MOSAIC search trajectory under a fixed 1M-token budget. The proposal agent quickly converges on a stable macro mixture, while the number of selected windows still varies because the three datasets have different average lengths.
Solution
XGuard
OrBench
IFEval
Interpretation
Base
2.7600
4.6667
3.4300
best no-tuning usability reference
Iter 2
4.4567
4.3300
3.7033
most balanced three-objective point
Iter 4
4.6700
4.4067
3.6533
strongest safety-oriented solution
Table 2: Non-dominated solutions found by MOSAIC.
Figure 3: Controlled MOSAIC comparison for iterations 2–4. The macro mixture is identical, so the performance differences are attributable to bucket weights and focus criteria.
Figure 4: Diagnostic evidence from the final MOSAIC Pareto solution. The L1 score can be decomposed into concrete slice-level and atomic failure signals that directly inform data actions.
Attack success rate (%) ↓
Over-refusal (%) ↓
Model
X-Teaming
ActorAttack
Crescendo
XSTest
OKTest
OR-Bench
PHTest
Base
76.00
51.00
91.00
7.20
4.00
4.67
11.00
Random LoRA
1.00
21.50
17.00
18.00
26.67
17.67
36.33
Pareto
15.00
26.00
63.00
5.20
5.33
4.33
9.33
Table 3: External validation of MOSAIC. Lower is better for attack success rate and over-refusal, higher is better for general capability.
Figure 5: Granular external validation of the final MOSAIC solution. The Pareto model avoids the over-refusal spike and instruction-following collapse observed in random static mixing.
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.
Jingyu Peng, Hongyu Chen, Jiancheng Dong +5
University of Science and Technology of China · City University of Hong Kong · Baidu Inc +1
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}.
Guoli Wang, Haonan Shi, Tu Ouyang +1
Case Western Reserve University, Cleveland, OH, USA
Conflicting objectives are general in RL alignment, and training on them data-efficiently is hard. Training a safety guard with RL means optimizing two objectives that conflict: catch real harm, and do not refuse benign prompts. Our finding is that over-refusal improves 22.4% to 12.8%, while under-refusal on adversarial attacks silently worsens 0.27 to 0.33. We present C-Guard, a constitution-grid instrument that generates the RL training data, and C-LIM, a per-cell learnability score that decides each cell's move: prune, densify, amend, expand. C-LIM flags the dead-weight data region before any training budget is spent: 187 untargeted rows had bought zero gain, and our method lifts the same region's learning impact 0.733 to 0.80. Code and the constitution are open-sourced.