AI 推論與安全
Speculative Probing Reuses MTP Draft Heads for Real-Time LLM Safety Monitoring with a Few Soft Prompts
Cornell Tech researchers repurpose MTP/Eagle3 modules originally designed for speculative decoding as sequence classifiers, without rerunning the main model. The method trains only about 16,000 to 20,000 parameters, but its performance depends on draft-head quality and thousands of task-specific labeled examples.
Real-time detection of prompt injection, harmful content, or reasoning loops typically involves a trade-off between inexpensive but information-limited hidden-state probes and Llama Guard–style models that require a separate model run. A Cornell Tech team proposes Speculative Probing, which instead performs classification using Qwen3.5’s multi-token prediction (MTP) head or MiniCPM4.1’s Eagle3 draft head. It appends one to five trainable soft prompts to the end of the input, recursively applies the single-layer draft head, and produces a binary classification through a linear sigmoid head. All weights in the main model and draft head remain frozen.
The key is that speculative decoding has already built the draft head’s KV cache for the full prefix. The monitor adds only a few queries. The authors estimate that the marginal inference cost of each additional classification task can be below 1%. Because the temporary positions are not written back to the production decoding cache, the approach does not affect the draft acceptance rate. The Qwen3.5-27B version with five soft prompts trains only about 18,000 parameters, making it lighter than MultiMax, which processes intermediate layers token by token.
Four evaluations cover instruction conflicts, repetitive reasoning, branching reasoning, and multilingual safety classification. The Qwen3.5-27B safety probe scores 78.2%, close to zero-shot GPT-5.4-mini at 78.1%, and outperforms Qwen3Guard and Llama Guard under the paper’s evaluation setup. In cross-dataset testing on WildGuardMix, SP-2 with Qwen3.5-9B retains a score of 79.0%, while both MultiMax and a standard MLP fall to 69.0%. However, this is not a training-free, general-purpose guardrail: each task still uses several thousand to 10,000 supervised examples, and models without a high-quality MTP or draft head cannot achieve the same cost advantage. Eagle3 also trails jointly pretrained MTP by roughly 4 to 12 percentage points. The next steps should focus on independent replication, real-world serving latency, and out-of-distribution attacks—not just the balanced classification datasets constructed by the authors.