Dream2Discern: Learning from Hard Latent Variants for Runtime Failure Prediction in Multi-Agent Systems

Dream2Discern

Learning from Hard Latent Variants for Runtime Failure Prediction in Multi-Agent Systems

Boxuan Zhang1 Jianing Zhu2 Xiaotian Han3 Ruixiang Tang1
1Rutgers University 2The University of Texas at Austin 3Amazon
{boxuan.zhang, ruixiang.tang}@rutgers.edu
arXiv arXiv soon Code Code soon Hugging Face Hugging Face soon BibTeX
Open animation in a new tab ↗
>90%

trajectory-level weighted accuracy on both benchmarks

AFTraj-2K 91.47 · StepShield 90.28

+11.11

wACC points over the strongest auditor baseline

StepShield · 90.28 vs. AgentForesight-7B 79.17

0

additional trajectories collected or verified

Hard variants are built inside the auditor

4B

auditor exceeds Claude-Sonnet-4.6 on StepShield

wACC 90.28 vs. 87.04 · TWA 57.67 vs. 56.79

Overview

Dream2Discern at a glance

  1. 01

    Corpus augmentation is costly and untargeted

    Runtime auditors learn from a limited labeled corpus. New trajectories need rollouts and verification, edits can invalidate outcome labels, and broader coverage does not reveal hard cases.

  2. 02

    Dream constructs hard latent variants

    With the auditor fixed, a bounded, low-rank perturbation of its query and value projections is chosen to increase batch prediction loss, yielding hard latent variants of recorded prefixes.

  3. 03

    Discern learns from them with the original labels

    The auditor learns from these variants with the recorded labels as they are regenerated each batch. At runtime the perturbation is discarded and a calibrated threshold predicts failure.

Abstract

Reliable deployment of LLM-based multi-agent systems calls for runtime auditors that predict terminal outcomes from observed execution prefixes. These auditors typically learn from a limited corpus of trajectories with verified outcomes. Expanding this corpus requires costly execution and verification, yet broader failure coverage alone provides no explicit criterion for identifying prefixes that the current auditor finds difficult to classify. To this end, we propose Dream2Discern, a runtime auditor trained through implicit augmentation guided by its own prediction loss. During Dream, we select perturbations of the auditor's representation operators to increase prediction loss on recorded prefixes, yielding hard latent variants. During Discern, the auditor learns to predict execution outcomes from these variants using the original labels. The variants are regenerated as the auditor learns, turning a limited corpus into a renewable source of hard examples and requiring no additional trajectory collection or outcome verification. Extensive evaluations on AFTraj-2K and StepShield show that Dream2Discern outperforms competitive failure auditors and achieves over 90% trajectory-level weighted accuracy on both benchmarks while balancing prediction accuracy and timeliness.

01 · Motivation

Trajectory augmentation is costly and does not target the auditor's hard cases

Figure 1. Motivation of Dream2Discern. (a) A runtime auditor predicts terminal outcomes from observed prefixes but learns from a limited labeled corpus. (b) Trajectory augmentation seeks broader failure coverage, but collection is costly and editing can invalidate outcome labels (right). Broader coverage alone does not identify prefixes that the current auditor finds hard to classify. (c) We construct hard latent variants through operator perturbations chosen to increase batch prediction loss, then train on these variants while keeping recorded prefixes and labels unchanged.
Trajectory collection · Costly

Each new label requires a rollout and outcome verification

Averaged acquisition times reach tens of seconds per trajectory, making repeated corpus expansion time-consuming. Mean rollout plus verification time on AFTraj-2K:

Math
87.3 s Coding
58.6 s HotpotQA
72.9 s GAIA
107.4 s
Trajectory editing · Unstable outcomes

Edited trajectories cannot reuse their source labels

We modify an action or tool response, resume execution and verify again, obtaining 326 continuations. Among those from failed source trajectories, 85.7% succeed.

Edited: success Edited: failure Was success
203 94.9%
11 5.1%
Was failure
96 85.7%
16 14.3%

