PaperlayerView exact revision on arXiv

Source-linked AI summary

Speculative Probing: LLM Monitoring at Speculative-Decoding Cost

Collin Zhang, Tingwei Zhang, Vitaly Shmatikov

arXiv:2608.28099v1cs.AIcs.CL

TL;DR

Real-time LLM monitoring requires classification under strict latency constraints, while existing probes trade context-awareness and accuracy against efficiency. Speculative Probing repurposes speculative-decoding heads using soft prompts and existing KV caches, achieving favorable accuracy-efficiency trade-offs with small probes.

  • Problem

    Real-time classification supports safety filtering and behavioral monitoring, but last-token probes can lack sequence-level expressiveness for patterns distributed across positions.

  • Method

    Speculative Probing freezes a speculative-decoding head, appends learned soft prompts, and trains a small task-specific classifier using the KV cache already computed during inference.

  • Results

    Speculative Probing consistently achieves a favorable accuracy-efficiency trade-off, with Qwen MTP probes outperforming MiniCPM Eagle3 probes by 4–12 pp across most tasks.

  • Takeaways & Limitations

    Speculative-decoding heads can serve as efficient sequence classifiers while requiring only a negligible number of task-specific parameters and no additional inference-pipeline architecture.

  • Takeaways & Limitations

    The approach depends on a well-trained speculative-decoding module and requires supervised training data at the scale of thousands of examples per task.

Abstract

from arXiv · show

Real-time classification during language model inference is valuable for safety filtering, behavioral analysis, and model monitoring, but current approaches force a trade-off between accuracy and efficiency. Hidden-state probes are fast but limited: they are either not context-aware: operating on a single vector and cannot model interactions across positions; or they are very costly: having dedicated classifier models (Llama Guard, Qwen Guard, LLM-as-judge) or performing computation on hidden states for all tokens and then pooling the results (MultiMax). This shows an intrinsic trade-off between efficiency and accuracy. However, we find that the speculative-decoding module in recent LLMs can be repurposed for efficient high-quality classification. By appending a trained soft prompt at the end of the target sequence, we can repurpose the speculative-decoding module into a sequence classifier. At inference time in a speculative-decoding pipeline, the KV cache is already in GPU memory, so classification adds negligible overhead. We evaluate on four classification tasks across four models (Qwen3.5-4B, 9B, 27B, MiniCPM4.1-8B). Our small probes consistently outperform zero-shot GPT-5.4-mini and, on multilingual prompt safety, match or beat specialized 8B safety classifiers (Qwen3Guard-Gen-8B, Llama-Guard-3-8B) without running a full LLM.

1 Introduction

Speculative Probing repurposes speculative-decoding heads for efficient sequence classification, addressing the accuracy–efficiency trade-off in real-time monitoring. Across four tasks and model families, small probes outperform zero-shot GPT-5.4-mini and can match dedicated safety classifiers.

  • Real-time classification supports safety filtering, behavioral monitoring, quality control, and user studies, but must meet strict latency constraints.
  • Current methods trade efficiency for context: hidden-state probes are cheap but limited, whereas separate classifiers achieve higher accuracy at substantially greater inference cost.
  • The method reuses speculative-decoding heads and their already-computed KV caches, adding a few low-cost forward passes without architectural changes to the inference pipeline.
  • Co-pretrained MTP probes outperform post-hoc Eagle3 probes by 4–12 pp across most tasks, although the comparison cannot isolate training recipe from base-model quality.
  • Speculative Probing freezes the base model and speculative-decoding head, training only recursive soft prompts and a linear head with approximately 16K–20K task-specific parameters.
  • Across four binary tasks and four models, probes on co-pretrained MTP and post-hoc Eagle3 heads beat zero-shot GPT-5.4-mini and match or beat dedicated 8B safety classifiers on multilingual prompt safety.

2 Related Work

