PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Aoxi Liu, Yupeng Chen, James Oldfield, Guanzhe Hong, Junchi Yu, Baoyuan Wu, Philip Torr, Adel Bibi

arXiv:2605.25893v1cs.AI

TL;DR

Safety monitoring for diffusion LLMs is underexplored despite their multi-step denoising trajectories. D2-Monitor uses hesitation-aware dynamic routing to achieve state-of-the-art monitoring with fewer than 0.85M parameters and the best effectiveness–efficiency trade-off across evaluated settings.

  • Problem

    Safety monitoring for D-LLMs remains underexplored, including whether their full denoising trajectories provide more safety-relevant information than single-step representations.

  • Method

    D2-Monitor uses an always-on lightweight probe to estimate hesitation and classify safety, routing high-hesitation trajectories to an advanced probe trained on hesitation windows.

  • Results

    Fewer than 0.85M parameters accompany state-of-the-art performance and the best effectiveness–efficiency trade-off across three datasets, four D-LLMs, and intra- and cross-dataset settings.

  • Takeaways & Limitations

    Hesitation-aware routing supports dynamic computational allocation and is particularly suited to resource-constrained D-LLM safety monitoring deployments.

  • Takeaways & Limitations

    The monitor instantiates its always-on probe as a linear probe, tying the approach to that lightweight probe choice.

Abstract

from arXiv · show

Despite the emergence of diffusion large language models (D-LLMs) as an alternative to autoregressive large language models (AR-LLMs), safety monitoring for D-LLMs remains largely unexplored. Unlike AR-LLMs, D-LLMs generate text through a multi-step denoising process, exposing intermediate hidden representations that may contain safety-relevant information unavailable in standard single-step monitoring setups. Motivated by the suitability of lightweight probes for always-on monitoring, we analyze which trajectory-level signals best indicate when such probes are likely to struggle. We find that the most informative signal is safety hesitation: intermediate hidden states repeatedly falling within a small margin of the probe's decision boundary. The number of such hesitation steps in D-LLM's trajectory predicts probe failure effectively, providing a proxy of sample difficulty. Building on this analysis, we propose $D^2$-Monitor, a bi-level safety monitor for D-LLMs. $D^2$-Monitor adopts a lightweight probe as an always-on monitor to jointly estimate hesitation and perform base classification. When the hesitation level exceeds a threshold, a more expressive but computationally heavier probe is activated. This dynamic routing mechanism allocates monitoring resources efficiently at test time. Evaluated on 3 datasets (WildguardMix, ToxicChat, OpenAI-Moderation) across 4 D-LLMs, $D^2$-Monitor achieves state-of-the-art performance with a compact parameter footprint ($\leq$ 0.85M parameters), and exhibits the best trade-off between effectiveness and efficiency relative to 8 baselines.

1 Introduction

Safety monitoring for diffusion LLMs is underexplored, despite alignment’s vulnerability to adversarial attacks and the richer signals exposed by multi-step denoising trajectories. D2-Monitor uses safety hesitation for dynamic routing between lightweight and advanced probes, achieving strong effectiveness-efficiency results with a compact parameter footprint.

  • Problem: Safety monitoring for D-LLMs remains underexplored, while alignment alone is insufficient because it remains vulnerable to adversarial attacks.The paper motivates external monitoring as critical for reducing risks from harmful model use.
  • Mechanistic discovery: Multi-step D-LLM trajectories provide richer safety-monitoring signals than single-step representations through hesitation steps in the safety probe space.Hesitation severity is characterized using probe margins and strongly correlates with linear probe performance.
  • Method: D2-Monitor combines a router, an always-on low-complexity base probe, and a high-complexity advanced probe activated when hesitation exceeds a threshold.The base probe jointly estimates hesitation and performs base-level safety classification; hesitation signals also curate advanced-probe training data.
  • Evaluation: 3 safety datasets and 4 D-LLMs are used to evaluate D2-Monitor under intra-dataset and cross-dataset settings.The evaluation covers WildGuardMix, ToxicChat, and OpenAI-Moderation.
  • Results: Fewer than 0.85M parameters, or 0.01% of an 8B model, accompany state-of-the-art performance and the best effectiveness-efficiency trade-off against 8 baselines.Additional analysis reports robustness across generation configurations and remasking strategies.

