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

Authors: ching-an  cheng
allen  nie
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 20:58:06

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

Main paper: 
 introduction: 
 Computational workflows that integrate large language models (LLMs), machine learning (ML) models, orchestration, retrievers, tools, etc., power many state-of-the-art AI applications [1] : from chatbots [2] , coding assistants [3] , robots [4] , to multi-agent systems [5] . However designing a computational workflow requires laborious engineering because many heterogeneous parameters (e.g. prompts, orchestration code, and ML hyper-parameters) are involved. Moreover, after deployment any erroneous behaviors of the workflow persist unless a developer manually updates it.We study a class of optimization problems motivated by automating the design and update of computational workflows. Computational workflows produce optimization problems with heterogeneous parameters, rich feedback (e.g. console output and user's verbal responses), and intricate objectives (beyond maximizing a score). Moreover, a workflow can have interdependent steps (e.g. adaptive orchestration, feedback control loops) and/or involve non-differentiable, semi-black-box, stochastic operations (e.g. ML models, simulations) whose behavior cannot be succinctly captured. As a result, the structure of the computation may change as the parameters and the inputs of the workflow vary.Due to its complexity, computational workflow optimization is usually framed as a black-box [6] or algorithm configuration [7] problem, and is tackled by general techniques like Bayesian Optimization [8] , Evolutionary Algorithms [9] , Reinforcement Learning (RL) [10] using scalar scores as feedback. But one observation of scalar feedback alone does not provide an improvement signal, so these algorithms are very inefficient when the parameter space is large (e.g. codes or natural language prompts). Recently LLM-based optimizers [11] [12] [13] [14] [15] [16] have been proposed as generative optimizers to improve efficiency, leveraging the prior of generative models learned from large pre-training corpora to optimize complex prompts and codes. In this paper, we investigate how generative optimization can be applied more broadly and systematically to optimize a general computational workflow end-to-end. Appendix B discusses related works in generative optimization.Inspired by back-propagation [17] , we take an end-to-end approach to computational workflow optimization. AutoDiff frameworks [18, 19] have scaled the back-propagation algorithm to optimize differentiable workflows (e.g., neural networks) with billions of parameters. We extend the idea of AutoDiff and design a new framework, Trace, for jointly optimizing all of the heterogeneous parameters in general computational workflows, which may not be differentiable.Trace treats a computational workflow as a graph like a neural network, where nodes are either inputs, parameters or the results of computation steps, and directional edges denote how nodes are created from others. But, instead of gradients, Trace propagates the execution trace of a workflow (which records the intermediate computed results and how they are used), via the notion of minimal subgraph (see Section 3.3) . We show that propagating the execution trace subsumes back-propagation for differentiable workflows, and remains applicable even for non-differentiable workflows. Viewing a workflow as a graph and using its execution trace is standard for software engineering; for instance, human developers use such traces to debug distributed systems [20] . Our novel insight is that execution traces also enable end-to-end generative optimization, because they provide the information needed to relate rich feedback to the parameters in general workflows.Trace is available as a Python library with an API inspired by PyTorch [19] . A user declares the parameters to be optimized using a trainable flag, decorates the workflow with node and bundle wrappers, and runs a Trace optimizer -just like how they would declare and train neural networks.Consider building an AI agent for the Battleship game (Fig. 1 ). The agent's policy (Fig. 2a ) has two components (reason and act), which are chained together to react to different board configurations. The Battleship environment provides feedback (binary reward in texts) to tell if the agent's action hit the hidden ships, and the goal is to hit all hidden ships as fast as possible. Consider how a human programmer might approach the problem. They may run the policy and change the code based on the def forward ( self , map ): plan = self . reason ( map ) output = self . act ( map , plan ) return output @trace . bundle ( trainable = True ) def reason ( self , map ) -> str : """ Given a Battleship map , analyze the game ... """ return [0 , 0] @trace . bundle ( trainable = True ) def act ( self , map , plan ):""" Given a map and plan , select a target coordinate ... """ return (a) We write a trainable policy in Python using Trace operators. 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 ". . . " and list them in Appendix J. Trace then builds a DAG as the workflow executes and updates the parameters (see Fig. 1 for the result).observed feedback. They may rewrite the code a few times to try different heuristics to solve this problem. They will fix any execution errors (e.g., out-of-bounds exceptions) by using stack traces.Our Trace framework accomplishes the programmer's goal automatically without adding complexity to the Python code. The user declares reason and act as trainable (Fig. 2a ) and then runs the agent in a PyTorch-like training loop (Fig. 2b ). During the execution, Trace records a directed acyclic graph (DAG) (Fig. 2c ) and uses it to compute the execution trace for optimization. Trace also automatically catches errors (e.g., syntax/semantic errors) and can use them as feedback. In Fig. 1 , we show what the agent learns as Trace optimizes2 its policy, where the learned policy is evaluated on new randomly generated games for generalization. The agent can quickly improve its performance and learn strategies that are increasingly complex. Remarkably, there is no mention of the specific Battleship environment API, nor details on how the functions reason and act should behave or adapt in Fig. 2a . The Trace optimizer figures out all the details through interactions as the computational graph unfolds and the feedback on the output is observed. Beyond code as parameters in this example, we also have experiments in Section 5 where prompts and other heterogenous parameters are optimized.The design of Trace is based on a new mathematical setup of iterative optimization, which we call Optimization with Trace Oracle (OPTO). In OPTO, an optimizer selects parameters and receives a computational graph as well as feedback on the computed output. Trace is a tool to efficiently convert the optimization of computational workflows into OPTO problems in practice.We argue that framing computational workflow optimization as OPTO can lead to faster convergence than a black-box approach. We present a constructive proof: We design a general-purpose efficient generative optimizer called OptoPrime, for OPTO. OptoPrime turns OPTO to a sequence of pseudoalgorithm problems. In each iteration of OPTO, we format the execution trace and output feedback as a pseudo-algorithm question and present it to an LLM for solution (GPT-4 using a ReAct-CoT prompt listed in Appendix G). In experiments, we apply OptoPrime unchanged to many disparate applications like prompt optimization, first-order numerical optimization, hyper-parameter tuning, and robot controller design. We find that the general purpose OptoPrime is competitive with specialized optimizers for each domain, e.g., achieving 10% higher accuracy on BigBenchHard [21] when optimizing a DSPy [22] program compared to their hand-designed optimizer.Trace, OPTO, and OptoPrime together provide the first tractable algorithm for optimizing general computational workflows end-to-end. The Trace framework a) leverages a workflow's structure and b) can incorporate rich feedback beyond scores, extending AutoDiff to complicated, non-differentiable computational workflows. With Trace, we conjecture that "training deep agent networks" (which fluidly mix computation of tensors, LLMs, and other programmable tools) will soon be possible. 
optimization with trace oracle: 
 OPTO is the foundation of Trace. In this section, we define this graph-based abstraction of iterative optimization and discuss how OPTO covers various computational workflow optimization problems.Preliminary We review the definition of a computational graph (see Fig. 2c ). A computational graph g is a DAG, where a node represents an object (such as tensors, strings, etc.) and an edge denotes an input-output relationship. We call a node without parents a root and a node without children a leaf, which are the inputs and outputs of the computational graph. In the context of optimization, some inputs are marked as trainable parameters, which are denoted as {X θ }. For a node X, its parents are the inputs to an operator that creates X. The descendants of node X are those that can be reached from X following the directed edges; the ancestors are defined conversely. Without loss of generality, we suppose that all computational operators have a unitary output 3 . In this way, we can associate the operator that creates the child node with the child node, and the full computation can be represented compactly as a DAG without explicitly representing the operators.The execution trace of a computational workflow is defined as the sequence of operations and their execution results invoked when computing the output from a set of inputs; execution traces can be expressed as a computational graph as defined above. OPTO is an abstract setup of iterative computational workflow optimization. An OPTO problem instance is defined by a tuple (Θ, ω, T ), where Θ is the parameter space, ω is the context of the problem, and T is a Trace Oracle. In each iteration, the optimizer selects a parameter θ ∈ Θ, which can be heterogeneous. Then the Trace Oracle T returns a trace feedback, denoted as τ = (f, g), where g is the execution trace represented as a DAG (the parameter is contained in the root nodes of g), and f is the feedback provided to exactly one of the output nodes of g. Finally, the optimizer uses the trace feedback τ to update the parameter according to the context ω and proceeds to the next iteration, as shown in Fig. 3 .Output FeedbackIn OPTO, the output feedback f is generic, such as scores, gradients, hints/explanations expressed in natural language, and console messages. The context ω provides invariant information to interpret the output feedback f as well as any known side-information, e.g., desired properties of the parameters. The context ω is fixed for an OPTO problem instance (similar to an instruction, or a problem definition), whereas the output feedback f can change with the parameter θ ∈ Θ and the resulting computation g. For example, ω may be "Minimize a loss function" and f is a loss. Alternatively, ω can be open-ended, like "Follow the feedback" and f describes how an output should be changed. In Section 3.2, we discuss how to define the context and output feedback when constructing OPTO problems in practice. In this paper, we focus on OPTO problems where f and ω can be expressed compactly in text. This covers a wide range of problems [23] , including those with scalar feedback.OPTO differs from a black-box setup in that the execution trace g shows the computational path toward the output, which provides information to construct a parameter update direction from f and ω. In the minimization example above, when the execution trace g is missing, it is unclear how the parameter can be improved given only a point evaluation of f . On the other hand, with g documenting the functions used to create the output, an update direction (e.g., a gradient) can be derived. We highlight that the structure of the computational graph g returned by the Trace Oracle T can be different in each iteration in the general case (as in Fig. 3 ) because some workflows can change with different inputs and parameters.To ground the OPTO setup, we show how OPTO is related to some existing problems with examples. We discuss other examples like hyperparameter tuning and multi-agent systems in Appendix C.Example 1 (Neural network with back-propagation). The parameters are the weights. g is the neural computational graph and f is the loss. An example context ω can be "Minimize loss". The back-propagation algorithm can be embedded in the OPTO optimizer, e.g., an OPTO optimizer can use τ to compute the propagated gradient at each parameter, and apply a gradient descent update.Example 2 (RL). The parameters are the policy. g is the trajectory (of states, actions, rewards) resulting from running the policy in a Markov decision process; i.e., g documents the graphical model of how a generated action is applied to the transition, which then returns the observation and reward. f can be a success flag. ω can be "Maximize return (cumulative rewards)" or "Maximize success". Example 3 (Prompt Optimization of an LLM Agent). The parameters are the prompt of an LLM workflow. g is the computational graph of the agent and f is the feedback about the agent's behavior (which can be scores or natural language). ω can be "Maximize score" or "Follow the feedback". 
