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

AI summary

Overview

Research area: Machine learning for reliable LLM-based multi-agent systems — specifically runtime failure prediction (early detection that a running agent trajectory will end in failure).

Technical level: Advanced. The paper assumes familiarity with LoRA adapters, Transformer internals (query/value projections), BCE loss, SVD truncation, and first-order Taylor analysis.

Scope: The paper proposes Dream2Discern, a runtime auditor trained by alternating a loss-guided perturbation search (“Dream”) with supervised learning on the resulting perturbed representations (“Discern”), evaluated on two agent-trajectory benchmarks.

What This Paper Is About

LLM-based multi-agent systems make reasoning and tool-use errors that can waste tool calls or modify external systems unexpectedly, so deployment needs an auditor that watches an execution prefix and predicts whether the task will ultimately fail. Existing learned auditors train on a small corpus of trajectories with verified outcomes, and the paper argues that simply adding more labelled trajectories or editing existing ones is both expensive and ineffective. Dream2Discern instead manufactures “hard” training examples internally, by perturbing the auditor’s own representation weights in directions that raise its prediction loss, and then training the auditor to classify those harder variants using the original labels.

Key Contributions

1. Evidence against trajectory-level augmentation. The authors measure the cost of acquiring labelled trajectories (averaging tens of seconds per trajectory over roughly 2K trajectories in AFTraj-2K) and show that, of 326 verified continuations obtained by editing and re-executing trajectories, 85.7% of those derived from failed source trajectories actually succeed — so source labels cannot be reused. Training a prediction head on original prefixes plus verified edits (“+Edited”) produced higher held-out class-balanced BCE than training on originals alone under a fixed budget.

2. The Dream2Discern method. A parameter-efficient auditor (frozen LLM backbone plus LoRA on query and value projections, plus a classification head) trained by alternating two steps: Dream, which searches for bounded, low-rank perturbations to the adapted query/value weights that maximize batch prediction loss with the head held fixed; and Discern, which updates the adapters and head on the perturbed forward pass using the recorded outcome labels.

3. A closed-form linearized perturbation rule. Proposition 1 shows the loss-maximizing perturbation within the relative-radius/rank budget is built from the top-r singular components of each scaled gradient, with relative magnitudes proportional to their retained Frobenius norms; the candidate is accepted only if it increases the measured nonlinear loss, with up to six successive halvings. Proposition 2 bounds the Discern loss as the unperturbed loss plus ρ·g_r, up to a second-order remainder.

4. Empirical gains on two benchmarks and two model families. Over 90% trajectory-level weighted accuracy (wACC) on both AFTraj-2K and StepShield, with up to 11.11% wACC improvement over the strongest auditor baseline, plus results on Llama backbones and ablations isolating the loss-guided perturbation from random or no perturbation.

Main Findings

• Headline accuracy: Dream2Discern-Qwen3-8B reaches 91.47 ± 0.17 wACC on AFTraj-2K (AUROC 95.82 ± 0.43, FPR95 31.95 ± 3.07, TWA 73.45 ± 0.71) and 87.50 ± 0.93 wACC on StepShield (AUROC 93.95 ± 0.63, FPR95 53.09 ± 8.40, TWA 55.69 ± 1.10).

• 4B variant: Dream2Discern-Qwen3-4B reaches 90.96 ± 0.60 wACC on AFTraj-2K and 90.28 ± 0.46 wACC on StepShield (TWA 57.67 ± 0.35), exceeding Claude-Sonnet-4.6 on StepShield wACC (87.04) and TWA (56.79).

• Against failure-auditor baselines: On AFTraj-2K, SAFE-Qwen3-8B scores 85.74 ± 1.22 wACC, Hide-and-Seek-Qwen3-4B 88.76 ± 1.14, TRACES-Qwen3-4B 88.35 ± 1.14, and AgentForesight-7B 77.11. The paper reports Dream2Discern outperforming SAFE, Hide-and-Seek, and TRACES on all four metrics averaged over three seeds; on StepShield, Claude-Sonnet-4.6 retains the best AUROC (95.58) and FPR95 (38.89), so Dream2Discern leads only wACC and TWA there.

• Token-uncertainty baselines are weak: Max/avg entropy and max/avg NLL under a Qwen3-8B proxy hover near chance on AUROC in several rows (e.g., Max NLL-8B: 45.20 AUROC on AFTraj-2K, 51.58 on StepShield).

• Gains transfer to Llama: Llama3.2-3B-Instruct with Dream2Discern reaches 91.27 ± 0.30 wACC on AFTraj-2K (vs TRACES 85.54 ± 4.44 and SAFE 85.54 ± 4.64) and 82.25 ± 4.20 on StepShield (vs TRACES 62.04 ± 2.45). On StepShield with Llama3.1-8B, Dream2Discern has the best wACC (79.48 ± 2.94) but a lower TWA (49.06 ± 2.67) than TRACES (56.01 ± 0.72).

• Loss-guided beats random and no perturbation, with exceptions: On AFTraj-2K with Qwen3-8B, Dream gives the best AUROC (96.21), wACC (91.57), and TWA (73.70) versus 94.60/89.46/69.80 for random and 93.65/90.06/71.04 for none. However, on AFTraj-2K with Qwen3-4B the no-perturbation variant has the best wACC (90.96 vs 90.36 for Dream), and on StepShield with Qwen3-8B random perturbation has the best wACC (89.81) and TWA (57.72), with Dream best only on AUROC (94.59).