The paper builds on probing, soft prompt tuning, speculative decoding, and inference-time monitoring. Existing work motivates lightweight probes and dedicated safety systems, while Speculative Probing combines these lines through speculative-decoding heads.

  • Probing classifiers train small models over frozen internal representations to test encoded information, with prior work finding useful transferable signal in intermediate layers.
  • Soft prompt tuning appends or prepends learned continuous vectors to frozen models, providing parameter-efficient adaptation alongside adapters and low-rank updates.
  • Speculative decoding accelerates generation by drafting candidate tokens and verifying them against the target, with lightweight drafter heads and multi-token prediction moving drafting into or upstream of the target model.
  • Production monitoring commonly uses dedicated classifiers for safety, while other systems inspect chain-of-thought traces, scan for PII, or train lightweight probes for hallucination and uncertainty estimation.

3 Method

Speculative Probing uses a frozen speculative-decoding head with recursively appended soft prompts to produce sequence-level classification features. It reuses cached computation, keeps decoding unchanged, and adds only a small marginal cost.

  • MTP and Eagle3 are speculative-decoding draft heads that use target-model hidden states to predict several future tokens.
  • MTP combines final-layer hidden states with shifted token embeddings, processes them through a decoder layer, and recursively feeds its previous output for later predictions.
  • Eagle3 instead fuses low-, mid-, and high-layer hidden states and is distilled post-hoc on a frozen base model.
  • Speculative Probing replaces predicted-token embeddings with learned soft prompts, chains k prompts recursively through the head, and applies a linear classification head.
  • Training freezes the base model and speculative-decoding head while optimizing only the soft prompts and linear classification head with binary cross-entropy.
  • SP-1, SP-2, and SP-5 denote configurations using one, two, and five recursive soft prompts, respectively.
  • With the prefix KV cache already resident, the added computation is O(k n d_h) attention + O(k d^2) feed-forward and requires no base-model recomputation.
  • Speculative Probing reads committed prefixes off the decoding path, so its scratch-buffer prompts do not alter the drafter cache or acceptance rate.

4 Last Token Hidden State Probing isn’t enough

A single pooled last-token representation struggles to detect information buried in long documents, whereas position-aware probes can localize such spans. In the PII experiment, Speculative Probing matched MultiMax with far fewer parameters and lower inference cost.

  • The motivating PII use case filters long user queries so only inputs without personally identifiable information enter model-improvement workflows.
  • The experiment injected PII into half of LongSafety documents, truncated examples to 2,048 tokens, and produced a roughly balanced dataset of approximately 1,100 documents.
  • SP-1 matched MultiMax’s accuracy while using 50× fewer trainable parameters and reusing computation already performed during speculative decoding.

5 Experimental Setup

The experiments compare lightweight probes and baselines across four models and four binary classification tasks. MultiMax uses token-level mid-layer representations with max pooling, while baselines include an MLP probe and dedicated or zero-shot classifiers.

  • Models and tasks: Four models span Qwen3.5-4B/9B/27B and MiniCPM4.1-8B, covering co-pretrained MTP and post-hoc-distilled Eagle3 heads.The Qwen3.5 models form a scaling series, while MiniCPM4.1 uses an Eagle3 layer.
  • Models and tasks: Four balanced binary tasks test instruction hierarchy compliance, repetitive chain-of-thought behavior, reasoning strategy, and multilingual prompt safety.Validation selects the best epoch, and the corresponding test accuracy is reported.
  • Baselines: The MLP probe trains a two-layer classifier on the last-token, last-layer hidden state with 5–8M trainable parameters.Its hidden dimension is 2,048, with ReLU activation and dropout 0.1.
  • Baselines: MultiMax projects each token independently, applies H learned heads, and hard-max-pools scores across the sequence.Each head uses a learned projection vector, and the total trainable parameter count is approximately 412K.
  • Baselines: Baselines include zero-shot GPT-5.4-mini on every task and dedicated Qwen3Guard-Gen-8B and Llama-Guard-3-8B classifiers for safety classification.GPT-5.4-mini receives one task-specific prompt template and returns a yes/no answer.
  • Evaluation protocol: Table 2 reports VerIH instruction-contradiction accuracy, bolds Speculative Probing columns, and underlines the best result in each row.Across methods, the best epoch is selected by validation accuracy before reporting test accuracy.

