Thank you for your insightful comments and questions.

Comparison with TextGrad

Please note that we submitted Trace to NeurIPS on May 15 2024, and TextGrad was uploaded to arXiv and Github on June 11 2024. So a direct comparison was impossible.

Narrative

We considered the narrative you suggested and decided that we needed a very strong motivation for an entirely new approach to optimization (Section 1.3). We found that showing the Trace software in action first was the fastest way to demonstrate that this kind of optimization (OPTO) is feasible, practical and powerful.

Limitation on Scalability

We agree that scalability is a limitation of OptoPrime, as discussed in Sec 6 and Sec 7. We wish to clarify that this limitation pertains specifically to OptoPrime and not to Trace itself. Trace is designed to scale efficiently to large graphs with minimal overhead (see line 242) and can handle non-textual nodes. We are inspired by the historical development of back-propagation, starting from optimizing networks with 10s of neurons (in the Rumelhart et al paper in 1986) to billions today. We anticipate better and scalable algorithms for OPTO developed in the future.

Definition

We will include a clear definition of an execution trace. The execution trace is defined as the sequence of operations and their execution results invoked when computing the output from a set of inputs. This execution trace can be expressed as a computational graph defined in Preliminary. We will clarify that the DAG g is the computational graph defined above.

Analyzing learned prompts

We will highlight the learned prompts in the revised Sec 5. One lesson from several automatic prompt engineering works is that human intuition is not a good guide for prompt engineering. Seemingly innocuous (and perhaps unreasonable) changes to a prompt template (like that in Figure A.12) can have large effects on an LLM’s behavior, thus motivating the need for their automatic optimization e.g. via Trace.

Bugs in Running DSpy baseline

Thank you for running our code. Unfortunately, the code uploaded to OpenReview included a stale version of the DSPy script. We apologize for the oversight. We provide a snippet of our correct setup in the one-page pdf. We will update the supplementary material.

Optimizing stateful agents

Optimizing stateful agents is similar to optimizing recurrent neural networks; we need to explicitly represent the state as inputs to the learned functions (discussed in line 193). Our experiments on Meta-World and Battleship are examples of stateful problems. In these problems, the environment is stateful and that state is returned as input to the functions that Trace’s learning; these graphs are similar to that of optimizing an agent with an internal state. We hope these experiments are sufficient to address the reviewer’s concern.

Battleship game

We independently found this error after the paper submission. The exact function definition used in the experiments is given in Appendix H, which is consistent with the code we submitted. You are right that Figure 2a over-simplifies important details about the function docstring and the current text in line 79 is misleading. We will change it to “Remarkably, there is no mention of Battleship environment APIs nor details on how the functions reason and act should behave or adapt in Fig. 2a”.

We agree that GPT4 likely has trained on code about the battleship game. But it does not know the API of the Battleship game (because we coded it up from scratch). If GPT4 did know how to solve the problem we presented, it could have solved it in the first iteration after one update, but that is not what we observed (Figure 1). This gap between the performance after one update and multiple updates indicate the sign and need of learning from feedback and interactions. We will use “excerpt” in the revision to clarify the code snippets in Figure 2.

LLM Prompt Optimization Example

An example of LLM agent as per line 767 can be found in a new experiment that we conducted for the virtualhome environment. This experiment is included in the one-page PDF for the rebuttal, including the code and figure. Virtualhome is a collaborative environment that requires two agents to work together to solve household tasks. Trace is asked to optimize and update a specific part of the prompt, which is the plan for future actions. Prior work (Guo et al., "Embodied LLM Agents Learn to Cooperate in Organized Teams", 2024) forces agents to have a round of conversation before they start the task. We show that Trace allows agents to have naturally emerging pro-social behaviors for some tasks (such as “putting plates into the dishwasher”), but not others (such as “reading a book”).

Graphs in Fig 3

Figure 3 is an illustration of a general OPTO problem setup. In the Battleship example, you are correct that the graph is the same in every iteration. But, for a general optimization problem, the graph structure can be different, e.g. different parameters changing program flow during the forward pass, or simply because the execution is stochastic. Meta-World is an example where the graph structure can be different across iterations. The graph is a chain describing the multi-step interactions with the environment, and each episode ends either when the robot successfully solves the problem or when timeout happens. Therefore, iterations with successful episodes can have a shorter chain than those that fail and time out.

Token Cost

The cost depends on the graph size and the tokens required to describe the problem. Running OptoPrime with GPT-4-Turbo for the most complex experiment in the paper, MetaWorld, costs <$30 USD for one task (over 10 seeds, 30 iterations). Costs of other experiments are a fraction of this.

