Enforcing LLM Safety through DMD-based Classification of Prompt-Response Embedding Dynamics
arXiv:2608.195792026-08-21
AI 답변이 안전한지, 사람처럼 읽지 않고 단어가 숫자 공간에서 움직이는 궤적으로 판단하는 방법
연구진은 원래 AI의 '환각'(사실이 아닌 그럴듯한 답변) 탐지용으로 만들어진 기법을 가져와 유해하거나 위험한 챗봇 답변을 잡아내는 데 응용했다. 사용자의 질문과 AI의 답변을 각각 숫자 벡터(임베딩) 흐름으로 바꾼 뒤, 안전한 예시로 학습한 예측 모델과 위험한 예시로 학습한 예측 모델 두 개를 만들어 어느 쪽이 다음 단어를 더 잘 맞히는지 비교한다. 답변만이 아니라 질문의 흐름까지 함께 고려하자 탐지 성능이 확실히 좋아졌으며, 특히 위험이 질문과 답변의 조합에서 생기는 경우 효과가 컸다.
무엇을 했나
LLM이 내놓는 각 단어를 숫자 벡터(임베딩)로 바꾸고, 답변 전체를 이 벡터들이 시간에 따라 움직이는 궤적으로 취급한다. 마치 움직이는 물체의 경로를 추적하듯 다루는 방식이다.
DMD(dynamic mode decomposition)라는 기법으로 Koopman operator라는 가벼운 예측 모델 두 개를 따로 학습시킨다. 하나는 '안전한' 단어 궤적이 어떻게 흘러가는지, 다른 하나는 '위험한' 궤적이 어떻게 흘러가는지 학습한다.
새로운 질문-답변 쌍이 들어오면 두 모델 중 어느 쪽이 그 궤적을 더 정확히 예측하는지 비교하는 '차등 잔차 점수(differential residual score)'를 계산하고, 기준값(threshold)을 넘으면 위험, 아니면 안전으로 분류한다.
Aegis, Synthetic CoT Safety, BeaverTails 세 개의 안전성 데이터셋과 Qwen3-Embed(0.6B), Mistral(7.2B), Llama-3(8B) 세 개의 임베딩 모델로 실험한 결과, 답변뿐 아니라 질문의 흐름까지 포함하면 탐지 성능이 꾸준히 향상됐다.
위험이 질문과 답변의 상호작용에서 나올 때는 Llama-3(원인-결과 방향으로만 정보를 보는 causal decoder 구조)가 가장 좋았고, 위험이 답변 자체에만 담겨 있을 때는 Qwen-Embed가 더 나았다. 이 방법은 별도의 파인튜닝이나 모델 내부 접근, 여러 번 샘플링이 필요 없으면서도 F1 점수 70%대 후반~80%대 중반을 기록해, 전용으로 파인튜닝된 가드 모델(80%대 중후반)에 근접했다.
((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.
왜 중요한가
이 방법은 모델 내부를 들여다볼 수 없는 상황, 즉 API로만 LLM을 쓰는 경우에도 겉으로 드러난 텍스트만으로 위험한 답변을 걸러낼 수 있는 길을 제시한다. 또한 물리 시스템을 다루던 동역학계(dynamical systems) 수학 도구를 AI 자체를 분석하는 데 거꾸로 적용한 새로운 접근이라는 점에서, AI 안전성 연구의 방향에 의미 있는 시사점을 준다.
((b)) Mistral
이 논문의 용어
임베딩(embedding) · 단어나 문장을 컴퓨터가 계산할 수 있는 숫자 벡터로 바꾼 표현
Koopman operator · 복잡하게 변화하는 시스템의 다음 상태를 예측하는 선형 수학 모델
DMD(dynamic mode decomposition) · 관측 데이터로부터 Koopman operator를 근사적으로 계산해내는 데이터 기반 기법
black-box 방식 · 모델 내부 정보 없이 겉으로 드러난 입력·출력 텍스트만으로 판단하는 방식
F1 score · 정밀도와 재현율을 함께 고려한 분류 성능 지표
본문에 싣지 못한 그림
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.