trace: the next autodiff: 
 We design a framework, Trace, to bring OPTO from an abstract concept to reality. Trace provides a lightweight Python4 tool to implement the Trace Oracle of OPTO for optimizing computational workflows. Through the OPTO framing, Trace separates the design of optimizers and domain-specific components so that general-purpose optimizers can be built that work across diverse domains.Trace is designed based on two primitives:• node is the wrapper of Python objects. When wrapped, a Python object is registered as a unique node in the computational graph of Trace. A node can be set trainable, which would make the node a parameter in OPTO. In addition, when using node to declare a parameter, one can also describe (in natural language) constraints that the parameter should obey. • bundle is the decorator to turn Python methods into operators. When a function is decorated, its docstring and source code are recorded as the definition of the operator, which infer how the output feedback should change the parameters. Moreover, functions decorated by bundle can be set trainable as well, which means that the code of the decorated method becomes a parameter. 5For any workflow, using Trace involves the following steps (see Fig. 2 ). First, the user declares the workflow's parameters using node and bundle, and defines the workflow's conceptual blocks as operators using bundle. Then the user creates an OPTO optimizer (such as OptoPrime in Section 4), and optionally provides the context ω for the problem. (A default context ω of OptoPrime is "Follow the feedback"). In addition, the user defines a mechanism to provide feedback to the computed result (e.g., scores, natural language suggestions, etc.), in analogy to defining a loss function in neural network training. After the declaration, optimization via Trace repeats the following: 1) Execute the decorated workflow. As it runs, a DAG is built in the backend, logging the computed results and their connections. 2) Initiate the propagation of the output feedback to the parameters by calling backward.(Any execution error is also treated as feedback; see Appendix D.) Internally, Trace extracts the minimal subgraph g connecting the parameters and the output and sends the OPTO optimizer the trace feedback τ = (f, g). 3) Call the OPTO optimizer's step method to update the parameters.There are many ways to represent a computational workflow as a computational graph, from abstracting the entire process as one big operator to listing all low-level steps as operators in the graph. In Trace, the level of abstraction is decided by how bundle is applied, as all steps underneath bundle are abstracted as one operator. The design of bundle allows tracing most Python codes, except for those modifying the content of an object reference in place. However, such a case can be avoided by first duplicating the object and then applying the modification to the copied object, similar to how a recurrent neural network is implemented.Different abstraction choices trade off the graph complexity and the description needed for each operator. Abstracting everything into a single operator makes a trivial graph but requires more descriptions to faithfully capture the workflow. On the other hand, not all details matter in optimization, so exposing every low-level operator can make the graph unnecessarily cluttered. The best representation depends on the application and OPTO optimizer at hand. This problem, we believe, is similar to the design of neural network architectures. Here, we suggest defining the operators to mimic the whiteboard system diagram of the computational workflow. This level of abstraction in our experiments strikes a good balance between the ease of documentation and the graph size.Apart from architecture design, another design question is what information goes into the context ω versus the description of each operator? For a single problem, there is no difference in principle; one can choose to provide details of all operators in g through the context ω. However, this will require manually crafting a context for every workflow. We suggest instead providing a description of the operators when they are defined using bundle. Then Trace will automatically generate the workflow-specific information while the same context ω is shared across many workflows. f ← fi // all fi are the same. 5: return {p : (f, g) for p in node.parents}Trace uses a recursive graph traversal algorithm (Algorithm 1) to propagate feedback in the reversed topological ordering. With different propagators, Algorithm 1 can implement various forward-backward schemes including back-propagation. 6 We propose a general propagator, Minimal Subgraph Propagator (MSP), in Algorithm 2. MSP propagates the trace feedback τ = (f, g), where g is implemented as a priority queue. Running Algorithm 1 with MSP (Algorithm 2) together implements the Trace Oracle of OPTO, which extracts the minimal subgraph between parameter nodes and output. 7 connecting the parameters and an output. Appendix E proves the following: Theorem 1. For a graph with N nodes and maximum degree W , Algorithms 1 and 2 have time complexity O(W N 2 log N ) and space complexity 8 O(W N ).By contrast, back-propagation has time and space complexities of O(W N d 2 ) and O(d) respectively, where d is the maximal dimension of tensors. The difference is because in the most general setting of computational graphs and feedback, the propagated feedback (no matter how it is represented) does not have a constant size and needs the full subgraph.Theorem 2. For generic computational graphs of N nodes, in the worst case, the propagated feedback needs a description length Ω(N ) to construct an improvement direction.Despite the worst case complexity of MSP, in practice the difference is negligible. Since MSP only involves merging priority queues of references, most actual computation happens in the forward pass (and also the optimizer's step method). However for very large problems with millions of nodes in the minimal subgraph, we anticipate that computational issues of MSP could arise. 
design of the first opto optimizer: 
 We introduce an LLM-based generative optimization algorithm OptoPrime for any text-based OPTO problem. We believe that this is one of many possible optimization algorithms for these problems and there is a large space to be explored for identifying efficient optimization methods for OPTO.Subgraph Representation One core challenge of designing an LLM-based OPTO optimizer is how to represent the execution trace subgraph g (which can involve various graph structures and heterogenous data) to LLMs, in a way that LLMs can understand and reason about the downstream effects of parameter updates. We leverage the LLMs' remarkable coding and debugging ability [3] . 6 In back-propagation, the message is the gradient ∇i and the propagate function returns J ⊤ i j ∇j to its ith parent, where Ji is the Jacobian to the ith parent and and the ∇j gradient received from the jth child. 7 The minimal subgraph gX→Y between nodes X and a node Y is defined as gX,Y := X {Y } {Z, P |Z ∈ ancestors(Y ), Z ∈ descendants(X), X ∈ X , P ∈ parents(Z), P ∈ parents(Y )}. 8 The space complexity refers to the extra space needed for the backward pass, not including the forward pass.# Code : a = bar ( x ) y = add (b , a ) z = mul (a , y ) # Definitions :[ mul ] This is a multiply operator .[ add ] This is an add operator .[ bar ] This is a method that does negative scaling . # Inputs : b =1.0 # Others : a =2.0 y =3.0 # Output z =6.0 # Variable x = -1.0 # Feedback : Output should be larger . We present the trace feedback (f, g) computed by Trace as a pseudo-algorithm problem: the subgraph g is expressed as a report of code execution with information about the computed values and descriptions of operators in g. Then we prompt the LLM to update the parameters in g based on feedback f given to the output. Fig. 4 shows an example. It is crucial to note that even though the lines look like an actual program, it is not the real program but the computational graph defined by bundle (see Section 3.2).Parameter Update We prompt the LLM with a ReAct-CoT style prompt (Appendix G.2) in one query, asking it to generate reasoning based on the graph, and a suggestion on the parameter changes. If the suggestion can be extracted from the response, we update the parameters.Optimization Memory OptoPrime optimizes most workflows reasonably well using just the traced graph and feedback, but it can run into issues when single feedback alone is not informative (e.g., the output feedback is rewards, but there is no description of how the rewards are generated). For robustness, we have a basic memory module in OptoPrime, which tracks the past parameter-feedback pairs as in-context examples. See Appendix G for details. 
