🔢 mathematics

Enforcing LLM Safety through DMD-based Classification of Prompt-Response Embedding Dynamics

This paper proposes a black-box safety classification framework for Large Language Models that leverages Koopman-based dynamical systems to analyze prompt-response embedding dynamics, demonstrating that incorporating prompt information significantly improves the detection of interaction-dependent harmful outputs.

Original authors: Mohamed Akrout, Olivera Kotevska, Dan Wilson

Published 2026-08-21

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Mohamed Akrout, Olivera Kotevska, Dan Wilson

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Large language models are the engines behind a new generation of artificial intelligence, capable of writing stories, solving problems, and holding conversations that feel remarkably human. Yet, beneath their fluent surface lies a persistent vulnerability: these systems can sometimes generate harmful, deceptive, or dangerous content, even when they have been trained to be helpful. This risk is not just a theoretical glitch; it is a practical barrier to using these tools in high-stakes situations like healthcare, legal advice, or customer service. For years, the standard way to catch these errors has been to look inside the model's code or to ask the model to generate many different answers and compare them. But in the real world, many powerful models are "black boxes," where the internal workings are hidden, and asking for multiple answers is too slow or expensive. This leaves a critical gap: how do we spot dangerous outputs quickly and reliably without needing to see the model's brain or wait for it to think twice?

A team of researchers at the University of Tennessee and Oak Ridge National Laboratory has proposed a different way to look at the problem. Instead of treating a language model as a static database of facts or a simple text generator, they treated it as a dynamic system, similar to how a physicist might study the flow of water or the movement of planets. In this view, every word a model generates is not just a standalone piece of text, but a step in a continuous, evolving trajectory. The researchers hypothesized that safe conversations and unsafe ones follow different patterns as they unfold, much like how a healthy heartbeat follows a different rhythm than an irregular one. By mapping the journey of words into a high-dimensional space and analyzing the shape of that journey, they developed a method to distinguish between safe and unsafe interactions without ever needing to open the model up.

The core of their approach involves watching how a model moves from one word to the next. When a user asks a question, the model does not just spit out an answer; it builds that answer one token, or word-piece, at a time. The researchers took these sequences of words and converted them into mathematical points in a vast, multi-dimensional space. They then used a technique called dynamic mode decomposition to fit a predictive model to these points. Think of this as drawing a line through a series of dots to see where the next dot is likely to fall. They built two separate sets of these predictive models: one trained on examples of safe conversations and another trained on examples of unsafe ones. When a new conversation happens, the system runs the sequence of words through both sets of models. If the sequence follows the path predicted by the safe model, it is likely safe. If it deviates sharply from the safe path and aligns more closely with the unsafe path, the system flags it as dangerous.

What makes this study particularly significant is that it does not look at the answer in isolation. The researchers realized that the safety of a response often depends entirely on the question that prompted it. A sentence like "I will give you a refund" might be perfectly safe in a customer service context but could be a dangerous hallucination if the user is asking about a package that was never shipped. To capture this nuance, the team designed their system to track the dynamics of both the user's prompt and the model's response together. They found that by watching the interaction between the two, rather than just the final output, the system could catch subtle forms of danger that other methods miss. This is especially true for models that generate text in a specific, causal order, where the relationship between the question and the answer is tightly woven.

The team tested their method on three different collections of data, ranging from general safety benchmarks to specific datasets designed to catch harmful advice and toxic speech. They used three different types of embedding models—tools that translate words into those mathematical points—to see which worked best. The results showed that their method was highly effective, particularly when it considered the prompt alongside the response. In one major test involving over 12,000 examples, the system achieved an AUC of over 0.8, correctly identifying unsafe content without needing any special training on that specific data. This is a strong performance for a "black-box" method that requires no access to the model's internal code.

The study also revealed a fascinating detail about how different types of AI models behave. The researchers found that the best tool for the job depended on the nature of the danger. For interactions where the risk comes from the specific relationship between a question and an answer, a model with a causal architecture—one that reads and writes in a strict forward sequence—performed best. However, for dangers that were clearly visible in the content of the answer itself, regardless of the question, a different type of model that focused on dense semantic meaning worked better. This suggests that there is no single "best" way to monitor AI safety; the right tool depends on the specific kind of risk being managed.

