D2-Monitor: Hesitation-Aware Routing을 이용한 Diffusion LLM의 다단계 안전 모니터링

Diffusion LLM은 다단계 denoising 경로에서 중간 은닉 상태에 안전 관련 정보를 담고 있어 단일-스텝 모니터링으로 포착하기 어려운 신호가 존재한다. 본 연구는 hesitation 지표를 활용한 bi-level 모니터(D2-Monitor)를 제안하고, 테스트 시점에 계산 자원을 효율적으로 배치하기 위해 hesitation severity가 임계치를 넘을 때만 고복잡도 프로브로 escalation하는 구조를 제시한다. 실험에서 3개 데이터셋, 4개 D-LLM에 대해 파라미터 풋프린터가 ≤0.85M인 경량 모니터가 최적의 효과-효율 트레이드오프를 달성한다.

1. 다단계 D-LLM은 각 denoising 스텝에서의 은닉 표현이 안전 판단에 중요한 정보를 담고 있다. 2) probe margin를 기반으로 한 hesitation 측정(d_s의 부호 변화 여부, nτ = Σ I[|d_s|<τ])은 샘플의 난이도를 정량화하는 효과적인 지표이다. 3) 이 지표를 이용해 쉬운 샘플은 base probe로, 어려운 샘플은 advanced probe로 전달하는 cascade 방식을 사용하면 계산 자원을 절약하면서도 안전성 성능을 유지할 수 있다. 4) Margin 기반 routing이 엔트로피/confidence 같은 probe-extrinsic 신호보다 안정적이고 강력한 난이도 예측을 제공하며, hesitation window를 이용한 advanced probe 학습이 성능을 보강한다. 5) 실험은 intra-/cross-데이터셋에서 일관된 성능 향상과 높은 효율을 보여준다.

Stage 1: OOF 스코어링으로 hesitation trajectories 수집. k-fold로 나누어진 학습 데이터에서 선형 프로브를 학습하고 각 샘플의 denoising step별 margin ds를 산출한다. |ds|<τ인 hesitation step을 식별하고, nτ를 계산한다. nτ>0인 경로를 advanced probe 학습 데이터로 선택한다. 패턴: 입력으로 ds를 받아 threshold τ에 따라 hesitation 여부를 판단하고, nτ를 합산해 결정한다. Stage 2: Base 및 Advanced Probes 학습. 기본 Linear Probe를 전체 학습 데이터에 대해 학습해 test 시에는 per-step Margin를 산출한다. Hesitation trajectory가 있는 샘플에 대해 hesitation window W를 정의하고, 그 윈도우에 있는_hidden_states_만 입력으로 advanced probe를 학습한다: 손실 L(y(i), f({h_s}_{s∈W(i)})). Stage 3: Cascade Detection. 테스트 시점에 각 denoising step에서 margin ds를 계산하고 hesitation step을 찾은 뒤, nτ가 λ를 초과하면 advanced probe로 escalates, 그렇지 않으면 base probe로 판단한다. Routing signal로 margin-based hesitance를 사용하며, τ와 λ는 검증 세트에서 최적화한다. Probe 아키텍처: base는 Linear Probe, advanced는 MLP(또는 TimeAttn)로 구성. 모델별 파라미터 수는 0.85M 이하로 제한되며, 객체는 사전 학습된 D-LLM의 denoising trajectory에서 템포럴 윈도우에 한정된 입력만 사용한다. 실험 구성은 WildGuardMix, ToxicChat, OpenAI-Moderation에서 intra- 및 cross-dataset 설정으로 진행된다.

메인 벤치마크: intra-dataset에서 D2-Monitor는 모든 baselines보다 우수한 F1를 달성하고, 파라미터 풋프린트가 작다. 예를 들어 D2-TimeAttn(Ours)은 LLaDA-8B-Instruct에서 F1이 89.4로 상위권에 위치하며, 0.54M의 파라미터로 운영된다. FLOPs 및 추론 시간은 Full-trajectory baselines 대비 현저히 낮고, D2-LP 단일 스텝 기반 방법에 근접한 효율성을 보인다(Table 5, Table 6). cross-dataset에서도 D2-MLP/TimeAttn은 ToxicChat, OpenAI-Moderation에서 baselines보다 우수한 F1을 기록하며 일반화 성능을 입증한다(Table 3). 라우팅 신호로 margin를 사용할 때 엔트로피/confidence 대비 일관된 성능 향상이 관찰된다(Fig. 4). 라우팅 임계 λ를 조정하면 고비용 고정합 모델의 비율을 조절할 수 있으며, 더 낮은 E[P]에서 더 나은 F1를 달성하는 경우가 많다(Fig. 1). 추론 시간 측면에서 D2-Monitor는 일부 샘플에만 고복잡도 프로브를 적용하고, hesitation 윈도우에 한정해 연산을 수행하므로 평균 FLOPs가 크게 감소한다(Table 5, Table 6). Adversarial 시나리오에서 hesitation이 강한 샘플일수록 routed 샘플의 adversarial 비율이 높아진다는 점이 실험으로 확인된다(Fig. 5).

전체 아키텍처는 base probe(LP) + router + advanced probe로 구성된다. 다단계 트래젝토리에서의 hesitation count nτ를 활용해 test-time routing 의 판단지를 만들고, nτ>λ인 경우에만 advanced probe로 escalation한다. hesitation Steps를 정의하는 신호 ds는 linear probe의 decision boundary에 대한 margin으로 측정되며, τ는 hesitation step을 정의하는 임계값이다. Stage 1에서 OOF scoring으로 hesitation trajectory를 수집하고, Stage 2에서 base와 advanced probes를 각각 학습한다. Stage 3에서 cascade로 inference를 수행하며, advanced probe 입력은 hesitation window 내의 숨겨진 상태 h_s들로 제한된다. TimeAttn/MLP 등의 고복잡도 프로브는 Swin 길이의 window에서 작동하므로 전체 시퀀스의 비용을 피한다. 실험 세부사항은 Appendix C에 기술되며, hyperparameters τ, λ은 검증 세트에서 최적화한다.

실험은 16B 파라미터 이하의 D-LLM에 한정되며, 더 큰 모델에서의 일반화 가능성은 확인되지 않았다. activation monitors의 안전성 취약성(적대적 공격에 의한 헛된 hesitation 조작 가능성) 가능성이 있으며, 향후 추가 보안 대책과 함께 평가할 필요가 있다.