6 Results

Across four monitoring tasks, Speculative Probing generally delivers a favorable accuracy–efficiency trade-off, often matching or surpassing more expensive methods. Its relative advantage depends on model, task, and context length.

  • 6.1 VerIH: Speculative Probing and MultiMax substantially outperform the last-token MLP on VerIH, while both surpass GPT-5.4-mini’s 80.3 accuracy on Qwen models.On Qwen3.5-4B, Speculative Probing exceeds MultiMax at lower cost; on 9B and 27B, MultiMax leads by around 2 points.
  • 6.2 CoT Repetition: Speculative Probing and MultiMax only slightly outperform MLP on CoT Repetition, indicating that last-token representations often encode repetitive-loop information.GPT-5.4-mini performs poorly on this task under the zero-shot setup.
  • 6.3 CoT Reasoning Strategy: 85.0 versus 81.5: Speculative Probing beats MultiMax on Qwen3.5-27B for CoT Reasoning Strategy, while remaining within 2–3 percentage points on Qwen3.5-4B and 9B.MULTIMAX-mid leads on 4B and 9B, whereas MiniCPM’s Eagle3 underperforms even MLP; the authors note that draft-head quality may matter.
  • 6.4 Nemotron Safety: Speculative Probing outperforms MLP on all Qwen models for Nemotron Safety and performs better than MultiMax on Qwen3.5-27B.For short prompts, last-token states can encode sufficient information, but Speculative Probing retains higher expressivity.
  • 6.5 Parameter Efficiency and Overall Comparison: Speculative Probing sits on the cost–accuracy Pareto frontier across all four tasks and strictly dominates MultiMax and dedicated baselines on three of them.Across 16 model–task cells, it matches or beats MultiMax on half and trails by only 2–4 percentage points on most remaining cells at a fraction of the cost.
  • 6.5 Parameter Efficiency and Overall Comparison: SP-1 is best on 7 of 16 model–task cells and within 2.5 percentage points on 14 of 16, making it a sensible default.SP-5 helps mainly on Qwen3.5-4B and long-context tasks where aggregating across more positions matters.

7 Conclusion

Speculative Probing repurposes auxiliary speculative-decoding heads as efficient sequence classifiers by combining frozen heads, soft prompts, and the KV cache already computed during decoding. The observed advantage of co-pretrained MTP probes over post-hoc Eagle3 probes is promising but not isolated from base-model or head-training differences.

  • Speculative Probing repurposes auxiliary speculative-decoding heads as highly efficient sequence classifiers.The approach keeps the auxiliary head frozen and appends a few soft prompts.
  • The method trains only a negligible number of task-specific parameters while achieving high classification accuracy.
  • Co-pretrained MTP probes outperform post-hoc distilled Eagle3 probes, mirroring their speculative-decoding acceptance rates.
  • The comparison pairs each head type with a single base model, so it cannot isolate the effect of the training recipe.

Limitations

Speculative Probing depends on native, well-trained speculative-decoding modules and substantial supervised data, limiting its applicability to architectures or tasks without those resources.

  • The method requires a well-trained speculative-decoding module, such as a co-pretrained MTP head.
  • Architectures without native auxiliary heads cannot currently use the near-zero marginal cost monitoring paradigm.
  • Probe quality is tightly coupled to the draft model’s pretraining budget.
  • Competitive accuracy requires supervised training data at the scale of thousands of examples per task.Novel, long-tail, or rapidly emerging scenarios may therefore require synthesized training data.

Potential Risks

Making behavioral tracking and safety filtering computationally trivial creates a risk of invasive surveillance or over-censorship without strict privacy safeguards.

  • Without strict privacy guardrails, computationally trivial monitoring could facilitate invasive user surveillance or over-censorship.