Perhaps the most surprising finding was that the system could sometimes detect danger just by looking at the user's prompt, even before the model generated a single word of the response. In tests using a dataset of human-written prompts, the system could distinguish between safe and unsafe questions based solely on the trajectory of the words in the question itself. This implies that the way humans phrase dangerous requests leaves a distinct mathematical signature that can be detected early. While the system was not perfect and did struggle slightly with very long sequences of text, the results demonstrate that viewing language generation as a dynamic process offers a powerful new lens for safety.

This work represents a shift in how we might approach the governance of artificial intelligence. Instead of relying solely on training models to be better or building complex filters to block bad words, this approach treats the generation process itself as a signal to be monitored. By fitting simple, linear models to the complex, non-linear dance of word generation, the researchers showed that it is possible to build a safety net that is both efficient and effective. The method does not require the heavy computational cost of retraining massive models or the privacy risks of exposing internal data. It simply watches the flow of conversation, looking for the moment the path diverges from safety. As these tools become more integrated into daily life, such methods offer a promising way to ensure that the conversation remains helpful, honest, and harmless.

Technical Summary: Enforcing LLM Safety through DMD-based Classification of Prompt-Response Embedding Dynamics

Problem Statement

Large Language Models (LLMs) are increasingly deployed in high-stakes applications, yet they remain prone to generating toxic, harmful, or policy-violating content, including hallucinations, hate speech, and instructions for illegal activities. While training-time alignment methods like Reinforcement Learning from Human Feedback (RLHF) exist, they are resource-intensive and do not eliminate the risk of unsafe outputs. Consequently, there is an urgent need for efficient, scalable, post-deployment monitoring methods.

Current detection approaches face significant limitations:

