Submodular Optimization for Minimal Augmentation in Robust Language Model Alignment

Safety alignment of large language models is fragile: even small fine-tuning perturbations elastically revert behaviors toward those of the pre-training, with degradation inversely proportional to the size of the alignment set. We ask how to achieve safety alignment with minimal augmentation. To this end, we model augmentation as a set of group actions on sequences and formalize robustness gains as a normalized, monotone submodular function over transformations. We then leverage submodular optimization to select minimal augmentations that provably improve robustness. Experiments confirm that our approach efficiently restores safety alignment while minimizing the overhead of augmentation.

Large language models can still give unsafe answers when users rephrase harmful requests or when the model is later fine-tuned on new data. Our work asks how to make safety training more robust without adding many extra examples or slowing the model down at inference time. We treat data augmentation like choosing tools from a toolbox: each transformation should cover a different weakness, and repeated or redundant transformations should be avoided. To do this, we use a mathematical selection method that greedily picks the most useful safety augmentations while keeping the training set small. We also study a simple way to spread safety cues across different positions in model outputs, making the learned refusal behavior harder to undo. Our theory gives guarantees for why this selection process should work, and our experiments show that it can substantially reduce harmful responses across several open-source language models. This provides a practical path toward safer language models that does not require expensive reasoning steps or extra checks every time the model answers.