• Variants are genuinely harder: With a trained Qwen3-8B fixed, the authors sample 50 batches of 32 prefixes from the AFTraj-2K training set; Dream perturbations increase batch BCE on every batch, while random perturbations cluster near zero.

• Sensitivity tracks errors: Using a fixed Qwen3-4B auditor and eight Dream perturbations from separate training batches, the paper reports that held-out prefixes with larger absolute score changes under training-selected perturbations are more often misclassified by the unperturbed auditor (analysis truncated in the provided text).

• Accuracy–timeliness trade-off: Sweeping the nominal false-alarm level α from 0.01 to 1.00, Dream2Discern shows higher wACC at operating points favouring early detection on AFTraj-2K and gains across a broad range of operating points on StepShield.

Methodology in Plain English

Start with a pre-trained LLM and freeze it. Attach small trainable adapters (LoRA) to the query and value weight matrices in the Transformer blocks, and put a simple classifier on top of the last-token hidden state; the classifier outputs a failure probability. Give it a short supervised warm start (100 steps), then alternate two steps for a 4-epoch budget.

In the Dream step, freeze everything and ask: what small change to the adapted query/value weights would make the current auditor worse at classifying this batch of already-labelled prefixes? The batch is class-balanced (equal successes and failures). The change is constrained to be small relative to the original weights and to have low rank, so it is a structured nudge rather than noise. Because the exact answer is hard, they linearize the loss, take the gradient with respect to each weight matrix, keep only its top singular directions, and scale by the allowed radius. They then check the true (nonlinear) loss, halving the step up to six times until the loss actually rises; if nothing works, they skip the perturbation.

In the Discern step, keep that perturbation fixed, run the prefixes through the perturbed network, and do a normal gradient update on the adapters and classifier, using each prefix’s original success/failure label. Gradients do not flow back into the perturbation search. Because the adapters change, next epoch’s Dream step finds new perturbations — so the same small corpus keeps generating fresh hard examples.

At inference time the perturbation is discarded entirely. Each prefix gets a failure score from one forward pass, compared to a step-dependent threshold fitted offline on successful training trajectories using an empirical-percentile / max-rank procedure; thresholds are calibrated at nominal α = 0.05 unless otherwise stated. Evaluation is trajectory-level: AUROC, FPR95, weighted accuracy (wACC), and time-weighted accuracy (TWA). AFTraj-2K is in-domain (test split from the same corpus as training); StepShield is out-of-domain at the template level (test trajectories come from incident templates never seen in training).

Why This Matters

Impact on research. The paper reframes data augmentation for runtime monitors: instead of collecting more trajectories or editing them (which breaks label validity — 85.7% of edited failures became successes), it treats the model’s own loss landscape as a renewable source of hard examples. It also supplies a concrete negative result about edited-prefix augmentation under a fixed training budget, and a linearized, provably bounded way to pick those hard variants.

Real-world applications:

• Guarding coding and software-development agents that can otherwise make unwanted file or repository changes.
• Safety monitoring for web-navigation and open-ended tool-use agents, where an unnecessary tool call can act on an external system.
• Early abort/rollback decisions in long multi-step agent pipelines, where the auditor must fire before the terminal step to be useful.
• Budget control for costly agent rollouts, flagging likely failures so compute is redirected.

Industry relevance. The method is parameter-efficient (LoRA only, frozen backbone) and adds one forward pass plus a threshold comparison at runtime, so it fits deployment constraints. A 4B-backbone auditor beating a frontier proprietary judge on StepShield wACC and TWA indicates that small, task-specialized auditors can substitute for expensive API-based judging. The paper’s own caveat — edits do not help under a fixed budget — is directly useful to teams planning trajectory-collection pipelines.

Future Directions

• Combine with, rather than replace, corpus expansion. The paper shows editing alone does not help under a fixed budget; whether Dream-style variants and verified new trajectories are complementary is not reported.
• Stress-test the exceptions. Dream perturbation loses to no-perturbation on AFTraj-2K wACC at Qwen3-4B and to random perturbation on StepShield wACC/TWA at Qwen3-8B, and Llama3.1-8B on StepShield trails TRACES on TWA — the conditions under which the search helps are not fully characterized.
• Tune and ablate the perturbation budget. The relative radius γ, rank cap r, and the choice of which Transformer blocks receive adapters are not reported as swept in the provided content, so their sensitivity is an open question.
• Extend beyond query/value projections and single-step search. The perturbation set is restricted to adapted query and value weights, and the linearized search is reset from zero each batch; whether broader operator sets or multi-step search yield harder variants is untested here.
• Scale the mechanism analysis. The variant-construction check uses 50 batches of 32 prefixes on a single benchmark with Qwen3-8B, and the sensitivity analysis eight perturbations; larger-scale validation is left open.

Target Audience

Researchers and engineers working on LLM agent reliability, runtime monitoring, and failure detection; practitioners who need a cheap, deployable auditor for multi-agent pipelines; and ML researchers interested in loss-guided data augmentation, adversarial or hard-example construction inside adapters, and parameter-efficient fine-tuning of frozen backbones. Readers without background in LoRA and Transformer internals will find the theory sections (Propositions 1 and 2) difficult, though the motivation and results tables are accessible.

Authors’ 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. Project page: https://zbox1005.github.io/dream2discern-project/