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

TL;DR

This work develops a general LLM-based generative optimizer called OptoPrime, 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.

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/.

Figures (27)

• Figure 1: Learning Example in Battleship. Trace automatically optimizes heterogeneous parameters (multiple codes) to train an agent to Battleship. Means and standard errors are computed over 10 random seeds.
• Figure 2: Python Code of the Battleship Example. To build a self-adapting agent with Trace, we only need to annotate some empty functions (reason, act) and set up an optimizer following PyTorch semantics. For space, we trim the docstrings of the empty functions with "$\dots$" and list them in \ref{['sec:optimized params examples']}. Trace then builds a DAG as the workflow executes and updates the parameters (see \ref{['fig:battleship-learning']} for the result).
• Figure 3:
• Figure 4: An example pseudo-code report generated by Trace for a program of x = Node(-1.0); z = bar(x) * (bar(x)+1) and the objective of $\max_x z$.
• Figure 5: Numerical Optimization and Traffic Optimization Results.

Theorems & Definitions (11)

• Example 1: Neural network with back-propagation
• Example 2: RL
• Example 3: Prompt Optimization of an LLM Agent
• Theorem 1
• Theorem 2
• Example 4: Neural network with back-propagation
• Example 5: Code Debugging
• Example 6: RL
• Example 7: Hyperparameter Tuning of ML Pipeline
• Example 8: Prompt Optimization of an LLM Agent