Source-linked AI summary
Trace is the Next AutoDiff: Generative Optimization with Rich Feedback, Execution Traces, and LLMs
Ching-An Cheng, Allen Nie, Adith Swaminathan
TL;DR
General computational workflows are difficult to optimize because they combine heterogeneous parameters, rich feedback, non-differentiable operations, and objectives beyond scalar scores. The paper formalizes Optimization with Trace Oracle, implements Trace to expose execution traces, and uses them in OptoPrime, which performs competitively across diverse optimization tasks. The framework remains limited by feedback-design requirements and workflow representations that exclude some stateful or distributed computations.
Problem
General computational workflows combine heterogeneous parameters, rich feedback, non-differentiable operations, and intricate objectives that standard AutoDiff does not cover directly.
Source: §1 · p. 1Method
The paper defines OPTO, builds Trace to convert workflows into OPTO instances, and represents trainable objects and operators with Python primitives.
Source: §2.1 · p. 4, §3.1 · p. 5Results
OptoPrime is competitive with specialized optimizers across prompt optimization, hyper-parameter tuning, robot controller design, numerical optimization, and code-related workflows.
Source: §5 · p. 7, §5.2 · p. 7, §5.3 · p. 8Takeaways & Limitations
Trace provides a platform for end-to-end generative optimization of computational workflows using execution traces and rich feedback.
Source: §1.1 · p. 2, §7 · p. 10OPTO requires a specified solution concept and feedback source, while Trace cannot directly represent some stateful in-place or distributed workflows.
Source: §2.1 · p. 4, §6 · p. 10
Abstract
from arXiv · showhide
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/.
1 Introduction
Trace extends end-to-end optimization beyond differentiable systems by using execution traces to connect rich feedback with heterogeneous workflow parameters. The paper formalizes this idea as OPTO, implements it in a PyTorch-like Python library, and introduces OptoPrime as a general generative optimizer.
1 Introduction
- Trace targets workflows with heterogeneous parameters, rich feedback, intricate objectives, interdependent steps, and non-differentiable or stochastic operations.These properties make the computational structure potentially change with parameters and inputs.Source: §1 · p. 1
1.1 Toward Efficient End-to-End Optimization of Computational Workflows
- Trace propagates execution traces through workflow graphs instead of gradients, retaining applicability to non-differentiable workflows while subsuming back-propagation for differentiable ones.The traces record intermediate results and how they are used, helping relate output feedback to parameters.Source: §1.1 · p. 2
1.2 Example of Trace in Action
- In the Battleship example, Trace marks code components as trainable, records an execution DAG, incorporates execution errors as feedback, and improves policies through interaction.The learned policy is evaluated on newly generated games, and the framework requires no explicit Battleship API or prescribed behavior for the trainable functions.Source: Figure 1 · p. 2, §1.2 · p. 3, Figure 2 · p. 3
1.3 A New World of Optimization
- OPTO has an optimizer receive a computational graph and output feedback, while Trace converts workflow optimization into OPTO instances using PyTorch-like syntax.This separates general-purpose optimizer design from domain-specific workflow components.Source: §1.3 · p. 3
- OptoPrime applies the same LLM-based optimization procedure across prompt, numerical, hyper-parameter, robot-controller, and other workflow problems, achieving 10% higher BigBenchHard accuracy than DSPy’s hand-designed optimizer.The reported comparison concerns optimization of a DSPy program.Source: §1.3 · p. 3
2 Optimization with Trace Oracle
OPTO models iterative workflow optimization through a trace oracle that returns both execution structure and feedback. This representation generalizes across differentiable, reinforcement-learning, and prompt-optimization settings while exposing information absent from black-box scores alone.
2.1 Problem Definition of OPTO
- OPTO represents an optimization problem as parameters Θ, context ω, and a trace oracle T that returns trace feedback τ = (f, g).The graph g contains the parameter roots and an output node receiving feedback f; the optimizer then updates parameters using ω.Source: §2.1 · p. 4
- The context ω supplies invariant instructions or side-information, while feedback f may be a score, gradient, explanation, hint, or console message expressed compactly in text.The framework focuses on OPTO problems whose feedback and context can be represented in text.Source: §2.1 · p. 4
- Unlike black-box optimization, OPTO exposes the computational path toward the output, allowing an update direction to be constructed from feedback and the trace.The graph may change across iterations when workflow behavior varies with inputs or parameters.Source: §2.1 · p. 4
- OPTO contains standard examples including neural-network back-propagation, reinforcement learning trajectories, and prompt optimization of LLM agents.Each example instantiates parameters, a computation graph, feedback, and an optimization context.Source: §2.1 · pp. 4–5
3 Trace: The Next AutoDiff
Trace operationalizes OPTO with Python primitives that build workflow DAGs, propagate feedback through minimal subgraphs, and update heterogeneous parameters. Its propagation method has explicit worst-case costs that exceed back-propagation but is usually negligible in practice.
3.1 Design of Trace
- Trace provides node and bundle primitives for registering trainable Python objects, defining operators, recording code and documentation, and specifying parameter constraints.Decorated functions can themselves be trainable, making their code a parameter.Source: §3.1 · p. 5, §3 · p. 5
- Each Trace iteration executes the decorated workflow, builds a DAG, propagates output feedback through the minimal parameter-to-output subgraph, and calls the optimizer’s step method.Execution errors can also be treated as feedback.Source: §3.1 · p. 5
3.2 Using Trace Primitives for Effective Execution Tracing
- The abstraction level set by bundle trades graph complexity against the descriptions required for operators, and in-place object mutation is the main Python tracing exception described.Duplicating an object before modification can avoid this limitation.Source: §3.2 · p. 5
3.3 Backward Feedback Propagation: Realizing the Trace Oracle of OPTO
- Trace’s backward message passing with the Minimal Subgraph Propagator combines feedback and graph references to implement the OPTO Trace Oracle.The propagator extracts the minimal subgraph connecting parameter nodes and the output.Source: §3.3 · p. 6
- Algorithms 1 and 2 require O(WN^2 log N) time and O(WN) extra space for graphs with N nodes and maximum degree W.Back-propagation has O(WNd^2) time and O(d) extra space, whereas generic trace feedback may require the full subgraph.Source: §3.3 · p. 6
- Although MSP has worst-case Ω(N) feedback-description complexity, its practical overhead is usually negligible because most computation occurs in the forward pass and optimizer step.The authors anticipate issues for minimal subgraphs containing millions of nodes.Source: §3.3 · p. 6
4 Design of the First OPTO Optimizer
OptoPrime is an LLM-based optimizer that turns traced execution graphs and output feedback into pseudo-algorithm problems. It represents traces as code-like reports, prompts an LLM for parameter updates, and supplements sparse feedback with optimization memory.
4 Design of the First OPTO Optimizer
- OptoPrime is an LLM-based generative optimizer designed for any text-based OPTO problem.The paper presents it as one possible optimizer within a larger space of OPTO methods.Source: §4 · p. 6
Subgraph Representation
- OptoPrime represents each trace subgraph as a pseudo-algorithm report containing execution values and operator descriptions, then asks an LLM to update parameters using output feedback.The report resembles a program but is not the original program; it represents the computational graph.Source: §4 · p. 7, Figure 4 · p. 7
Parameter Update
- A ReAct-CoT-style prompt asks the LLM to reason over the graph and suggest parameter changes, which are applied when the suggestion can be extracted.This makes the optimizer’s update step depend on an interpretable text response.Source: §4 · p. 7
Optimization Memory
- OptoPrime includes a memory module that stores past parameter-feedback pairs for cases where a single feedback signal is not informative.The module provides in-context examples when rewards lack information about how they were generated.Source: §4 · p. 7
5 Experiments
Trace and OptoPrime are evaluated across numerical optimization, traffic control, joint prompt-and-code optimization, robot control, and comparison with TextGrad. Across these settings, execution traces support competitive optimization, while memory and trace access materially affect performance.
5.1 Validating with Numerical Optimization
- Trace matches Adam on synthetic numerical optimization, whereas masking the computational graph causes the feedback-only variant to struggle finding x∗.The task uses randomly generated computational graphs and 30 trials.Source: §5.1 · p. 7
5.2 Tuning Hyperparameters to Orchestrate Complex Systems
- Trace quickly becomes competitive with SCATS for traffic control, while OPRO does not; memory is crucial, but graph materialization and longer prompts add overhead.GP and PSO appear poor because 50 iterations are insufficient for convergence.Source: §5.2 · pp. 7–8
5.3 Unifying Prompts and Functions Optimization
- Trace optimizes a DSPy agent jointly across its meta-prompt, prompt-construction code, and answer-extraction code, outperforming COPRO especially on algorithmic tasks.The evaluation uses Big-Bench Hard with 15 training examples, 5 validation examples, and the remaining examples for testing.Source: §5.3 · p. 8, Table 1 · p. 8
5.4 Long-Horizon Robot Manipulator Control
- OptoPrime is the top-performing robot-controller optimizer, especially with memory, while masking execution traces significantly reduces performance and stability.The experiment uses simulated Sawyer manipulation tasks with traced 10-step rollouts and held-out initial conditions.Source: §5.4 · p. 9, Figure 6 · p. 9
5.5 Comparison with TextGrad
- Trace applies directly to the TextGrad evaluation pipeline, and OptoPrime, TextGrad, and its reimplementation achieve similar success rates across the reported reasoning experiments.The comparison uses the same LLM APIs around the same time and evaluates Solution and Prompt optimization tasks.Source: §5.5 · p. 9, Table 2 · p. 10
- OptoPrime is about 3x faster wall-clock time than TextGrad because it makes one LLM call per optimization step rather than calls linear in graph size.The experiment-time comparison includes training, validation, and testing in TextGrad’s original pipeline.Source: §5.5 · p. 9, Table 2 · p. 10
6 Limitations
Trace, OPTO, and OptoPrime remain preliminary, with limitations in representable workflows, optimizer guarantees, token efficiency, and practical problem scale.
6 Limitations
- Trace cannot currently represent stateful in-place functions or distributed and parallel workflows without modification.Asynchronous execution is theoretically possible when the overall graph remains acyclic.Source: §6 · p. 10
- OptoPrime is preliminary, lacks a provable optimality guarantee, and uses more tokens than OPRO despite favorable reported performance.The authors call for more token-efficient generative optimization algorithms.Source: §6 · p. 10
- Practical problem scale is constrained by the LLM’s debugging ability and context limits, while OPTO still requires a specified solution concept and feedback source.The paper identifies feedback design and optimality notions as open research questions.Source: §6 · p. 10
7 Conclusion and Future Work
The paper presents Trace and OptoPrime as an initial framework for end-to-end generative optimization, while identifying several directions for improving optimizer reasoning, scalability, and modality support.
- Trace converts computational workflow optimization into OPTO, and OptoPrime provides a tractable optimizer for this formulation.Source: §7 · p. 10
- OptoPrime connects optimization to LLM reasoning, motivating the use of techniques such as chain-of-thought, tool use, and multi-agent workflows.Source: §7 · p. 10
- Future work includes hybridizing LLMs with search, improving counterfactual reasoning and graph representations, and studying OPTO’s theoretical properties.Source: §7 · p. 10
- The current framework focuses on textually representable contexts and feedback, while future Trace applications may address richer non-textual inputs using modalities such as vision.Source: §7 · p. 10
A Perspective: Deep Agent Workflows
Trace frames automatic optimization of heterogeneous, interactive computational workflows as a broader analogue of differentiable neural-network optimization, while connecting OPTO to related workflow, prompt, hyperparameter, and feedback-learning settings.
A Perspective: Deep Agent Workflows
- Deep Agent Workflows are presented as computational workflows that could be automatically and scalably optimized rather than manually engineered.The paper compares this transition with the shift from hand-coded neural-network weights to AutoDiff-based optimization.Source: A Perspective: Deep Agent Workflows · p. 14
- Trace uses heterogeneous parameters and rich feedback, extending beyond AutoDiff’s numerical parameters and black-box methods’ limited scalar feedback.The paper reports greater efficiency than black-box methods in experiments by using a generalization of back-propagation.Source: A Perspective: Deep Agent Workflows · p. 14
- Trace currently excludes recursive bundle operators and distributed or parallel workflows, and tracing conditional branches involves representation trade-offs.The choice of what to trace and how to represent it remains an ongoing design problem.Source: A Perspective: Deep Agent Workflows · p. 14
- OPTO can use natural-language, image, or video feedback, but the current OptoPrime design does not handle non-textual parameters, inputs, or feedback.Effective optimization also depends on operator descriptions and an informative output-feedback source.Source: A Perspective: Deep Agent Workflows · p. 14
B Related Work
- Unlike workflow frameworks that trace for profiling or debugging, Trace uses a framework-agnostic graph representation for automatic joint optimization.Trace can be applied to workflows combining frameworks such as AutoGen, LangChain, and DSPy.Source: B Related Work · p. 15
- Compared with prompt- and code-focused optimizers, Trace targets nontrivial workflows by jointly optimizing prompts, hyperparameters, and codes with rich feedback.Its OPTO abstraction supports instantiated optimizers across problems instead of requiring problem-specific optimization prompts.Source: B Related Work · pp. 15–16
- OPTO generalizes partial monitoring and can be viewed as a language-feedback setup whose observations include execution traces.The paper uses Trace for policy learning in language-feedback problems grounded in applications.Source: B Related Work · p. 16
- The paper relates Trace to back-propagation while leaving forward-mode, checkpointing, and truncated-back-propagation analogues for future work.Source: B Related Work · p. 16
D Trace Handles Error in Execution as Feedback
Trace turns execution errors into optimization feedback by recording them as exception nodes and propagating their messages to parameters.
- Execution errors truncate the computational graph, create a special exception node, and provide feedback that an OPTO optimizer can propagate to parameters.Source: D Trace Handles Error in Execution as… · p. 17
- This mechanism supports automatic debugging and early constraint enforcement during workflow optimization.Invalid code or intermediate results can generate feedback before the workflow completes, helping avoid repeated errors.Source: D Trace Handles Error in Execution as… · p. 17
E.1 Proof of Complexity
The complexity analysis characterizes minimal-subgraph propagation and information requirements, while the broader OPTO framework raises open questions about solvability, verification, and solution existence.
E.1 Proof of Complexity
- O(WN) space and an additional WN log N time factor arise because MSP merges W propagated subgraphs of size O(N).Source: §E.1 · p. 17
E.2 Proof of Lower bounds
- In a k-digit matching problem, the minimal subgraph conveys the k bits needed to interpret feedback, avoiding 2^k possibilities without that information.Source: §E.2 · pp. 17–18
F When is OPTO Efficiently Solvable?
- OPTO’s broad contexts and feedback make its general properties and efficient solvability fundamental questions beyond this paper’s scope.Source: F When is OPTO Efficiently Solvable? · p. 18
F.1 What is a solution?
- OPTO solution verification may use a human, machine-learning model, or polynomial-time algorithm, and depends on context and output feedback rather than execution traces.Source: §F.1 · p. 18
F.2 Does a solution exist?
- The framework assumes at least one parameter can be verified as a solution; contradictory feedback under a follow-feedback context can violate this assumption.Source: §F.2 · p. 18
F.3 Can OPTO be efficiently solved?
OPTO can be difficult to solve without execution traces, while trace information may expose corrective search directions that make some instances efficiently solvable.
F.3 Can OPTO be efficiently solved?
- Without execution traces, some OPTO instances cannot be efficiently solved from context and output feedback alone.The paper frames OPTO as a search problem that can be NP-hard.Source: §F.3 · pp. 18–19
- Execution traces enable polynomial-time approximate tabular reinforcement learning and gradient computation for neural-network training.The trace reveals Markovian structure and trajectories in tabular RL and supports back-propagation for neural networks.Source: §F.3 · p. 19
- The authors conjecture that OPTO is efficiently solvable when context and trace feedback provide enough information to construct a corrective search direction.When they do not, OPTO reduces to black-box optimization; identifying efficiently solvable subsets remains open.Source: §F.3 · p. 19
G.1 Backward Step of Trace
Trace represents workflow optimization through a minimal computational subgraph, while OptoPrime’s prompts turn the resulting OPTO update into an LLM-solvable pseudo-algorithm problem.
G.1 Backward Step of Trace
- The minimal subgraph extracts the portion of the full workflow graph relevant to the optimization step.The appendix visualizes this extraction using the example in Fig. 4.Source: §G.1 · p. 19
- Trace’s optimization representation builds a computation graph and debugging report from program execution, distinct from the original program.The example maximizes z by changing x through intermediate computations.Source: Figure A.1 · p. 19
G.2 Prompts used in OptoPrime
- OptoPrime combines a representation prompt with a ReAct+CoT output prompt that requests reasoning, an answer, and a suggestion.The system prompt concatenates these two components.Source: §G.2 · p. 19, Figure A.2 · p. 20, Figure A.3 · p. 20
- Its user prompt has versions with or without memory and is populated from a template instructing the LLM to change variables according to feedback.The template defines the OPTO context instruction used to guide updates.Source: §G.2 · p. 20, Figure A.4 · p. 21, Figure A.5 · p. 21
H Comparison between Trace and TextGrad
Across heterogeneous workflow tasks, Trace uses execution-graph information and joint optimization to support broader representations than TextGrad and more informed optimization than feedback-only baselines.
H Comparison between Trace and TextGrad
- Trace jointly optimizes heterogeneous parameters and can represent numbers, floats, strings, code, prompts, and hyper-parameters, unlike string-restricted alternatives.This supports combined numerical optimization and code-plus-prompt optimization.Source: H Comparison between Trace and TextGrad · p. 22
- Trace separates the optimizer-independent TraceGraph infrastructure from OptoPrime, enabling other generative optimizers to target a shared workflow representation.The paper compares this modularity with the absence of independent computation-graph objects in TextGrad.Source: H Comparison between Trace and TextGrad · p. 22
- Trace supports arbitrary user-defined functions through bundle decorators, whereas TextGrad relies on special functions such as FormattedLLMCall.The difference concerns how workflow computations are represented to an LLM.Source: H Comparison between Trace and TextGrad · p. 22
- Trace’s scalability advantage over one-LLM-call-per-node strategies may diminish for very large graphs, motivating hybrid divide-and-conquer optimizers.The paper notes that not all computation nodes have equal optimization difficulty.Source: H Comparison between Trace and TextGrad · p. 22
I.1 Token Counts
- OptoPrime consumes significantly more tokens than OPRO, yet OPRO remains worse even with 7–10 times more iterations matched for token cost.The paper attributes OPRO’s plateau to missing trace information rather than a token limit.Source: §I.1 · p. 23
I.3 Numerical Optimization
- Trace matches Adam, the best-in-class first-order optimizer, on numerical optimization while OPRO eventually finds solutions but struggles as computation graphs become more complex.The comparison uses 30 trials over different graphs and absolute error as the metric.Source: §I.3 · p. 24
I.4 Traffic Control
- In traffic control, Trace quickly becomes competitive with SCATS while OPRO does not, although memory is crucial and Trace adds graph-materialization overhead.GP and PSO also perform poorly under the reported iteration budget, while black-box methods can improve with enough iterations.Source: §I.4 · p. 24
I.5 BigBench-Hard
- Trace jointly optimizes prompts and code on BigBench-Hard, outperforming DSPy’s COPRO by 10% on algorithmic tasks.This avoids relying on extensive output post-processing to satisfy automatic evaluation formats.Source: §I.5 · pp. 25–26, Table A.2 · p. 26
I.6 LLFBench Meta-World
Trace optimizes robot controller code over long-horizon simulator rollouts, using execution traces to improve policies across training and held-out initial conditions.
I.6 LLFBench Meta-World
- The experiments use deterministic repeated initial conditions during optimization, 30 rollout-update iterations, and evaluation on 10 held-out initial conditions.Each iteration consists of one episode rollout and one update.Source: §I.6 · p. 27
- OptoPrime is the top-performing optimizer in the Sawyer Meta-World experiments, especially when equipped with memory.The results report training and held-out testing success rates over 10 seeds.Source: §I.6 · p. 27, Figure A.8 · p. 27
- Masking execution-trace information causes significant degradation in performance and stability, supporting trace-based search directions over output feedback alone.The ablation masks inputs, intermediate values, code, and definitions from the LLM context.Source: §I.6 · pp. 27–28
- Trace learns complex controller logic in dozens of interactions and generalizes from the training initial condition to held-out testing conditions.The rollout graph contains about 30 operations and resembles back-propagation through time.Source: §I.6 · pp. 27–28
- Success rates vary substantially across random seeds, and informative task-related context is necessary for the LLM-based optimizers to work efficiently.Except for Reach, runs often find either a policy near 1.0 or one near 0.0 success rate.Source: §I.6 · p. 28
J Examples of the Optimized Parameters in the Experiments
This section illustrates optimized parameters from the Battleship, BigBench, and Meta-World experiments through learned code, prompts, and LLM responses, while documenting reproducibility and ethics-related checklist items.
J Examples of the Optimized Parameters in the Experiments
- The experiments provide examples of learned Battleship code, BigBench Predict modules and prompt templates, and Meta-World Pick-Place code and LLM responses.The figures show both starting and learned implementations, with functions sharing names learned across different iterations or trials.Source: Figure A.9 · p. 29, Figure A.10 · p. 30, Figure A.11 · p. 31
2. Limitations
- The paper provides complete proofs for its Trace complexity theorems and information-theoretic lower bound on propagated feedback.These results are proved in Appendix E.Source: §2 · p. 42
- The main experiments are described in Section 5 and Appendix I, with exact supplementary code requiring only an OpenAI API key for reproduction.The checklist states that all Trace and OptoPrime code and all reported experiments are supplied.Source: §2 · pp. 42–43
- All experiments include standard error bars computed from at least 10 replications.The checklist also notes that the compute resources are listed at the start of Section 5.Source: §2 · pp. 43–44
- The paper reports no crowdsourcing or human-subject studies and releases code for Trace and OptoPrime, but no datasets or models.The released repository includes licenses and tutorial notebooks documenting the code functionality.Source: §2 · pp. 45–46
- IRB approval or equivalent may be required for human-subjects research, depending on where the research is conducted.Source: §2 · p. 47
- Human-subjects research procedures can vary substantially across institutions and locations, requiring adherence to institutional guidelines and the NeurIPS Code of Ethics.Source: §2 · p. 47
- Initial submissions should omit information that could break anonymity, when applicable.Source: §2 · p. 47