Technical Details

We currently require Python>=3.8, and we will update the setup.py to add the missing dependencies. Thank you for your feedback and running the code!

Thank you for your insightful comments and questions.

Limited Scalability

We agree that scalability is a limitation of OptoPrime, as discussed in Section 6 (Limitations) and Section 7 (Conclusion) where we note the current focus on textualizable problems. However, we wish to clarify that this limitation pertains specifically to OptoPrime and not to Trace itself. Trace is designed to scale to large graphs with minimal overhead (see line 242) and can handle non-textual nodes. OptoPrime's limitation arises from converting the propagated Trace graph into a single query that fits in the context window for a current-gen LLM. As its name suggests, OptoPrime is a first step in solving OPTO problems. We are inspired by the historical development of back-propagation, starting from optimizing networks with 10s of neurons (in the Rumelhart et al paper in 1986) to billions today. We anticipate better and scalable algorithms for OPTO developed in the future, such as using multi-agent workflows to scale to large graphs and VLMs to interpret multi-modal parameters and feedback.

Lacking an automated method

We wish to clarify that Trace constructs the computational graph automatically (line 183) while the program executes (dynamically like in PyTorch), rather than requiring users to pre-define it or pre-compile it (statically like in Tensorflow 1). This means that the created graph changes automatically based on the workflow process, with different inputs potentially resulting in different graph structures. In Section 3.1, we discuss how to abstract the workflow using @bundle operators. This abstraction is different from manually defining the graph, which may be the source of the confusion. The @bundle decorator is an optional feature that allows users to simplify the workflow’s text representation (e.g. to reduce input token costs), enabling LLMs to better understand and optimize the workflow.

Limited Experimental Improvement

We respectfully disagree on the comment on the performance gap between OptoPrime and OPRO. In almost all experiments (Figure 1, 5b, 6a, 6b), OptoPrime is a significant improvement over OPRO, e.g. 2x-4x improvement in success or rewards. Only the experiment in Figure 6c shows the different algorithms’ performance is within the error margin.

Difference between Node and Graph Optimization

Consider the following example that shows why optimizing over a graph is better than optimizing only individual nodes.

@bundle()
def function1(x):
    return x > 0
@bundle()
def function2(y):
    return y % 2 == 0
def xor_test(x, y):
    return function1(x).neq(function2(y))
input1 = node(3, trainable=True); input2 = node(4, trainable=True)
xor_test(input1, input2).backward(feedback=”Find a set of inputs to make the return True.”)

When we optimize an individual node, we only see that node input, the function output, and the feedback. We do not see the other inputs and how they can affect the outcome. When we optimize for the full graph however, we can see that each input only partially affects the outcome, and we need to jointly optimize both inputs to achieve a desired outcome.

Discrepancy between Fig 5 and Fig 6’s Conclusions

In all the experiments, OptoPrime with memory (denoted as Trace) performs better than OptoPrime without memory (denoted as Trace NoMem) and OptoPrime with memory but with the execution trace info removed from the prompt (denoted as Trace Masked). We noticed that the current writing is not clear about what each method means, which may cause confusion. Across the different ablations and Figure 5 and 6, we consistently see that memory improves performance, and masking the execution trace information hurts performance. We will better clarify the ablations in the revision.

Token Efficiency

Thank you for the excellent suggestion, we will include token counts for the OPRO and OptoPrime prompts in the paper. Here are the statistics for the prompts at the first iteration of optimization (note OPRO's token usages grows 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.

Thank you for your insightful comments and question.

Yes, in principle, Trace can indeed be adapted for use with other programming languages. The core design of Trace is based on the primitives node and @bundle, which define the nodes and operators, respectively, for the directed acyclic graph (DAG) abstraction of the traced computational process. Once the DAG is created in any programming environment, the algorithms used in Trace can be applied.

In our current implementation, we overloaded Python’s magic methods to seamlessly integrate with existing Python code. This approach may not be feasible in other programming languages due to the limitations of operator overloading, which could result in a less clean interface. Nonetheless, by building a set of operators using the idea of @bundle, we can create DAGs to abstract the computational process in other languages. One example demonstrating the feasibility of such an adaptation is the C++ versions of AutoDiff libraries like PyTorch (which are also based on DAGs).

Therefore, we believe that the DAG-based design that Trace employs can be effectively adapted to other programming languages. However, we acknowledge that developing Trace libraries for other languages can require non-trivial engineering. We hope that the impressive results demonstrated in our paper will inspire future development to adapt Trace to a broader range of programming environments.