2 Related Work

The paper situates D-LLMs as diffusion-based, iterative denoising alternatives to left-to-right AR-LLMs and motivates external safety guardrails by LLM vulnerabilities to adversarial attacks. It distinguishes guardrails that deploy additional LLMs as safety monitors.

  • Autoregressive and diffusion language models: AR-LLMs generate text strictly left to right through next-token prediction, whereas D-LLMs use iterative denoising to extend diffusion modeling to discrete text.D-LLMs are built upon masked diffusion models and adapt diffusion-based generative modeling from continuous domains to text.
  • Safety guardrails: Despite extensive safety training, LLMs remain vulnerable to adversarial attacks, motivating external safety guardrails for industry-deployed models.The passage links this need particularly to legal and regulatory requirements.
  • Safety guardrails: One guardrail category deploys an additional LLM trained as a safety classifier to filter model inputs and outputs.Llama-Guard is identified as a representative model fine-tuned on safety tasks.

3 Exploring Safety Monitoring in D-LLMs

D-LLMs expose safety-relevant information across multi-step denoising trajectories, making full-trajectory probing more effective than relying only on the final step. Probe-intrinsic hesitation severity, measured by the number of near-boundary steps, provides a finer-grained indicator of probe difficulty than binary hesitation.

  • Diffusion Large Language Models: D-LLMs generate sequences through forward noising and reverse multi-step denoising, with intermediate hidden states available throughout the trajectory.The reverse process updates masked positions and may re-mask predicted positions using strategies such as random or low-confidence re-masking.
  • Problem Setup: Mean pooling over tokens produces a step-wise trajectory H of aggregated hidden states, reflecting safety information distributed across bidirectional representations.The hidden representation has dimensions D × L × S before aggregation, where D, L, and S denote hidden dimension, sequence length, and denoising steps.
  • Problem Setup: A linear probe is used because its lightweight, interpretable design supports always-on monitoring.The probe is trained by minimizing empirical cross-entropy loss.
  • Trajectory Utilization: Both multi-step readouts achieve higher Acc and F1 than the single-step baseline on most models, indicating that intermediate denoising steps contain safety-relevant information.The compared variants are LP (Last Step), LP (Mean), and LP (MV), using controlled training conditions and either temporal-mean or majority-vote trajectory readouts.
  • Hesitation Characterization: Hesitation can be measured extrinsically from entropy or confidence, or intrinsically by flagging steps whose probe margin satisfies |d_s| < τ.A trajectory is binary-flagged when it contains at least one hesitation step.
  • Hesitation Severity: n_τ counts hesitation steps and stratifies difficulty more finely than the binary τ criterion, with probe F1 gaps reaching up to ∼0.30 between extreme buckets.Probe F1 generally decreases monotonically as n_τ increases, while trajectories with small n_τ can perform close to the stable n_τ = 0 subset.

4 Method

D2-Monitor is a hesitation-aware safety monitoring framework that dynamically allocates test-time computation according to estimated sample difficulty. It uses a lightweight linear probe for always-on scoring and routes sufficiently hesitant examples to an advanced probe using hesitation-window representations.

  • Framework: D2-Monitor dynamically allocates test-time compute based on estimated sample difficulty from hesitation steps in D-LLM representation trajectories.The framework is motivated by the finding that hesitation-step counts estimate linear-probe classification difficulty.
  • Probe Design: The linear probe serves as the low-complexity base probe because it achieves approximately 0.90 F1 on low-hesitation samples and identifies estimation difficulty effectively.The framework also considers two high-complexity variants: an MLP probe and a temporal attention probe (TimeAttn).
  • Stage 1: Out-of-Fold Scoring and Hesitation Trajectories Collection: Out-of-fold scoring produces leakage-free signed margins for training examples, which are used to identify hesitation steps for advanced-probe training.Each fold is scored using a linear probe trained on the remaining k−1 folds.
  • Stage 2: Base and Advanced Probes Training: The advanced probe is trained exclusively on trajectories containing hesitation, using hidden states within each trajectory’s minimal contiguous hesitation window.The base linear probe is trained on the full training set and computes test-time step-wise margins and hesitation severity n_τ.
  • Stage 3: Cascade Detection: At test time, the base probe counts hesitant denoising steps as hesitation severity n_τ, and the router compares n_τ with λ to select the base or advanced probe.Examples with n_τ ≤ λ use the low-complexity probe; examples with n_τ > λ are passed to the advanced probe using the extracted hesitation window.

