SafeCtrl-RL: Inference-Time Adaptive Behaviour Control for LLM Dialogue via RL-Driven Prompt Optimisation

Image generated by AI

TLDR

• What: An inference-time behavioral control framework (SafeCtrl-RL) that uses a Reinforcement Learning (DQN) agent to dynamically select from 11 distinct prompt-refinement strategies based on real-time safety and quality feedback.
• Who's at risk: Production LLM dialogue applications (e.g., educational chatbots, RAG pipelines, and code assistants built on models like DeepSeek-R1, Llama 3, or customized uncensored architectures) vulnerable to adversarial jailbreaks and toxic generation.
• Key number: SafeCtrl-RL improves the macro safety-quality score (Macro-$P_{\text{Score}}$) from a baseline of 0.305 to 0.818 ($\Delta P_{\text{Macro}} = +0.513$) while outperforming state-of-the-art prompt optimization methods like OPRO by up to 139% in safety-critical settings.

Ensuring that Large Language Models (LLMs) behave safely and contextually during real-time interactions remains an active battleground for AI safety engineers. Standard offline alignment techniques—such as Supervised Fine-Tuning (SFT) and Reinforcement Learning from Human Feedback (RLHF)—frequently fail against creative, out-of-distribution jailbreak prompts. This vulnerability is especially critical today for customer-facing systems, RAG pipelines in Cursor, or Google’s AI Overviews, where safety violations can result in immediate brand damage and regulatory non-compliance.

To address these vulnerabilities, Orme et al. (2026) introduce SafeCtrl-RL, an inference-time behavioral control framework. Instead of editing weights or relying on brittle, static system prompts, SafeCtrl-RL formulates dialogue safety as a sequential decision process. It deploys an external Reinforcement Learning agent to dynamically rewrite prompts on the fly, systematically driving the target LLM toward high-quality, harmless responses.

Threat Model & Safeguarding Context

To understand the defensive positioning of SafeCtrl-RL, we must define the interaction boundaries between the adversarial user, the target model, and the safeguard framework.

Background & Related Work

Existing methods for controlling LLM safety and behavior fall into three distinct paradigms. Unlike typical adversarial defense literature that relies on static input filters, SafeCtrl-RL treats prompt modification as a dynamic control system.

Comparison of LLM Behavior Control Paradigms

Unlike PoisonedRAG (Zou et al., CCS 2024), which demonstrates how adversaries exploit retrieval context to bypass safety, SafeCtrl-RL acts directly at the system prompt layer. This allows the system to suppress harmful behaviors dynamically, functioning as a real-time behavioral unlearning safeguard.

Methodology

SafeCtrl-RL models the dialogue safeguarding task as a closed-loop control system. As detailed in Section 3, the architecture consists of three core components:

[ User Input ] ---> [ LLM Response ] ---> [ Safety-Quality Evaluator ]
       ^                      |                        |
       |                      v                        v
[ System Prompt ] <--- [ DQN Agent ] <--- [ 36-D State Vector Encoder ]

1. The 36-Dimensional State Vector

Rather than feeding raw text histories to the RL controller, a State Encoder compresses the interaction trace into a 36-dimensional normalized vector (described in Table 4). This state representation is constructed from:

• Training & Iteration Progress (Features 1–5): Tracks metrics like Category Progress and Performance History Fullness.
• Performance Metrics (Features 6–11): Includes the Recent Mean Performance (up to last 10 scores), Recent Volatility, and Performance Consistency.
• Strategy Performance (Features 12–15): Records the historical success of each available refinement action.
• Risk and Safety Signals (Features 25–27): Captures the Unsafe Category Score (mapped via a 12-category safety taxonomy) and Consecutive Error Rates.
• Prompt characteristics (Features 31–36): Measures syntactic features and a Locality-Sensitive Hash (LSH) fingerprint of the user prompt to recognize structurally similar inputs.

2. The 11-Action Discrete Space