White-box methods require access to internal states (hidden states, attention maps), which is often impossible with commercial APIs.
Gray-box methods rely on token-level probabilities or logit entropy, which are not always exposed.
Black-box methods typically rely on sampling multiple responses to check for consistency or using external knowledge graphs, which can be computationally expensive or require external retrieval.
Interaction Blindness: Many existing black-box approaches analyze only the response text, failing to detect safety violations that arise specifically from the interaction between a safe-looking prompt and a harmful response (e.g., a model incorrectly granting a refund based on a user's query).

Methodology

The paper proposes a black-box safety classification framework that treats the LLM's token-by-token generation process as a discrete-time dynamical system. The method extends a previously proposed hallucination detection framework to safety classification by incorporating both prompt and response dynamics.

1. Dynamical Systems Framework

The core assumption is that the evolution of token embeddings in an LLM can be modeled as a dynamical system. The authors utilize Koopman operator theory to linearize these dynamics.

Embedding: Both the prompt and the response are projected into high-dimensional embedding spaces using pre-trained embedding models (e.g., Qwen, Mistral, Llama-3).
Koopman Operator Estimation: Using Extended Dynamic Mode Decomposition (EDMD), the authors fit separate linear predictive models (Koopman operators) for "safe" and "unsafe" regimes.
- Four operators are estimated: $A^{(p)}_s, A^{(p)}_u$ for safe/unsafe prompts, and $A^{(r)}_s, A^{(r)}_u$ for safe/unsafe responses.
- These operators are learned from training data consisting of token embedding sequences.

2. Differential Residual Score

To classify a new interaction, the method computes a differential residual score ( $\Delta E$ ) that compares the prediction errors of the safe and unsafe models against the observed token evolution.

For a prompt with $P$ tokens and a response with $L$ tokens, the prediction error $\epsilon$ is calculated for each token step $k$ using the respective Koopman operators.
The score $\Delta E$ is defined as the difference between the cumulative prediction errors of the unsafe models and the safe models:
$\Delta E = \left( \sum_{j=1}^{L-1} (\epsilon^{(r)}_{u,j})^2 + \sum_{j=1}^{P-1} (\epsilon^{(p)}_{u,j})^2 \right)^{1/2} - \left( \sum_{j=1}^{L-1} (\epsilon^{(r)}_{s,j})^2 + \sum_{j=1}^{P-1} (\epsilon^{(p)}_{s,j})^2 \right)^{1/2}$
Classification Logic: If the output is unsafe, the unsafe models ( $A_u$ ) should predict the trajectory better, resulting in a lower error for the unsafe term and a negative $\Delta E$ . Conversely, safe outputs yield a positive $\Delta E$ .
A binary decision is made by comparing $\Delta E$ to a threshold $\eta$ .

Key Contributions

Prompt-Aware Safety Classification: The primary contribution is the extension of DMD-based classification to include prompt dynamics. The authors demonstrate that safety is often interaction-dependent; a response that appears safe in isolation may be unsafe given a specific prompt. By fitting separate dynamical systems for prompts and responses, the method captures these causal patterns.
Black-Box Efficiency: The method operates entirely on output text embeddings, requiring no access to internal model weights, token probabilities, or multiple stochastic samples.
Empirical Analysis of Embedding Architectures: The study systematically evaluates how different embedding model architectures (dense decoder-only vs. causal decoder) interact with different types of safety violations.

Experimental Results

The method was evaluated on three safety benchmarks using three embedding models of varying scales: Qwen3-Embed (0.6B), Mistral (7.2B), and Llama-3 (8.0B).

Aegis AI Content Safety Dataset (Interaction-Dependent):
- This dataset features violations arising from the interplay between user intent and model compliance.
- Result: Incorporating prompt embeddings significantly improved performance. Llama-3 emerged as the dominant model (F1 = 77.0% at $L \ge 150$ ), outperforming Qwen-Embed and Mistral.
- Insight: Llama-3's causal decoder architecture effectively captures the temporal dynamics of the prompt-response interaction, reducing tracking error for safe samples and increasing it for unsafe ones when prompt context is included.
Synthetic CoT Safety Dataset:
- This dataset involves structured reasoning and refusal patterns.
- Result: All models achieved high accuracy (>88%), though F1 scores were lower due to class imbalance. Adding prompt embeddings provided a modest but consistent boost across models.
BeaverTails Dataset (Response-Only Violations):
- This dataset contains human-written prompts and responses where violations are often intrinsic to the response content (e.g., hate speech, violence) regardless of the prompt.
- Result: Qwen-Embed (a dense semantic encoder) outperformed the others, achieving an F1 of 86.7% at $L \ge 150$ when prompt embeddings were included.
- Insight: For violations directly encoded in the response tokens, dense semantic representations excel. Interestingly, prompt-only classification on this dataset also achieved high F1 scores (83.4%), suggesting human-written unsafe prompts carry distinct dynamical signatures, though response dynamics remain the primary signal.

Comparison to Baselines: On the Aegis dataset, the proposed black-box method achieved an F1 of 77.0%, compared to 86.8% for a fine-tuned white-box model (AegisGuard). The authors note this is a significant achievement for a method requiring no task-specific training.

Significance and Claims

The paper claims to advance the paradigm of "Dynamical Systems for AI" rather than the more common "AI for Dynamical Systems." By treating LLM generation as a dynamical system, the authors demonstrate that Koopman-based prediction can effectively distinguish safe from unsafe outputs.

The study highlights a nuanced relationship between embedding model architecture and the nature of safety violations:

Causal Decoders (e.g., Llama-3) are superior for interaction-dependent violations where the prompt-response relationship is critical.
Dense Semantic Encoders (e.g., Qwen-Embed) are superior for content-isolated violations where the unsafe nature is intrinsic to the response tokens.

The authors conclude that this approach opens the door for using dynamical systems theory to analyze and control AI systems, offering a scalable, black-box alternative to resource-intensive training-time alignment or complex multi-sample inference strategies. Future work suggested includes extending the method to multi-class harm categorization and exploring ensemble strategies of different embedding models.