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

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

Paper

References (51)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer GPUE8/10 · confidence 4/52024-06-28

Summary

The paper proposes a problem specification (OPTO) and a solution (optimizer OptoPrime) as well as a software framework (Trace) for agentic program optimization. The paper demonstrates high performance of the bundle above in solving 5 benchmarks. The paper demonstrates the ability to optimize 3 types of parameters: numeric, text prompts, and python code. The paper proposes Trace as an analogue of imperative-style AutoDiff engines like PyTorch for agentic programs relying on (black box) LLM inferences. The evaluated benchmarks cover a reinforcement learning setting (Battleship, Traffic control and Robot manipulation), scalar regression (Numerical optimization) and question answering (classification and more in BigBenchHard).

Strengths

[1] The authors do a great job in demonstrating the operational performance of Trace on 5 examples: Battleship, Numerical optimization, Traffic control, BigBenchHard, and Robot manipulation. [2] Great to see Trace compared with both traditional optimizers and the recent LLM-based (OPRO) in Figure 5b. [3] It is nice to see that the proposed method is designed to generalize to several types of payload (demonstrated 3: numeric values, text and code). [4] Experiments in “5.3 Unifying Prompts and Functions Optimization” neatly showcase joint optimization of prompts and code which is a clear novelty, even though one may argue that a text prompt and a text of a piece of code are both textual information. Good to see that the optimization works even without the history like Yang 2024 (OPRO) or population like Pryzant 2023 (ProTeGi). [5] The learning seems very sample-efficient with the reported scores achieved in just several single-sample iterations. [6] After running exp/run_prompt_bigbench.py I confirm the reproducibility of the results on BBH.

Weaknesses