5 Experiment

Experiments evaluate D^2-Monitor across three safety datasets and four D-LLMs against eight baselines, measuring accuracy, F1, and expected test-time parameter usage. Across datasets, decoding configurations, and remasking strategies, D^2-Monitor consistently delivers strong effectiveness-efficiency performance and validates margin-based routing.

  • Experimental Setup: Experiments use WildGuardMix, ToxicChat, and OpenAI-Moderation with four open-source D-LLMs: LLaDA-8B-Base, LLaDA-8B-Instruct, LLaDA-1.5-8B, and LLaDA-2.0-mini-16B.WildGuardMix has 86.8k training and 1.7k test prompts; ToxicChat has 5.08k training and 5.08k test prompts.
  • Evaluation Metrics: The evaluation reports accuracy, F1, and E[P], where E[P] captures expected parameters used per example under cascade routing.For D^2-Monitor, E[P] = |θLP| + ρ · |θg|.
  • Main Results: D^2-Monitor consistently outperforms all baselines in accuracy and F1 while achieving the best overall effectiveness-efficiency trade-off.The baselines include single-step, mean-based, voting, TimeAttn, and LSTM approaches.
  • Robustness to Generation Length and Step Length: D^2-Monitor consistently outperforms all baselines when step length varies from 1 to 8 or generation length varies from 16 to 128 without retraining.Methods are trained with generation length 128 and step length 4, then tested across both decoding axes.
  • Robustness to Remasking Strategy: D^2-Monitor maintains consistent superiority under low-confidence, entropy, and random remasking strategies.All methods are trained under low-confidence remasking and tested under the three strategies.
  • Ablation on Routing Signal: Margin-based routing consistently outperforms entropy- and confidence-based routing and is more robust to threshold selection.The ablation evaluates D2-MLP and D2-TimeAttn on ToxicChat at thresholds covering 30%-70% hesitation samples.

6 Conclusion

D2-Monitor is a dynamic bi-level safety monitor for D-LLMs that uses safety hesitation along denoising trajectories to estimate sample difficulty and route difficult cases to an advanced probe. A lightweight linear probe remains always on, while a heavier probe is activated when hesitation exceeds a predefined threshold.

  • 6 Conclusion: D2-Monitor defines hesitation steps as denoising steps whose hidden states yield low probe margins, and uses their count as a proxy for sample difficulty.Hesitation trajectories identified from training data are used to train the advanced probe.
  • 6 Conclusion: A lightweight linear probe jointly evaluates hesitation and performs base safety classification as D2-Monitor’s always-on monitor.This design provides continuous monitoring before any second-stage activation.
  • 6 Conclusion: When hesitation steps exceed a predefined threshold, D2-Monitor activates a more expressive but computationally heavier probe for second-stage classification.The hesitation-aware routing mechanism dynamically allocates computational resources.

A Limitation · B Broader Impacts, Safeguards, and Licenses

D2-Monitor shows strong performance with a compact parameter footprint, but experiments are limited to D-LLMs with up to 16B parameters. The paper emphasizes low-cost safety monitoring, defensive artifacts, and compliance with licenses for datasets and models.

  • A Limitation: <0.85M parameters: D2-Monitor achieves superior performance in intra-dataset and cross-dataset settings.Experiments cover a variety of D-LLM models.
  • A Limitation: 16B parameters: computational constraints limit experiments to D-LLMs up to this size, despite an expected extension to larger models.The passage also notes that activation monitors are vulnerable to adversaries.
  • B Broader Impacts, Safeguards, and Licenses: D2-Monitor provides a lightweight, always-on mechanism for detecting harmful or adversarial inputs at low computational cost.This is particularly valuable for resource-constrained or edge deployment where heavyweight LLM-as-monitor solutions are infeasible.
  • B Broader Impacts, Safeguards, and Licenses: The work aims to reduce misuse risk by improving the practicality of safety monitoring for D-LLMs.The stated societal impact concerns harmful content generation.
  • B Broader Impacts, Safeguards, and Licenses: ≤0.85M parameters: released artifacts are lightweight safety probes trained on publicly available benchmarks and intended as defensive mechanisms.The paper does not release high-risk pretrained models, generative models, or scraped datasets.
  • B Broader Impacts, Safeguards, and Licenses: All existing assets are properly cited and used in accordance with their licenses.This licensing statement applies to the assets used in the work.
  • B Broader Impacts, Safeguards, and Licenses: WildGuardMix uses ODC-BY, ToxicChat uses CC-BY-NC-4.0 for non-commercial research, and OpenAI-Moderation uses MIT.These are the licenses specified for the three datasets.
  • B Broader Impacts, Safeguards, and Licenses: LLaDA-8B-Base, LLaDA-8B-Instruct, and LLaDA-1.5 use MIT, while LLaDA-2.0-mini uses Apache-2.0.These are the licenses specified for the listed base D-LLMs.