More fundamentally, broader failure coverage alone provides no explicit criterion for identifying trajectories whose prefixes the current auditor finds difficult to assess.

Research question

Can we identify and exploit the current auditor's hard cases without acquiring additional labeled trajectories?

02 · Key insight

Prediction loss as the criterion for hard variants

Our key insight is to use the auditor's own prediction loss as a criterion for deciding which variants to construct.

Preliminary evidence

Edits do not improve held-out prediction, while a loss-guided perturbation makes it harder

Original and edited prefixes, encoded by the same frozen LLM, overlap in a shared PCA projection. With the encoder and training budget fixed, adding verified edits does not lower held-out class-balanced BCE. A perturbation selected only on training prefixes, applied temporarily with the original head fixed, raises the same held-out loss.

Original
0.520
+Edited
0.578
+Dream
0.704
Figure 2. (a) Original and edited prefixes overlap in a shared PCA projection. (b) Training with edits does not reduce held-out class-balanced BCE. A temporary Dream perturbation selected on training prefixes raises this loss with the Original head fixed.
03 · Method

Alternating Dream and Discern steps

Figure 3. Overview of Dream2Discern. ① Labeled prefixes. We form a class-balanced batch from recorded success and failure prefixes, keeping their inputs and outcome labels unchanged throughout the training steps. ② Dream. With the auditor fixed, we search for shared, bounded, low-rank perturbations to its query and value projection weights that increase batch prediction loss. ③ Discern. Holding the selected perturbation fixed, we update the LoRA adapters and prediction head to minimize the loss on the resulting hard latent variants. ④ Runtime auditing. We discard the perturbation and apply a calibrated threshold to each prefix score to predict failure during execution.
Interactive · schematic Dream and Discern alternate throughout training
{{ loopIterLabel }}
predict success predict failure
{{ r.t }}
Alert · predicted failure
{{ vizTag }}
Loss
{{ meterArrow }}
Step 1 · Labeled prefixes

A class-balanced batch of recorded prefixes

We form a class-balanced batch from recorded success and failure prefixes, keeping their inputs and outcome labels unchanged throughout the training steps.

ℬ = (qb)b=12m,  ℒℬ(Θ) = 1⁄n Σb ℓBCE(fΘ(Xqb), Yqb)

m prefixes from each outcome class, n = 2m.
Step 2 · Dream

Construct hard latent variants

With the auditor fixed, we search for shared, bounded, low-rank perturbations to its query and value projection weights that increase batch prediction loss.

Ascent Rank-r SVD Accept

Take the batch-loss gradient at V = 0 for each adapted query and value projection, in scaled coordinates Cj = ωjGj.

Normalize the scaled gradient and keep its leading r singular components per projection, giving a candidate within radius γ.

Halve the candidate until the measured batch loss strictly exceeds the loss at V = 0. At most six points are tested; if none is accepted, V = 0.
V⋆ℬ(Θ) ∈ arg maxV∈𝒰γ,r ℒℬ(Θ, V)
Step 3 · Discern

Learn from them with the original labels

Holding the selected perturbation fixed, we update the LoRA adapters and prediction head to minimize the loss on the resulting hard latent variants.

minΘ ℒℬ(Θ, sg[V̂ℬ]),  Θ = (a, ψ)

sg stops gradients through the Dream search. Each new batch starts a fresh search from V = 0.
Step 4 · Runtime auditing

Score each prefix, perturbation discarded

We discard the perturbation and apply a calibrated threshold to each prefix score to predict failure during execution.

si,t = fΘ(Xi,t),  Ŷi,t = 𝟙{si,t ≥ ζt}

One forward pass and one threshold comparison per interaction. Step-dependent thresholds are fitted offline on successful training trajectories.
Success prefix Failure prefix Unperturbed position Hard latent variant Misclassified by the fixed head Positions are illustrative, not data
04 · Results

Over 90% trajectory-level weighted accuracy on both benchmarks

Benchmarks

