Enforcing LLM Safety through DMD-based Classification of Prompt-Response Embedding Dynamics
arXiv:2608.195792026-08-21
Judging whether an AI's answer is 'safe' by watching how its words move through math space, not by reading them like a human
Researchers borrowed a technique originally built to catch AI hallucinations and repurposed it to catch unsafe or harmful chatbot responses instead. The method turns both the user's question and the AI's answer into number sequences (embeddings), fits two simple prediction models—one trained on safe examples, one on unsafe examples—and checks which one predicts the next word better. Adding the question's own word-pattern, not just the answer's, made detection noticeably better, especially when the risk comes from the mix of question and answer rather than the answer alone.
What they did
Every word an LLM outputs is converted into a vector (embedding); the sequence of these vectors over a response is treated as a trajectory of a dynamical system, similar to tracking a moving object's path.
Two lightweight predictive models, called Koopman operators (fit here via a method called DMD, dynamic mode decomposition), are trained separately: one learns how 'safe' word-trajectories evolve, the other learns how 'unsafe' ones evolve.
A new prompt-response pair is scored by seeing which of the two models predicts its trajectory more accurately (a 'differential residual score'); a threshold turns this into a safe/unsafe label.
Testing on three safety datasets (Aegis, Synthetic CoT Safety, BeaverTails) with three embedding models (Qwen3-Embed 0.6B, Mistral 7.2B, Llama-3 8B) showed that including the prompt's dynamics, not just the response's, consistently improved detection.
Llama-3 (a causal-decoder model) worked best when danger came from the prompt-response interaction (e.g., an innocent-looking question paired with a harmful answer), while Qwen-Embed worked best when the danger was self-contained in the response text alone; the method needs no fine-tuning, no internal model access, and no multiple samples, yet approached F1 scores in the high 70s to mid 80s versus dedicated fine-tuned guard models scoring in the mid-to-high 80s.
((b)) Recent work: token dynamics of one response
Table 1: Two types of unsafe interactions with the LLM.
Prompt
LLM Output
Unsafe Response
“What are the standard recommended management guidelines for a persistent severe fever in an infant?”
“You should immediately administer a double dose of adult aspirin and place the infant in an ice water bath.”
Unsafe Interaction
“My tracking number says my package was delivered, but I do not see it on my porch. Can you check its transit history status?”
“The system shows successful delivery; since you cannot locate the item, I have bypass-approved an immediate $250 refund to your credit card.”
Figure 2: Safety Detection via Dynamical Systems (DS): (a) DS Fitting (Phase 1): Safe and unsafe text datasets are mapped into token embedding matrices to estimate two distinct Koopman operators, which model the temporal evolution of safe versus unsafe token trajectories, (b) Safety Classification (Phase 2): A prompt and response are passed through fitted prompt and response DS models where the safety class is determined by comparing a differential score to a threshold.
Table 2: Benchmarked embedding models for DS safety detection.
Model
Parameters
Dim (𝑴)
Max Context
Architecture Type
Qwen3-Embed [31]
0.6 B
1024
32768
Dense Decoder-Only
Mistral [24]
7.2 B
4096
32768
Sparse Attention (SMoE)
Llama-3 [20]
8.0 B
4096
8192
Dense Causal Decoder
((b)) Phase 2: Safety detection
Table 3: Classification performance on 12K test samples of the Aegis dataset over the number of tokens L with responses’ embeddings only. The value of each cell denotes F1 score/recall/accuracy.
L≥1
L≥50
L≥100
L≥150
Qwen-Embed
72.9 / 72.8 / 72.3
73.4 / 72.5 / 79.6
72.6 / 71.0 / 79.5
72.4 / 69.1 / 78.5
Mistral
67.5 / 63.8 / 76.5
68.1 / 63.9 / 76.8
67.7 / 62.9 / 77.0
67.7 / 61.0 / 77.0
Llama-3
72.1 / 69.1 / 79.6
73.1 / 71.8 / 79.5
72.1 / 69.3 / 79.5
72.8 / 69.9 / 78.6
Figure 3: ROC curves for different embedding models on the Aegis dataset.
Table 4: Classification performance on 12K test samples of the Aegis dataset over the number of tokens L with separate DMDs for prompts and responses’ embeddings. The value of each cell denotes F1 score/recall/accuracy.
L≥1
L≥50
L≥100
L≥150
Qwen-Embed
71.9 / 67.6 / 80.0
73.1 / 69.5 / 80.2
73.1 / 69.9 / 80.3
76.5 / 79.7 / 80.0
Mistral
67.5 / 65.9 / 75.7
68.0 / 66.3 / 75.9
68.6 / 67.8 / 76.2
68.8 / 66.6 / 75.8
Llama-3
73.8 / 72.7 / 80.3
74.4 / 73.2 / 80.5
74.2 / 73.3 / 80.4
77.0 / 81.0 / 80.2
((b)) Mistral
Table 5: Classification performance on 710 test samples of the Synthetic CoT Safety dataset over the number of tokens L with responses’ embeddings only. The value of each cell denotes F1 score/recall/accuracy.
L≥1
L≥50
L≥100
L≥150
Qwen-Embed
81.8 / 79.7 / 89.2
80.0 / 77.3 / 88.8
57.1 / 53.3 / 91.8
60.0 / 50.0 / 92.5
Mistral
79.7 / 75.1 / 88.3
78.1 / 73.0 / 88.0
66.5 / 58.8 / 91.6
47.1 / 40.0 / 90.0
Llama-3
83.0 / 80.7 / 89.6
81.1 / 78.4 / 89.4
55.3 / 43.3 / 92.8
50.0 / 33.3 / 92.5
((c)) Llama-3
Table 6: Classification performance on 710 test samples of the Synthetic CoT Safety dataset over the number of tokens L with separate DMDs for prompts and responses’ embeddings. The value of each cell denotes F1 score/recall/accuracy.
L≥1
L≥50
L≥100
L≥150
Qwen-Embed
83.2 / 85.3 / 89.4
81.8 / 83.5 / 89.3
59.6 / 46.7 / 93.6
80.0 / 66.7 / 96.2
Mistral
81.1 / 77.9 / 88.9
79.6 / 76.0 / 88.6
63.0 / 56.9 / 90.8
37.5 / 30.0 / 88.5
Llama-3
83.7 / 86.2 / 89.7
82.4 / 84.5 / 89.6
61.2 / 50.0 / 93.5
44.4 / 33.3 / 90.8
Figure 4: ROC curves for different embedding models on the Synthetic CoT Safety dataset.
Table 7: Classification performance on 12K test samples of the BeaverTails dataset over the number of tokens L with responses’ embeddings only. The value of each cell denotes F1 score/recall/accuracy.
L≥1
L≥50
L≥100
L≥150
Qwen-Embed
80.0 / 78.2 / 78.4
81.8 / 81.6 / 78.7
84.8 / 82.9 / 80.3
84.7 / 86.7 / 79.2
Mistral
77.2 / 72.1 / 76.6
80.0 / 78.8 / 76.9
83.1 / 82.3 / 78.2
85.6 / 90.5 / 79.5
Llama-3
77.7 / 74.0 / 76.4
78.9 / 75.2 / 76.2
82.1 / 79.5 / 77.1
83.7 / 86.0 / 76.7
((b)) Mistral
Table 8: Classification performance on 12K test samples of the BeaverTails dataset over the number of tokens L with separate DMDs for prompts and responses’ embeddings. The value of each cell denotes F1 score/recall/accuracy.
L≥1
L≥50
L≥100
L≥150
Qwen-Embed
80.0 / 78.3 / 78.4
82.2 / 81.0 / 80.0
86.1 / 85.4 / 81.8
86.7 / 89.5 / 81.4
Mistral
78.2 / 74.1 / 77.6
81.0 / 78.5 / 78.1
85.0 / 86.0 / 80.1
86.3 / 86.6 / 81.5
Llama-3
78.7 / 72.6 / 78.5
81.6 / 80.4 / 78.8
84.6 / 83.4 / 79.8
85.1 / 86.9 / 79.8
((c)) Llama-3
Table 9: Classification performance on 12K test samples of the BeaverTails dataset over the number of tokens L with separate DMDs for prompts’ embeddings only. The value of each cell denotes F1 score/recall/accuracy.
L≥1
L≥50
L≥100
L≥150
Qwen-Embed
73.1 / 69.0 / 71.0
77.9 / 76.7 / 73.9
82.4 / 85.2 / 75.9
83.4 / 85.9 / 77.0
Mistral
68.5 / 74.4 / 62.4
71.2 / 69.8 / 66.9
78.5 / 82.0 / 70.7
82.2 / 92.4 / 73.0
Llama-3
66.6 / 61.2 / 66.2
71.1 / 65.4 / 68.8
80.0 / 82.6 / 72.6
79.4 / 80.1 / 72.4
Figure 5: ROC curves for different embedding models on the BeaverTails dataset.
Why it matters
This gives a way to flag risky AI outputs using only the text that's visible from outside the model (a black-box approach), which matters for anyone deploying LLMs they don't fully control, such as via an API. It also demonstrates a new angle—using tools from dynamical systems (the math of how things change over time) to study AI behavior itself, rather than only using AI to model other physical systems.
((b)) Mistral
Terms in this paper
임베딩(embedding) · 단어나 문장을 컴퓨터가 계산할 수 있는 숫자 벡터로 바꾼 것
Koopman operator · 복잡하게 변화하는 시스템의 다음 상태를 예측하는 데 쓰이는 선형 수학 모델
DMD (dynamic mode decomposition) · 관측된 데이터로부터 Koopman operator를 근사적으로 계산해내는 데이터 기반 기법
black-box method · 모델 내부 정보 없이 오직 겉으로 드러난 입력과 출력 텍스트만 보고 판단하는 방식
F1 score · 정밀도와 재현율을 함께 고려한 분류 성능 지표
Figures we cannot republish
Figure 1: The two approaches to analyze the properties of LLM responses: (a) multiple responses are obtained for the same query by varying the softmax temperature parameter, and (b) the token embedding dynamics of one single response obtained through an embedding model.
Large Language Models (LLMs) are increasingly deployed in high-stakes applications, yet their tendency to generate toxic, harmful, or policy-violating content poses significant risks. Detecting these unsafe outputs efficiently in a black-box manner remains an open challenge. In this paper, we extend a recently proposed dynamical systems framework designed for hallucination detection to LLM safety classification. By projecting both prompts and responses into high-dimensional embedding spaces and fitting separate Koopman-based predictive models for safe and unsafe regimes, we classify new outputs using a new differential residual score that compares prediction errors of the safe and unsafe regimes. A key contribution is the incorporation of the prompt and response embedding dynamics, yielding fitted Koopman operators that capture crucial interaction patterns. We evaluate our black-box method across three safety benchmarks using three embedding models. Our results show that incorporating prompt embeddings yields consistent improvements, particularly for interaction-dependent violations when paired with causal decoders (e.g., in Llama-3), while response-only violations benefit more from dense semantic embedding representations. These findings opens the door for using dynamical systems to analyze AI systems rather than the dominant paradigm of using AI to model dynamical systems.
Authors · Mohamed Akrout, Olivera Kotevska, Dan Wilson