C Experiment Details · C.1 Hyperparameter Tuning

Hyperparameters are selected using validation data while controlling probe capacity for fair architectural comparisons. The protocol specifies separate tuning procedures for baselines, the proposed method, test-time routing, and cross-dataset transfer.

  • C.1 Hyperparameter Tuning: All hyperparameters are tuned on a validation subset formed by splitting the original training set 4:1, with key probe dimensions fixed for fair comparison.Architectural dimensions such as hidden size or hidden dimension are controlled rather than extensively tuned.
  • C.1 Hyperparameter Tuning: Baseline hyperparameters are selected by validation-set grid search, after which the best configuration is retrained on the full training set before test evaluation.No test data is used during model selection.
  • C.1 Hyperparameter Tuning: For OOF scoring, the method fixes the linear probe at learning rate 1e−3 and weight decay 1e−4, while tuning MLP and TimeAttn on validation hesitation examples.The fixed linear-probe configuration is also used for the method’s linear probe component.
  • C.1 Hyperparameter Tuning: The method’s test-time hyperparameters are the hesitation threshold τ and routing parameter λ, with τ chosen for approximately 50% hesitation-example coverage.Larger τ increases coverage but may add stable steps, whereas smaller τ produces cleaner signals with fewer examples.
  • C.1 Hyperparameter Tuning: λ is fixed to the value achieving the best validation F1 score, while τ and λ are retuned per target dataset under cross-dataset distribution shift.For ToxicChat, 20% of training data is reserved for validation and τ and λ are selected by grid search; OpenAI-Moderation lacks a training split.

C.2 Probe Architectures · D Additional Results

The paper specifies normalization and pooling choices for trajectory-based probes, then defines linear, MLP, attention, and recurrent architectures for binary classification. These probes map denoising trajectories to scalar logits using configurations detailed below.

  • C.2 Probe Architectures: All probes are trained with Adam for 50 epochs using batch size 256 on an NVIDIA A40 GPU with 48GB VRAM.The supplied passage states these as shared training conditions across the probe baselines.
  • C.2 Probe Architectures: Normalization uses per-feature statistics for full-trajectory probes and per-step statistics for single-step probes.Full-trajectory statistics have shape [D], whereas single-step statistics have shape [S, D].
  • C.2 Probe Architectures: Fixed-length probes reduce the step dimension after normalization using mean pooling, last-step pooling, or majority vote.Last-step uses h_1, while majority vote classifies each step independently at test time.
  • C.2 Probe Architectures: Linear Probe computes the scalar logit as s = w^T h̄, yielding LP (Mean), LP (Last Step), and LP (MV).The weight vector satisfies w ∈ R^D.
  • C.2 Probe Architectures: The two-layer MLP uses the same three pooling variants as the linear probe, with hidden dimension K = 256.The variants are MLP (Mean), MLP (Last Step), and MLP (MV).
  • C.2 Probe Architectures: TimeAttn processes the full trajectory without pooling by applying layer normalization and additive attention over denoising steps.Its attended representation is classified with a two-layer MLP and layer normalization.
  • C.2 Probe Architectures: LSTM projects each step with layer normalization and GELU, then processes the sequence using a 2-layer unidirectional LSTM.The projection dimension is d_p = 512, the hidden size is d_h = 128, and the final hidden state feeds a layer-normalized linear head.

D.1 More Evaluation Metrics