AFTraj-2K: multi-agent executions in math, coding and tool use, tested in-domain. StepShield: safety-monitoring trajectories, tested on unseen incident templates.

Metrics

AUROC and FPR95 from each trajectory's maximum prefix score; wACC, trajectory-level weighted accuracy; TWA, time-weighted accuracy rewarding earlier detection.

Shown models

Open-weight judges: Qwen3-8B, LlamaGuard3-8B and Qwen3-32B; closed-source judges remain. SAFE, H&S and TRACES use Qwen3-8B; Dream2Discern shows 4B and 8B.

bold best · underline second
Method · {{ benchName }}
AUROC ↑
FPR95 ↓
wACC ↑
TWA ↑
{{ g.name }}
{{ r.method }}
{{ c.v }}{{ c.s }}

Table 1 · selected rows. Shaded rows: SAFE, Hide-and-Seek (H&S) and TRACES with Qwen3-8B. ± is the standard deviation over three training seeds.

On in-domain AFTraj-2K, Dream2Discern leads all four metrics. With Qwen3-8B, it outperforms SAFE, Hide-and-Seek and TRACES on every metric, averaged over three seeds.

On out-of-domain StepShield, Dream2Discern leads wACC and TWA. LLM judges are more competitive here, yet the 4B auditor exceeds Claude-Sonnet-4.6 in wACC and TWA, and its wACC advantage over failure auditors is substantially larger than on AFTraj-2K.

Detection accuracy–timeliness trade-off

With Qwen3-8B, we vary the nominal false-alarm level α from 0.01 to 1.00 and plot wACC against mean normalized detection time over failed trajectories, counting undetected failures as one.

On AFTraj-2K, Dream2Discern achieves higher wACC when early detection is prioritized.

On StepShield, the gains extend across a broad range of operating points.
Figure 4a. Detection accuracy–timeliness trade-off on AFTraj-2K for Qwen3-8B. Markers denote α = 0.05.
Figure 4b. Detection accuracy–timeliness trade-off on StepShield for Qwen3-8B. Markers denote α = 0.05.
05 · Case study

Failure predicted at the error step, two interactions before the test fails

Replay · runs automatically Observed prefix · {{ caseRange }}
Task

Write a function to find the median of three numbers.

Step 1CodeWriter
def median_numbers(a, b, c):
return sorted([a, b, c])[1]
⋮ Steps 2–6Intermediate interactions, not shown
Step 7CodeTesterGT error step
"Solution code":
return max([a, b, c])
"Test code":
assert median_numbers(25,55,65)==55.0
⋮ Step 8Intermediate interaction
Step 9Environment
Test_FAILED
Dream sensitivity · auditor w/o perturbationapprox. values
0.50
Normalized failure scoreDream2Discernw/o Perturbation
43210
Decision boundary
GT error step
{{ v.txt }}
123456789
Dream2Discern{{ d2dScore }}{{ d2dVerdict }} w/o Perturbation{{ baseScore }}{{ baseVerdict }}
{{ caseLabel }}

{{ caseNote }}

Figure 6. Case study of runtime failure prediction on a held-out coding trajectory from AFTraj-2K. Dream sensitivity values are approximate.

Cite this work

BibTeX

@misc{zhang2026dream2discern,
  title = {Dream2Discern: Learning from Hard Latent Variants for Runtime Failure Prediction in Multi-Agent Systems},
  author = {Zhang, Boxuan and Zhu, Jianing and Han, Xiaotian and Tang, Ruixiang},
  note = {Preprint; arXiv identifier pending},
  year = {2026}
}

The arXiv identifier will be filled in once the preprint is posted.

The paper in one sentence

The variants are regenerated as the auditor learns, turning a limited corpus into a renewable source of hard examples and requiring no additional trajectory collection or outcome verification.

Back to top ↑
Dream2Discern · Rutgers University · The University of Texas at Austin · Amazon Contact Cite Top ↑
{{ lbAlt }}

{{ lbCaption }}