Step representations improve failure tracing in AI agent systems
ReCast: Attribution-Oriented Step Representation Learning for LLM-Based Agent Systems
Artificial Intelligence
Summary
When AI systems follow a series of steps to complete tasks, mistakes made early on can cause problems later, but it’s hard to find where things first went wrong. The authors studied how AI models represent these steps and found the usual method doesn’t clearly show the earliest errors. They created a new way called ReCast that changes step information to better highlight which step caused a failure. Their approach works better than existing ones at identifying the original cause of mistakes in several tests.
What this means in practice
- •For ai system developers: Pinpoint early mistakes in complex AI agent workflows to improve debugging efficiency.
- •For software quality engineers: Analyze failure origins in multi-step AI processes to enhance system reliability and error handling.
Authors
Weilin Jin, Mingyu Wang, Taiyu Zhu, Ziqi Zhou, Wenbo Li, Haoyang Huang, Nan Duan, Yifan Wu, Ying Li, Zhonghai Wu
Abstract
In LLM-based agent systems, failures can originate from early steps whose effects propagate through subsequent interactions, making their origins difficult to identify. To trace such failures back to their origin, failure attribution has been formulated as the task of identifying the earliest step responsible for the failure. Recent methods leverage LLM internal signals for failure attribution, typically using hidden states as step representations. We therefore conduct an empirical study to evaluate how effectively these representations distinguish root-cause steps from other steps and find limited separation. Motivated by this observation, we propose ReCast, a step representation learning method that transforms hidden states from a frozen LLM into attribution-oriented step representations. ReCast first selects attribution-relevant layers, then constructs complementary pattern and deviation features, and finally learns contextualized step representations through an encoder trained with contrastive and ranking objectives. We also introduce ReCast-2K, a training dataset for failure attribution. ReCast achieves the best Hit@1 across four benchmarks, surpassing the strongest baseline by 5.65 and 9.19 pp on Who&When Algorithm and Handcrafted, respectively. Code is available at https://anonymous.4open.science/r/ReCast-5FB6 .