Additional metrics confirm that D2-MLP and D2-TimeAttn preserve strong precision, recall balance, and false-refusal performance. Efficiency analysis shows hesitation-aware routing substantially reduces inference time and FLOPs while retaining leading F1 results.

  • Additional Metrics: D2-MLP and D2-TimeAttn achieve the best or second-best performance across nearly all six reported metrics and models.The additional metrics show that gains in accuracy and F1 do not come at the expense of precision, recall balance, or false refusal rate.
  • Inference Time: 2.4×–6.6× speedup: D2-MLP reduces inference cost versus MLP (Mean) across different LLaDA models.Hesitation-aware routing processes only hesitant samples and their localized hesitation windows.
  • FLOPs: 2SD FLOPs per sample: the cascade applies its base linear probe at every denoising step, then conditionally invokes an expert on hesitant samples within a minimal hesitation window.The expert is activated for fraction pesc of samples and uses an average hesitation-window length Swin ≤ S.
  • FLOPs: 0.7–1.0 MFLOPs per sample: D2-MLP is 2–3× cheaper than MLP (Mean) and 35–150× cheaper than sequence-based baselines while delivering the highest F1 scores in Table 1.D2-TimeAttn remains 4–5× cheaper than running TimeAttn on the full trajectory despite its heavier expert.

D.2 Robustness to Random Seeds · E Additional Analysis

The authors test robustness to random initialization by retraining all methods on LLaDA-8B-Instruct across five seeds. Results remain consistent, with low standard deviations, and the proposed method retains the best performance among baselines.

  • D.2 Robustness to Random Seeds: Main results use a fixed random seed of 2026.
  • D.2 Robustness to Random Seeds: The robustness check retrains all methods on LLaDA-8B-Instruct.
  • D.2 Robustness to Random Seeds: Five random seeds—0, 1, 2, 3, and 4—are used for retraining.
  • D.2 Robustness to Random Seeds: Results are consistent across the five random seeds.
  • D.2 Robustness to Random Seeds: All methods exhibit low standard deviations across seeds.
  • D.2 Robustness to Random Seeds: The proposed method continues to achieve the best performance among all baselines.

E.1 Analysis of Hesitation Dynamics · E.2 Margin Outperforms Entropy and Confidence as Step-Count Signal · E.3 Hesitation Severity Captures Adversarial Inputs

The analysis shows that hesitation is temporally persistent, predicts probe difficulty, and is associated with adversarial inputs. Among the evaluated step-count signals, margin-, entropy-, and confidence-based counts all qualitatively track difficulty, while adversarial prevalence rises with hesitation severity.

  • E.1 Analysis of Hesitation Dynamics: The analysis measures local prediction instability through the probability that the signed margin crosses the decision boundary at the next step as a function of |d_s|.Cross-boundary probability is evaluated across all four LLaDA models.
  • E.1 Analysis of Hesitation Dynamics: Low-margin steps tend to persist across multiple subsequent denoising steps, with persistence probabilities decaying gradually as the step horizon increases.This indicates that hesitation is temporally structured rather than purely local.
  • E.2 Margin Outperforms Entropy and Confidence as Step-Count Signal: The step-count construction underlying n_τ generalizes beyond margin to probe-extrinsic entropy- and confidence-based hesitation signals.The signals are instantiated using per-step entropy E_s and confidence C_s with respective thresholds τ_E and τ_C.
  • E.2 Margin Outperforms Entropy and Confidence as Step-Count Signal: All three step-count signals—probe margin, step-wise entropy, and step-wise confidence—qualitatively predict difficulty because probe F1 decreases monotonically as hesitation counts increase.The comparison is conducted on LLaDA-8B-Instruct across n_signal buckets and five threshold settings parameterized by hesitant ratio.
  • E.3 Hesitation Severity Captures Adversarial Inputs: Adversarial samples constitute an increasing fraction of higher n_τ buckets across three LLaDA variants, linking greater hesitation severity to adversarially designed prompts.WildGuardMix’s adversarial prompts are described as the harder benchmark portion because they are constructed to evade safety classifiers.
  • E.3 Hesitation Severity Captures Adversarial Inputs: 47% is the dataset-wide adversarial baseline used to contextualize the routed subset’s adversarial fraction in WildGuardMix.Routing thresholds are selected on a held-out validation set per method.
Loading 2605.25893v1…