A Task Samples and Judge Prompts

The appendix provides representative task examples, judge prompts, dataset-labeling details, and demonstrations for instruction compliance, chain-of-thought behavior, reasoning structure, safety, and PII filtering.

  • A Task Samples and Judge Prompts: Two representative examples per task are provided, with one example for each class.
  • A Task Samples and Judge Prompts: Long-CoT examples abbreviate traces with [...], while full traces typically contain 1,000–6,000 tokens.
  • A.1 VerIH examples: VerIH labels come from per-instance LLM-as-judge grading, illustrated with correct and contradictory instruction examples.
  • A.3 CoT Repetition examples: The repetition examples contrast a looping trace that repeatedly rechecks a settled answer with a short productive trace that reaches an answer once.
  • A.3 CoT Repetition examples: The CoT Repetition task identifies traces that revisit the same reasoning without adding information, distinguishing loops from productive reasoning.
  • A.4 CoT Reasoning Strategy examples: CoT Reasoning Strategy classifies reasoning shape as linear or branched, independently of whether the solution is correct.
  • A.5 Nemotron Safety Examples: Nemotron Safety labels are taken directly from the dataset’s prompt_label field, whose samples are translated from Arabic originals.
  • A.5 Nemotron Safety Examples: The appendix also reproduces the GPT-5.4 judge prompt used for PII filtering, including examples of identifying contactable private information.

B Training Details

Speculative Probing uses a sigmoid classifier with a trained soft prompt and task-dependent training schedules, while MLP and MULTIMAX use separate probe configurations. Model selection is based on validation accuracy, and experiments compare training costs across methods.

  • Training configuration: Speculative Probing uses a sigmoid head with 0.2 dropout on the pooled hidden state and soft-prompt lengths k ∈{1, 2, 5}.Training uses AdamW with learning rate 3×10−4, cosine decay to 10−4, weight decay 0.01, and batch size 1.
  • Training configuration: Speculative Probing trains for 10 epochs on verih and nemotron_safety, and 4 epochs on CoT Repetition and CoT Reasoning Strategy.
  • Baseline configurations: MLP uses AdamW with learning rate 5×10−4, batch size 32, and 10 epochs, while MULTIMAX uses the same optimizer schedule.
  • Evaluation protocol: For every method, the best epoch is selected by validation accuracy, and the corresponding test accuracy is reported.
  • Compute and timing: Experiments run on 2×NVIDIA RTX PRO 6000 Blackwell GPUs; Speculative Probing takes minutes to hours, versus seconds to minutes for MLP and MULTIMAX.

C Zero-Shot GPT-5.4-mini Baseline Prompts

The zero-shot GPT-5.4-mini baseline uses one fixed yes/no prompt template per task, embeds each test example, and parses the first yes/no token as the prediction. Templates define task-specific criteria for reasoning traces, safety, and instruction conflicts.

  • Baseline procedure: GPT-5.4-mini receives one fixed prompt template per task, with the test example embedded at the {{TEXT}} placeholder.The templates explain the task, request a one-word yes/no answer, and provide no in-context examples.
  • Baseline procedure: The baseline parses the first yes/no token from each response into a binary prediction.
  • Task prompts: For verih, the prompt asks whether system and user instructions conflict because following both simultaneously is impossible or contradictory.Compatible instructions receive the no label.
  • Task prompts: For CoT Repetition, the prompt labels traces repetitive when they revisit points, restate ideas, or loop without making progress.Traces that progress without obvious repetition receive the no label.
  • Task prompts: For CoT Reasoning Strategy, the prompt labels reasoning as branched when it considers multiple alternative approaches or hypotheses before committing.A single linear chain is labeled not branched.
  • Task prompts: For nemotron_safety, the prompt labels harmful, dangerous, illegal, sexual, violent, self-harm, hate, or otherwise disallowed requests as unsafe.Benign or safe prompts receive the no label.
Loading 2608.28099v1…