The DQN controller does not generate text directly. Instead, at each iteration $k$, it selects an action $a_k \in \mathcal{A}$ containing 11 prompt engineering strategies (Section 3.3.2):

• Direct Strategies: Minimal (raw prompt only) and Raw History.
• Summarisation-Based: AI Summary Only, AI Enhanced, Progressive Summary, and Hybrid.
• Performance-Aware: Best-Worst-Recent, Performance Tiered, Trajectory Focused, Contrast Learning, and Adaptive Performance.

3. The Objective Function and Safety Gate

The agent is trained to maximize a combined reward function. To prevent the agent from optimizing conversational quality at the expense of safety, Orme et al. (2026) define the reward using an exponential weighted product combined with a hard safety gate:

Where:

• $q \in [0, 1]$ represents the linguistic quality score (coherence, relevance).
• $s \in [0, 1]$ represents the safety score (absence of toxic content).
• $\alpha = 0.6$ balances the quality/safety trade-off.
• $\beta = 10.0$ shapes reward sharpness, heavily penalizing minor deviations.
• $\theta = 0.8$ acts as a hard safety threshold. Any response that fails critical safety checks ($\min(M_{\text{crit}}) < \theta$) receives a reward of exactly $0.0$.

DQN Agent Architecture

The controller is a Deep Q-Network (DQN) with a Multi-Layer Perceptron (MLP) architecture configured with layer dimensions [512, 512, 256]. It is trained over 5,000,000 timesteps with a learning rate of $1.0 \times 10^{-4}$ and a replay buffer size of 100,000 transitions (as specified in Table 6).

Key Results

The framework was evaluated using a highly adversarial, unified corpus of over 1,000,000 prompts merging datasets from PKU-SafeRLHF, TOXIC-DPO, BeaverTails, and DarkSide DPO across 12 distinct harm categories (Section 4.4).

Performance Across Diverse LLM Architectures

Table 1 highlights how SafeCtrl-RL compares against handcrafted prompt strategies and state-of-the-art prompt optimization algorithms. Performance is measured using $P_{\text{Score}}$ (the final joint safety-quality reward) and $\Delta P$ (improvement over the unprotected "Plain" baseline).

Key Takeaways from the Evaluation

1. Consistent Generalization: Unlike handcrafted strategies, which display high variance across models (e.g., ai_enhanced scores $0.869 on BlackSheep but drops to \0.497$ on DialoGPT), SafeCtrl-RL delivers stable safety gains across all evaluated architectures.
2. Failure of Search-Based Optimization: Offline optimization methods like OPRO and TextGradient perform poorly. OPRO achieved a Macro-$P_{\text{Score}}$ of just $0.342$. Because these methods are optimized offline, they struggle to adapt to the highly dynamic context of real-time dialogue.
3. High Efficiency: Figure 2 demonstrates that SafeCtrl-RL forms a clear upper efficiency envelope. It converges to highly optimal, safe outputs in an average of 3.7 iterations, whereas fixed, handcrafted strategies require $4.5 to \5.2$ iterations to reach comparable safety levels.

Behavioral Retention (Unlearning Persistence)

A key finding in Section 5.4 is the persistence of safe behaviors even after removing the dynamic controller. By providing a minimal "Consistency Instruction" during post-safeguard evaluation (Section D.2), the researchers measured the model's behavioral retention:

• Evil-Alpaca retained 94.50% of the safer behavior patterns ($\Delta_{\text{mean}} = +0.451$).
• BlackSheep retained 90.50% ($\Delta_{\text{mean}} = +0.483$).
• DeepSeek-R1-Distill retained 72.10% ($\Delta_{\text{mean}} = +0.381$).

This shows that the iterative refinement process dynamically steers the model's activation states, reinforcing safe generation patterns that persist even when the safety guardian is removed.

Limitations & Open Questions

