Training LLMs to Enforce Multi-Level Instruction Hierarchies via Gravity-Weighted Direct Preference Optimization

Technical Summary: Training LLMs to Enforce Multi-Level Instruction Hierarchies via Gravity-Weighted Direct Preference Optimization

Problem Statement
Production Large Language Models (LLMs) currently process instructions from sources with vastly different trust levels—ranging from immutable platform safety policies and developer system prompts to end-user queries and potentially adversarial tool outputs. Architecturally, these inputs are concatenated into a single token sequence where every token receives uniform attention, a state described as executing "in kernel mode." This structural uniformity creates a vulnerability to prompt injection, where malicious inputs hijack model behavior, and leaves models without a principled mechanism to resolve conflicts between legitimate but competing instructions from different hierarchy levels. While prior work has proposed training-based defenses, existing approaches typically formalize hierarchies of only three or four levels, treat all violations as equally severe, and fail to evaluate the full set of pairwise level interactions required by modern deployment architectures (e.g., OpenAI's five-tier stakeholder model).

Methodology
The authors formalize a $k$-level instruction hierarchy problem, instantiating it for $k=5$ levels:

1. L0 (Platform Governance): Immutable safety rules.
2. L1 (Developer System Prompt): Persona and scope.
3. L2 (Per-User Configuration): Session preferences.
4. L3 (User Message): Runtime input.
5. L4 (Data/Tool Outputs): External content.

This structure yields ten distinct pairwise priority relationships that a compliant model must enforce. To train models to respect these relationships, the paper introduces Gravity-Weighted Direct Preference Optimization (GW-DPO).

• Loss Function: GW-DPO adapts Offset-DPO (ODPO) by introducing a per-sample offset $\delta(i, j)$ that scales with the structural distance between conflicting levels. The loss is defined as:
  $$L_{GW-DPO} = -\mathbb{E} \left[ \log \sigma \left( \beta \cdot \Delta r_\theta(x, y_w, y_l) - \delta(i, j) \right) \right]$$
  where $\Delta r_\theta$ is the implicit reward margin.
• Schedules: Two schedules for $\delta(i, j)$ are evaluated:
  • Linear: Depends only on the privilege gap ($j-i$).
  • Bilateral: Factors in both the privilege gap and the privilege of the victim level ($k-1-i$), effectively weighting the severity of a violation by the "gravity" of the compromised level.
• Architectural Modifications: The training pipeline integrates two components from prior research:
  • Special Delimiter Tokens: Distinct start/end tokens for each hierarchy level to mark boundaries.
  • Instructional Segment Embeddings (ISE): Learnable per-token embeddings indexed by hierarchy level, added to token embeddings.
• Training Pipeline: The process involves a two-phase approach:
  1. SFT (Supervised Fine-Tuning): Teaches the model to recognize delimiters and integrate ISE information using a dataset of aligned, partial, and trivially misaligned prompts.
  2. GW-DPO: Uses a curriculum learning strategy (easy-to-hard based on hierarchy gap) to optimize the preference objective on a dataset containing explicit conflicts between levels.

Key Contributions

1. Formalization of $k=5$ Hierarchy: The paper defines a five-level instruction hierarchy grounded in production architectures, generating a custom evaluation suite that tests all ten pairwise priority relationships.
2. Gravity-Weighted DPO: The introduction of GW-DPO, specifically the bilateral schedule, which scales the training margin based on both the distance between conflicting levels and the privilege of the victim level. This approach Pareto-improves over standard DPO and the linear variant.
3. Component Isolation via Ablations: The study isolates the contributions of ISE and hierarchy depth:
   • ISE is identified primarily as a refusal-threshold calibrator. Removing ISE leaves conflict-resolution accuracy largely intact but causes an order-of-magnitude spike in over-refusal (refusing benign prompts), suggesting ISE helps the model distinguish between structured adversarial spans and benign structured spans without altering underlying conflict reasoning.
   • Five-level vs. Three-level Training: The comparison reveals a generality–specialization tradeoff. While three-level training (collapsing L0-L2 into a "System" tier) achieves higher scores on specific conflict pairs via surface-form recognition, it regresses on generalization to flat-text prompts and fails to resolve intra-system conflicts (e.g., L1 vs. L2) that lack surface cues.

Results
Evaluated on Llama-3.1-8B-Instruct, the GW-DPO (bilateral) configuration achieved the following:

• Conflict Resolution: It attained the highest Macro Pairwise Priority Adherence ($PPA_{macro} = 0.838$) and Weighted Hierarchy Score ($WHS = 0.885$), outperforming standard DPO and the linear variant.
• Over-Refusal: It maintained an Over-Refusal Rate (ORR) of 0.057, roughly half the rate of standard DPO (0.120) and significantly lower than the No-ISE configuration (0.576).
• Generalization: On flat-text reference splits (stripping delimiters), GW-DPO (bilateral) maintained strong performance ($PPA_{ref} = 0.793$), though slightly lower than the linear variant, indicating robust internalization of hierarchy semantics.
• External Benchmarks:
  • XSTest: Showed improved calibration on safe prompts compared to standard DPO, though with a trade-off in handling "safe-but-unusual" prompts.
  • SEP (Instruction-Data Separation): Achieved a probe-execution rate of 0.118, a ~5x improvement over the base model, demonstrating effective separation of instructions from data.
  • TensorTrust: Achieved a robustness rate of 0.818, significantly outperforming the base model (0.212), though slightly lower than standard DPO (0.906), reflecting the calibration trade-off.
  • MT-Bench: Preserved general chat-assistant utility, with GW-DPO (bilateral) scoring 7.425, comparable to or exceeding the off-the-shelf anchor.

Significance and Claims
The paper claims that structural margins, specifically the bilateral schedule in GW-DPO, serve as a useful and complementary lever for hierarchy compliance. The authors assert that their recipe is ready to scale to higher $k$ as deployment architectures become more layered.

Crucially, the paper frames the choice between three-level and five-level training not as a simple ordering question but as a generality–specialization tradeoff. Five-level training preserves resolution capability on plain-text rewrites and keeps over-refusal calibrated, whereas three-level training delivers narrower mastery of specific subsets via surface correlations. The authors emphasize that the bilateral schedule is the only configuration that simultaneously improves conflict resolution and reduces over-refusal compared to the linear schedule, aligning with the "no write up" constraints of classical security models (Biba, 1977).

The work concludes that while the recipe effectively defends against prompt injection and enforces hierarchy, it is not a complete adversarial defense against motivated attackers and requires careful policy consideration regarding the suppression of user concerns in favor of organizational priorities.