Trace is the Next AutoDiff: Generative Optimization with Rich Feedback, Execution Traces, and LLMs

Abstract: We study a class of optimization problems motivated by automating the design and update of AI systems like coding assistants, robots, and copilots. AutoDiff frameworks, like PyTorch, enable efficient end-to-end optimization of differentiable systems. However, general computational workflows can be non-differentiable and involve rich feedback (e.g. console output or user's responses), heterogeneous parameters (e.g. prompts, codes), and intricate objectives (beyond maximizing a score). We investigate end-to-end generative optimization -- using generative models such as LLMs within the optimizer for automatic updating of general computational workflows. We discover that workflow execution traces are akin to back-propagated gradients in AutoDiff and can provide key information to interpret feedback for efficient optimization. Formally, we frame a new mathematical setup, Optimization with Trace Oracle (OPTO). In OPTO, an optimizer receives an execution trace along with feedback on the computed output and updates parameters iteratively. We provide a Python library, Trace, that efficiently converts a workflow optimization problem into an OPTO instance using PyTorch-like syntax. Using Trace, we develop a general LLM-based generative optimizer called OptoPrime. In empirical studies, we find that OptoPrime is capable of first-order numerical optimization, prompt optimization, hyper-parameter tuning, robot controller design, code debugging, etc., and is often competitive with specialized optimizers for each domain. We envision Trace as an open research platform for devising novel generative optimizers and developing the next generation of interactive learning agents. Website: https://microsoft.github.io/Trace/.

• The paper introduces the Trace framework and OPTO paradigm, extending auto differentiation to non-differentiable computations via rich execution traces.
• It details how computational workflows are transformed into directed acyclic graphs to enable effective back-propagation of varied feedback.
• Empirical results in domains such as prompt tuning and robotic control demonstrate significant improvements in optimization efficiency.

Overview of the Paper "Trace is the New AutoDiff — Unlocking Efficient Optimization of Computational Workflows"

This paper presents a framework named Trace aimed at the automatic optimization of computational workflows, promoting the notion of “Trace is the New AutoDiff.” The core contribution of the paper is the introduction of a novel framework, Trace, and an associated optimization paradigm named Optimization with Trace Oracle (OPTO), implemented through a general-purpose optimizer called OptoPrime. The authors argue that this new abstraction allows for efficient and rich feedback-driven optimization across various domains by extending the idea of back-propagation beyond differentiable computations.

The Trace Framework and Its Design

The Trace framework is designed to automatically convert computational workflows into optimization problems, allowing the use of a general-purpose optimizer. Specifically, Trace treats the computational workflow as a directed acyclic graph (DAG) where nodes represent parameters (including, but not limited to, prompts, hyperparameters, and code) and directional edges denote the computational relationships among these nodes.

In implementing Trace, two primary constructs are utilized:

1. Node: Wraps Python objects and logs them as unique entities within the global graph, allowing them to be flagged as trainable parameters.
2. Bundle: Decorates Python methods, transforming them into operators that are traceable and optimizable, thereby allowing Trace to back-propagate execution feedback effectively.

Trace ensures that the entire computational graph involved in workflow execution is captured and the dependencies between nodes are utilized to propagate feedback from the output back to the parameters effectively. This provides an automatic means to incorporate rich and varied forms of feedback, such as natural language and error messages, far beyond the scalar feedback used in traditional black-box optimization techniques.

Optimization with Trace Oracle (OPTO)

The OPTO framework abstracts the iterative optimization process, defining a new optimization problem class where the optimizer receives an execution trace (captured graph) along with rich feedback. The setup generalizes beyond simple gradient calculations to accommodate non-differentiable computations and more complex feedback signals. This abstraction allows the optimizer to derive meaningful update directions, facilitating efficient optimization.

Implementation and Examples

Trace's practical applications are demonstrated through various examples, ranging from traditional numerical optimization, prompt optimization for LLMs, and hyperparameter tuning to robotic controller design. Empirical studies showcase how Trace converts these diverse optimization tasks into OPTO problems, leveraging the execution trace to significantly improve optimization efficiency compared to state-of-the-art algorithms for each specific domain.

For instance, in a robot controller design task within a simulation environment, Trace effectively utilizes language feedback to guide the iterative improvement of the controller, demonstrating its capability to handle intricate dependencies and dynamically changing computational graphs. Similarly, in the context of optimizing prompts for LLM workflows, Trace jointly optimizes prompts and associated processing codes, achieving improved performance over domain-specific optimizers.

Implications and Future Directions

The work posits several implications for both practical and theoretical advancements in optimization of AI workflows:

• Practical Implications: The ability of Trace to handle various forms of feedback and non-differentiable parameters opens new possibilities for automating the adaptation and optimization of complex AI systems, ranging from coding assistants and chatbots to multi-agent environments and robotic systems.
• Theoretical Implications: The introduction of OPTO as a foundational framework for optimization suggests new lines of inquiry into efficient abstractions for handling rich feedback and dynamically changing computational structures.

Looking forward, future research may explore the design of specialized propagators within the Trace framework, enhancing its scalability and extending its applicability to even more complex and large-scale workflows. Moreover, advancements in LLMs' capabilities, and their integration into the optimization process, could further improve the efficacy and flexibility of OPTO-based methods.

Conclusions

The paper delineates a new paradigm, positioning Trace and OPTO at the confluence of computational graph-based optimization and LLM utilization, thereby extending the principles of AutoDiff to broader, more heterogenous AI workflows. Through concrete examples and rigorous empirical validation, it demonstrates how this unified approach can lead to substantial improvements in the efficiency and quality of optimization across various domains.