experiments: 
 We evaluate the Trace framework with OptoPrime. We implement the state-of-the-art LLM optimizer OPRO [13] as a baseline; in comparison with OptoPrime, OPRO does not use the execution trace but relies on the memory of parameter and feedback pairs. For these experiments, we use GPT-4-0125-Preview. We run the experiments on a standard PC with 16 GB RAM, and Trace introduces no measurable overhead on executing the workflow. We also conduct experiments to compare Trace and OptoPrime with a concurrent AutoDiff-like framework, TextGrad [24] , which was released after Trace was submitted to NeurIPS. We show that TextGrad can be easily implemented as an optimizer in Trace, and OptoPrime achieves similar or better performance than TextGrad while using much less computation time. In the rest of this section, we will denote Trace+OptoPrime simply as Trace. We report the token usages of all approaches in all experiments in Appendix I.1.First, we want to validate if OptoPrime can solve classical differentiable optimization problems, since they are a special case of OPTO. Consider the problem of min x |h(x) -y * | for a target y * . We construct a synthetic task environment that randomly creates y * and the computational graph of h with arbitrarily complex connections between numerical variables (see Appendix I.3 for details). We evaluate OptoPrime (denoted as Trace) and a variant that does not see the graph (Trace Masked); both the optimizers do not use memory. The output feedback is "The output should be <larger/smaller>" (this feedback has the same information as the gradient w.r.t. h). We compare also the performance of Adam optimizer [25] . We run 30 trials over different randomly generated problems. All methods see the same randomness. Trace is able to match the best-in-class Adam; on the other hand, without access to the full computational graph, the feedback-alone optimizer struggles to find x * (Figure 5a ).We tested Trace in a traffic control problem, which is an instance of hyper-parameter tuning. We used UXSim [26] to simulate traffic at a four-way intersection, where the trainable parameters are 2 integers in [15, 90] , which are the green light duration for each direction of traffic flow. The feedback is the estimated delay experienced by all vehicles due to intersections, and the goal of an optimizer is to minimize the delay using the fewest number of traffic simulations. To this end, this optimizer must find the right trade-off for temporally distributed and variable demands. In Fig. 5 we report the performance of a SOTA heuristic from the traffic control literature, SCATS [27] as well as two black-box optimization techniques: Gaussian Process Minimization (GP) [8] and Particle Swarm Optimization (PSO) [28] . All methods use the same starting parameters. Trace denotes OptoPrime using memory, and Trace NoMem denotes OptoPrime without memory. We report further details in Appendix I.4. GP and PSO appear bad because 50 iterations are insufficient for their convergence; given enough iterations, both will eventually perform well. Trace is quickly competitive with the SCATS heuristic, whereas OPRO is not. Moreover, we find that memory is crucial for Trace to perform well for this task. But we note that Trace consumes extra overhead compared to other methods, since Trace has to materialize the resulting computation graph and query an LLM with effectively a longer prompt than that of OPRO.Many LLM agents today, e.g., specified by LangChain [29] and DSPy [22] , have many components. These libraries provide optimization tools to optimize a small portion of their workflows, predominantly the prompt that goes into an LLM call. However, for building self-adapting agents that can modify their own behavior, only allowing the change to one part of a workflow but not others can be limiting. In this experiment, we test Trace's ability in joint prompt optimization and code generation. Specifically, we optimize a given DSPy-based LLM agent and tune its three components: the meta-prompt prompt_template, a function create_prompt that modifies the prompt with the current question, and a function extract_answer that post-processes the output of an LLM call.We set up an end-to-end prompt-and-code optimization pipeline. We use an automatic evaluation function to compare the LLM's output with the ground truth, which requires the LLM agent to generate outputs not only with the correct answer but also in the correct format (following the guidelines of [30] ). We use the Big-Bench Hard dataset [21] (15 examples for training, 5 for validation, and the rest for testing). We compare Trace with DSPy's COPRO module (which optimizes the meta-prompt).In Table We test the ability of Trace to optimize long-horizon workflows with complex dependencies and to "back-propagate through time". We use Trace to train the controller code (in Python) for a simulated Sawyer robot manipulator. We use the Meta-World environment from LLF-Bench [23] as the simulator and consider three tasks: Reach, Pick-place, and Push. For each task, LLF-Bench provides a task instruction and meaning of the action space, which we use as the context ω of the OPTO problem. The observation is a Python dictionary of vectors, indicating the end-effector position, the goal position, the gripper status, etc. The action space is a 4-dimensional vector to control the relative position of the end-effector and the gripper state. In each time step, the LLF-Bench Meta-World simulator returns the observation and natural language feedback to guide the robot. An episode ends if the robot successfully solves the problem or because of time-out. We consider an episodic training setting. The initial condition for all iterations in training is the same. We evaluate the learned policy in terms of success, starting from 10 held-out initial conditions. The task horizon is 10 steps, which is sufficient for task completion, and each training iteration has one rollout. The output feedback in OPTO is a textual representation of task success. In addition to the controller code, we also decorated the reset and step functions of the gym environment so that the entire rollout can be traced end-to-end. We compare Trace with OPRO; because of the streaming OPTO setting, our OPRO implementation only proposes one candidate in each iteration, which is then evaluated and provided with the output feedback.The experimental results are summarized in Fig. 6 . Trace denotes OptoPrime using memory, and Trace NoMem denotes OptoPrime without memory. We show learned code in Appendix J. OptoPrime is clearly the top-performing optimizer, especially the version with memory. OPRO is able to solve Reach at the start, but its performance degraded over iterations (this instability was mentioned in [13] ) and gets a similar performance as OptoPrime (without memory) in Push. To validate that the performance of OptoPrime is indeed due to using the execution trace, we include an ablation where we mask out the execution trace, which leads to a significant decline in performance and stability. This experiment features the most complex graph structures among all the experiments. The experimental results here are quite impressive, showing that Trace is able to learn a sophisticated control logic in dozens of interactions, not only working on the training initial conditions but also on held-out testing ones too. We discuss some limitations in Appendix I.6.After the submission of our work to NeurIPS, another AutoDiff-like framework, TextGrad [24] , was released, which shares the same goal of end-to-end optimizing AI workflows as Trace. In comparison, TextGrad propagates text feedback, whereas Trace propagates minimal subgraphs (see Section 3.3). The graph-based design of Trace, which separates the tracing infrastructure and optimization algorithms, makes it more flexible. In fact, we easily implemented TextGrad as an optimizer in the Trace framework, but the reverse is not possible (because TextGrad couples the infrastructure and the optimization algorithm together). In addition, unlike TextGrad, Trace supports jointly optimizing heterogeneous parameters and can be applied to directly trace a given computational workflow without the need to rewrite the workflow using pre-defined templates. Please see Appendix H for more discussion comparing the two frameworks.In this experiment, we apply Trace to directly decorate the evaluation code released with the TextGrad library and optimize the parameters following their training/evaluation pipeline line-by-line. This experimental design makes the comparison fair by allowing each optimizer to access the same LLM APIs around the same time, and showcases the flexibility of Trace framework to optimize any computational workflow. We pick the Solution optimization [24, Table 2 ] and Prompt optimization [24, Table 3 ] for the reasoning tasks experiments. Please see [24] for details on the exact setup. We compare OptoPrime, TextGrad9 , and a reimplementation10 of TextGrad as an optimizer in Trace.We find that all these algorithms achieve similar success rates in these experiments. One noticeable difference is that OptoPrime is about 3x faster wall-clock time than TextGrad since OptoPrime makes a single call to LLM in each optimization step, whereas TextGrad calls linear to the graph's size.  
limitations: 
 We highlight that Trace, OPTO and OptoPrime are a first step towards end-to-end generative optimization and building self-adapting workflows. They have limitations in their current form. OPTO captures rich feedback, but it is important to specify a solution concept as well as the feedback source. We provide guidance for feedback design in Section 3.2 and discuss notions of optimality in Appendix F. We believe designing feedback will be as important as designing loss function in deep learning, both of which are open research questions. Also, Trace cannot convert all computational workflows into OPTO problems, e.g., stateful functions that modify their state in place cannot be represented as a DAG without modification, and distributed/parallel computing workflows are incompatible with the current implementation (though in theory Trace can run in an asynchronous way so long as the overall graph does not end up with cycles). Finally, while Trace is designed to be generic and future-proof, the OptoPrime optimizer is preliminary. Although we demonstrated that OptoPrime could work well with moderate-size graphs, it is not a provably optimal algorithm and uses more tokens than OPRO, though, in our experiments, OPRO's performance does not improve even when given a large token budget. The debugging ability and context limits of the LLM used in OptoPrime crucially determine the scale of problems that we can practically address today. Consequently, more research is needed for designing token-efficient generative optimization algorithms. 
