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

TL;DR

This paper introduces a novel framework called OPTO that uses execution traces and large language models to optimize complex, non-differentiable workflows, enabling applications like code debugging, robot control, and hyper-parameter tuning.

Contribution

It formalizes the OPTO framework, develops the Trace library for workflow optimization, and demonstrates the effectiveness of the LLM-based optimizer OptoPrime across various domains.

Findings

OptoPrime effectively performs first-order numerical optimization.

OptoPrime successfully optimizes prompts, hyper-parameters, and control policies.

The framework often matches or exceeds specialized optimizers in diverse tasks.

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