From Noisy Traces to Root Causes: Structural Trajectory Analysis and Causal Extraction for Agent Optimization

Summary§

STRACE (Structural Trajectory Analysis and Causal Extraction) addresses the challenge of optimizing long-horizon LLM agents by reducing noise in execution traces. It operates at two levels: batch-level failure pattern mining to select representative failures, and per-trace causal localization over a textual dependency graph to isolate root causes. This enables more precise optimization signals, leading to significant performance gains.

Core Methodology§

Batch-Level Failure Pattern Mining§

STRACE clusters similar failure traces using embeddings of error messages or step outcomes. Redundant traces within a cluster are pruned, preserving only diverse failures that cover distinct failure modes. This prevents overfitting to low-value failures and reduces computational cost.

Per-Trace Causal Localization§

For each selected trace, STRACE constructs a textual dependency graph where nodes are agent steps (observations, actions, thought) and edges denote causal dependencies (e.g., a thought leading to an action). The graph is built using a lightweight dependency parser or an LLM-based relation extractor. Then, it performs root cause localization by identifying the minimal set of nodes that are most influential on the final failure. This is formalized as: find the subset of nodes $C$ that maximizes the causal effect on the failure outcome $F$, subject to a budget $k$:

$$ C^* = \argmax_{C, |C| \leq k} I(C; F) $$

where $I(C;F)$ is the mutual information approximated by the reduction in uncertainty when observing $C$. In practice, STRACE uses a greedy algorithm: remove non-causal steps by iteratively pruning nodes that have no influence on the failure, as measured by a learned causal mask.

Optimization Context Construction§

The filtered traces and extracted root-cause modules are concatenated into a high-signal optimization context. This context is fed to an LLM optimizer (e.g., GPT-4) to generate improved agent policies. The optimization objective is:

$$ \theta^* = \argmax_\theta \mathbb{E}_{\tau \sim \mathcal{T}, c \sim \text{STRACE}(\tau)} [R(\pi_\theta(c))] $$

where $\tau$ is a trajectory, $c$ is the extracted context, and $R$ is task success.

Implementation Details§

• Dependency graph construction: Use BERT-based relation extraction to link steps. Edges are weighted by confidence.
• Root cause localization: Train a small MLP that predicts the failure probability given a subset of steps. Use Shapley values to score each step's contribution.
• Batch clustering: Use K-means on sentence embeddings of step summaries. Cluster number determined by silhouette score.

Code Snippet (Illustrative)§

import torch
from transformers import AutoModel, AutoTokenizer

def build_dependency_graph(steps):
    tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
    model = AutoModel.from_pretrained("bert-base-uncased")
    # Simplified: embed each step and compute pairwise attention
    embeddings = [model(torch.tensor(tokenizer.encode(s))).last_hidden_state.mean(0) for s in steps]
    adj = torch.zeros(len(steps), len(steps))
    for i in range(len(steps)):
        for j in range(len(steps)):
            adj[i,j] = torch.cosine_similarity(embeddings[i], embeddings[j], dim=0)
    return adj

def causal_localization(steps, adj, failure_label):
    # Greedy removal of low-influence nodes
    remaining = set(range(len(steps)))
    while len(remaining) > 1:
        scores = []
        for node in remaining:
            subset = remaining - {node}
            # Predict failure with subset (simplified)
            score = predict_failure(steps[list(subset)])
            scores.append(score)
        if max(scores) - min(scores) < 0.1:
            break
        worst = min(scores, key=lambda x: scores[x])
        remaining.remove(worst)
    return list(remaining)

Results§

On VeruSAGE-Bench (formal verification), STRACE improved success rate from 42.5% to 58.5% (1.4×). Ablations confirm both batch filtering and causal localization are essential.