While SafeCtrl-RL represents a notable advancement in runtime alignment, practitioners should consider several operational constraints before deploying it:

1. Inference Latency: As shown in Table 11, the average end-to-end execution time per safeguarded query is 89.3 seconds when using cloud-based evaluators like Gemini 2.0 Flash (which alone incurs a ~20.4-second bottleneck per call). While acceptable for offline content moderation, this latency is too high for interactive, real-time user chats.
2. Evaluator Dependency: The system relies heavily on the accuracy of the external evaluator (DeepEval). If an attacker successfully jailbreaks or poisons the evaluation LLM itself, the entire closed-loop control system breaks down.
3. Evaluation of Highly Aligned Foundation Models: The paper's evaluations focused on under-aligned or uncensored models (such as Evil-Alpaca and BlackSheep). How SafeCtrl-RL interacts with highly aligned, closed-source models (such as GPT-4o or Claude 3.5 Sonnet) remains an open research question.

What Practitioners Should Do

If you are deploying LLMs in production environments where safety and brand protection are paramount, you can adapt the core design principles of SafeCtrl-RL with the following steps:

1. Shift from Hard Filters to Closed-Loop Refinement

Rather than instantly blocking a user when a safety violation is detected, implement a multi-pass regeneration loop. Use a fast local evaluator to score intermediate responses, and instruct your system to self-correct when scores fall below your safety threshold.

2. Implement Lightweight Local Evaluators to Cut Latency

To bypass the ~20-second latency of commercial API calls, deploy a highly optimized local classifier. Models like Llama-Guard-3-8B or specialized, task-specific BERT classifiers can evaluate safety in milliseconds.

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

# Load a fast local safety checker
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-Guard-3-8B")
model = AutoModelForSequenceClassification.from_pretrained("meta-llama/Llama-Guard-3-8B")

def check_safety(prompt, response):
    inputs = tokenizer(f"User: {prompt}\nAgent: {response}", return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    logits = outputs.logits
    # Extract binary classification (safe vs unsafe)
    is_safe = torch.argmax(logits, dim=-1).item() == 0
    return is_safe

3. Log State Transitions to Monitor Policy Drift

Maintain a rolling history of conversation metrics (such as the 36-D state vector used in this paper). Track sliding averages of response safety scores, generation lengths, and system error rates. This structured logging makes it easy to flag when system behavior begins to drift or when an attacker is actively attempting to exploit the system.

The Takeaway

SafeCtrl-RL demonstrates that LLM safety does not require expensive weight fine-tuning or rigid input filtering. By framing dialogue safety as an online, reinforcement learning-driven control problem, we can dynamically guide black-box models toward safe, high-quality responses in real time. For security teams and ML engineers, this framework offers a highly flexible blueprint for building adaptive, context-aware guardrails that keep pace with evolving conversational risks.

Den's Take

What excites me about SafeCtrl-RL is the shift away from brittle, static system prompts toward dynamic, inference-time adaptive control. Weight alignment is costly and struggles with drift; real-time orchestration is where the industry is heading.

However, as a practitioner, I am highly skeptical of the operational overhead. An average of 3.7 extra iterations per response is a massive latency penalty. In a high-throughput production environment—like Klarna, where latency spikes can directly threaten a $30M daily checkout volume—this overhead is a complete non-starter. More critically, introducing a stateful DQN controller creates an entirely new attack surface. An adversary could easily craft multi-turn dialogue payloads designed to pollute the feedback signals, tricking the RL controller into choosing weaker prompt strategies over time.

This circular vulnerability is why we must treat feedback loops with extreme caution. In our prior paper, IterInject: Indirect Prompt Injection Against LLM Agents via Feedback-Guided Iterative Optimization, we demonstrated how iterative, feedback-guided optimization loops can be systematically exploited by adversaries to bypass safeguards. SafeCtrl-RL is a clever defense, but without hardening the DQN controller against adversarial input, we are simply moving the goalposts of the exploit.