conclusion and future work: 
 We created Trace that can convert a computational workflow optimization problem into an OPTO problem, and we demonstrated a tractable OPTO optimizer, OptoPrime. This is just a first step towards a new paradigm of end-to-end generative optimization, with exciting avenues for future work. We discuss a few selected ones below. Please see Appendix A for a longer discussion.In OptoPrime, we connect optimization to an LLM's reasoning capability. Techniques that have been proposed to improve LLM reasoning, e.g. Chain-of-Thought [31] , Few-Shot Prompting [32] , Tool Use [33] , and Multi-Agent Workflows [5] could also help improve OptoPrime or design new OPTO optimizers. We conjecture that a hybrid workflow of LLM and search algorithms, can enable a truly general-purpose OPTO optimizer. Along the way, we must settle on how to delineate the agent vs. the optimizer. How to trade off the generality of optimizer vs. crafting side-information in the context ω to achieve task-specific performance is an open question.In Trace, we chose a specific propagator (MSP), which maximally preserves information for a general computation graph. We can instead specialize it for specific computations, e.g. to accommodate very large graphs with a hierarchical graph representation. Going a step beyond the basic memory module we experimented with in OptoPrime, we anticipate that an optimizer that can reason about how a workflow will behave under counterfactual parameter settings can be more efficient than OptoPrime and can enable a divide-and-conquer approach to OPTO. More research is needed to study the theoretical properties of OPTO (such as optimization landscape and complexity). We hope our preliminary effort in Appendix F can provide some guidance. Finally, in this paper, we focused on output feedback and context that can be compactly textualized. We anticipate that computational workflows with rich non-textual contexts and output feedback will also benefit from automatic generative optimization through appropriate applications of Trace (e.g., with VLMs). 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes]Justification: All the experiments use the same compute resources (LLM API for OptoPrime and machine to run Trace), which are listed at the start of Section 5.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: In Sections 6 and 7 we discuss both the aspirational goals and their broader impacts.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: There are no models, data or APIs released along with the paper. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: Datasets used in the experiments (e.g. BigBench-Hard) are properly attributed. Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: We posit that the current practice of manually engineering computational workflows to build AI systems is analogous to programmers in the early 2000s hand-coding neural network weights to create engaging AI characters in video games [34] . Just like AutoDiff enabled the automatic and scalable optimization of deep neural networks with billions of parameters, we believe that Trace is the first step towards automatic and scalable optimization of "Deep Agent Workflows" to power even more capable AI systems. However, there are several limitations of the current implementation of Trace that need to be addressed to build Deep Agent Workflows.When designing interactive AI systems that learn from their interactions, we need to define the parameters and feedback of the system. Parameters are the internal attributes that can be updated by the learning algorithm employed by the system. Feedback are the things observed and recorded by the system as a product of its interactions, and that provide signal for learning. Trace enables the development of new learning algorithms (e.g. through OptoPrime) that incorporate rich feedback to update heterogenous parameters. In contrast, AutoDiff for deep neural networks uses numerical feedback (e.g. rewards or loss functions) to optimize numerical parameters (e.g. tensors). Black-box optimization techniques (e.g. Reinforcement Learning) can use numerical feedback to optimize heterogenous parameters (e.g. codes, hyper-parameters as well as tensors), though they are inefficient.In the experiments, we saw that Trace was more efficient than black-box methods by using a generalization of back-propagation. Finally, Trace can use rich feedback (e.g. language) to extract more signal for learning.What can be traced? Trace cannot convert all computational workflows into OPTO problems.Workflows with recursive bundle operators or those requiring distributed/parallel computing are not compatible with the current implementation. A future work would be to expand the Trace implementation to support these scenarios. In addition, the current implementation of Trace does not trace the execution within an operator defined by bundle, though in principle this is possible.There are also ambiguities in how an existing workflow can be traced and represented as a DAG. One example is when there is some sub-workflow following an if condition and another one following else. One can choose to wrap the entire code, including if and else, by bundle as a single operator. On the other hand, one can also just wrap the sub-workflows and not trace the if condition nor represent it as part of the DAG. (That is, suppose the if condition is true; from the DAG, one cannot see the alternate path under else). The latter choice has a flavor of applying "stop-gradient" on the boolean condition, whereas the first choice enables back-propagation through also the logical condition. We summarized some design considerations in Section 3.2. We foresee the choice of what to trace and how to trace in building Deep Agent Workflows will be an on-going research problem, similar to neural network architecture design.Where do we get rich feedback? The OPTO framework captures an abstraction of rich feedback, called trace feedback (the execution trace and the output feedback), but it requires specification of operator descriptions and output feedback source to guide the optimization effectively. OPTO can be more efficiently solvable than black-box problems only when the trace feedback provides information beyond reward signals; otherwise, information-wise, OPTO is no easier than black-box problems. The Trace framework automates the generation of the execution trace in OPTO, when users of Trace decorate the workflow end-to-end. Currently, OptoPrime uses the docstring of operators in the execution trace to understand the operators; nonetheless, Trace logs also the source code of the decorated methods, which can also be used in the future to design optimizers that uses more details of the operators. In our experiments, we focus on output feedback that is automatically generated. In other contexts, e.g. users interacting with chatbots, we can natively gather natural language feedback or synthesize feedback [16] based on raw observations. We also anticipate that feedback in the form of images (e.g. users' gestures) or videos (e.g. videogame player showing a desired correction to agent's behavior) will be readily available and can be used as the output feedback in the OPTO framework. While Trace can handle this feedback, the current design of OptoPrime does not handle non-textual content (either in parameters, inputs or feedbacks). We anticipate future work along the lines of [16] on enhancing feedback design and developing guidelines for designing more informative and directive optimizer and feedback mechanisms.How to design adaptive optimizer for general OPTO problems? The proposed OptoPrime optimizer shows the possibility of designing a single optimization algorithm to solve a range OPTO problems from diverse domains. However, we remark that OptoPrime is just the first step; it is akin to the vanilla gradient descent algorithm, which shows a proof of concept but is not scalable for large problems. The current design of OptoPrime has several scalability limitations due to its summary approach that updates parameters through one call to an LLM. While we show in Section 5. Framework for Computational Workflows Frameworks such as LangChain [29] , Semantic Kernels, AutoGen [5] , DSPy [22] allow composing computational workflows and provide handengineered optimizers to tune an LLM's context (i.e. prompt templates, few shot examples, or tool libraries) using scalar feedback with black-box search techniques. They support tracing of the workflow to aid in profiling, debugging and visualization. In contrast, Trace uses tracing for automatic optimization, and constructs a different representation of the computational graph which is suited for that purpose. Moreover, Trace is designed to be general-purpose and agnostic to the underlying frameworks of computational workflows users choose. In principle, one can apply Trace to decorate and tune a workflow based on a mix of Autogen, LangChain, DSPy codes. In fact, our experiments in Appendix I use workflows declared using both AutoGen and DSPy.There are multiple efforts to optimize the computational graph of LLM workflows, which is a special case of the OPTO problem. These algorithms focus on optimizing prompts. SAMMO [14] is an example for prompts that uses additional graph structure to make the optimizer efficient. SAMMO represents the prompt parameter itself as a program so as to enable more efficient black-box search through the space of programs. DSPy [22] can optimize directly the prompts or the few-shot examples to include using scalar reward feedback. Retroformer [35] uses another small language model (LM) to provide suggestions/feedback (i.e. changing prompts) to improve the behavior of an actor LLM, where the small LM is tuned by offline RL. Deep Language Networks [36] view all of the prompts in an LLM workflow as tunable parameters and jointly optimizes them. They discovered that optimizing each parameter in isolation instead produces subpar results. [37] frames LLM systems as graph where nodes are operations and edges are messages/connections. (Note that this is different from the DAG used in Trace; here nodes are messages and edges are input-output of operators) and optimizes for the connection on edges (binary variables) by REINFORCE using scalar reward feedback and prompts by LLMs . They optimize each component separately without considering each other; for example, the prompts are optimized individually without considering the graph topology or how they are used down the road. We suspect this approach can be less stable. Their prompt optimization part also does not take output feedback, but simply use an LLM to self-check whether the prompt meets the need of generating desired functions the user specified. In contrast to these works, through the OPTO framing, Trace supports joint optimization of all parameters (prompts, hyperparameters, codes) with rich feedback, and is agnostic to graph structures (e.g., changing these parameters can dynamically change the graph structure and connections between nodes). Users of Trace are free to specify which parameters they want to automatically optimize via online interactions.There is a huge and fast growing literature on using LLMs as optimizers to improve prompts [12, 11, [38] [39] [40] or codes [41, 15, 42, 38] . Different from the works mentioned above, here the focus has been on an isolated problem (e.g., changing the behavior of a single LLM or improving the code generation in the question-answering format) rather than considering a non-trivial workflow or agent with multiple components like above. They do not consider optimizing prompts or codes as one component of a bigger workflow (e.g. implementing an autonomous agent), which is harder and requires the right credit assignment. In addition, these LLM-based optimizers, including OPRO [13] , often propose only principles of how prompts should be designed and requires crafting problem specific prompts (as opposed to a single optimization prompt that can be applied to different problems). For adapting them to new problems, users need to design new prompts. Trace can also be applied to optimize trivial OPTO problems where the returned graph has just a single node of the parameter (which are the scenarios considered by these works). Nonetheless, the main focus of this paper is to study how optimization can be done efficiently as the graph becomes nontrivial and for diverse applications. Trace achieves this by using the abstract OPTO problem framing. Since OPTO encapsulates domain specific info in the graph, it enables designing fully instantiated optimizers that can be applied to multiple problems, rather than just principles which then requires hand crafting prompts for individual problems like previous works.Recent works like [43, 44] use LLMs to optimize numerical hyperparameters, as an alternate to Bayesian optimization. Here in the experiments we show that Trace + OptoPrime also can effectively learn hyperparamters, faster than Bayesian optimization.The main difference between Trace and the aforementioned work is the representation of the problem. In Trace, we provide the graph to the LLM-based optimization (through the pseudo-algorithm representation), and we consume rich language feedbacks on the output, both of which accelerates hyper-parameter optimization.OPTO Related Setups OPTO is a generalization of partial monitoring games [45] . If there exists a latent loss function that the feedback f adheres to (e.g. as in [23] ), those OPTO instances can be written as partial monitoring game. However OPTO admits a more general notion of feedback f , and we discuss solution concepts for them in Appendix F. On the other hand, OPTO can be also viewed as a special case of Learning from Language Feedback (LLF) setup defined in [23] with observations as the trace feedback. This is a framing of a meta LLF problem. In the LLFBench Meta-World experiments of this paper (Section 5), we show Trace can be used to learn policy for LLF problems grounded to an application too.AutoDiff and Back-propagation Back-propagation has been shown to be a very effective tool in optimizing differential computational workflows. Our design of Trace is inspired by back-propagation and the ease of use of the AutoDiff framework PyTorch [19] . Nonetheless, we highlight that backpropagation (Backward Mode Differentiation) is not the only AutoDiff algorithm. For example, the gradient can be computed in a forward mode (Forward Mode Differentiation) as well, and there are also techniques of Checkpointing [46] and Truncated Back-Propagation approximation [47] for efficiency. What are the equivalent ideas of these methods for general computational workflows? We think this is an interesting future research direction.To ground the OPTO setup, we show how OPTO is related to some existing problems with examples.Example 4 (Neural network with back-propagation). The parameters are the weights. g is the neural computational graph and f is the loss. An example context ω can be "Minimize loss". The back-propagation algorithm, in view of the OPTO formulation, is embedded in the OPTO optimizer. For example, an OPTO optimizer here is a composition of back-propagation and gradient descent, where back-propagation takes τ to compute the propagated gradient at the parameter.Example 5 (Code Debugging). The parameters are the codes. g denotes the stacked trace and f is the error message returned by a compiler. ω can be "Make no error".Example 6 (RL). The parameters are the policy. g is the trajectory (of states, actions, rewards) resulting from running the policy in a Markov decision process; that is, g documents the graphical model of how an action generated by the policy, applied to the transition dynamics which then returns the observation and reward, etc. f can be the termination signal or a success flag. ω can be "Maximize return" or "Maximize success". Example 7 (Hyperparameter Tuning of ML Pipeline). The parameters are e.g. learning rates and architectures. g describes the stages of the ML pipeline and the evaluation on the validation set, and f is the validation loss. ω can be "Minimize validation error". Example 8 (Prompt Optimization of an LLM Agent). The parameters are the prompt of an LLM workflow. g is the computational graph of the agent and f is the feedback about the agent's behavior (which can be scores or natural language). ω can be "Maximize score" or "Follow the feedback".Example 9 (Multi-Agent Collaboration). The parameters are each agent's prompts. g describes the entire conversation flow between agents, and f is the feedback about whether the task is successful after each agent performs their action. ω can be "A group of agents coordinate to finish a task.".As mentioned, the computational graph g returned by the Trace Oracle T may have different graph structures. The length of the execution trace, e.g., in the debugging example above depends on how far the code executes. Similarly, the rollout length of in the RL problem can be randomly determined.The formulation of the Trace Oracle abstracts the details of a computational workflow, so problems from different domains can be framed in the unified framework. This abstraction allows us to design the computational tool Trace for various applications.It is worth mentioning that execution error can be directly used as feedback to optimize parameters in Trace. When execution error happens within a method decorated by bundle, Trace would adds a special exception node to the global computational graph and throw an TraceExecutionError to stop the computation. The computational graph ends at where the execution error happens. This exception node becomes the new output of the inputs to the decorated method (since the original method raises an error) and is the output of the truncated computational graph. Messages in TraceExecutionError can then be used as the feedback f in OPTO and propagated from the exception node to the parameters. By calling an OPTO optimizer, the parameters can be updated to avoid causing the same execution error. See the exception handling code in Fig. 2 .We find that this error handling mechanism has two convenient usages. First, this allows using Trace to automatically debug issues in the workflow due to incorrect parameter settings. Such errors can happen frequently especially when codes are parameters, as during optimization codes not satisfying syntax or downstream API requirements can happen. The second usage is to enforce constraints the workflow has to satisfy at different stags of computation. With Trace, if an intermediate computed result does not satisfy the constraint, we can simply throw an exception which states the desired constraint. This error signal would be caught by Trace and can then provide early feedback to efficiently improve the parameters, since the graph is truncated at the error. Consider an OPTO problem whose goal is to find a parameter matching a k-digit binary number. The computation checks each digit against a reference number in an arbitrary order. The feedback is either "N th check failed" or "All checks succeeded". Propagated feedback must communicate k bits of information to interpret the feedback correctly; and the minimal subgraph conveys exactly that information. Updating the parameter using the minimal subgraph is trivial, whereas without it there are 2 k possibilities to check.F When is OPTO Efficiently Solvable?We show that OPTO covers a wide range of complicated optimization problems. This shows that if OPTO can be efficiently solved, then many complex workflows can be efficiently optimized. However, the generality of OPTO also raises some fundamental questions, such as if OPTO is well defined and when OPTO can be efficiently solved. These questions stem from its generality of the context ω and the output feedback f in OPTO, since e.g. they can be anything descriable texts. This flexibility makes the scope of OPTO go beyond standard mathematical optimization problems, where a setup has a fixed context ω (e.g., "First-order optimization") and a fixed type of output feedback f (a descent direction). Fully characterize the properties of OPTO, due to its generality, is beyond the scope of this paper and would require years of future research to come. Nonetheless, here we attempt to provide some preliminary answers and point out some research questions.F.1 What is a solution?Classical mathematical optimization problems have a problem definition which itself is the solution concept. For example, in a minimization problem, it is clear we want to find the minimum of an objective function; even for problems as abstract and general as an equilibrium problem, the problem setup clearly states the solution concept of finding a point/set satisfying an equilibrium inequality [48] .One common pattern of these problems is that the solution concept is something that can be described as conditions on feedback that the parameter should satisfy.By contrast, in a OPTO problem (Θ, ω, T ), by varying the context ω, the desired parameter can change from one extreme to another. For example ω may state "Follow the feedback" or "The feedback is adversarial.". Therefore, we need define the solution concept of OPTO differently, rather than just using the feedback. We need to also consider the context ω appropriately. Below we make an attempt to give an axiom of OPTO for its solution to be well defined.There is an verification oracle (a human, a machine learning model, or a polynomial-time algorithm) when given (θ, ω, f ) can verify whether θ is a solution or not.Notice the verification oracle in Axiom 1 is not limited to just algorithms. This is intentional because we currently do not have algorithms that are intelligent enough to process the wide range of contexts and feedback that OPTO allows. Therefore, we include human judgement or the use of LLMs or other AI systems as part of the definition, while acknowledging the impreciseness of the statement due to OPTO's soft computing nature. Lastly we note the verifiability is only defined with respect to the context ω and the output feedback f , not the execution trace g. That is, the verification of a solution depends only on the output of computation.Under Axiom 1, we can start to ask the basic question of whether a solution to an OPTO problem exists or not. There are clearly problems where no solution exists (that is, no parameter in Θ can be verified by the verification oracle). For example if the feedback f is contradicting and yet the context ω is "Follow the feedback.", then there would be no solution that is satisfactory. On the other hand, if ω is "Ignore the feedback", all parameters can be solutions. In the following, we assume solutions of OPTO under consider exist. This assumption would rule out problems, e.g., where the feedback is adversary to the context, and makes solving OPTO is a well-defined search problem. Assumption 1. For an OPTO problem (θ, ω, T ), we assume there is at least a parameter θ ∈ Θ such that it can be verified as a solution by the verification oracle.So far our discussion establishes OPTO as a well-defined search problem, based on qualification on the context ω and the output feedback f . However search problems can be NP-hard. In other words, we know that, without the execution trace, there are search problem instances modeled by some ω and oracle giving f that cannot be efficiently solved. Take RL for a tabular MDP as an example of OPTO problem. Without the execution trace (i.e., not seeing the Markovian structure and trajectories), the problem has an exponential complexity (due to the size of the policy space) and we know by using the execution trace here, tabular RL can be solved approximately in polynomial time [49] . Another example is training of neural networks. Without the execution trace, we have a complex black-box optimization with a loss value, without gradients, whereas an execution trace allows implementation of back-propagation to compute the gradients at the parameters.More broadly speaking, if we consider a "human" as an optimizer for OPTO, we see that (expert) engineers/researchers, when equpped with additional computational tools, can efficiently solve a broad range of OPTO problems (such as by using the execution trace. From these observations, we conjecture using information in the execution trace is the key to unlock efficient OPTO. More precisely, we conjecture that OPTO is efficiently solvable when the context and the trace feedback need to provide information to construct a corrective search direction. For example, when the output feedback back is just a scalar loss, and yet the context + execution trace feedback does not provide enough information to compute a descending direction then OPTO reduces back to a black box problem. (See the problem instance in Appendix E.2). Nonetheless, identifying which subsets of OPTO are efficiently solvable is a big open research question.The MSP extracts the minimal subgrpah of the full computational graph of the workflow. Here we show a visualization using the example in Fig. 4 . [ mul ] This is a multiply operator [ add ] This is an add operator .[ bar ] This is a method that does negative scaling . (b) We create a succinct summary of the computation graph using a language that mimics a program. OptoPrime is an LLM-based optimizer. Its prompt is composed of the following parts. where + denotes concatenation. We list the prompt templates of different components below.You ' re tasked to solve a coding / algorithm problem . You will see the instruction, the code, the documentation of each function used in the code, and the feedback about the execution result .Specifically , a problem will be composed of the following parts :-# Instruction : the instruction which describes the things you need to do or the question you should answer .-# Code : the code defined in the problem .-# Documentation : the documentation of each function used in # Code . The explanation might be incomplete and just contain high -level description . You can use the values in # Others to help infer how those functions work .-# Variables : the input variables that you can change .-# Constraints : the constraints or descriptions of the variables in # Variables .-# Inputs : the values of other inputs to the code, which are not changeable .-# Others : the intermediate values created through the code execution .-# Outputs : the result of the code output .-# Feedback : the feedback about the code's execution result .In # Variables , # Inputs , # Outputs , and # Others , the format is :< data_type > < variable_name > = < value >If < type > is (code), it means < value > is the source code of a python code, which may include docstring and definitions . Beyond the empirical results described in Section 5.5, there are several first principles reasons to prefer the Trace design to build future LLM-powered generative optimizers.Joint vs. Individual Parameter Optimization: When there are multiple parameters to be optimized in a workflow, Trace takes a "joint optimization" view (constructing the minimal subgraph involving all of the parameters described in the initialization of OptoPrime). This flexible design also allows an alternative approach of "individual optimization" akin to co-ordinate descent if desired, i.e., fix all but one parameter and individually optimize each parameter. Such an optimization heuristic can be accomplished, e.g., by instantiating an OptoPrime instance per parameter and iterating through their updates; however, for many problems (like the toy example below), this strategy, which TextGrad employs, can be sub-optimal because a consistent optimizer must additionally maintain state about any suggested updates to all other parameters when reasoning about an individual parameter update.1 @bundle () 2 def function1 ( x ) :3 return x > 0 4 @bundle () 5 def function2 ( y ) : Heterogenous Parameters: Trace encapsulates all Python datatypes as node, including numbers, floats, strings, etc. In contrast, many generative optimization libraries like TextGrad restrict their "variables" (i.e. parameters and intermediate values) to only be strings; consequently running numerical optimization experiments, or optimizing code alongwith prompts and hyper-parameters can be challenging to set up.Modularity: Trace provides a clean separation between the infrastructure (a platform akin to PyTorch) and a generative optimizer (OptoPrime akin to the Adam optimizer). Trace provides an object "TraceGraph" that describes the computational graph. This graph is independent of optimizers and LLMs. Such independent computation graph objects do not exist in TextGrad. We hope the Trace-provided interface can spur the development of many other generative optimizers, like how the OpenAI Gym API allowed the development of several Deep RL algorithms targeting a unified representation of RL problems.Flexible Representation: The bundle functionality of Trace is more flexible in representing computations to an LLM than pre-defined templates used in other optimization frameworks. TextGrad relies on special functions to chain variables together, such as using "FormattedLLMCall", while Trace supports any user-defined functions through decorators like bundle.The strategy taken by OptoPrime can scale better to larger graphs (both in terms of error accumulation across multiple LLM calls, and the costs of multiple LLM calls) compared to requiring one LLM call per computation node (and we see with bundle that not all computation nodes present equal difficulty for an optimizer). The difference in terms of computation costs between the two strategies can be large, especially when not all operators in the graph are as complex as querying LLM. However there may be even better optimizers that blend OptoPrime and TextGrad so as to divide-and-conquer huge computation graphs, we are excited about this avenue for future research.We include token counts for the OPRO and OptoPrime prompts used across all our experiments, at the first iteration of optimization (note that OPRO's token usages grows linearly with iterations). We can see that indeed OptoPrime consumes significantly more tokens than OPRO. However, we observe consistently that even allowing 7-10x more iterations of OPRO so as to equalize token costs, the OPRO performance plateaus to a worse level than OPTOPrime (e.g. Figure 1 : OPRO at Iter 7 vs. OptoPrime at Iter 2; Figure 5b : OPRO at Iter 50 vs. OptoPrime at Iter 5; Figure 6b : OPRO at Iter 30 vs. OptoPrime at Iter 10, etc.). OPRO is suboptimal not due to a token limit but instead a lack of information, which is captured and represented using Trace. We implement a simple battleship game board in Python. The exact code is in the supplement. The game offers a string-based visualization of the board. It randomly places different types of ships on a 2-dimensional board with pre-specified width and height when it initializes. The agent does not see the ship location and has to select a coordinate on the board to hit next. One additional rule of this game is that the agent can go again if their previous coordinate selection (fire) is a hit, not counting as the finish of a turn. In Figure 1 , we ran 10 trials, where in each trial, we ran 20 iterations of training. We measure the reward as % of ship squares hit (over all squares occupied by ships). The reward plateaued at 60% because the game has a chance element (heuristics and strategies can only go so far -strategy is only in effect if a hit happens. Otherwise, there is no information about where ships might be).Any classical numerical optimization problem can be framed as an OPTO problem. Consider h(x) and a target y * , in a context ω finding the y * by changing x; we know the most useful corrective f feedback to change x is the gradient ∇ h x. Similar to Trace, AutoDiff packages like PyTorch's AutoGrad have implemented dynamic graph construction with special classes like torch.Tensor. We want to validate whether it is possible to rely on binary text feedback, a graph automatically constructed by Trace, and OptoPrime to update x in the context of minimizing |y -y * |.We constructed a synthetic task environment where we can create a complex computation graph with arbitrarily complex connections between numerical variables. The focus of this environment is on the complexity of the graph, not on the complexity of the numerical operators. Therefore, we only use one-dimensional input and basic arithmetic operators to create a numerical optimization problem solvable by a first-order optimizer. This environment constructs a computational graph by sampling a number of times. At each time, it will either use a previously computed variable or sample a new variable, and an operation will be sampled to combine them. The optimization task is, for a fixed number of steps, an optimizer needs to output x that minimizes y.We evaluate the following baseline methods. Basic Agent: a basic LLM agent that simply stores past information of (x t-1 , y t-1 ) in context before choosing the next x t . OPRO Agent: a basic LLM agent but we implement the state-of-the-art LLM optimizer OPRO [13] , which updates the meta-prompt of the basic LLM agent. Torch + Adam: the problem we construct is end-to-end differentiable. Therefore, we simply pass in torch.Tensor(x) as input and use Adam optimizer to update. We tune the learning rate slightly and found 1e-1 to work well. We compare two kinds of Trace-based optimizers: Trace, where we allow OptoPrime to read in the entire computation graph before updating x, or Trace Masked, where we hide the computation graph.We run 30 trials over different computation graphs and start all methods with the same initial x, y * . We compute the absolute error, which is |y -y * |. On average, Trace is able to match the best-in-class first-order gradient optimizer Adam [25] . It is not entirely surprising that all the other baselines are performing worse due to a lack of access to the computation graph. To our surprise, OPRO, by only accessing the history of input and output, as well as changing the meta-prompt, is able to eventually discover the correct solution. This confirms why there were early signs of success using LLMs for black-box optimization in a simple plug-and-play style. However, OPRO is not an efficient optimizer because it lacks access to the Trace oracle. We show OPRO struggles even more when the computation graph gets more complex.We tested OptoPrime in a traffic control problem which is an instance of hyper-parameter tuning. We used UXSim [26] to simulate traffic at a four-way intersection, where the tunable parameters are the duration of the green lights for each direction of traffic flow. The feedback is a scalar loss calculated by monitoring the flow of a pseudo-random sequence of vehicles arriving at the intersection over a period of 30 minutes. The loss computes an estimate of the delay experienced by all vehicles due to the intersection, as well as variability in this estimate for every link in the network; lower values are better. The goal of an optimizer is to identify values for all of the green light duration so as to minimize the loss using the fewest number of traffic simulations. If the green light duration for a given traffic flow direction is set too low, then vehicles will queue up over time and experience delays, thereby lowering the score for the intersection. However, if the green light duration for a given direction is set too high, vehicles in other directions will queue up and experience delays, thereby lowering the score for the intersection. Hence an optimizer must find the right trade-off for temporally distributed and variable demands.In Figure 5 we report the performance of a SOTA heuristic from the traffic control literature, SCATS [27] (adapted to this toy setting) as well as two black-box optimization techniques: Gaussian Process Minimization (GP) [8] and Particle Swarm Optimization (PSO) [28] . All methods are initialized to evaluate the same starting parameter. GP and PSO further evaluate 5 random parameters; moreover, if they query a previously evaluated point, that query is replaced by a randomly sampled parameter. GP constructs a surrogate model to mimic the black-box traffic simulation function which maps from parameters to observed score. Then it minimizes a utility function (e.g. the lower confidence bound) using the surrogate model to pick the next parameter to evaluate. PSO on the other hand maintains 5 particles in parameter space, each with a position and velocity. At each iteration of PSO, particles update their positions according to their previous positions and velocity, evaluate the function at the updated positions, and update the velocities of all particles using the observed values.Although GP and PSO are both black-box methods, GP can be thought to replace Trace oracle with instead a smooth differentiable surrogate function; whereas PSO is very different and maintains a candidate set of parameters (can be thought of as conceptually related to OptoPrime with memory).GP appears to be bad because even when it converged, the exploration heuristic randomly samples parameters rather than pick the converged parameter. PSO appears bad because 10 iterations is insufficient for its convergence. Note that given enough number of iterations, black-box approaches will eventually perform well. Trace is quickly competitive with the SCATS heuristic, whereas OPRO is not. Moreover, we find that memory is crucial for OptoPrime to perform well at this task. Finally, Trace consumes additional overhead compared to black-box methods; beyond the space and time complexity for running the traffic simulation, Trace additionally materializes the computation graph per iteration. Thus it can also be more expensive per LLM call compared to OPRO.Perhaps more surprisingly, there are many components that a workflow needs to learn. Some of these components can be the prompt to generate output from an LLM, while other components can be code that needs to further process these outputs. In many workflows today, enabled by LangChain [29] and DSPy [22] , only a small part of this workflow, predominantly, the input to an LLM API call, is optimized. These libraries optimize input to an LLM, and human engineers process that input and integrate it into other systems. Indeed, both libraries can enable robust and swift large-scale engineering efforts to build LLM-based software. However, if our goal is to develop self-adapting agents that can modify their own behavior, we should not ignore one of LLM's greatest strengths: code generation. Trace allows us to unify prompt optimization and code generation, which enables the creation of agents capable of fast learning.In this example of an LLM-based workflow (Figure A.7), there are three parameters that are flagged as trainable for the optimizer: prompt_template, create_prompt, and extract_answer. Note that two of them require the LLM to generate Python code, and one of them requires the LLM to modify a text. Trace abstracted away the different data types and enabled direct update and optimization of them. Furthermore, a human engineer is often tasked with writing an error-free extract_answer.The output of an LLM can be highly stochastic and can often change over time; the code that is used to extract the response of an LLM has to be extremely robust and, therefore, arduous to create. Whenever a major distribution shift happens in the LLM output, this code needs to be rewritten by a human engineer, and it is hard for humans to anticipate all of LLM's output patterns.We set up the task of end-to-end workflow optimization. Unlike a typical LLM benchmark evaluation, where a lot of effort went into creating the perfect evaluate(answer, target) method so that all kinds of LLM outputs were post-processed, cleaned, and formatted to match the ground truth, we choose a simple evaluation function (that extracts a segment or does exact string matching) and place the burden on the workflow itself to figure out how to create the right answer to satisfy the evaluation metric. We choose Big-Bench Hard [21] as our task because it has 23 subtasks and contains both language and algorithmic tasks.We numbers. Even though DSPy's meta-prompt optimization is trained on each task individually, the output of LLM to the evaluation method is still not post-processed, resulting in low performances of these tasks. However, Trace can optimize code and LLM prompt jointly to successfully deliver the response expected by an automatic evaluation method.We test the ability of Trace to optimize long-horizon workflows with complex dependencies. We experiment with using Trace to train controller (python code) for a simulated Sawyer robot manipulator. We use the Meta-World environment of LLF-Bench [23] as the simulator and consider three tasks reach, pick-place and push. LLF-Bench is a simulated benchmark with gym interface for testing an agent's ability to learn from language feedback. In these LLF-Bench Meta-World tasks, the observation is a dictionary where each field denotes a feature of the state and has a vector value (e.g., the end-effector position, the goal position, the gripper status, etc.). The keys of the observation dictionary can differ for each task. The action space is 4-dimensional, which controls the relative position of the end-effector and the state of the gripper. In each time step, the LLF-Bench Meta-World simulator returns the observation dictionary and natural language feedback to guide the robot (we use the 'a' mode of LLF-Bench, with which the language feedback would contain information about the current performance, explanation of past successes and failures, and suggestions for the next step). An episode ends if the robot successfully solve the problem or because of time-out. For each task, LLF-Bench also provides a task instruction explaining that the task is about controlling a Sawyer robot arm and the meaning of the action space (see [23] ). We use that as the context ω of the OPTO problem. We consider an episodic setting. For each experiment (a random seed), we randomly sample an initial configuration. Then for each iteration of optimization, we reset the simulator to that sampled initial configuration and run the robot policy for 1011 steps or until the episode termination due to success. We compute the sum of rewards and gives the output feedback f in texts in the format of "Success: <true/false> Return: <score>". Note that the initial condition for all iterations within an experiment is the same so that the optimization problem is deterministic. To evaluate the learned policy's performance, for each experiment, we additionally run the learned policy starting from 10 held-out initial conditions, different from the fixed training initial condition. For each training algorithm discussed, we run it with 30 iterations, where each iteration consists of one episode rollout and one update.To optimize the controller with Trace, we declare the control code as the parameter using the bundle decorator with trainable set to True; the initial control code simply outputs a zero vector [0, 0, 0, 0]. We decorate also the reset and the step function of the gym environment, so that the entire rollout of an episode can be traced end-to-end. In our implementation, a prototypical rollout would create a graph with around 30 operations where the controller code parameter is used multiple times. This graph structure is similar to that of running a recurrent neural network. For Trace, we experiment with OptoPrime with and without a memory of size 10. In addition to Trace, we implement the state-of-the-art LLM optimizer OPRO [13] as a baseline. Compared with Trace, OPRO does not use the execution trace information but rely on just memory of parameter and feedback pairs 12 To make the comparison with Trace fair, we append the feedback observation from LLF-Bench given at each time step to the final feedback received by OPRO; on the other hand, Trace uses the simple final feedback of success and return and has to read the per-step feedback from the execution trace. To run OPRO in the OPTO setting, our implementation only proposes a single candidate in each iteration, which is then evaluated and provided with the output feedback. Since in [13] OPRO generates about 10 samples per iteration, so one iteration in [13] is roughly equivalent to 10 iterations here.The experimental results are summarized in Fig. A .8, where we show the success rates at both the training initial condition as well as the held-out testing initial conditions over 10 seeds. OptoPrime is clearly the top-performing optimizer, especially the version with memory. OPRO is able to solve Reach at the start but its performance degraded over iterations (this instability was observed in [13] ) and gets similar performance as OptoPrime (without memory) in Push. To validate that the performance of OptoPrime is indeed due to using the execution trace, we include an ablation where we mask out information in #Inputs, #Others, #Code, #Definition in the LLM context (see Fig.A.1b),which lead to significant degrade in performance and stability. This ablation shows that additionally using the execution trace provides more informed search direction compared with just using just the output feedback, which agrees with our hypothesis. This experiment features the most complex graph structures, and using Trace for optimization here is similar to back-propagation over time. The experimental results here are quite impressive, showing that Trace is able to learn a complex control logic in a dozens of interactions, not only working on the training initial conditions but also on the held-out testing ones too. Nonetheless, we want to point out some limitations in the current experimental results. We find that the success rate of the learned policy varies largely across random seeds. Except for Reach (the simplest task), in a seed, often either it finds a policy close to 1.0 success rate or 0.0 success rate. Therefore, the plots can roughly be interpreted as how long it takes to find a working policy. In addition, in these experiments, we find that providing task-related context is necessary. We find the context needs to be informative enough for humans to understand the problem 13 ; otherwise, the optimization can be solved efficiently with the time scale considered here. Nonetheless, this requirement is reasonable, as there is no free lunch. Use code to process the response and find where the answer is .Can use self . call_llm (" Return the answer from this text : " + response ) again to refine the answer if necessary .prompt_template : The prompt that was used to query LLM to get the response question : Question has a text describing the question but also " Options " response : LLM returned a string response Process it and return the answer in the exact format that the evaluator wants to see .Be mindful of the type of answer you need to produce .It can be ( A ) /( B ) , a number like 8 , or a string , or Yes / No .""" answer = response . split ( " Answer : " ) [1] . strip () return answer ## Iteration > 0 def extract_answer ( self , response ) :""" Need to read in the response , which can contain additional thought , deliberation and an answer .Use code to process the response and find where the answer is . Uses code to meticulously parse the response to locate the answer section .Optionally refines the answer by querying the LLM again if necessary .response : string from LLM , expected format contains ' Answer : ' followed by the answer . question : Description of the question being addressed , may include ' Options '""" answer = response . split ( ' Answer : ') [ ## Iteration 0 ( initialization ) prompt_template = """ Given the fields ' question ' , produce the fields ' answer '.---Follow the following format .---Question : {} Answer : """ ## Iteration > 0 prompt_template = """ Given the fields ' question ' , produce the fields ' answer '.---Follow the following format .Question : Answer : Please respond precisely in the format required .---Question : {} """ " reasoning ": " The task involves controlling a robotic arm to perform a pick -place operation by interpreting observations and generating appropriate control actions (4 D vectors ) to move the arm towards the goal . The objective is to adjust the controller ' s code ( __code0 ) based on feedback received at each step of the process . The feedback within # Others suggests whether the movements made by the arm are leading it closer to or farther from the goal , along with advice on what coordinates to aim for next ( e . g . , [ -0.07 0.68 0.12 0. ]) . The controller ' s code currently returns a static action [0 , 0 , 0 , 0] , which does not adjust based on observations and hence does not efficiently move the robotic arm towards the goal . Given the repetitive feedback suggesting to go to [ -0.07 0.68 0.12 0. ] , it appears that the controller needs to dynamically compute actions based on the difference between the current positions ( hand_pos , puck_pos ) and the goal_pos , along with possibly manipulating the gripper state for pick and place operation ." , " answer ": "" , " suggestion ": {...} } # Iteration 5 {" reasoning ": " The instruction involves creating a controller for a Sawyer robot arm to solve a pick -and -place task by providing sequential control actions based on observations of the environment and state of the robot . The feedback indicates that despite the actions taken , the success criteria have not been met , with the latest attempt achieving a return of 0.9574427394250022. The code provided for __code0 defines a controller function using the observed state to compute control actions aimed at grasping a puck and moving it to a goal position while adjusting the gripper based on its distance from the puck . Feedback throughout the execution indicates that the controller adapts its strategy based on proximity to the puck and whether it has been picked up . Feedback also suggests that the gripper ' s opening and closing actions are central to the task ' s success , with corrections being made to close the gripper when necessary . The controller ' s logic for determining movements toward the puck or goal and adjusting the gripper state appears reasonable , but the feedback implies that adjustments might be required to ensure task completion . This could include fine -tuning the conditions under which the gripper is closed or adjusting the thresholds for puck pickup and delivery to the goal ." , " answer ": "" , " suggestion ": {} # Iteration 10 { " reasoning ": " The task involves programming a robot arm to pick up a puck and place it at a goal position by providing sequential 4 -dimensional action vectors . The action vectors represent target positions and gripper states ( open or close ) based on the robot ' s and puck ' s current positions . The feedback indicates that despite optimizing the controller logic to close the gripper when near the puck , move towards the puck when it ' s not grasped , and head towards the goal once the puck is grasped , the actions have not resulted in task completion . Specifically , repeated suggestions to close the gripper ( indicating it remained open when it should have been closed ) and guidance on where to move next suggest that the action selection logic might not be effectively responding to the environment ' s state or might not be properly engaging the gripper at the right times . The latest code includes a condition to normalize the direction of movement and checks for gripper position , but there seems to be a misunderstanding or oversight in adjusting the robot ' s behavior based on the scenario to successfully complete the pick -place task ." , " answer ": "" , " suggestion ": {...} } } " reasoning ": " In this task , the controller function for the Sawyer robot arm manipulates the robot ' s actions based on the observed environment to achieve a pick -and -place task .The function correctly defines actions based on the robot ' s hand position , the puck ' s position , the goal position , and the gripper ' s state . However , the feedback after each step in the documentation indicates that although the robot is making progress towards the goal , there is an issue with the gripper ' s action state not being switched to ' close ' (1) at the crucial moment when the robot is close enough to the puck . Most feedback comments suggest that the robot is moving in the right direction but fails to grasp the puck due to the gripper not closing . This inefficiency appears to arise from the conditions set for closing the gripper and possibly the normalized direction towards the goal after the puck has been grasped . The condition for closing the gripper ( when close enough to the puck ) seems to be met , but the action selected afterwards does not effectively lead to puck manipulation or goal achievement in the most optimized manner , as indicated by the ' Feedback ' messages suggesting movements that imply needing adjustments in action selection ." , " answer ": "" , " suggestion ": {...} } # Iteration 13 ( Success ) LLM response : { " reasoning ": " The instruction requests improvements based on feedback for a Sawyer robot arm ' s controller function , aiming to accomplish a pick -and -place task . The controller function accepts observations about the robot and the environment and computes an action to move the robot accordingly . \ n \ nThe feedback provided indicates the task was successfully completed with a suggestion that the final command should have been '[0.04 , 0.88 , 0.24 , 1] ' , which probably indicates the precise position and gripper state ( closed ) needed to achieve the task goal . The feedback suggests the need for fine -tuning actions especially towards the end of the task to ensure the robot ' s gripper is appropriately positioned and the gripper state is correct for completing the task . Since the feedback for the current code positively notes ' Success : True ' with a high return value , no modifications are recommended for ' __code0 '. \ n \ nPrevious feedback instances with lower success metrics and recommendations for change highlight the iterative improvements to the controller logic , specifically related to grip control and movement precision . The aim was to adjust the controller ' s response to observation inputs to accurately position the gripper and handle the puck . The current code evidently manages these aspects well by normalizing directions and managing the gripper ' s open -close state effectively under varying circumstances ( proximity to puck or goal and whether the puck is already grasped ) . These adjustments have likely led to the successful output in the described problem ." , " answer ": " No changes are recommended as the feedback indicates successful task completion and suggests the controller ' s current logic adequately addresses the task ' s requirements ." , " suggestion ": {} } • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.Justification: The contributions of this paper (a framework analogous to PyTorch to generalize the backpropagation algorithm) do not have direct safety or security implications.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).We use a new general-purpose LLM-based optimizer OptoPrime that we detail in Section 4.A multi-output operator can always be modeled by a single-output operator and single-output indexers.The current implementation is in Python. The DAG-based design of Trace can be adapted to other programming languages also, just like AutoDiff libraries that are available in C++.This would add an extra parent (i.e., the trainable code) to the computed child node.The numbers in the original paper cannot be reproduced exactly despite using the released TextGrad code.All algorithms in Trace see the same graph, which is slightly different from the graph of TextGrad.We set the problem horizon to be 10 steps, as we find the expert policies implemented in LLF-Bench can solve these problems within 10 steps.The original version of OPRO uses parameter-score pairs. Since we're interested in the more general setup of OPTO, we extend it to use parameter-feedback pairs.The original instructions in the v2 environments of LLF-Bench does not contain task specific background, but only the task name. We find this task name alone, except for Reach, is too vague even for humans to understand how to read the observation dictionary is related to the problem, so we added additionally one line to explain the task. Pick-place: "The goal of the task is to pick up a puck and put it to a goal position.". Push: "The goal of the task is to push a puck to a goal position.". Without this extra information, none of the LLM-based optimizers works in the experiments. 


Decision: 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. As a case study, this paper also developed a general-purpose optimizer, OptoPrime that can effectively solve OPTO problems,such as  first-order numerical optimization, prompt optimization, hyper-parameter tuning, robot controller design, code debugging, etc.,

The limitations of the works lie in the scalability such as the context length and graph size.  Despite some concerns about scalability and the need for clearer experimental results, the strengths of the paper—especially its novelty, versatility, and demonstrated performance—outweigh these weaknesses. Therefore, I recommend acceptance, with the suggestion that the authors consider revising the narrative structure and enhancing the clarity of experimental results in the final version.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: As someone outside of this field, I find this paper to be quite impressive.  I did not identify any specific weaknesses.
questions: questions: 1. Could Trace be adapted for use with other programming languages, and what would such an adaptation entail in terms of architectural changes?
limitations: limitations: Yes
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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.
limitations: limitations: Yes
paper_summary: paper_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).
strength_weakness: strength_weakness: [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: 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.
limitations: 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.

Comments: comment: Thank you for addressing the weaknesses and questions. I intend to keep my score.
Thank you for your detailed response and congratulations on your excellent work!!