Trace is a generative optimization framework that trains AI systems end to end with general feedback: numerical rewards or losses, natural-language critique, or compiler errors. It generalizes back-propagation by capturing an AI system's execution trace as a graph and passing that trace to an LLM-based optimizer, which updates the trainable parts in place.
It is a PyTorch-like Python library. You write the agent as ordinary Python, mark the parts that may change with two primitives (node for values and bundle for functions), and call backward and step on an optimizer, the same way you would train a neural network. Prompts, code, numbers, and objects can all be trainable parameters. The library accompanies the Trace paper (NeurIPS 2024) and includes the code to reproduce its experiments.
| Optimizer approach | Computation graph | Code as functions | Library support | Large graphs |
|---|---|---|---|---|
| OPRO | No | No | No | Yes |
| TextGrad | Yes | No | Yes | Yes |
| Trace (OPRO, OptoPrime, TextGrad) | Yes | Yes | Yes | Yes |
OptoPrime, the project's own algorithm, places the whole computation graph in context, which can limit it on graphs with more than a few hundred operations. TextGrad runs about 2–3 times slower than OptoPrime. LLM calls go through LiteLLM by default, with AutoGen kept as an alternative backend. The project labels itself a research-purpose beta.
Features
- Two primitives: node wraps values and bundle wraps Python functions so their execution is traced and optionally trainable
- Agent models: a model decorator turns a Python class into an agent whose methods and prompts can be optimized
- Optimizers: OptoPrime, OPRO, and TextGrad share one interface and can be swapped without changing the agent
- General feedback: scores, natural-language text, and execution errors all serve as the training signal
- Graph visualization: built-in rendering of the traced computation graph during backward
- Custom optimizers: new algorithms can work directly on the TraceGraph abstraction
- LLM backends: LiteLLM by default, with AutoGen as an alternative for caching and key management
- Worked examples: tutorials cover code optimization, self-improving agents, multi-agent collaboration in VirtualHome, BigBench-Hard prompt optimization, and robotic arm control