[1] The paper does not compare with “TEXTGRAD: Automatic Differentiation via Text” (arXiv:2406.07496v1) which is a concurrent work. [2] The narrative goes from describing a piece of software (Trace) to the problem specification (OPTO) and then to the algorithmic solution (OptoPrime). I would expect the problem specification and the proposed algorithm to be explained first while the software details are discussed later or even in the Appendix. [3] The experiment in 5.3 confirms that the proposed method works in principle, however due to the small scale of the optimized program (3 parameters), there is no clear signal that Trace will scale up for 10s or 100s of trainable parameters. [4] The definition of a trace (line 126) is not given. I am familiar with the verb “to trace” with respect to a symbolic program, however it is unclear what “a trace” means. From a random web page (https://www.ituonline.com/tech-definitions/what-is-an-execution-trace/): Definition: Execution Trace An execution trace is a record of the sequence of operations executed by a program during its run. This trace includes function calls, variable values, and the flow of control among other details. This definition is not mathematically clear. [5] The work lacks the analysis of Figure A.12 where the learned prompt template does not make much sense. [6] Running the baseline exp/run_prompt_bigbench_dspy.py fails. [7] The paper does not demonstrate the optimization success in agentic programs with memory, i.e. when the optimized class like Predict below has internal state that changes from call to call: ``` @trace_class class Predict(LLMCallable): ```

Questions

[1] On lines 79-80 you claim that “Remarkably, there is 80 no mention of Battleship nor details on how the functions reason and act should behave or adapt 81 in Fig. 2a”. However in Appendix H, Iteration 0 (Initialization) you provide relatively detailed instructions: “Given a map, analyze the board in a game. On map, O denotes misses, X denotes successes, and . denotes unknown positions.”. Since the Battleship game is a very well known one, ChatGPT could reproduce the code memorized from the public repositories. How do you assess the risk of this type of leak? [2] Where can an example of prompt optimization of an LLM agent as per line 767 be found? [3] I appreciate quoting “Complete” on Figure 2a as the example is not clear without `__call__()` and `select_coordinate()` that can only be found in the supplementary materials in exp/battleship_exp.py -> Policy2. I suggest replacing it with “An excerpt from”. [4] In Figure 3 it is not clear how can g1, g2 and g3 be different since the connectivity of the graph is defined by the optimized torch-style module, specifically a chain of 2 optimizable nodes Reason and Act in the example. [5] It would be valuable to know the USD cost of OpenAI API usage per experiment. [6] On the technical side, python version is not mentioned, datasets and ray are not on the requirements list. However I appreciate providing OAI_CONFIG_LIST_sample.

Rating

8

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

[1] The authors honestly admit that one of the main limitations of the current approach is LLM context length. Indeed, the scaling of the proposed algorithm that packs the entire trace into a single context for LLM inference inside OptoPrime. [2] Scaling of the proposed framework to more learnable parameters and more sophisticated agentic programs is yet to be demonstrated.

Reviewer aQXt5/10 · confidence 3/52024-07-12

Summary

This paper proposes an end-to-end optimization framework, Trace, for the automatic design and updating of artificial intelligence systems. Trace is based on Optimization with Trace Oracle (OPTO), treating the computational workflow of AI systems as a graph of neural networks, which can be updated via backpropagation. Additionally, this paper introduces OptoPrime, a general optimizer based on large language models (LLM), as a specific implementation of Trace. In the experimental section, the paper compares Trace with the state-of-the-art LLM optimizer OPRO across various tasks, including numerical optimization, traffic control, and robotic control. The results demonstrate the superior performance of Trace. The primary contribution of this work lies in modeling computational workflows as OPTO problems and designing Trace + OptoPrime to address these issues.

Strengths

Strengths: 1.Clarity of Writing: The structure of the article is clear, the language is fluent, and it is easy to understand. 2.Novelty: This paper introduces a novel end-to-end optimization framework, Trace, which views the computational process as a graph and utilizes execution trace information to optimize parameters. Compared to traditional black-box optimization methods, this approach offers higher efficiency and greater interpretability.

Weaknesses

Weaknesses: 1.Limited Scalability: The OptoPrime optimizer mentioned in the paper has some scalability limitations, such as difficulty in handling parameters or nodes that cannot be represented textually, and challenges in dealing with computational workflows that contain a large number of nodes. 2.The graph for Trace requires manual design, lacking automated methods. 3.Limited Experimental Improvement: The experimental section shows limited improvement, with the differences between Trace and OPRO not being particularly significant.

Questions

1.Difference Between Graph Optimization and Individual Node Optimization: Graph optimization considers the entire computational workflow as an interconnected system, optimizing parameters in a holistic manner, which can lead to more coordinated and efficient results. In contrast, optimizing individual nodes treats each component in isolation, potentially missing out on interactions between nodes. However, the paper lacks experimental evidence to support the effectiveness of Trace in performing graph optimization over individual node optimization. 2.Discrepancy Between Ablation Studies in Figure 5 and Figure 6: The conclusions drawn from the ablation studies in Figure 5 and Figure 6 differ, raising questions about consistency. 3.Length of Trace's Prompts Compared to OPRO and Token Efficiency: The prompts used by Trace are longer than those used by OPRO. It is important to quantify this difference in length to assess its impact on token efficiency.

Rating

5

Confidence

3

Soundness

2

Presentation

2

Contribution

2

Limitations

Yes

Reviewer YuBN7/10 · confidence 2/52024-07-13

Summary

The paper introduces Trace, a novel optimization framework that instances the concept of Optimization with Trace Oracle (OPTO). In Trace, the computational workflows is treated as dynamic graphs and rich information, including intermediate results, processing details and computational graph, are used as feedback for optimization instead of traditional gradients. The framework includes a general-purpose optimizer, OptPrime, to solve the OPTO problem. In the experiment section, Trace is shown to be comparable to the first-order gradient optimizer on a numerical optimization task, and outperform baseline LLM-based methods across a wide range of tasks including , hyper-parameter tuning, robot controller design, etc. Additionally , Trace offers a Python interface that can integrate seamlessly with PyTorch.

Strengths

1.I think using execution traces instead of gradients for optimization is quite innovative. This allows for the optimization of workflows that are non-differentiable or have dynamic computation graphs. 2.This framework can be applied to a wide range of tasks, including robotics and AI systems. 3.Trace has shown superior performance compared to other LLM-based optimizers and has demonstrated results comparable to those of traditional gradient-based optimization methods. 4.The Python interface for Trace simplifies integration with existing codebases.

Weaknesses

As someone outside of this field, I find this paper to be quite impressive. I did not identify any specific weaknesses.

Questions

1. Could Trace be adapted for use with other programming languages, and what would such an adaptation entail in terms of architectural changes?

Rating

7

Confidence

2

Soundness

4

Presentation

4

Contribution

3

Limitations

Yes

Reviewer GPUE2024-08-12

Thank you for addressing the weaknesses and questions. I intend to keep my score.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC