SoFA: Shielded On-the-fly Alignment via Priority Rule Following

TL;DR

SoFA reframes AI alignment as on-the-fly, priority-driven rule following and introduces PriorityDistill, a semi-automated pipeline that extracts rule-following signals from LLM simulations. By distilling (r,i,y) triplets and using a reference-guided loss, the approach strengthens both the integration and maintenance of rules, enabling shielded, adaptive responses across unseen regulations. Empirical results across multiple benchmarks show improved depth and breadth of alignment with modest alignment tax, and a PriorityRules dataset of over 20k rules offers a valuable resource for future work. The work advances practical, scalable on-the-fly alignment, with implications for safer and more controllable LLM deployments.

Abstract

The alignment problem in Large Language Models (LLMs) involves adapting them to the broad spectrum of human values. This requirement challenges existing alignment methods due to diversity of preferences and regulatory standards. This paper introduces a novel alignment paradigm, priority rule following, which defines rules as the primary control mechanism in each dialog, prioritizing them over user instructions. Our preliminary analysis reveals that even the advanced LLMs, such as GPT-4, exhibit shortcomings in understanding and prioritizing the rules. Therefore, we present PriorityDistill, a semi-automated approach for distilling priority following signals from LLM simulations to ensure robust rule integration and adherence. Our experiments show that this method not only effectively minimizes misalignments utilizing only one general rule but also adapts smoothly to various unseen rules, ensuring they are shielded from hijacking and that the model responds appropriately.

Figures (16)

• Figure 1: Learning-based alignment v.s. on-the-fly alignment via priority rule following. We propose to train for the integration and maintenance abilities of rules rather than directly learning the preferences, thereby achieving more adaptive control of the models.
• Figure 2: Example of test scenarios designed to challenge the integration and maintenance capabilities of LLMs. These scenarios require LLMs to accurately infer implicit knowledge behind rules (e.g., the irrelevance of photosynthesis to human diet and nutrition advice) and to handle conflicting instructions effectively.
• Figure 3: Pass rates of different models in the preliminary study. All evaluated models exhibit limited ability to integrate the rule effectively, especially when encountered with conflicting instructions.
• Figure 4: Overview of our Simulation Pipeline. The pipeline starts with extending seed instruction and rules set, then automatically identifies key rule-instruction pairs. Ultimately, it steers the model's response through a CoT process to ensuring that the model correctly applies the rules and maintains the relative priority. This CoT process is then distilled into the model parameters through direct learning of the $\left(r,i,y\right)$ triplet. The details are in Appendix \ref{['sec:pipeline_detail']}.
• Figure 5: Llama-2-chat system message.