# D^2-Monitor: Dynamic Safety Monitoring for Diffusion LLMs via Hesitation-Aware Routing

> D²-Monitor leverages hesitation in a diffusion LLM's hidden states to dynamically route samples between lightweight and heavyweight safety probes, achieving state-of-the-art safety detection with minimal parameters.

- **Source:** [arXiv](https://arxiv.org/abs/2605.25893)
- **Published:** 2026-05-28
- **Permalink:** https://picx.dev/p/Q15j91
- **Whiteboard:** https://picx.dev/p/Q15j91/image

## Summary

# D²-Monitor: Dynamic Safety Monitoring for Diffusion LLMs via Hesitation-Aware Routing - Summary

## Summary (Overview)
*   **Problem & Discovery:** Safety monitoring for Diffusion Large Language Models (D-LLMs) is underexplored. The paper identifies **safety hesitation**—intermediate hidden states repeatedly falling near a probe's decision boundary—as a key signal that predicts when lightweight safety probes will fail, serving as a proxy for sample difficulty.
*   **Proposed Solution:** **D²-Monitor**, a dynamic bi-level safety monitor. It uses a lightweight linear probe as an always-on monitor to estimate hesitation and perform base classification. When hesitation severity exceeds a threshold, a router activates a more expressive (but heavier) advanced probe for second-stage classification.
*   **Key Results:** Evaluated on 3 safety datasets (WildGuardMix, ToxicChat, OpenAI-Moderation) across 4 D-LLMs, D²-Monitor achieves **state-of-the-art performance** with an extremely compact parameter footprint (≤ 0.85M parameters).
*   **Efficiency- Effectiveness Trade-off:** D²-Monitor exhibits the **best trade-off** between effectiveness (F1 score) and efficiency (expected test-time parameters) compared to 8 baselines, making it suitable for resource  -constrained deployment.
*   **Generalization:** The method demonstrates strong performance in both intra-dataset and cross-dataset evaluation settings, and is robust to variations in generation length, step length, and re-masking strategies.

## Introduction and Theoretical Foundation
Despite the emergence of **Diffusion Large Language Models (D-LLMs)** as a promising alternative to autoregressive LLMs (AR-LLMs) due to their parallel decoding and iterative refinement, **safety monitoring for D-LLMs remains largely unexplored**. Unlike AR-LLMs, D-LLMs generate text through a multi-step denoising process, exposing a trajectory of intermediate hidden states. This trajectory may contain safety-relevant information unavailable in standard single-step monitoring setups used for AR-LLMs.

Existing safety monitoring literature focuses on AR-LLMs and falls into two categories: **LLM-as-monitors** (using additional LLMs as classifiers) and **probe-based monitors** (lightweight classifiers on internal representations). Probe-based monitors are particularly suited for always-on, low-cost deployment.

The paper is motivated by recent findings that intermediate D-LLM outputs can oscillate between correct and incorrect answers during reasoning tasks. The authors hypothesize that analogous **instability or "hesitation"** occurs in the safety probe space for D-LLMs. They aim to characterize this hesitation and leverage it to build a more efficient and effective safety monitor.

The **theoretical foundation** is built on the analysis of probe **margins** (distance to the decision boundary) across the denoising trajectory. The core insight is that trajectories with many steps where the margin is small (hesitation steps) are harder for probes to classify correctly. This establishes **hesitation severity** as an effective, probe-intrinsic signal for estimating sample difficulty.

## Methodology

### 3.1 Preliminary & Problem Setup
Diffusion LLMs define a discrete diffusion process. Given a prompt, the reverse denoising process starts from a partially masked state $\tilde{x}^{(S)}$ and produces $S$ denoising steps. For a dataset of $I$ prompts with safety labels $y^{(i)} \in \{0, 1\}$, the D-LLM produces a hidden representation $H^{(i)}_{\text{raw}} \in \mathbb{R}^{D \times L \times S}$ at a particular layer. After mean-pooling over the sequence dimension, we get a step-wise representation matrix $H^{(i)} = [h^{(i)}_1, ..., h^{(i)}_S] \in \mathbb{R}^{D \times S}$. The dataset is $\mathcal{D} = \{H^{(i)}, y^{(i)}\}_{i=1}^I$. A safety probe $f$ is learned by minimizing the empirical cross-entropy loss:
$$
\min_f \frac{1}{I} \sum_{i=1}^I \mathcal{L}\left(y^{(i)}, f(H^{(i)})\right).
$$

### 3.2 & 3.3: Analysis of Useful Signals
*   **Multi-step vs. Single-step:** The paper first establishes that the full denoising trajectory $H$ carries more safety-relevant signal than a single-step representation (e.g., the final step $h_1$). Multi-step readout strategies (Mean pooling and Majority Vote) outperform single-step probing.
*   **Hesitation Characterization:** A **hesitation step** is defined as a step whose hidden state yields a signed probe margin $d_s$ with $|d_s| < \tau$ for a threshold $\tau$. The **hesitation severity** for a trajectory is defined as:
    $$n_\tau = \sum_{s=1}^S \mathbb{I}[|d_s| < \tau],$$
    which counts the number of hesitation steps.
*   **Key Finding:** Hesitation severity $n_\tau$ is strongly correlated with linear probe performance (F1 score). Trajectories with higher $n_\tau$ are significantly harder for the probe to classify correctly. This signal is more predictive of difficulty than probe-extrinsic signals like token distribution entropy or confidence.

### 4.1 & 4.2: D²-Monitor Design and Implementation
D²-Monitor is a hesitation-aware, bi-level safety monitoring framework with three components:
1.  **Low-complexity Base Probe:** A linear probe that serves as an always-on monitor.
2.  **Router:** Uses hesitation severity $n_\tau$ to decide whether to escalate a sample.
3.  **High-complexity Advanced Probe:** A more expressive probe (MLP or Temporal Attention) activated for hard samples.

**Training and Inference Pipeline:**
1.  **Stage 1 - Out-of-Fold (OOF) Scoring:** The training set is split into $k$ folds. A linear probe trained on $k-1$ folds scores the held-out fold to compute leakage-free margins $\{d^{(i)}_s\}_{s=1}^S$ and hesitation severity $n^{(i)}_\tau$ for each training example.
2.  **Stage 2 - Probe Training:**
    *   **Base Probe:** Trained on the full training set.
    *   **Advanced Probe:** Trained **only** on trajectories with $n^{(i)}_\tau > 0$ (hesitation trajectories). For each such trajectory, a **hesitation window** $\mathcal{W}^{(i)}$ (minimal contiguous span containing all hesitation steps) is constructed. The advanced probe $f$ is trained using:
    $$\min_f \sum_{i: n^{(i)}_\tau > 0} \mathcal{L}\left(y^{(i)}, f(\{h^{(i)}_s\}_{s \in \mathcal{W}^{(i)}})\right).$$
3.  **Stage 3 - Cascade Detection (Inference):**
    *   For a test example, the base probe computes margins $d_s$ for all steps and the hesitation severity $n_\tau$.
    *   The **router** compares $n_\tau$ to a threshold $\lambda$.
    *   If $n_\tau \leq \lambda$, the sample is classified by the base probe via majority vote: $\hat{y} = \text{sign}\left(\sum_{s=1}^S \text{sign}(d_s)\right)$.
    *   If $n_\tau > \lambda$, the hesitation window $\mathcal{W}$ is extracted and passed to the advanced probe for final prediction.

## Empirical Validation / Results

### 5.1 Experiment Setup
*   **Datasets:** WildGuardMix (adversarial prompts), ToxicChat (real user-AI chats), OpenAI-Moderation.
*   **Models:** Four open-source D-LLMs: LLaDA-8B-Base, LLaDA-8B-Instruct, LLaDA-1.5-8B, LLaDA-2.0-mini-16B.
*   **Baselines (8):** Single-step methods (LP/MLP on last step) and full-trajectory methods (LP/MLP with Mean/MV readout, TimeAttn, LSTM).
*   **Metrics:** Accuracy (Acc), F1 score, and **Expected Parameters at test time** $\mathbb{E}[P] = |\theta_{LP}| + \rho \cdot |\theta_g|$, where $\rho$ is the fraction of examples routed to the advanced probe.

### 5.2 Main Results
**Tables 1 & 2** show intra-dataset performance on WildGuardMix and ToxicChat. **Table 3** shows cross-dataset generalization (train on WildGuardMix, test on ToxicChat and OpenAI-Moderation).

**Table 1: Intra-dataset performance on WildGuardMix test set. Best in bold, second-best underlined.**
| Method | LLaDA-8B-Base | LLaDA-8B-Instruct | LLaDA-1.5 | LLaDA-2.0-mini |
| :--- | :--- | :--- | :--- | :--- |
| | **E[P]** | **Acc** | **F1** | **Acc** | **F1** | **Acc** | **F1** | **E[P]** | **Acc** | **F1** |
| *Single-step* | | | | | | | | | | |
| LP (Last Step) | 4e-3M | 84.6 | 84.1 | 87.4 | 87.0 | 88.2 | 87.9 | 2e-3M | 77.6 | 76.3 |
| MLP (Last Step) | 1.05M | 85.8 | 85.4 | 87.1 | 86.8 | 87.9 | 87.5 | 0.52M | 79.0 | 78.0 |
| *Full-trajectory* | | | | | | | | | | |
| LP (MV) | 4e-3M | 86.7 | 86.2 | 88.2 | 87.9 | 87.8 | 87.5 | 2e-3M | 80.0 | 79.0 |
| LP (Mean) | 4e-3M | 86.9 | 86.5 | 88.2 | 88.0 | 87.9 | 87.7 | 2e-3M | 80.5 | 79.4 |
| MLP (MV) | 1.05M | 86.9 | 86.6 | 87.9 | 87.6 | 88.3 | 88.0 | 0.52M | 81.3 | 80.5 |
| MLP (Mean) | 1.05M | 87.4 | 87.0 | 87.7 | 87.4 | 88.3 | 88.0 | 0.52M | 81.4 | 80.6 |
| TimeAttn | 1.59M | 87.4 | 86.9 | 87.9 | 87.5 | 88.0 | 87.7 | 0.80M | 81.7 | 80.8 |
| LSTM | 2.57M | 87.1 | 86.6 | 87.8 | 87.4 | 88.1 | 87.8 | 1.51M | 81.7 | 80.9 |
| **D²-MLP (Ours)** | **≤0.36M** | **88.1** | **87.8** | **89.9** | **89.7** | **89.3** | **89.1** | **0.17M** | **83.7** | **82.9** |
| **D²-TimeAttn (Ours)** | **≤0.54M** | **88.6** | **88.3** | **89.6** | **89.4** | **89.3** | **89.0** | **0.26M** | **83.7** | **83.0** |

**Key Findings:**
*   D²-Monitor variants (**D²-MLP** and **D²-TimeAttn**) achieve the highest Accuracy and F1 scores across all models and datasets.
*   They do this with a significantly lower expected parameter count $\mathbb{E}[P]$ than most non-linear baselines (MLP, TimeAttn, LSTM).
*   The method shows strong **cross-dataset generalization** (Table 3), outperforming baselines when trained on one dataset and tested on another.

### 5.3 Analysis
*   **Efficiency-Effectiveness Trade-off:** By adjusting the routing threshold $\lambda$, D²-Monitor provides a Pareto frontier that dominates other methods, offering the best F1 for a given computational budget (see Figure 1 in the paper).
*   **Robustness:** D²-Monitor maintains superior performance when tested under different generation lengths, step lengths, and re-masking strategies without retraining (Figure 3).
*   **Ablation on Routing Signal:** Margin-based hesitation routing consistently outperforms routing based on entropy or confidence scores from the D-LLM's token distribution (Figure 4). The margin signal more precisely isolates samples that genuinely benefit from the advanced probe.
*   **Hesitation Captures Adversarial Inputs:** Analysis shows that trajectories with high hesitation severity $n_\tau$ are disproportionately drawn from the **adversarial split** of WildGuardMix. This means D²-Monitor's router effectively channels adversarial/hard samples to the advanced probe.

## Theoretical and Practical Implications
*   **Theoretical:** The paper provides a novel characterization of **safety hesitation** in D-LLMs, linking the dynamics of probe margins across the denoising trajectory to classification difficulty. It demonstrates that probe-intrinsic signals are more informative for monitoring than model-centric uncertainty metrics.
*   **Practical:** D²-Monitor offers a **practical, lightweight safety solution** for deploying D-LLMs. Its dynamic routing mechanism allows for efficient resource allocation: easy samples are processed cheaply, while computational resources are focused on hard (often adversarial) samples. With ≤ 0.85M parameters, it is feasible for edge deployment and provides a state-of-the-art balance of safety and efficiency.

## Conclusion
The paper introduces **D²-Monitor**, a dynamic safety monitoring framework for Diffusion LLMs that leverages intrinsic hesitation signals from the multi-step denoising trajectory. The core contributions are:
1.  The discovery that **hesitation severity** ($n_\tau$) in the probe margin space is a strong predictor of sample difficulty for safety probes.
2.  A **bi-level monitor design** that uses this signal for both curating advanced probe training data and for test-time routing.
3.  **State-of-the-art results** on multiple safety benchmarks across several D-LLMs, achieved with a compact parameter footprint and the best effectiveness-efficiency trade-off.

The insights from D²-Monitor provide a promising direction for developing reliable and efficient safety monitors tailored to the unique generative dynamics of D-LLMs. Future work should explore the robustness of such hesitation-aware monitors against adaptive adversaries and their scaling to even larger D-LLMs.

---

_Markdown view of https://picx.dev/p/Q15j91, served by PicX — AI-generated visual whiteboard summaries of research papers._